From 59169b55b58870cd0a1ee8316585c91522536beb Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 14 Jul 2025 11:45:41 +0200 Subject: [PATCH 001/292] add (the idea) of a new Variables class --- src/struphy/models/base.py | 20 +++++++++++++++++++- src/struphy/models/variables.py | 10 ++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) create mode 100644 src/struphy/models/variables.py diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 221e5247f..7f1a00819 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1686,6 +1686,8 @@ def _init_variable_dicts(self): Initialize em-fields, fluid and kinetic dictionaries for information on the model variables. """ + # from struphy.models.variables import Variable + # electromagnetic fields, fluid and/or kinetic species self._em_fields = {} self._fluid = {} @@ -1710,17 +1712,33 @@ def _init_variable_dicts(self): # initial conditions if "background" in self.params["em_fields"]: + # background= self.params["em_fields"]["background"].get(var_name) self._em_fields[var_name]["background"] = self.params["em_fields"]["background"].get(var_name) + # else: + # background = None + if "perturbation" in self.params["em_fields"]: + # perturbation = self.params["em_fields"]["perturbation"].get(var_name) self._em_fields[var_name]["perturbation"] = self.params["em_fields"]["perturbation"].get(var_name) + # else: + # perturbation = None # which components to save if "save_data" in self.params["em_fields"]: + # save_data = self.params["em_fields"]["save_data"]["comps"][var_name] self._em_fields[var_name]["save_data"] = self.params["em_fields"]["save_data"]["comps"][var_name] else: self._em_fields[var_name]["save_data"] = True - + # save_data = True + + # self._em_fields[var_name] = Variable(name=var_name, + # space=space, + # background=background, + # perturbation=perturbation, + # save_data=save_data,) + # overall parameters + # print(f'{self._em_fields = }') self._em_fields["params"] = self.params["em_fields"] for var_name, space in self.species()["fluid"].items(): diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py new file mode 100644 index 000000000..15606e76c --- /dev/null +++ b/src/struphy/models/variables.py @@ -0,0 +1,10 @@ +class Variable: + def __init__(self, + name : str, + space : str, + background, + perturbation, + save_data + ): + self._name = name + self._space = space \ No newline at end of file From 1995b6cd2ab76075581e181823ce392a8b566cc3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 14 Jul 2025 11:47:55 +0200 Subject: [PATCH 002/292] add new folder topology/ with file grids.py --- src/struphy/topology/__init__.py | 0 src/struphy/topology/grids.py | 112 +++++++++++++++++++++++++++++++ 2 files changed, 112 insertions(+) create mode 100644 src/struphy/topology/__init__.py create mode 100644 src/struphy/topology/grids.py diff --git a/src/struphy/topology/__init__.py b/src/struphy/topology/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/struphy/topology/grids.py b/src/struphy/topology/grids.py new file mode 100644 index 000000000..fb2798dc2 --- /dev/null +++ b/src/struphy/topology/grids.py @@ -0,0 +1,112 @@ +import numpy as np + + +class TensorProductGrid: + """Grid as a tensor product of 1d grids. + + Parameters + ---------- + Nel : tuple[int] + Number of elements in each direction. + + p : tuple[int] + Spline degree in each direction. + + spl_kind : tuple[bool] + Kind of spline in each direction (True=periodic, False=clamped). + + dirichlet_bc : tuple[tuple[bool]] + Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). + + nquads : tuple[int] + Number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials). + + nq_pr : tuple[int] + Number of Gauss-Legendre quadrature points in each direction for geometric projectors (default = p+1, leads to exact integration of degree 2p+1 polynomials). + + mpi_dims_mask: Tuple of bool + True if the dimension is to be used in the domain decomposition (=default for each dimension). + If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. + """ + + def __init__(self, + Nel: tuple, + p: tuple, + spl_kind: tuple, + *, + dirichlet_bc: tuple = None, + nquads: tuple = None, + nq_pr: tuple = None, + mpi_dims_mask: tuple = None,): + + assert len(Nel) == len(p) == len(spl_kind) + + self._Nel = Nel + self._p = p + self._spl_kind = spl_kind + + # boundary conditions at eta=0 and eta=1 in each direction (None for periodic, 'd' for homogeneous Dirichlet) + if dirichlet_bc is not None: + assert len(dirichlet_bc) == len(Nel) + # make sure that boundary conditions are compatible with spline space + assert np.all([bc == [False, False] for i, bc in enumerate(dirichlet_bc) if spl_kind[i]]) + + self._dirichlet_bc = dirichlet_bc + + # default p: exact integration of degree 2p+1 polynomials + if nquads is None: + self._nquads = tuple([pi + 1 for pi in p]) + else: + assert len(nquads) == len(Nel) + self._nquads = nquads + + # default p + 1 : exact integration of degree 2p+1 polynomials + if nq_pr is None: + self._nq_pr = tuple([pi + 1 for pi in p]) + else: + assert len(nq_pr) == len(Nel) + self._nq_pr = nq_pr + + # mpi domain decomposition directions + if mpi_dims_mask is None: + self._mpi_dims_mask = (True,)*len(Nel) + else: + assert len(mpi_dims_mask) == len(Nel) + self._mpi_dims_mask = mpi_dims_mask + + @property + def Nel(self): + """Tuple of number of elements (=cells) in each direction.""" + return self._Nel + + @property + def p(self): + """Tuple of B-spline degrees in each direction.""" + return self._p + + @property + def spl_kind(self): + """Tuple of spline type (periodic=True or clamped=False) in each direction.""" + return self._spl_kind + + @property + def dirichlet_bc(self): + """None, or Tuple of boundary conditions in each direction. + Each entry is a list with two entries (left and right boundary), "d" (hom. Dirichlet) or None (periodic). + """ + return self._dirichlet_bc + + @property + def nquads(self): + """Tuple of number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials).""" + return self._nquads + + @property + def nq_pr(self): + """Tuple of number of Gauss-Legendre quadrature points in histopolation (default = p + 1) in each direction.""" + return self._nq_pr + + @property + def mpi_dims_mask(self): + """Tuple of bool; whether to use direction in domain decomposition.""" + return self._mpi_dims_mask \ No newline at end of file From b2e8de1f9bf7884793ef59c1dbb55cbbe6e16b11 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 14 Jul 2025 15:46:46 +0200 Subject: [PATCH 003/292] also search for .py files when scanning for parameters files --- src/struphy/console/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index 66d448f16..c14ca54aa 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -227,7 +227,7 @@ def struphy(): def get_params_files(i_path): if os.path.exists(i_path) and os.path.isdir(i_path): - params_files = recursive_get_files(i_path) + params_files = recursive_get_files(i_path, contains=(".yml", ".yaml", ".py")) else: print("Path to input files missing! Set it with `struphy --set-i PATH`") params_files = [] From aaab3bb54210f4ac6ed5b59428298025f3ef90b5 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 14 Jul 2025 15:56:13 +0200 Subject: [PATCH 004/292] new files struphy/io/inp/parameters.py and strupyh/io/inp/options.py containing the classes for the new input handling --- src/struphy/io/options.py | 62 +++++++++++++++++++++++++ src/struphy/io/parameters.py | 88 ++++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+) create mode 100644 src/struphy/io/options.py create mode 100644 src/struphy/io/parameters.py diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py new file mode 100644 index 000000000..51cc8c36f --- /dev/null +++ b/src/struphy/io/options.py @@ -0,0 +1,62 @@ +from dataclasses import dataclass +from typing import Literal, get_args + +# needed for import in StruphyParameters +from struphy.geometry import domains +from struphy.topology import grids +from struphy.fields_background import equils + +SplitAlgos = Literal["LieTrotter", "Strang"] +PolarRegularity = Literal[-1, 1] + + +@dataclass +class Units: + """... + + Parameters + ---------- + x : float + Unit of length in m. + """ + x: float = 1.0 + B: float = 1.0 + n: float = 1.0 + kBT: float = None + + +@dataclass +class Time: + """... + + Parameters + ---------- + x : float + Unit of length in m. + """ + + dt: float = 1.0 + Tend: float = 1.0 + split_algo: SplitAlgos = "LieTrotter" + + def __post_init__(self): + options = get_args(SplitAlgos) + assert self.split_algo in options, f"'{self.split_algo}' is not in {options}" + + +@dataclass +class DerhamOptions: + """... + + Parameters + ---------- + x : float + Unit of length in m. + """ + + polar_ck: int = -1 + local_projectors: bool = False + + def __post_init__(self): + options = get_args(PolarRegularity) + assert self.polar_ck in options, f"'{self.polar_ck}' is not in {options}" \ No newline at end of file diff --git a/src/struphy/io/parameters.py b/src/struphy/io/parameters.py new file mode 100644 index 000000000..dfadd9931 --- /dev/null +++ b/src/struphy/io/parameters.py @@ -0,0 +1,88 @@ +from struphy.geometry.base import Domain +from struphy.topology import grids +from struphy.fields_background.base import FluidEquilibrium +from struphy.io import options + + +class StruphyParameters(): + """Wrapper around Struphy simulation parameters.""" + + def __init__(self, + model: str = None, + domain: Domain = None, + grid: grids.TensorProductGrid = None, + equil: FluidEquilibrium = None, + units: options.Units = None, + time: options.Time = None, + derham: options.DerhamOptions = None, + em_fields=None, + fluid=None, + kinetic=None, + diagnostic_fields=None, + ): + + self._options = options + + self._model = model + self._domain = domain + self._grid = grid + self._equil = equil + self._units = units + self._time = time + self._derham = derham + self._em_fields = em_fields + self._fluid = fluid + self._kinetic = kinetic + self._diagnostic_fields = diagnostic_fields + + ## possible inputs + + @property + def options(self): + return self._options + + ## parameters + + @property + def model(self): + return self._model + + @property + def domain(self): + return self._domain + + @property + def grid(self): + return self._grid + + @property + def equil(self): + return self._equil + + @property + def units(self): + return self._units + + @property + def time(self): + return self._time + + @property + def derham(self): + return self._derham + + @property + def em_fields(self): + return self._em_fields + + @property + def fluid(self): + return self._fluid + + @property + def kinetic(self): + return self._kinetic + + @property + def diagnostic_fields(self): + return self._diagnostic_fields \ No newline at end of file From b4872e79dba88ba5a73ce51ff5310427eb84f8f7 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 14 Jul 2025 15:59:16 +0200 Subject: [PATCH 005/292] use new class StruphyParameters to set up a Maxwell simulation --- src/struphy/console/main.py | 2 +- src/struphy/io/inp/params_Maxwell.py | 55 +++++++++ src/struphy/io/setup.py | 162 +++++++++++++++++++-------- src/struphy/main.py | 29 ++--- src/struphy/models/base.py | 84 +++++++------- src/struphy/models/toy.py | 4 +- 6 files changed, 231 insertions(+), 105 deletions(-) create mode 100644 src/struphy/io/inp/params_Maxwell.py diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index c14ca54aa..51c0e3c59 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -16,7 +16,7 @@ # struphy path import struphy -import struphy.utils.utils as utils +from struphy.utils import utils libpath = struphy.__path__[0] __version__ = importlib.metadata.version("struphy") diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py new file mode 100644 index 000000000..a739764ba --- /dev/null +++ b/src/struphy/io/inp/params_Maxwell.py @@ -0,0 +1,55 @@ +from struphy.io import options + +# model +model = "Maxwell" + +# units +units = options.Units(x=1.0, + B=1.0, + n=1.0, + kBT=1.0,) + +# time +time = options.Time(split_algo="LieTrotter") + +# geometry +domain = options.domains.Cuboid() + +# grid +grid = options.grids.TensorProductGrid( + Nel=(12, 14, 1), + p=(2, 3, 1), + spl_kind=(False, True, True), +) + +# derham options +derham = options.DerhamOptions() + +# fluid equilibrium +equil = options.equils.HomogenSlab() + +# FOR NOW: initial conditions and options +em_fields = {} +em_fields["background"] = {} +em_fields["perturbation"] = {} +em_fields["options"] = {} + +em_fields["background"]["e_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} +em_fields["background"]["b_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} + +em_fields["perturbation"]["e_field"] = {} +em_fields["perturbation"]["b_field"] = {} +em_fields["perturbation"]["e_field"]["TorusModesCos"] = {"given_in_basis": [None, "v", None], + "ms": [[None], [1, 3], [None]]} +em_fields["perturbation"]["b_field"]["TorusModesCos"] = {"given_in_basis": [None, "v", None], + "ms": [[None], [1, 3], [None]]} + +solver = {"type": ["pcg", "MassMatrixPreconditioner"], + "tol": 1.0e-08, + "maxiter": 3000, + "info": False, + "verbose": False, + "recycle": True,} + +em_fields["options"]["Maxwell"] = {"algo": "implicit", + "solver": solver} diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index ff7323f45..d6b70b12b 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -1,9 +1,12 @@ -from dataclasses import dataclass - import numpy as np from mpi4py import MPI +import importlib.util +import sys -from struphy.utils.utils import dict_to_yaml +from struphy.utils.utils import dict_to_yaml, read_state +from struphy.topology.grids import TensorProductGrid +from struphy.io.parameters import StruphyParameters +from struphy.io.options import DerhamOptions def derive_units( @@ -126,7 +129,7 @@ def derive_units( return units -def setup_domain_and_equil(params: dict, units: dict = None): +def setup_domain_and_equil(params: StruphyParameters): """ Creates the domain object and equilibrium for a given parameter file. @@ -135,9 +138,6 @@ def setup_domain_and_equil(params: dict, units: dict = None): params : dict The full simulation parameter dictionary. - units : dict - All Struphy units. - Returns ------- domain : Domain @@ -147,51 +147,56 @@ def setup_domain_and_equil(params: dict, units: dict = None): The equilibrium object. """ - from struphy.fields_background import equils + # from struphy.fields_background import equils from struphy.fields_background.base import ( NumericalFluidEquilibrium, NumericalFluidEquilibriumWithB, NumericalMHDequilibrium, ) - from struphy.geometry import domains - - if "fluid_background" in params: - for eq_type, eq_params in params["fluid_background"].items(): - eq_class = getattr(equils, eq_type) - if eq_type in ("EQDSKequilibrium", "GVECequilibrium", "DESCequilibrium"): - equil = eq_class(**eq_params, units=units) - else: - equil = eq_class(**eq_params) + # from struphy.geometry import domains + + if params.equil is not None: + # for eq_type, eq_params in params["fluid_background"].items(): + # eq_class = getattr(equils, eq_type) + # if eq_type in ("EQDSKequilibrium", "GVECequilibrium", "DESCequilibrium"): + # equil = eq_class(**eq_params, units=units) + # else: + # equil = eq_class(**eq_params) + + equil = params.equil # for numerical equilibria, the domain comes from the equilibrium if isinstance(equil, (NumericalMHDequilibrium, NumericalFluidEquilibrium, NumericalFluidEquilibriumWithB)): domain = equil.domain # for all other equilibria, the domain can be chosen idependently else: - dom_type = params["geometry"]["type"] - dom_class = getattr(domains, dom_type) + # dom_type = params["geometry"]["type"] + # dom_class = getattr(domains, dom_type) - if dom_type == "Tokamak": - domain = dom_class(**params["geometry"][dom_type], equilibrium=equil) - else: - domain = dom_class(**params["geometry"][dom_type]) + # if dom_type == "Tokamak": + # domain = dom_class(**params["geometry"][dom_type], equilibrium=equil) + # else: + # domain = dom_class(**params["geometry"][dom_type]) + domain = params.domain # set domain attribute in mhd object - equil.domain = domain + equil.domain = params.domain # no equilibrium (just load domain) else: - dom_type = params["geometry"]["type"] - dom_class = getattr(domains, dom_type) - domain = dom_class(**params["geometry"][dom_type]) + # dom_type = params["geometry"]["type"] + # dom_class = getattr(domains, dom_type) + # domain = dom_class(**params["geometry"][dom_type]) + domain = params.domain equil = None return domain, equil def setup_derham( - params_grid, + grid: TensorProductGrid, + options: DerhamOptions, comm=None, domain=None, mpi_dims_mask=None, @@ -202,8 +207,8 @@ def setup_derham( Parameters ---------- - params_grid : dict - Grid parameters dictionary. + grid : TensorProductGrid + The FEEC grid. comm: Intracomm MPI communicator (sub_comm if clones are used). @@ -227,28 +232,28 @@ def setup_derham( from struphy.feec.psydac_derham import Derham # number of grid cells - Nel = params_grid["Nel"] + Nel = grid.Nel # spline degrees - p = params_grid["p"] + p = grid.p # spline types (clamped vs. periodic) - spl_kind = params_grid["spl_kind"] + spl_kind = grid.spl_kind # boundary conditions (Homogeneous Dirichlet or None) - dirichlet_bc = params_grid["dirichlet_bc"] + dirichlet_bc = grid.dirichlet_bc # Number of quadrature points per histopolation cell - nq_pr = params_grid["nq_pr"] + nq_pr = grid.nq_pr # Number of quadrature points per grid cell for L^2 - nq_el = params_grid["nq_el"] + nquads = grid.nquads # C^k smoothness at eta_1=0 for polar domains - polar_ck = params_grid["polar_ck"] + polar_ck = options.polar_ck # local commuting projectors - local_projectors = params_grid["local_projectors"] + local_projectors = options.local_projectors derham = Derham( Nel, p, spl_kind, dirichlet_bc=dirichlet_bc, - nquads=nq_el, + nquads=nquads, nq_pr=nq_pr, comm=comm, mpi_dims_mask=mpi_dims_mask, @@ -264,7 +269,7 @@ def setup_derham( print(f"spline degrees:".ljust(25), p) print(f"periodic bcs:".ljust(25), spl_kind) print(f"hom. Dirichlet bc:".ljust(25), dirichlet_bc) - print(f"GL quad pts (L2):".ljust(25), nq_el) + print(f"GL quad pts (L2):".ljust(25), nquads) print(f"GL quad pts (hist):".ljust(25), nq_pr) print( "MPI proc. per dir.:".ljust(25), @@ -407,12 +412,81 @@ def pre_processing( else: parameters_path = parameters - with open(parameters) as file: - params = yaml.load(file, Loader=yaml.FullLoader) + if ".yml" in parameters or ".yaml" in parameters: + with open(parameters) as file: + params = yaml.load(file, Loader=yaml.FullLoader) + elif ".py" in parameters: + # print(f'{parameters = }') + # Read struphy state file + state = read_state() + i_path = state["i_path"] + # load parameter.py + spec = importlib.util.spec_from_file_location("parameters", parameters) + params_in = importlib.util.module_from_spec(spec) + sys.modules["parameters"] = params_in + spec.loader.exec_module(params_in) + + if not hasattr(params_in, "model"): + params_in.model = None + + if not hasattr(params_in, "domain"): + params_in.domain = None + + if not hasattr(params_in, "grid"): + params_in.grid = None + + if not hasattr(params_in, "equil"): + params_in.equil = None + + if not hasattr(params_in, "units"): + params_in.units = None + + if not hasattr(params_in, "time"): + params_in.time = None + + if not hasattr(params_in, "derham"): + params_in.derham = None + + if not hasattr(params_in, "em_fields"): + params_in.em_fields = None + + if not hasattr(params_in, "fluid"): + params_in.fluid = None + + if not hasattr(params_in, "kinetic"): + params_in.kinetic = None + + if not hasattr(params_in, "diagnostic_fields"): + params_in.diagnostic_fields = None + + print(f'{params_in = }') + print(f'{params_in.model = }') + print(f'{params_in.domain = }') + print(f'{params_in.grid = }') + print(f'{params_in.equil = }') + print(f'{params_in.units = }') + print(f'{params_in.time = }') + print(f'{params_in.derham = }') + print(f'{params_in.em_fields = }') + print(f'{params_in.fluid = }') + print(f'{params_in.kinetic = }') + print(f'{params_in.diagnostic_fields = }') + + params = StruphyParameters(model=params_in.model, + domain=params_in.domain, + grid=params_in.grid, + equil=params_in.equil, + units=params_in.units, + time=params_in.time, + derham=params_in.derham, + em_fields=params_in.em_fields, + fluid=params_in.fluid, + kinetic=params_in.kinetic, + diagnostic_fields=params_in.diagnostic_fields,) if model_name is None: - assert "model" in params, "If model is not specified, then model: MODEL must be specified in the params!" - model_name = params["model"] + assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" + model_name = params.model if mpi_rank == 0: # copy parameter file to output folder diff --git a/src/struphy/main.py b/src/struphy/main.py index 433e85b6f..38745ab0f 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -50,7 +50,6 @@ def main( Number of domain clones (default=1) """ - import copy import os import time @@ -97,8 +96,8 @@ def main( ) if model_name is None: - assert "model" in params, "If model is not specified, then model: MODEL must be specified in the params!" - model_name = params["model"] + assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" + model_name = params.model if rank < 32: print(f"Rank {rank}: calling struphy/main.py for model {model_name} ...") @@ -176,13 +175,16 @@ def main( data.add_data({key_time: val}) data.add_data({key_time_restart: val}) - time_params = params["time"] + # retrieve time parameters + dt = params.time.dt + Tend = params.time.Tend + split_algo = params.time.split_algo # set initial conditions for all variables if not restart: model.initialize_from_params() - total_steps = str(int(round(time_params["Tend"] / time_params["dt"]))) + total_steps = str(int(round(Tend / dt))) else: model.initialize_from_restart(data) @@ -191,7 +193,7 @@ def main( time_state["value_sec"][0] = data.file["restart/time/value_sec"][-1] time_state["index"][0] = data.file["restart/time/index"][-1] - total_steps = str(int(round((time_params["Tend"] - time_state["value"][0]) / time_params["dt"]))) + total_steps = str(int(round((Tend - time_state["value"][0]) / dt))) # compute initial scalars and kinetic data, pass time state to all propagators model.update_scalar_quantities() @@ -208,7 +210,7 @@ def main( print("\nINITIAL SCALAR QUANTITIES:") model.print_scalar_quantities() - split_algo = time_params["split_algo"] + split_algo = split_algo print(f"\nSTART TIME STEPPING WITH '{split_algo}' SPLITTING:") # time loop @@ -217,7 +219,7 @@ def main( comm.Barrier() # stop time loop? - break_cond_1 = time_state["value"][0] >= time_params["Tend"] + break_cond_1 = time_state["value"][0] >= Tend break_cond_2 = run_time_now > runtime if break_cond_1 or break_cond_2: @@ -249,12 +251,12 @@ def main( # perform one time step dt t0 = time.time() with ProfileManager.profile_region("model.integrate"): - model.integrate(time_params["dt"], time_params["split_algo"]) + model.integrate(dt, split_algo) t1 = time.time() # update time and index (round time to 10 decimals for a clean time grid!) - time_state["value"][0] = round(time_state["value"][0] + time_params["dt"], 10) - time_state["value_sec"][0] = round(time_state["value_sec"][0] + time_params["dt"] * model.units["t"], 10) + time_state["value"][0] = round(time_state["value"][0] + dt, 10) + time_state["value_sec"][0] = round(time_state["value_sec"][0] + dt * model.units["t"], 10) time_state["index"][0] += 1 run_time_now = (time.time() - start_simulation) / 60 @@ -297,9 +299,9 @@ def main( step = str(time_state["index"][0]).zfill(len(total_steps)) message = "time step: " + step + "/" + str(total_steps) - message += " | " + "time: {0:10.5f}/{1:10.5f}".format(time_state["value"][0], time_params["Tend"]) + message += " | " + "time: {0:10.5f}/{1:10.5f}".format(time_state["value"][0], Tend) message += " | " + "phys. time [s]: {0:12.10f}/{1:12.10f}".format( - time_state["value_sec"][0], time_params["Tend"] * model.units["t"] + time_state["value_sec"][0], Tend * model.units["t"] ) message += " | " + "wall clock [s]: {0:8.4f} | last step duration [s]: {1:8.4f}".format( run_time_now * 60, t1 - t0 @@ -337,7 +339,6 @@ def main( # Read struphy state file state = utils.read_state() - o_path = state["o_path"] parser = argparse.ArgumentParser(description="Run an Struphy model.") diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 7f1a00819..23d6a238f 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -22,6 +22,7 @@ from struphy.propagators.base import Propagator from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml +from struphy.io.parameters import StruphyParameters class StruphyModel(metaclass=ABCMeta): @@ -30,8 +31,8 @@ class StruphyModel(metaclass=ABCMeta): Parameters ---------- - params : dict - Simulation parameters, see from :ref:`params_yml`. + params : StruphyParameters + Simulation parameters. comm : mpi4py.MPI.Intracomm MPI communicator for parallel runs. @@ -47,7 +48,7 @@ class StruphyModel(metaclass=ABCMeta): def __init__( self, - params: dict, + params: StruphyParameters = None, comm: MPI.Intracomm = None, clone_config: CloneConfig = None, ): @@ -87,28 +88,25 @@ def __init__( ) # create domain, equilibrium - self._domain, self._equil = setup_domain_and_equil( - params, - units=self.units, - ) + self._domain, self._equil = setup_domain_and_equil(params) if comm.Get_rank() == 0 and self.verbose: print("\nTIME:") print( f"time step:".ljust(25), "{0} ({1:4.2e} s)".format( - params["time"]["dt"], - params["time"]["dt"] * self.units["t"], + params.time.dt, + params.time.dt * self.units["t"], ), ) print( f"final time:".ljust(25), "{0} ({1:4.2e} s)".format( - params["time"]["Tend"], - params["time"]["Tend"] * self.units["t"], + params.time.Tend, + params.time.Tend * self.units["t"], ), ) - print(f"splitting algo:".ljust(25), params["time"]["split_algo"]) + print(f"splitting algo:".ljust(25), params.time.split_algo) print("\nDOMAIN:") print(f"type:".ljust(25), self.domain.__class__.__name__) @@ -117,7 +115,7 @@ def __init__( print((key + ":").ljust(25), val) print("\nFLUID BACKGROUND:") - if "fluid_background" in params: + if params.equil is not None: print("type:".ljust(25), self.equil.__class__.__name__) for key, val in self.equil.params.items(): print((key + ":").ljust(25), val) @@ -125,8 +123,8 @@ def __init__( print("None.") # create discrete derham sequence - if "grid" in params: - dims_mask = params["grid"]["dims_mask"] + if params.grid is not None: + dims_mask = params.grid.mpi_dims_mask if dims_mask is None: dims_mask = [True] * 3 @@ -136,7 +134,8 @@ def __init__( derham_comm = clone_config.sub_comm self._derham = setup_derham( - params["grid"], + params.grid, + params.derham, comm=derham_comm, domain=self.domain, mpi_dims_mask=dims_mask, @@ -1265,13 +1264,13 @@ def initialize_data_output(self, data, size): ################### @classmethod - def model_units(cls, params, verbose=False, comm=None): + def model_units(cls, params: StruphyParameters, verbose: bool=False, comm: MPI.Intracomm=None,): """ Return model units and print them to screen. Parameters ---------- - params : dict + params : StruphyParameters model parameters. verbose : bool, optional @@ -1299,27 +1298,24 @@ def model_units(cls, params, verbose=False, comm=None): # look for bulk species in fluid OR kinetic parameter dictionaries Z_bulk = None A_bulk = None - if "fluid" in params: + if params.fluid is not None: if cls.bulk_species() in params["fluid"]: Z_bulk = params["fluid"][cls.bulk_species()]["phys_params"]["Z"] A_bulk = params["fluid"][cls.bulk_species()]["phys_params"]["A"] - if "kinetic" in params: + if params.kinetic is not None: if cls.bulk_species() in params["kinetic"]: Z_bulk = params["kinetic"][cls.bulk_species()]["phys_params"]["Z"] A_bulk = params["kinetic"][cls.bulk_species()]["phys_params"]["A"] # compute model units - if "kBT" in params["units"]: - kBT = params["units"]["kBT"] - else: - kBT = None + kBT = params.units.kBT units = derive_units( Z_bulk=Z_bulk, A_bulk=A_bulk, - x=params["units"]["x"], - B=params["units"]["B"], - n=params["units"]["n"], + x=params.units.x, + B=params.units.B, + n=params.units.n, kBT=kBT, velocity_scale=cls.velocity_scale(), ) @@ -1368,7 +1364,7 @@ def model_units(cls, params, verbose=False, comm=None): eps0 = 8.8541878128e-12 # vacuum permittivity (F/m) equation_params = {} - if "fluid" in params: + if params.fluid is not None: for species in params["fluid"]: Z = params["fluid"][species]["phys_params"]["Z"] A = params["fluid"][species]["phys_params"]["A"] @@ -1387,7 +1383,7 @@ def model_units(cls, params, verbose=False, comm=None): for key, val in equation_params[species].items(): print((key + ":").ljust(25), "{:4.3e}".format(val)) - if "kinetic" in params: + if params.kinetic is not None: for species in params["kinetic"]: Z = params["kinetic"][species]["phys_params"]["Z"] A = params["kinetic"][species]["phys_params"]["A"] @@ -1700,7 +1696,7 @@ def _init_variable_dicts(self): # create dictionaries for each em-field/species and fill in space/class name and parameters for var_name, space in self.species()["em_fields"].items(): assert space in {"H1", "Hcurl", "Hdiv", "L2", "H1vec"} - assert "em_fields" in self.params, 'Top-level key "em_fields" is missing in parameter file.' + assert self.params.em_fields is not None, '"em_fields" is missing in parameter file.' if self.rank_world == 0 and self.verbose: print("em_field:".ljust(25), f'"{var_name}" ({space})') @@ -1711,22 +1707,22 @@ def _init_variable_dicts(self): self._em_fields[var_name]["space"] = space # initial conditions - if "background" in self.params["em_fields"]: - # background= self.params["em_fields"]["background"].get(var_name) - self._em_fields[var_name]["background"] = self.params["em_fields"]["background"].get(var_name) + if "background" in self.params.em_fields: + # background= self.params.em_fields["background"].get(var_name) + self._em_fields[var_name]["background"] = self.params.em_fields["background"].get(var_name) # else: # background = None - if "perturbation" in self.params["em_fields"]: - # perturbation = self.params["em_fields"]["perturbation"].get(var_name) - self._em_fields[var_name]["perturbation"] = self.params["em_fields"]["perturbation"].get(var_name) + if "perturbation" in self.params.em_fields: + # perturbation = self.params.em_fields["perturbation"].get(var_name) + self._em_fields[var_name]["perturbation"] = self.params.em_fields["perturbation"].get(var_name) # else: # perturbation = None # which components to save - if "save_data" in self.params["em_fields"]: - # save_data = self.params["em_fields"]["save_data"]["comps"][var_name] - self._em_fields[var_name]["save_data"] = self.params["em_fields"]["save_data"]["comps"][var_name] + if "save_data" in self.params.em_fields: + # save_data = self.params.em_fields["save_data"]["comps"][var_name] + self._em_fields[var_name]["save_data"] = self.params.em_fields["save_data"]["comps"][var_name] else: self._em_fields[var_name]["save_data"] = True # save_data = True @@ -1739,7 +1735,7 @@ def _init_variable_dicts(self): # overall parameters # print(f'{self._em_fields = }') - self._em_fields["params"] = self.params["em_fields"] + self._em_fields["params"] = self.params.em_fields for var_name, space in self.species()["fluid"].items(): assert isinstance(space, dict) @@ -1819,7 +1815,7 @@ def _allocate_variables(self): from struphy.pic.base import Particles # allocate memory for FE coeffs of electromagnetic fields/potentials - if "em_fields" in self.params: + if self.params.em_fields is not None: for variable, dct in self.em_fields.items(): if "params" in variable: continue @@ -1834,7 +1830,7 @@ def _allocate_variables(self): self._pointer[variable] = dct["obj"].vector # allocate memory for FE coeffs of fluid variables - if "fluid" in self.params: + if self.params.fluid is not None: for species, dct in self.fluid.items(): for variable, subdct in dct.items(): if "params" in variable: @@ -1850,7 +1846,7 @@ def _allocate_variables(self): self._pointer[species + "_" + variable] = subdct["obj"].vector # marker arrays and plasma parameters of kinetic species - if "kinetic" in self.params: + if self.params.kinetic is not None: for species, val in self.kinetic.items(): assert any([key in val["params"]["markers"] for key in ["Np", "ppc", "ppb"]]) @@ -1952,7 +1948,7 @@ def _allocate_variables(self): # TODO # allocate memory for FE coeffs of diagnostics - if "diagnostics" in self.params: + if self.params.diagnostic_fields is not None: for key, val in self.diagnostics.items(): if "params" in key: continue diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 6fd5b17fb..fd195cb23 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -60,8 +60,8 @@ def __init__(self, params, comm, clone_config=None): super().__init__(params, comm=comm, clone_config=clone_config) # extract necessary parameters - algo = params["em_fields"]["options"]["Maxwell"]["algo"] - solver = params["em_fields"]["options"]["Maxwell"]["solver"] + algo = params.em_fields["options"]["Maxwell"]["algo"] + solver = params.em_fields["options"]["Maxwell"]["solver"] # set keyword arguments for propagators self._kwargs[propagators_fields.Maxwell] = { From c76a6e5922f1c9daca190c060950fd3434d1ea57 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 14 Jul 2025 16:00:16 +0200 Subject: [PATCH 006/292] formatting --- src/struphy/io/inp/params_Maxwell.py | 55 +++++++++-------- src/struphy/io/options.py | 28 +++++---- src/struphy/io/parameters.py | 72 +++++++++++----------- src/struphy/io/setup.py | 89 ++++++++++++++-------------- src/struphy/models/base.py | 15 +++-- src/struphy/models/variables.py | 10 +--- src/struphy/topology/grids.py | 37 ++++++------ 7 files changed, 159 insertions(+), 147 deletions(-) diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index a739764ba..73cf77da6 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -4,10 +4,12 @@ model = "Maxwell" # units -units = options.Units(x=1.0, - B=1.0, - n=1.0, - kBT=1.0,) +units = options.Units( + x=1.0, + B=1.0, + n=1.0, + kBT=1.0, +) # time time = options.Time(split_algo="LieTrotter") @@ -30,26 +32,31 @@ # FOR NOW: initial conditions and options em_fields = {} -em_fields["background"] = {} -em_fields["perturbation"] = {} -em_fields["options"] = {} +em_fields["background"] = {} +em_fields["perturbation"] = {} +em_fields["options"] = {} -em_fields["background"]["e_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} +em_fields["background"]["e_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} em_fields["background"]["b_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} -em_fields["perturbation"]["e_field"] = {} -em_fields["perturbation"]["b_field"] = {} -em_fields["perturbation"]["e_field"]["TorusModesCos"] = {"given_in_basis": [None, "v", None], - "ms": [[None], [1, 3], [None]]} -em_fields["perturbation"]["b_field"]["TorusModesCos"] = {"given_in_basis": [None, "v", None], - "ms": [[None], [1, 3], [None]]} - -solver = {"type": ["pcg", "MassMatrixPreconditioner"], - "tol": 1.0e-08, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True,} - -em_fields["options"]["Maxwell"] = {"algo": "implicit", - "solver": solver} +em_fields["perturbation"]["e_field"] = {} +em_fields["perturbation"]["b_field"] = {} +em_fields["perturbation"]["e_field"]["TorusModesCos"] = { + "given_in_basis": [None, "v", None], + "ms": [[None], [1, 3], [None]], +} +em_fields["perturbation"]["b_field"]["TorusModesCos"] = { + "given_in_basis": [None, "v", None], + "ms": [[None], [1, 3], [None]], +} + +solver = { + "type": ["pcg", "MassMatrixPreconditioner"], + "tol": 1.0e-08, + "maxiter": 3000, + "info": False, + "verbose": False, + "recycle": True, +} + +em_fields["options"]["Maxwell"] = {"algo": "implicit", "solver": solver} diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 51cc8c36f..0d0fd549d 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -1,10 +1,11 @@ from dataclasses import dataclass from typing import Literal, get_args +from struphy.fields_background import equils + # needed for import in StruphyParameters from struphy.geometry import domains from struphy.topology import grids -from struphy.fields_background import equils SplitAlgos = Literal["LieTrotter", "Strang"] PolarRegularity = Literal[-1, 1] @@ -13,50 +14,51 @@ @dataclass class Units: """... - + Parameters ---------- x : float Unit of length in m. """ + x: float = 1.0 B: float = 1.0 n: float = 1.0 kBT: float = None - - + + @dataclass class Time: """... - + Parameters ---------- x : float Unit of length in m. """ - + dt: float = 1.0 Tend: float = 1.0 split_algo: SplitAlgos = "LieTrotter" - + def __post_init__(self): options = get_args(SplitAlgos) assert self.split_algo in options, f"'{self.split_algo}' is not in {options}" - - + + @dataclass class DerhamOptions: """... - + Parameters ---------- x : float Unit of length in m. """ - + polar_ck: int = -1 local_projectors: bool = False - + def __post_init__(self): options = get_args(PolarRegularity) - assert self.polar_ck in options, f"'{self.polar_ck}' is not in {options}" \ No newline at end of file + assert self.polar_ck in options, f"'{self.polar_ck}' is not in {options}" diff --git a/src/struphy/io/parameters.py b/src/struphy/io/parameters.py index dfadd9931..ba2484bf7 100644 --- a/src/struphy/io/parameters.py +++ b/src/struphy/io/parameters.py @@ -1,28 +1,28 @@ -from struphy.geometry.base import Domain -from struphy.topology import grids from struphy.fields_background.base import FluidEquilibrium +from struphy.geometry.base import Domain from struphy.io import options - - -class StruphyParameters(): +from struphy.topology import grids + + +class StruphyParameters: """Wrapper around Struphy simulation parameters.""" - - def __init__(self, - model: str = None, - domain: Domain = None, - grid: grids.TensorProductGrid = None, - equil: FluidEquilibrium = None, - units: options.Units = None, - time: options.Time = None, - derham: options.DerhamOptions = None, - em_fields=None, - fluid=None, - kinetic=None, - diagnostic_fields=None, - ): - + + def __init__( + self, + model: str = None, + domain: Domain = None, + grid: grids.TensorProductGrid = None, + equil: FluidEquilibrium = None, + units: options.Units = None, + time: options.Time = None, + derham: options.DerhamOptions = None, + em_fields=None, + fluid=None, + kinetic=None, + diagnostic_fields=None, + ): self._options = options - + self._model = model self._domain = domain self._grid = grid @@ -34,55 +34,55 @@ def __init__(self, self._fluid = fluid self._kinetic = kinetic self._diagnostic_fields = diagnostic_fields - + ## possible inputs - + @property def options(self): return self._options - + ## parameters - + @property def model(self): return self._model - + @property def domain(self): return self._domain - + @property def grid(self): return self._grid - + @property def equil(self): return self._equil - + @property def units(self): return self._units - + @property def time(self): return self._time - + @property def derham(self): return self._derham - + @property def em_fields(self): return self._em_fields - + @property def fluid(self): return self._fluid - + @property def kinetic(self): return self._kinetic - + @property def diagnostic_fields(self): - return self._diagnostic_fields \ No newline at end of file + return self._diagnostic_fields diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index d6b70b12b..6af443fa6 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -1,12 +1,13 @@ -import numpy as np -from mpi4py import MPI import importlib.util import sys -from struphy.utils.utils import dict_to_yaml, read_state -from struphy.topology.grids import TensorProductGrid -from struphy.io.parameters import StruphyParameters +import numpy as np +from mpi4py import MPI + from struphy.io.options import DerhamOptions +from struphy.io.parameters import StruphyParameters +from struphy.topology.grids import TensorProductGrid +from struphy.utils.utils import dict_to_yaml, read_state def derive_units( @@ -162,7 +163,7 @@ def setup_domain_and_equil(params: StruphyParameters): # equil = eq_class(**eq_params, units=units) # else: # equil = eq_class(**eq_params) - + equil = params.equil # for numerical equilibria, the domain comes from the equilibrium @@ -196,7 +197,7 @@ def setup_domain_and_equil(params: StruphyParameters): def setup_derham( grid: TensorProductGrid, - options: DerhamOptions, + options: DerhamOptions, comm=None, domain=None, mpi_dims_mask=None, @@ -425,64 +426,66 @@ def pre_processing( params_in = importlib.util.module_from_spec(spec) sys.modules["parameters"] = params_in spec.loader.exec_module(params_in) - + if not hasattr(params_in, "model"): params_in.model = None - + if not hasattr(params_in, "domain"): params_in.domain = None - + if not hasattr(params_in, "grid"): params_in.grid = None - + if not hasattr(params_in, "equil"): params_in.equil = None - + if not hasattr(params_in, "units"): params_in.units = None - + if not hasattr(params_in, "time"): params_in.time = None - + if not hasattr(params_in, "derham"): params_in.derham = None - + if not hasattr(params_in, "em_fields"): params_in.em_fields = None - + if not hasattr(params_in, "fluid"): params_in.fluid = None - + if not hasattr(params_in, "kinetic"): params_in.kinetic = None - + if not hasattr(params_in, "diagnostic_fields"): params_in.diagnostic_fields = None - - print(f'{params_in = }') - print(f'{params_in.model = }') - print(f'{params_in.domain = }') - print(f'{params_in.grid = }') - print(f'{params_in.equil = }') - print(f'{params_in.units = }') - print(f'{params_in.time = }') - print(f'{params_in.derham = }') - print(f'{params_in.em_fields = }') - print(f'{params_in.fluid = }') - print(f'{params_in.kinetic = }') - print(f'{params_in.diagnostic_fields = }') - - params = StruphyParameters(model=params_in.model, - domain=params_in.domain, - grid=params_in.grid, - equil=params_in.equil, - units=params_in.units, - time=params_in.time, - derham=params_in.derham, - em_fields=params_in.em_fields, - fluid=params_in.fluid, - kinetic=params_in.kinetic, - diagnostic_fields=params_in.diagnostic_fields,) + + print(f"{params_in = }") + print(f"{params_in.model = }") + print(f"{params_in.domain = }") + print(f"{params_in.grid = }") + print(f"{params_in.equil = }") + print(f"{params_in.units = }") + print(f"{params_in.time = }") + print(f"{params_in.derham = }") + print(f"{params_in.em_fields = }") + print(f"{params_in.fluid = }") + print(f"{params_in.kinetic = }") + print(f"{params_in.diagnostic_fields = }") + + params = StruphyParameters( + model=params_in.model, + domain=params_in.domain, + grid=params_in.grid, + equil=params_in.equil, + units=params_in.units, + time=params_in.time, + derham=params_in.derham, + em_fields=params_in.em_fields, + fluid=params_in.fluid, + kinetic=params_in.kinetic, + diagnostic_fields=params_in.diagnostic_fields, + ) if model_name is None: assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 23d6a238f..2374f1191 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -17,12 +17,12 @@ ProjectedFluidEquilibriumWithB, ProjectedMHDequilibrium, ) +from struphy.io.parameters import StruphyParameters from struphy.io.setup import setup_derham, setup_domain_and_equil from struphy.profiling.profiling import ProfileManager from struphy.propagators.base import Propagator from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml -from struphy.io.parameters import StruphyParameters class StruphyModel(metaclass=ABCMeta): @@ -1264,7 +1264,12 @@ def initialize_data_output(self, data, size): ################### @classmethod - def model_units(cls, params: StruphyParameters, verbose: bool=False, comm: MPI.Intracomm=None,): + def model_units( + cls, + params: StruphyParameters, + verbose: bool = False, + comm: MPI.Intracomm = None, + ): """ Return model units and print them to screen. @@ -1712,7 +1717,7 @@ def _init_variable_dicts(self): self._em_fields[var_name]["background"] = self.params.em_fields["background"].get(var_name) # else: # background = None - + if "perturbation" in self.params.em_fields: # perturbation = self.params.em_fields["perturbation"].get(var_name) self._em_fields[var_name]["perturbation"] = self.params.em_fields["perturbation"].get(var_name) @@ -1727,12 +1732,12 @@ def _init_variable_dicts(self): self._em_fields[var_name]["save_data"] = True # save_data = True - # self._em_fields[var_name] = Variable(name=var_name, + # self._em_fields[var_name] = Variable(name=var_name, # space=space, # background=background, # perturbation=perturbation, # save_data=save_data,) - + # overall parameters # print(f'{self._em_fields = }') self._em_fields["params"] = self.params.em_fields diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 15606e76c..df27c1168 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -1,10 +1,4 @@ class Variable: - def __init__(self, - name : str, - space : str, - background, - perturbation, - save_data - ): + def __init__(self, name: str, space: str, background, perturbation, save_data): self._name = name - self._space = space \ No newline at end of file + self._space = space diff --git a/src/struphy/topology/grids.py b/src/struphy/topology/grids.py index fb2798dc2..cda18c021 100644 --- a/src/struphy/topology/grids.py +++ b/src/struphy/topology/grids.py @@ -3,7 +3,7 @@ class TensorProductGrid: """Grid as a tensor product of 1d grids. - + Parameters ---------- Nel : tuple[int] @@ -23,28 +23,29 @@ class TensorProductGrid: nq_pr : tuple[int] Number of Gauss-Legendre quadrature points in each direction for geometric projectors (default = p+1, leads to exact integration of degree 2p+1 polynomials). - + mpi_dims_mask: Tuple of bool True if the dimension is to be used in the domain decomposition (=default for each dimension). If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. """ - - def __init__(self, - Nel: tuple, - p: tuple, - spl_kind: tuple, - *, - dirichlet_bc: tuple = None, - nquads: tuple = None, - nq_pr: tuple = None, - mpi_dims_mask: tuple = None,): - + + def __init__( + self, + Nel: tuple, + p: tuple, + spl_kind: tuple, + *, + dirichlet_bc: tuple = None, + nquads: tuple = None, + nq_pr: tuple = None, + mpi_dims_mask: tuple = None, + ): assert len(Nel) == len(p) == len(spl_kind) - + self._Nel = Nel self._p = p self._spl_kind = spl_kind - + # boundary conditions at eta=0 and eta=1 in each direction (None for periodic, 'd' for homogeneous Dirichlet) if dirichlet_bc is not None: assert len(dirichlet_bc) == len(Nel) @@ -69,7 +70,7 @@ def __init__(self, # mpi domain decomposition directions if mpi_dims_mask is None: - self._mpi_dims_mask = (True,)*len(Nel) + self._mpi_dims_mask = (True,) * len(Nel) else: assert len(mpi_dims_mask) == len(Nel) self._mpi_dims_mask = mpi_dims_mask @@ -105,8 +106,8 @@ def nquads(self): def nq_pr(self): """Tuple of number of Gauss-Legendre quadrature points in histopolation (default = p + 1) in each direction.""" return self._nq_pr - + @property def mpi_dims_mask(self): """Tuple of bool; whether to use direction in domain decomposition.""" - return self._mpi_dims_mask \ No newline at end of file + return self._mpi_dims_mask From a397d4848034282965767a207694cc26d5fd8f8d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 15 Jul 2025 08:02:45 +0200 Subject: [PATCH 007/292] some cleanup --- src/struphy/io/options.py | 3 +-- src/struphy/io/parameters.py | 17 ++++------------- src/struphy/main.py | 1 - 3 files changed, 5 insertions(+), 16 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 0d0fd549d..b5a0fe81b 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -1,9 +1,8 @@ from dataclasses import dataclass from typing import Literal, get_args -from struphy.fields_background import equils - # needed for import in StruphyParameters +from struphy.fields_background import equils from struphy.geometry import domains from struphy.topology import grids diff --git a/src/struphy/io/parameters.py b/src/struphy/io/parameters.py index ba2484bf7..35b12d5ab 100644 --- a/src/struphy/io/parameters.py +++ b/src/struphy/io/parameters.py @@ -1,7 +1,7 @@ from struphy.fields_background.base import FluidEquilibrium from struphy.geometry.base import Domain -from struphy.io import options from struphy.topology import grids +from struphy.io.options import Units, Time, DerhamOptions class StruphyParameters: @@ -13,15 +13,14 @@ def __init__( domain: Domain = None, grid: grids.TensorProductGrid = None, equil: FluidEquilibrium = None, - units: options.Units = None, - time: options.Time = None, - derham: options.DerhamOptions = None, + units: Units = None, + time: Time = None, + derham: DerhamOptions = None, em_fields=None, fluid=None, kinetic=None, diagnostic_fields=None, ): - self._options = options self._model = model self._domain = domain @@ -35,14 +34,6 @@ def __init__( self._kinetic = kinetic self._diagnostic_fields = diagnostic_fields - ## possible inputs - - @property - def options(self): - return self._options - - ## parameters - @property def model(self): return self._model diff --git a/src/struphy/main.py b/src/struphy/main.py index 38745ab0f..a5a2cf3b0 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -210,7 +210,6 @@ def main( print("\nINITIAL SCALAR QUANTITIES:") model.print_scalar_quantities() - split_algo = split_algo print(f"\nSTART TIME STEPPING WITH '{split_algo}' SPLITTING:") # time loop From eeecd245612a3f56c004efc2af8ad0b8efe1bf1d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 15 Jul 2025 13:23:17 +0200 Subject: [PATCH 008/292] new options class FieldsBackground --- src/struphy/io/options.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index b5a0fe81b..d8bec3fba 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -4,10 +4,13 @@ # needed for import in StruphyParameters from struphy.fields_background import equils from struphy.geometry import domains +from struphy.initial import perturbations +from struphy.kinetic_background import maxwellians from struphy.topology import grids SplitAlgos = Literal["LieTrotter", "Strang"] PolarRegularity = Literal[-1, 1] +BackgroundOpts = Literal["LogicalConst", "FluidEquilibrium"] @dataclass @@ -61,3 +64,22 @@ class DerhamOptions: def __post_init__(self): options = get_args(PolarRegularity) assert self.polar_ck in options, f"'{self.polar_ck}' is not in {options}" + + +@dataclass +class FieldsBackground: + """... + + Parameters + ---------- + x : float + Unit of length in m. + """ + + kind: str = "LogicalConst" + values: tuple = (1.5,) + variable: str = None + + def __post_init__(self): + options = get_args(BackgroundOpts) + assert self.kind in options, f"'{self.kind}' is not in {options}" From 677c238727ce80ae9f13b789cd144fe5e4dc820e Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 15 Jul 2025 13:25:01 +0200 Subject: [PATCH 009/292] new classes `Variables` and `Variable` for handling model variables --- src/struphy/models/variables.py | 119 +++++++++++++++++++++++++++++++- 1 file changed, 118 insertions(+), 1 deletion(-) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index df27c1168..6952ec3c2 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -1,4 +1,121 @@ +from typing import Callable + +from struphy.fields_background.base import FluidEquilibrium +from struphy.kinetic_background.base import KineticBackground + + +class Variables: + """Class for all Variable objects of a model.""" + def __init__(self, **species_dct): + + assert "em_fields" in species_dct + assert "fluid" in species_dct + assert "kinetic" in species_dct + + for name, space in species_dct["em_fields"].items(): + setattr(self, name, Variable(name, space)) + + for name, space in species_dct["fluid"].items(): + setattr(self, name, Variable(name, space)) + + for name, space in species_dct["kinetic"].items(): + setattr(self, name, Variable(name, space)) + + @property + def all(self): + return self.__dict__ + + def add_background(self, variable: str, background,): + assert variable in self.all, f'Variable {variable} is not part of model variables {self.all = }' + var = getattr(self, variable) + assert isinstance(var, Variable) + var.add_background(background) + + def add_perturbation(self, variable: str, perturbation: Callable, given_in_basis: tuple = None,): + assert variable in self.all, f'Variable {variable} is not part of model variables {self.all = }' + var = getattr(self, variable) + assert isinstance(var, Variable) + var.add_perturbation(perturbation=perturbation, + given_in_basis=given_in_basis,) + + class Variable: - def __init__(self, name: str, space: str, background, perturbation, save_data): + """Single variable (unknown) of a StruphyModel. + """ + def __init__(self, name: str, space: str, save_data: bool = True): + self._name = name self._space = space + self._save_data = save_data + + self._background = [] + self._perturbation = [] + + ## attributes + + @property + def name(self): + return self._name + + @property + def space(self): + return self._space + + @property + def save_data(self): + return self._save_data + + @property + def background(self): + return self._background + + # @background.setter + # def background(self, new): + # assert isinstance(new, FluidEquilibrium | KineticBackground) + # self._background = new + + @property + def perturbation(self): + return self._perturbation + + # @perturbation.setter + # def perturbation(self, new): + # assert isinstance(new, Perturbation) + # self._perturbation = new + + @property + def initial_condition(self): + if not hasattr(self, "_initial_condition"): + self.set_initial_condition() + return self._initial_condition + + ## methods + + def add_background(self, background): + # assert isinstance(...) + self._background += [background] + + def add_perturbation(self, + perturbation: Callable = None, + given_in_basis: tuple = None,): + # assert isinstance(...) + self._perturbation += [(perturbation, given_in_basis)] + + def set_initial_condition(self): + self._initial_condition = InitialCondition(self.background, + self.perturbation,) + + def eval_initial_condition(self, eta1, eta2, eta3, *v): + """Callable initial condition as sum of background + perturbation.""" + return self.initial_condition(eta1, eta2, eta3, *v) + + +class InitialCondition: + """Callable initial condition as sum of background + perturbation.""" + def __init__(self, + background: FluidEquilibrium | KineticBackground = None, + perturbation : Callable = None,): + + self._background = background + self._perturbation = perturbation + From 0988014d94abb74d2e1cfc89723a45a54e79cac5 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 16 Jul 2025 08:03:16 +0200 Subject: [PATCH 010/292] new classes: `Species`, `MultiSpecies` and `SubSpeceis`; these will replace the species dictionary in models. --- src/struphy/io/inp/params_Maxwell_lw.py | 88 +++++++++++++++++++ .../models/{variables.py => species.py} | 0 2 files changed, 88 insertions(+) create mode 100644 src/struphy/io/inp/params_Maxwell_lw.py rename src/struphy/models/{variables.py => species.py} (100%) diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py new file mode 100644 index 000000000..68bf38067 --- /dev/null +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -0,0 +1,88 @@ +from struphy.models import (toy, fluid, kinetic, hybrid,) +from struphy.io import options + +# model class (do not instantiate) +Model = toy.Maxwell + +# units +units = options.Units( + x=1.0, + B=1.0, + n=1.0, + kBT=1.0, +) + +# time +time = options.Time(split_algo="LieTrotter") + +# geometry +domain = options.domains.Cuboid() + +# grid +grid = options.grids.TensorProductGrid( + Nel=(12, 14, 1), + p=(2, 3, 1), + spl_kind=(False, True, True), +) + +# derham options +derham = options.DerhamOptions() + +# fluid equilibrium (can be used as part of initial conditions) +equil = options.equils.HomogenSlab() + +# light-weight instance of model +model = Model() +print(f'{model.__em_fields__ = }') + +# species parameters +species = model.species() +print(f'{species.em_fields = }') +print(f'{species.fluid = }') +print(f'{species.kinetic = }') +# species.em_fields.set_options(prop, prop.options()) +# model.fluid.set_phys_params("mhd", options.PhysParams()) +# model.fluid.set_propagator_options("mhd", prop, prop.options()) +# model.kinetic.set_phys_params("mhd", options.PhysParams()) + +# initial conditions for model variables (background + perturbation) +species.em_fields.add_background("e_ield", options.FieldsBackground(kind="LogicalConst", + values=(0.3, 0.15, None),),) +species.em_fields.add_perturbation("e_field", options.perturbations.TorusModesCos(ms=[[None], [1, 3], [None]],), + given_in_basis=(None, "v", None),) + +species.em_fields.add_background("b_field", options.FieldsBackground(kind="LogicalConst", + values=(0.3, 0.15, None),),) +species.em_fields.add_perturbation("b_field", options.perturbations.TorusModesCos(ms=[[None], [1, 3], [None]],), + given_in_basis=(None, "v", None),) + +# FOR NOW: initial conditions and options +em_fields = {} +em_fields["background"] = {} +em_fields["perturbation"] = {} +em_fields["options"] = {} + +em_fields["background"]["e_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} +em_fields["background"]["b_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} + +em_fields["perturbation"]["e_field"] = {} +em_fields["perturbation"]["b_field"] = {} +em_fields["perturbation"]["e_field"]["TorusModesCos"] = { + "given_in_basis": [None, "v", None], + "ms": [[None], [1, 3], [None]], +} +em_fields["perturbation"]["b_field"]["TorusModesCos"] = { + "given_in_basis": [None, "v", None], + "ms": [[None], [1, 3], [None]], +} + +solver = { + "type": ["pcg", "MassMatrixPreconditioner"], + "tol": 1.0e-08, + "maxiter": 3000, + "info": False, + "verbose": False, + "recycle": True, +} + +em_fields["options"]["Maxwell"] = {"algo": "implicit", "solver": solver} diff --git a/src/struphy/models/variables.py b/src/struphy/models/species.py similarity index 100% rename from src/struphy/models/variables.py rename to src/struphy/models/species.py From eb0fa98afec73ca3b8501b5dce1f782643ee9b31 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 16 Jul 2025 08:04:06 +0200 Subject: [PATCH 011/292] belongs to previous commit --- src/struphy/models/species.py | 192 +++++++++++++++++++++------------- 1 file changed, 121 insertions(+), 71 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 6952ec3c2..e518513f8 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,121 +1,171 @@ +from copy import deepcopy from typing import Callable from struphy.fields_background.base import FluidEquilibrium +from struphy.initial.base import InitialCondition from struphy.kinetic_background.base import KineticBackground -class Variables: - """Class for all Variable objects of a model.""" - def __init__(self, **species_dct): - - assert "em_fields" in species_dct - assert "fluid" in species_dct - assert "kinetic" in species_dct - - for name, space in species_dct["em_fields"].items(): - setattr(self, name, Variable(name, space)) - - for name, space in species_dct["fluid"].items(): - setattr(self, name, Variable(name, space)) - - for name, space in species_dct["kinetic"].items(): - setattr(self, name, Variable(name, space)) - +class Species: + """Handles the three species types in StruphyModel: em_fields, fluid, kinetic.""" + + def __init__( + self, + fluid: tuple = None, + kinetic: tuple = None, + em_fields: bool = True, + ): + # fluid species + if fluid is None: + self._fluid = None + else: + self._fluid = MultiSpecies() + for f in fluid: + self._fluid.add_species(name=f) + + # kinetic species + if kinetic is None: + self._kinetic = None + else: + self._kinetic = MultiSpecies() + for k in kinetic: + self._kinetic.add_species(name=k) + + # electromagnetic fields + if em_fields: + self._em_fields = SubSpecies(name="em_fields") + + @property + def fluid(self): + return self._fluid + + @property + def kinetic(self): + return self._kinetic + + @property + def em_fields(self): + return self._em_fields + + +class MultiSpecies: + """Handles multiple fluid or kinetic species.""" + + def __init__( + self, + ): + pass + @property def all(self): return self.__dict__ - - def add_background(self, variable: str, background,): - assert variable in self.all, f'Variable {variable} is not part of model variables {self.all = }' + + def add_species(self, name: str = "mhd"): + setattr(self, name, SubSpecies(name)) + + +class SubSpecies: + """Handles the three species types in StruphyModel: em_fields, fluid, kinetic.""" + + def __init__(self, name: str = "mhd"): + self._name = name + + @property + def name(self): + return self._name + + @property + def all(self): + out = deepcopy(self.__dict__) + out.pop("_name") + return out + + def add_variable(self, name: str = "velocity", space: str = "Hdiv"): + setattr(self, name, Variable(name, space)) + + def add_background( + self, + variable: str, + background, + ): + assert variable in self.all, f"Variable {variable} is not part of model variables {self.all.keys()}" var = getattr(self, variable) assert isinstance(var, Variable) var.add_background(background) - - def add_perturbation(self, variable: str, perturbation: Callable, given_in_basis: tuple = None,): - assert variable in self.all, f'Variable {variable} is not part of model variables {self.all = }' + + def add_perturbation( + self, + variable: str, + perturbation: Callable, + given_in_basis: tuple = None, + ): + assert variable in self.all, f"Variable {variable} is not part of model variables {self.all.keys()}" var = getattr(self, variable) assert isinstance(var, Variable) - var.add_perturbation(perturbation=perturbation, - given_in_basis=given_in_basis,) + var.add_perturbation( + perturbation=perturbation, + given_in_basis=given_in_basis, + ) + + def set_options(self, propagator): + pass class Variable: - """Single variable (unknown) of a StruphyModel. - """ + """Single variable (unknown) of a StruphyModel.""" + def __init__(self, name: str, space: str, save_data: bool = True): - self._name = name self._space = space self._save_data = save_data - + self._background = [] self._perturbation = [] - + ## attributes - + @property def name(self): return self._name - + @property def space(self): return self._space - + @property def save_data(self): return self._save_data - + @property def background(self): return self._background - - # @background.setter - # def background(self, new): - # assert isinstance(new, FluidEquilibrium | KineticBackground) - # self._background = new - + @property def perturbation(self): return self._perturbation - - # @perturbation.setter - # def perturbation(self, new): - # assert isinstance(new, Perturbation) - # self._perturbation = new - + @property - def initial_condition(self): + def initial_condition(self): if not hasattr(self, "_initial_condition"): self.set_initial_condition() return self._initial_condition - + ## methods - + def add_background(self, background): # assert isinstance(...) self._background += [background] - - def add_perturbation(self, - perturbation: Callable = None, - given_in_basis: tuple = None,): + + def add_perturbation( + self, + perturbation: Callable = None, + given_in_basis: tuple = None, + ): # assert isinstance(...) self._perturbation += [(perturbation, given_in_basis)] - - def set_initial_condition(self): - self._initial_condition = InitialCondition(self.background, - self.perturbation,) - - def eval_initial_condition(self, eta1, eta2, eta3, *v): - """Callable initial condition as sum of background + perturbation.""" - return self.initial_condition(eta1, eta2, eta3, *v) - - -class InitialCondition: - """Callable initial condition as sum of background + perturbation.""" - def __init__(self, - background: FluidEquilibrium | KineticBackground = None, - perturbation : Callable = None,): - - self._background = background - self._perturbation = perturbation + def set_initial_condition(self): + self._initial_condition = InitialCondition( + background=self.background, + perturbation=self.perturbation, + ) From 24b76bfd2325c7b2fd4f33ca6aa1a3f87355f47c Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 16 Jul 2025 08:05:32 +0200 Subject: [PATCH 012/292] formatting and cleaning of new params file --- src/struphy/io/inp/params_Maxwell_lw.py | 53 ++++++++++++++++++------- src/struphy/io/options.py | 17 ++++++-- src/struphy/io/parameters.py | 3 +- 3 files changed, 53 insertions(+), 20 deletions(-) diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index 68bf38067..b554a8eb8 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -1,5 +1,10 @@ -from struphy.models import (toy, fluid, kinetic, hybrid,) from struphy.io import options +from struphy.models import ( + fluid, + hybrid, + kinetic, + toy, +) # model class (do not instantiate) Model = toy.Maxwell @@ -33,28 +38,48 @@ # light-weight instance of model model = Model() -print(f'{model.__em_fields__ = }') +print(f"{model.__em_fields__=}") # species parameters species = model.species() -print(f'{species.em_fields = }') -print(f'{species.fluid = }') -print(f'{species.kinetic = }') +print(f"{species.em_fields=}") +print(f"{species.fluid=}") +print(f"{species.kinetic=}") # species.em_fields.set_options(prop, prop.options()) # model.fluid.set_phys_params("mhd", options.PhysParams()) # model.fluid.set_propagator_options("mhd", prop, prop.options()) # model.kinetic.set_phys_params("mhd", options.PhysParams()) # initial conditions for model variables (background + perturbation) -species.em_fields.add_background("e_ield", options.FieldsBackground(kind="LogicalConst", - values=(0.3, 0.15, None),),) -species.em_fields.add_perturbation("e_field", options.perturbations.TorusModesCos(ms=[[None], [1, 3], [None]],), - given_in_basis=(None, "v", None),) - -species.em_fields.add_background("b_field", options.FieldsBackground(kind="LogicalConst", - values=(0.3, 0.15, None),),) -species.em_fields.add_perturbation("b_field", options.perturbations.TorusModesCos(ms=[[None], [1, 3], [None]],), - given_in_basis=(None, "v", None),) +species.em_fields.add_background( + "e_field", + options.FieldsBackground( + kind="LogicalConst", + values=(0.3, 0.15, None), + ), +) +species.em_fields.add_perturbation( + "e_field", + options.perturbations.TorusModesCos( + ms=[[None], [1, 3], [None]], + ), + given_in_basis=(None, "v", None), +) + +species.em_fields.add_background( + "b_field", + options.FieldsBackground( + kind="LogicalConst", + values=(0.3, 0.15, None), + ), +) +species.em_fields.add_perturbation( + "b_field", + options.perturbations.TorusModesCos( + ms=[[None], [1, 3], [None]], + ), + given_in_basis=(None, "v", None), +) # FOR NOW: initial conditions and options em_fields = {} diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index d8bec3fba..bb4ae3d54 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -8,9 +8,9 @@ from struphy.kinetic_background import maxwellians from struphy.topology import grids +## generic options + SplitAlgos = Literal["LieTrotter", "Strang"] -PolarRegularity = Literal[-1, 1] -BackgroundOpts = Literal["LogicalConst", "FluidEquilibrium"] @dataclass @@ -48,6 +48,12 @@ def __post_init__(self): assert self.split_algo in options, f"'{self.split_algo}' is not in {options}" +## field options + +PolarRegularity = Literal[-1, 1] +BackgroundOpts = Literal["LogicalConst", "FluidEquilibrium"] + + @dataclass class DerhamOptions: """... @@ -64,8 +70,8 @@ class DerhamOptions: def __post_init__(self): options = get_args(PolarRegularity) assert self.polar_ck in options, f"'{self.polar_ck}' is not in {options}" - - + + @dataclass class FieldsBackground: """... @@ -83,3 +89,6 @@ class FieldsBackground: def __post_init__(self): options = get_args(BackgroundOpts) assert self.kind in options, f"'{self.kind}' is not in {options}" + + +## kinetic options diff --git a/src/struphy/io/parameters.py b/src/struphy/io/parameters.py index 35b12d5ab..6d695a326 100644 --- a/src/struphy/io/parameters.py +++ b/src/struphy/io/parameters.py @@ -1,7 +1,7 @@ from struphy.fields_background.base import FluidEquilibrium from struphy.geometry.base import Domain +from struphy.io.options import DerhamOptions, Time, Units from struphy.topology import grids -from struphy.io.options import Units, Time, DerhamOptions class StruphyParameters: @@ -21,7 +21,6 @@ def __init__( kinetic=None, diagnostic_fields=None, ): - self._model = model self._domain = domain self._grid = grid From f248d0246395b7ec1c21dccda82d33e876a9f105 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 16 Jul 2025 14:34:36 +0200 Subject: [PATCH 013/292] a lot of tryout --- src/struphy/initial/base.py | 30 +++++ src/struphy/io/inp/params_Maxwell_lw.py | 34 ++++-- src/struphy/linear_algebra/solver.py | 12 ++ src/struphy/models/base.py | 112 ++++++++---------- src/struphy/models/species.py | 9 +- src/struphy/models/toy.py | 71 +++++------ src/struphy/ode/utils.py | 2 +- src/struphy/parameters/__init__.py | 0 src/struphy/parameters/derham_params.py | 21 ++++ src/struphy/parameters/time_params.py | 22 ++++ src/struphy/parameters/units_params.py | 15 +++ src/struphy/propagators/base.py | 94 ++++++++------- src/struphy/propagators/propagators_fields.py | 52 ++++---- 13 files changed, 300 insertions(+), 174 deletions(-) create mode 100644 src/struphy/initial/base.py create mode 100644 src/struphy/linear_algebra/solver.py create mode 100644 src/struphy/parameters/__init__.py create mode 100644 src/struphy/parameters/derham_params.py create mode 100644 src/struphy/parameters/time_params.py create mode 100644 src/struphy/parameters/units_params.py diff --git a/src/struphy/initial/base.py b/src/struphy/initial/base.py new file mode 100644 index 000000000..0899ee8b6 --- /dev/null +++ b/src/struphy/initial/base.py @@ -0,0 +1,30 @@ +from typing import Callable + +from struphy.fields_background.base import FluidEquilibrium +from struphy.io.options import FieldsBackground +from struphy.kinetic_background.base import KineticBackground + + +class InitialCondition: + """Callable initial condition as sum of background + perturbation.""" + def __init__(self, + background: list = None, + perturbation: list = None, + equil: FluidEquilibrium = None,): + + for b in background: + assert isinstance(b, (FieldsBackground, KineticBackground)) + self._background = background + + for p in perturbation: + assert isinstance(p, tuple) + assert len(p) == 2 + assert isinstance(p[0], Callable) + assert isinstance(p[1], tuple) + self._perturbation = perturbation + + self._equil = equil + + def __call__(eta1, eta2, eta3, *v): + return 1 + \ No newline at end of file diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index b554a8eb8..eb53c3508 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -1,13 +1,13 @@ from struphy.io import options -from struphy.models import ( - fluid, - hybrid, - kinetic, - toy, -) +# from struphy.models import ( +# fluid, +# hybrid, +# kinetic, +# toy, +# ) -# model class (do not instantiate) -Model = toy.Maxwell +# import model +from struphy.models.toy import Maxwell as Model # units units = options.Units( @@ -38,13 +38,13 @@ # light-weight instance of model model = Model() -print(f"{model.__em_fields__=}") +species = model.species +propagators = model.propagators -# species parameters -species = model.species() print(f"{species.em_fields=}") print(f"{species.fluid=}") print(f"{species.kinetic=}") + # species.em_fields.set_options(prop, prop.options()) # model.fluid.set_phys_params("mhd", options.PhysParams()) # model.fluid.set_propagator_options("mhd", prop, prop.options()) @@ -81,6 +81,18 @@ given_in_basis=(None, "v", None), ) +# propagator options +print(f'{model.propagators.Maxwell = }') +print(f'{model.propagators.Maxwell.set_options = }') + + + + + + + + + # FOR NOW: initial conditions and options em_fields = {} em_fields["background"] = {} diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py new file mode 100644 index 000000000..be337f01e --- /dev/null +++ b/src/struphy/linear_algebra/solver.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass + +@dataclass +class SolverParameters: + """ + Parameters for psydac solvers.""" + + tol: float = 1e-8 + maxiter: int = 3000 + info: bool = False + verbose: bool = False + recycle: bool = True \ No newline at end of file diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 2374f1191..8f4742e15 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -2,6 +2,7 @@ import operator from abc import ABCMeta, abstractmethod from functools import reduce +from typing import Callable import numpy as np import yaml @@ -23,6 +24,8 @@ from struphy.propagators.base import Propagator from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml +from struphy.models.species import Species, SubSpecies +from struphy.propagators.base import Propagators class StruphyModel(metaclass=ABCMeta): @@ -48,41 +51,44 @@ class StruphyModel(metaclass=ABCMeta): def __init__( self, - params: StruphyParameters = None, + units=None, + time=None, + domain=None, + grid=None, + derham=None, + equil=None, comm: MPI.Intracomm = None, clone_config: CloneConfig = None, ): - assert "em_fields" in self.species() - assert "fluid" in self.species() - assert "kinetic" in self.species() - - assert "em_fields" in self.options() - assert "fluid" in self.options() - assert "kinetic" in self.options() - - if params is None: - params = self.generate_default_parameter_file( - save=False, - prompt=False, - ) + # assert "em_fields" in self.species() + # assert "fluid" in self.species() + # assert "kinetic" in self.species() + + # assert "em_fields" in self.options() + # assert "fluid" in self.options() + # assert "kinetic" in self.options() + # defaults + if units is None: + pass + + # mpi config self._comm_world = comm self._clone_config = clone_config - self._params = params - - # get rank and size if self.comm_world is None: self._rank_world = 0 else: self._rank_world = self.comm_world.Get_rank() - # initialize model variable dictionaries - self._init_variable_dicts() + # initialize static version of species and propagators + self.init_species() + self.init_propagators() + return # compute model units self._units, self._equation_params = self.model_units( - self.params, + units, verbose=self.verbose, comm=self.comm_world, ) @@ -206,58 +212,49 @@ def __init__( return params - @staticmethod - @abstractmethod - def species(): - """Species dictionary of the form {'em_fields': {}, 'fluid': {}, 'kinetic': {}}. - - The dynamical fields and kinetic species of the model. + ## abstract methods - Keys of the three sub-dicts are either: - - a) the electromagnetic field/potential names (b_field, e_field) - b) the fluid species names (e.g. mhd) - c) the names of the kinetic species (e.g. electrons, energetic_ions) - - Corresponding values are: - - a) a space ID ("H1", "Hcurl", "Hdiv", "L2" or "H1vec"), - b) a dict with key=variable_name (e.g. n, U, p, ...) and value=space ID ("H1", "Hcurl", "Hdiv", "L2" or "H1vec"), - c) the type of particles ("Particles6D", "Particles5D", ...).""" - pass + @abstractmethod + def init_species() -> Species: + """Static version of the species of a model (no runtime parameters).""" + @abstractmethod + def init_propagators() -> Propagators: + """Static version of the propagators of a model (no runtime parameters).""" + @staticmethod @abstractmethod - def bulk_species(): - """Name of the bulk species of the plasma. Must be a key of self.fluid or self.kinetic, or None.""" - pass + def bulk_species() -> SubSpecies: + """Bulk species of the plasma. Must be an attribute of species_static().""" @staticmethod @abstractmethod - def velocity_scale(): - """String that sets the velocity scale unit of the model. - Must be one of "alfvén", "cyclotron" or "light".""" - pass + def velocity_scale() -> str: + """Velocity unit scale of the model. + Must be one of "alfvén", "cyclotron", "light" or "thermal".""" @staticmethod def diagnostics_dct(): """Diagnostics dictionary. Model specific variables (FemField) which is going to be saved during the simulation. """ - pass - - @staticmethod - @abstractmethod - def propagators_dct(cls): - """Dictionary holding the propagators of the model in the sequence they should be called. - Keys are the propagator classes and values are lists holding variable names (str) updated by the propagator.""" - pass - + @abstractmethod def update_scalar_quantities(self): """Specify an update rule for each item in ``scalar_quantities`` using :meth:`update_scalar`.""" pass + ## basic properties + + @property + def species(self) -> Species: + return self._species + + @property + def propagators(self) -> Propagators: + """A list of propagator instances for the model.""" + return self._propagators + @property def params(self): """Model parameters from :code:`parameters.yml`.""" @@ -361,11 +358,6 @@ def prop_markers(self): """Module :mod:`struphy.propagators.propagators_markers`.""" return self._prop_markers - @property - def propagators(self): - """A list of propagator instances for the model.""" - return self._propagators - @property def kwargs(self): """Dictionary holding the keyword arguments for each propagator specified in :attr:`~propagators_cls`. @@ -1687,8 +1679,6 @@ def _init_variable_dicts(self): Initialize em-fields, fluid and kinetic dictionaries for information on the model variables. """ - # from struphy.models.variables import Variable - # electromagnetic fields, fluid and/or kinetic species self._em_fields = {} self._fluid = {} diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index e518513f8..ca1cd9faa 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -121,7 +121,10 @@ def __init__(self, name: str, space: str, save_data: bool = True): self._background = [] self._perturbation = [] - + + self._has_particles = False + if "Particles" in space: + self._has_particles = True ## attributes @property @@ -143,6 +146,10 @@ def background(self): @property def perturbation(self): return self._perturbation + + @property + def has_particles(self): + return self._has_particles @property def initial_condition(self): diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index fd195cb23..f92b2d38c 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,7 +1,9 @@ import numpy as np from struphy.models.base import StruphyModel +from struphy.models.species import Species from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers +from struphy.propagators.base import Propagators class Maxwell(StruphyModel): @@ -28,14 +30,17 @@ class Maxwell(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["em_fields"]["e_field"] = "Hcurl" - dct["em_fields"]["b_field"] = "Hdiv" - return dct - + def init_species(self): + self._species = Species(em_fields=True) + self._species.em_fields.add_variable(name="e_field", space="Hcurl") + self._species.em_fields.add_variable(name="b_field", space="Hdiv") + + def init_propagators(self): + self._propagators = Propagators() + self._propagators.add(propagators_fields.Maxwell, + self.species.em_fields.e_field, + self.species.em_fields.b_field,) + @staticmethod def bulk_species(): return None @@ -44,38 +49,34 @@ def bulk_species(): def velocity_scale(): return "light" - @staticmethod - def propagators_dct(): - return {propagators_fields.Maxwell: ["e_field", "b_field"]} + # __em_fields__ = [(v.name, v.space) for k, v in species_static().em_fields.all.items() if k != "_name"] + # __fluid_species__ = species()["fluid"] + # __kinetic_species__ = species()["kinetic"] + # __bulk_species__ = bulk_species() + # __velocity_scale__ = velocity_scale() + # __propagators__ = [prop.__name__ for prop in propagators_dct()] - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): + def __init__(self, params=None, comm=None, clone_config=None): # initialize base class super().__init__(params, comm=comm, clone_config=clone_config) # extract necessary parameters - algo = params.em_fields["options"]["Maxwell"]["algo"] - solver = params.em_fields["options"]["Maxwell"]["solver"] - - # set keyword arguments for propagators - self._kwargs[propagators_fields.Maxwell] = { - "algo": algo, - "solver": solver, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation - self.add_scalar("electric energy") - self.add_scalar("magnetic energy") - self.add_scalar("total energy") + # algo = params.em_fields["options"]["Maxwell"]["algo"] + # solver = params.em_fields["options"]["Maxwell"]["solver"] + + # # set keyword arguments for propagators + # self._kwargs[propagators_fields.Maxwell] = { + # "algo": algo, + # "solver": solver, + # } + + # # Initialize propagators used in splitting substeps + # self.init_propagators() + + # # Scalar variables to be saved during simulation + # self.add_scalar("electric energy") + # self.add_scalar("magnetic energy") + # self.add_scalar("total energy") def update_scalar_quantities(self): en_E = 0.5 * self.mass_ops.M1.dot_inner(self.pointer["e_field"], self.pointer["e_field"]) diff --git a/src/struphy/ode/utils.py b/src/struphy/ode/utils.py index 4f306a60d..c3fd06bd5 100644 --- a/src/struphy/ode/utils.py +++ b/src/struphy/ode/utils.py @@ -18,7 +18,7 @@ class ButcherTableau: """ @staticmethod - def available_methods(): + def available_methods() -> list: meth_avail = [ "rk4", "forward_euler", diff --git a/src/struphy/parameters/__init__.py b/src/struphy/parameters/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/struphy/parameters/derham_params.py b/src/struphy/parameters/derham_params.py new file mode 100644 index 000000000..8af2dc735 --- /dev/null +++ b/src/struphy/parameters/derham_params.py @@ -0,0 +1,21 @@ +from dataclasses import dataclass +from typing import Literal, get_args + +polar_options = Literal[-1, 1] + +@dataclass +class DerhamOptions: + """... + + Parameters + ---------- + x : float + Unit of length in m. + """ + + polar_ck: int = -1 + local_projectors: bool = False + + def __post_init__(self): + options = get_args(polar_options) + assert self.polar_ck in options, f"'{self.polar_ck}' is not in {options}" \ No newline at end of file diff --git a/src/struphy/parameters/time_params.py b/src/struphy/parameters/time_params.py new file mode 100644 index 000000000..8baf172d3 --- /dev/null +++ b/src/struphy/parameters/time_params.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass +from typing import Literal, get_args + +SplitAlgos = Literal["LieTrotter", "Strang"] + +@dataclass +class Time: + """... + + Parameters + ---------- + x : float + Unit of length in m. + """ + + dt: float = 1.0 + Tend: float = 1.0 + split_algo: SplitAlgos = "LieTrotter" + + def __post_init__(self): + options = get_args(SplitAlgos) + assert self.split_algo in options, f"'{self.split_algo}' is not in {options}" \ No newline at end of file diff --git a/src/struphy/parameters/units_params.py b/src/struphy/parameters/units_params.py new file mode 100644 index 000000000..b0f30857f --- /dev/null +++ b/src/struphy/parameters/units_params.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass + +@dataclass +class Units: + """... + + Parameters + ---------- + x : float + Unit of length in m. + """ + x: float = 1.0 + B: float = 1.0 + n: float = 1.0 + kBT: float = 1.0 \ No newline at end of file diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index ac7887ea8..83546ecb3 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -8,10 +8,11 @@ from struphy.feec.mass import WeightedMassOperators from struphy.feec.psydac_derham import Derham from struphy.geometry.base import Domain +from struphy.models.species import Variable class Propagator(metaclass=ABCMeta): - """Base class for Struphy propagators used in Struphy models. + """Base class for propagators used in StruphyModels. Note ---- @@ -25,50 +26,32 @@ def __init__(self, *vars): Parameters ---------- - vars : Vector or Particles - :attr:`struphy.models.base.StruphyModel.pointer` of variables to be updated. + vars : Variable + Variables to be updated. """ - from psydac.linalg.basic import Vector - - from struphy.pic.particles import Particles - - self._feec_vars = [] - self._particles = [] + comm = None for var in vars: - if isinstance(var, Vector): - self._feec_vars += [var] - elif isinstance(var, Particles): - self._particles += [var] + assert isinstance(var, Variable) + if var.has_particles: + # comm = var.obj.mpi_comm + pass else: - ValueError( - f'Variable {var} must be of type "Vector" or "Particles".', - ) + # comm = var.obj.comm + pass + self._vars = vars # for iterative particle push self._init_kernels = [] self._eval_kernels = [] - # mpi comm - if self.particles: - comm = self.particles[0].mpi_comm - else: - comm = self.derham.comm self._rank = comm.Get_rank() if comm is not None else 0 @property - def feec_vars(self): - """List of FEEC variables (not particles) to be updated by the propagator. - Contains FE coefficients from :attr:`struphy.feec.SplineFunction.vector`. - """ - return self._feec_vars - - @property - def particles(self): - """List of kinetic variables (not FEEC) to be updated by the propagator. - Contains :class:`struphy.pic.particles.Particles`. + def vars(self): + """List of Variables to be updated by the propagator. """ - return self._particles + return self._vars @property def init_kernels(self): @@ -92,9 +75,19 @@ def rank(self): """MPI rank, is 0 if no communicator.""" return self._rank + @abstractmethod + def set_options(self, **opts): + """Set the dynamical options of the propagator (kwargs). + """ + + @abstractmethod + def allocate(): + """Allocate all data/objects for an instance. + """ + @abstractmethod def __call__(self, dt): - """Update from t -> t + dt. + """Update variables from t -> t + dt. Use ``Propagators.feec_vars_update`` to write to FEEC variables to ``Propagator.feec_vars``. Parameters @@ -102,13 +95,6 @@ def __call__(self, dt): dt : float Time step size. """ - pass - - @staticmethod - @abstractmethod - def options(): - """Dictionary of available propagator options, as appearing under species/options in the parameter file.""" - pass @property def derham(self): @@ -310,3 +296,31 @@ def add_eval_kernel( args_eval, ) ] + + +class Propagators: + """Handels the Propagators of a StruphyModel.""" + def __init__(self): + pass + + @property + def all(self): + return self.__dict__ + + def add(self, prop: Propagator, *vars): + print(f'{prop = }') + print(f'{prop.__name__ = }') + for var in vars: + print(var) + setattr(self, prop.__name__, prop(*vars)) + + def set_options( + self, + name: str, + **opts, + ): + print(f'{self.all = }') + assert name in self.all, f"Propagator {name} is not part of model propagators {self.all.keys()}" + prop = getattr(self, name) + assert isinstance(prop, Propagator) + prop.set_options(**opts) \ No newline at end of file diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 6be69243a..f996d7b3c 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -2,6 +2,7 @@ from collections.abc import Callable from copy import deepcopy +from typing import Literal, get_args import numpy as np import scipy as sc @@ -42,6 +43,8 @@ from struphy.pic.particles import Particles5D, Particles6D from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator +from struphy.models.species import Variable +from struphy.linear_algebra.solver import SolverParameters class Maxwell(Propagator): @@ -57,38 +60,37 @@ class Maxwell(Propagator): :ref:`time_discret`: Crank-Nicolson (implicit mid-point). System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`. """ - @staticmethod - def options(default=False): - dct = {} - dct["algo"] = ["implicit"] + ButcherTableau.available_methods() - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, [], verbose=False) - - return dct - def __init__( self, - e: BlockVector, - b: BlockVector, - *, - algo: dict = options(default=True)["algo"], - solver: dict = options(default=True)["solver"], + e: Variable, + b: Variable, ): super().__init__(e, b) + OPTS_ALGO = Literal["implicit", *ButcherTableau.available_methods(),] + OPTS_SOLVER_TYPE = Literal[("pcg", "MassMatrixPreconditioner"), ("cg", None),] + + def set_options(self, + algo: OPTS_ALGO = "implicit", + solver_type: OPTS_SOLVER_TYPE = ("pcg", "MassMatrixPreconditioner"), + solver_params: SolverParameters = None, + ): + options = get_args(self.OPTS_ALGO) + assert algo in options, f"'{algo}' is not in {options}" + options = get_args(self.OPTS_SOLVER_TYPE) + assert solver_type in options, f"'{solver_type}' is not in {options}" + self._algo = algo + self._solver = solver_type + + if solver_params is None: + solver_params = SolverParameters() + self._solver_params = solver_params + + def allocate(self): + solver = {} + # obtain needed matrices M1 = self.mass_ops.M1 M2 = self.mass_ops.M2 From 3a0e5879cd3c526a9cb0e574112c5ae1779fc6cd Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 17 Jul 2025 14:24:01 +0200 Subject: [PATCH 014/292] add method `derive_units` to Units class; remove dataclass decorator --- src/struphy/io/options.py | 134 ++++++++++++++++++++++-- src/struphy/parameters/__init__.py | 0 src/struphy/parameters/derham_params.py | 21 ---- src/struphy/parameters/time_params.py | 22 ---- src/struphy/parameters/units_params.py | 15 --- 5 files changed, 127 insertions(+), 65 deletions(-) delete mode 100644 src/struphy/parameters/__init__.py delete mode 100644 src/struphy/parameters/derham_params.py delete mode 100644 src/struphy/parameters/time_params.py delete mode 100644 src/struphy/parameters/units_params.py diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index bb4ae3d54..0bd9a21ca 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -1,5 +1,6 @@ from dataclasses import dataclass from typing import Literal, get_args +import numpy as np # needed for import in StruphyParameters from struphy.fields_background import equils @@ -7,26 +8,145 @@ from struphy.initial import perturbations from struphy.kinetic_background import maxwellians from struphy.topology import grids +from struphy.physics.physics import ConstantsOfNature ## generic options SplitAlgos = Literal["LieTrotter", "Strang"] -@dataclass class Units: - """... + """ + Base units are passed to __init__, other units derive from these. Parameters ---------- x : float - Unit of length in m. + Unit of length in meters. + + B : float + Unit of magnetic field in Tesla. + + n : float + Unit of particle number density in 1e20/m^3. + + kBT : float, optional + Unit of internal energy in keV. + Only in effect if the velocity scale is set to 'thermal'. """ - x: float = 1.0 - B: float = 1.0 - n: float = 1.0 - kBT: float = None + def __init__(self, + x: float = 1.0, + B: float = 1.0, + n: float = 1.0, + kBT: float = None,): + + self._x = x + self._B = B + self._n = n * 1e20 + self._kBT = kBT + + @property + def x(self): + return self._x + + @property + def B(self): + return self._B + + @property + def n(self): + """Unit of particle number density in 1/m^3.""" + return self._n + + @property + def kBT(self): + return self._kBT + + @property + def v(self): + """Unit of velocity in m/s.""" + return self._v + + @property + def t(self): + """Unit of time in s.""" + return self._t + + @property + def p(self): + """Unit of pressure in Pa, equal to B^2/mu0 if velocity_scale='alfvén'.""" + return self._p + + @property + def rho(self): + """Unit of mass density in kg/m^3.""" + return self._rho + + @property + def j(self): + """Unit of current density in A/m^2.""" + return self._j + + 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.""" + + from mpi4py import MPI + + con = ConstantsOfNature() + + # velocity (m/s) + if velocity_scale is None: + self._v = 1.0 + + elif velocity_scale == "light": + self._v = con.c + + elif velocity_scale == "alfvén": + assert A_bulk is not None, 'Need bulk species to choose velocity scale "alfvén".' + self._v = self.B / np.sqrt(self.n * A_bulk * con.mH * con.mu0) + + elif velocity_scale == "cyclotron": + assert Z_bulk is not None, 'Need bulk species to choose velocity scale "cyclotron".' + assert A_bulk is not None, 'Need bulk species to choose velocity scale "cyclotron".' + self._v = Z_bulk * con.e * self.B / (A_bulk * con.mH) * self.x + + elif velocity_scale == "thermal": + assert A_bulk is not None, 'Need bulk species to choose velocity scale "thermal".' + assert self.kBT is not None + self._v = np.sqrt(self.kBT * 1000 * con.e / (con.mH * A_bulk)) + + # time (s) + self._t = self.x / self.v + + # return if no bulk is present + if A_bulk is None: + self._p = None + self._rho = None + self._j = None + else: + # pressure (Pa), equal to B^2/mu0 if velocity_scale='alfvén' + self._p = A_bulk * con.mH * self.n * self.v ** 2 + + # mass density (kg/m^3) + self._rho = A_bulk * con.mH * self.n + + # current density (A/m^2) + self._j = con.e * self.n * self.v + + # print to screen + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + units_used = (" m", " T", " m⁻³", "keV", " m/s", " s", " bar", " kg/m³", " A/m²",) + print("\nUNITS:") + for (k, v), u in zip(self.__dict__.items(), units_used): + if v is None: + print(f"Unit of {k[1:]} not specified.") + else: + print( + f"Unit of {k[1:]}:".ljust(25), + "{:4.3e}".format(v) + u, + ) @dataclass diff --git a/src/struphy/parameters/__init__.py b/src/struphy/parameters/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/struphy/parameters/derham_params.py b/src/struphy/parameters/derham_params.py deleted file mode 100644 index 8af2dc735..000000000 --- a/src/struphy/parameters/derham_params.py +++ /dev/null @@ -1,21 +0,0 @@ -from dataclasses import dataclass -from typing import Literal, get_args - -polar_options = Literal[-1, 1] - -@dataclass -class DerhamOptions: - """... - - Parameters - ---------- - x : float - Unit of length in m. - """ - - polar_ck: int = -1 - local_projectors: bool = False - - def __post_init__(self): - options = get_args(polar_options) - assert self.polar_ck in options, f"'{self.polar_ck}' is not in {options}" \ No newline at end of file diff --git a/src/struphy/parameters/time_params.py b/src/struphy/parameters/time_params.py deleted file mode 100644 index 8baf172d3..000000000 --- a/src/struphy/parameters/time_params.py +++ /dev/null @@ -1,22 +0,0 @@ -from dataclasses import dataclass -from typing import Literal, get_args - -SplitAlgos = Literal["LieTrotter", "Strang"] - -@dataclass -class Time: - """... - - Parameters - ---------- - x : float - Unit of length in m. - """ - - dt: float = 1.0 - Tend: float = 1.0 - split_algo: SplitAlgos = "LieTrotter" - - def __post_init__(self): - options = get_args(SplitAlgos) - assert self.split_algo in options, f"'{self.split_algo}' is not in {options}" \ No newline at end of file diff --git a/src/struphy/parameters/units_params.py b/src/struphy/parameters/units_params.py deleted file mode 100644 index b0f30857f..000000000 --- a/src/struphy/parameters/units_params.py +++ /dev/null @@ -1,15 +0,0 @@ -from dataclasses import dataclass - -@dataclass -class Units: - """... - - Parameters - ---------- - x : float - Unit of length in m. - """ - x: float = 1.0 - B: float = 1.0 - n: float = 1.0 - kBT: float = 1.0 \ No newline at end of file From ce3862a46c6abf8719e4986c64fb055ecba92714 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 17 Jul 2025 16:05:47 +0200 Subject: [PATCH 015/292] add function `check_option` to options.py --- src/struphy/io/options.py | 49 +++++++++++++++++++-------------------- 1 file changed, 24 insertions(+), 25 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 0bd9a21ca..f2fc61c54 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -10,10 +10,32 @@ from struphy.topology import grids from struphy.physics.physics import ConstantsOfNature + +def check_option(opt, options): + opts = get_args(options) + assert opt in opts, f"Option '{opt}' is not in {opts}." + ## generic options SplitAlgos = Literal["LieTrotter", "Strang"] +@dataclass +class Time: + """... + + Parameters + ---------- + x : float + Unit of length in m. + """ + + dt: float = 1.0 + Tend: float = 1.0 + split_algo: SplitAlgos = "LieTrotter" + + def __post_init__(self): + check_option(self.split_algo, SplitAlgos) + class Units: """ @@ -138,7 +160,6 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk # print to screen if verbose and MPI.COMM_WORLD.Get_rank() == 0: units_used = (" m", " T", " m⁻³", "keV", " m/s", " s", " bar", " kg/m³", " A/m²",) - print("\nUNITS:") for (k, v), u in zip(self.__dict__.items(), units_used): if v is None: print(f"Unit of {k[1:]} not specified.") @@ -149,31 +170,11 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk ) -@dataclass -class Time: - """... - - Parameters - ---------- - x : float - Unit of length in m. - """ - - dt: float = 1.0 - Tend: float = 1.0 - split_algo: SplitAlgos = "LieTrotter" - - def __post_init__(self): - options = get_args(SplitAlgos) - assert self.split_algo in options, f"'{self.split_algo}' is not in {options}" - - ## field options PolarRegularity = Literal[-1, 1] BackgroundOpts = Literal["LogicalConst", "FluidEquilibrium"] - @dataclass class DerhamOptions: """... @@ -188,8 +189,7 @@ class DerhamOptions: local_projectors: bool = False def __post_init__(self): - options = get_args(PolarRegularity) - assert self.polar_ck in options, f"'{self.polar_ck}' is not in {options}" + check_option(self.polar_ck, PolarRegularity) @dataclass @@ -207,8 +207,7 @@ class FieldsBackground: variable: str = None def __post_init__(self): - options = get_args(BackgroundOpts) - assert self.kind in options, f"'{self.kind}' is not in {options}" + check_option(self.kind, BackgroundOpts) ## kinetic options From 9c97d94aae9755e1c9bac04fb3f32a2a794e1878 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 17 Jul 2025 16:07:41 +0200 Subject: [PATCH 016/292] finished light-weight init of models base class --- src/struphy/io/inp/params_Maxwell_lw.py | 86 +++-------- src/struphy/models/base.py | 189 +++++++++++++----------- 2 files changed, 120 insertions(+), 155 deletions(-) diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index eb53c3508..fe89741c9 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -1,28 +1,21 @@ from struphy.io import options -# from struphy.models import ( -# fluid, -# hybrid, -# kinetic, -# toy, -# ) # import model from struphy.models.toy import Maxwell as Model +verbose = True # units -units = options.Units( - x=1.0, - B=1.0, - n=1.0, - kBT=1.0, -) - -# time -time = options.Time(split_algo="LieTrotter") +units = options.Units(x=2.0,) # geometry domain = options.domains.Cuboid() +# fluid equilibrium (can be used as part of initial conditions) +equil = options.equils.HomogenSlab() + +# time +time = options.Time() + # grid grid = options.grids.TensorProductGrid( Nel=(12, 14, 1), @@ -33,23 +26,17 @@ # derham options derham = options.DerhamOptions() -# fluid equilibrium (can be used as part of initial conditions) -equil = options.equils.HomogenSlab() - # light-weight instance of model -model = Model() +model = Model(verbose=verbose) species = model.species propagators = model.propagators -print(f"{species.em_fields=}") -print(f"{species.fluid=}") -print(f"{species.kinetic=}") - -# species.em_fields.set_options(prop, prop.options()) # model.fluid.set_phys_params("mhd", options.PhysParams()) -# model.fluid.set_propagator_options("mhd", prop, prop.options()) # model.kinetic.set_phys_params("mhd", options.PhysParams()) +# propagator options +propagators.Maxwell.set_options(verbose=verbose) + # initial conditions for model variables (background + perturbation) species.em_fields.add_background( "e_field", @@ -57,6 +44,7 @@ kind="LogicalConst", values=(0.3, 0.15, None), ), + verbose=verbose, ) species.em_fields.add_perturbation( "e_field", @@ -64,6 +52,7 @@ ms=[[None], [1, 3], [None]], ), given_in_basis=(None, "v", None), + verbose=verbose, ) species.em_fields.add_background( @@ -72,6 +61,7 @@ kind="LogicalConst", values=(0.3, 0.15, None), ), + verbose=verbose, ) species.em_fields.add_perturbation( "b_field", @@ -79,47 +69,5 @@ ms=[[None], [1, 3], [None]], ), given_in_basis=(None, "v", None), -) - -# propagator options -print(f'{model.propagators.Maxwell = }') -print(f'{model.propagators.Maxwell.set_options = }') - - - - - - - - - -# FOR NOW: initial conditions and options -em_fields = {} -em_fields["background"] = {} -em_fields["perturbation"] = {} -em_fields["options"] = {} - -em_fields["background"]["e_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} -em_fields["background"]["b_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} - -em_fields["perturbation"]["e_field"] = {} -em_fields["perturbation"]["b_field"] = {} -em_fields["perturbation"]["e_field"]["TorusModesCos"] = { - "given_in_basis": [None, "v", None], - "ms": [[None], [1, 3], [None]], -} -em_fields["perturbation"]["b_field"]["TorusModesCos"] = { - "given_in_basis": [None, "v", None], - "ms": [[None], [1, 3], [None]], -} - -solver = { - "type": ["pcg", "MassMatrixPreconditioner"], - "tol": 1.0e-08, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, -} - -em_fields["options"]["Maxwell"] = {"algo": "implicit", "solver": solver} + verbose=verbose, +) \ No newline at end of file diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 8f4742e15..2c9ca183c 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -22,10 +22,15 @@ from struphy.io.setup import setup_derham, setup_domain_and_equil from struphy.profiling.profiling import ProfileManager from struphy.propagators.base import Propagator +from struphy.propagators.hub import Propagators from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml -from struphy.models.species import Species, SubSpecies -from struphy.propagators.base import Propagators +from struphy.models.species import Species, SubSpecies +from struphy.io.setup import derive_units +from struphy.io.options import Units +from struphy.geometry.base import Domain +from struphy.geometry.domains import Cuboid +from struphy.fields_background.equils import HomogenSlab class StruphyModel(metaclass=ABCMeta): @@ -51,26 +56,23 @@ class StruphyModel(metaclass=ABCMeta): def __init__( self, - units=None, - time=None, - domain=None, - grid=None, - derham=None, - equil=None, + units: Units = None, + domain: Domain = None, + equil: FluidEquilibrium = None, comm: MPI.Intracomm = None, clone_config: CloneConfig = None, + verbose: bool = False, ): - # assert "em_fields" in self.species() - # assert "fluid" in self.species() - # assert "kinetic" in self.species() - - # assert "em_fields" in self.options() - # assert "fluid" in self.options() - # assert "kinetic" in self.options() # defaults if units is None: - pass + units = Units() + + if domain is None: + domain = Cuboid() + + if equil is None: + equil = HomogenSlab() # mpi config self._comm_world = comm @@ -84,17 +86,54 @@ def __init__( # initialize static version of species and propagators self.init_species() self.init_propagators() - return - # compute model units - self._units, self._equation_params = self.model_units( - units, - verbose=self.verbose, - comm=self.comm_world, - ) + # other light-weight inits + self._units = units + self.units.derive_units(verbose=verbose) + self.init_equation_params(units=self.units, verbose=verbose) + + self.config_domain_and_equil(domain, equil) + + def init_equation_params(self, units: Units, verbose=False): + # set equation parameters for each species + if self.species.fluid is not None: + for _, species in self.species.fluid.all.items(): + assert isinstance(species, SubSpecies) + species.set_equation_params(units=units, verbose=verbose) + + if self.species.kinetic is not None: + for name, species in self.species.kinetic.all.items(): + assert isinstance(species, SubSpecies) + species.set_equation_params(units=units, verbose=verbose) + + def config_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): + """If a numerical equilibirum is used, the domain is taken from this equilibirum. """ + if equil is not None: + self._equil = equil + if "Numerical" in self.equil.__class__.__name__: + self._domain = self.equil.domain + else: + self._domain = domain + self._equil.domain = domain + else: + self._domain = domain + self._equil = None + + def allocate(self, + grid=None, + derham_params=None, + time=None, + ): + + self.init_derham(grid, derham_params) + + self.allocate_species() + self.allocate_propagators() + + + + - # create domain, equilibrium - self._domain, self._equil = setup_domain_and_equil(params) if comm.Get_rank() == 0 and self.verbose: print("\nTIME:") @@ -221,13 +260,13 @@ def init_species() -> Species: @abstractmethod def init_propagators() -> Propagators: """Static version of the propagators of a model (no runtime parameters).""" - - @staticmethod + + @property @abstractmethod def bulk_species() -> SubSpecies: """Bulk species of the plasma. Must be an attribute of species_static().""" - @staticmethod + @property @abstractmethod def velocity_scale() -> str: """Velocity unit scale of the model. @@ -1255,27 +1294,15 @@ def initialize_data_output(self, data, size): # Class methods : ################### - @classmethod def model_units( - cls, - params: StruphyParameters, + self, + units: Units, verbose: bool = False, comm: MPI.Intracomm = None, ): """ Return model units and print them to screen. - Parameters - ---------- - params : StruphyParameters - model parameters. - - verbose : bool, optional - print model units to screen. - - comm : obj - MPI communicator. - Returns ------- units_basic : dict @@ -1285,36 +1312,26 @@ def model_units( Derived units for velocity, pressure, mass density and particle density. """ - from struphy.io.setup import derive_units - - if comm is None: - rank = 0 - else: - rank = comm.Get_rank() + print(f'{self.species.em_fields.all = }') # look for bulk species in fluid OR kinetic parameter dictionaries Z_bulk = None A_bulk = None - if params.fluid is not None: - if cls.bulk_species() in params["fluid"]: - Z_bulk = params["fluid"][cls.bulk_species()]["phys_params"]["Z"] - A_bulk = params["fluid"][cls.bulk_species()]["phys_params"]["A"] - if params.kinetic is not None: - if cls.bulk_species() in params["kinetic"]: - Z_bulk = params["kinetic"][cls.bulk_species()]["phys_params"]["Z"] - A_bulk = params["kinetic"][cls.bulk_species()]["phys_params"]["A"] + if self.bulk_species is not None: + Z_bulk = self.bulk_species.Z + A_bulk = self.bulk_species.A # compute model units - kBT = params.units.kBT + kBT = units.kBT units = derive_units( Z_bulk=Z_bulk, A_bulk=A_bulk, - x=params.units.x, - B=params.units.B, - n=params.units.n, + x=units.x, + B=units.B, + n=units.n, kBT=kBT, - velocity_scale=cls.velocity_scale(), + velocity_scale=self.velocity_scale, ) # print to screen @@ -1361,43 +1378,43 @@ def model_units( eps0 = 8.8541878128e-12 # vacuum permittivity (F/m) equation_params = {} - if params.fluid is not None: - for species in params["fluid"]: - Z = params["fluid"][species]["phys_params"]["Z"] - A = params["fluid"][species]["phys_params"]["A"] + if self.species.fluid is not None: + for name, species in self.species.fluid.all.items(): + Z = species.Z + A = species.A # compute equation parameters - om_p = np.sqrt(units["n"] * (Z * e) ** 2 / (eps0 * A * mH)) - om_c = Z * e * units["B"] / (A * mH) - equation_params[species] = {} - equation_params[species]["alpha"] = om_p / om_c - equation_params[species]["epsilon"] = 1.0 / (om_c * units["t"]) - equation_params[species]["kappa"] = om_p * units["t"] + om_p = np.sqrt(units.n * (Z * e) ** 2 / (eps0 * A * mH)) + om_c = Z * e * units.B / (A * mH) + equation_params[name] = {} + equation_params[name]["alpha"] = om_p / om_c + equation_params[name]["epsilon"] = 1.0 / (om_c * units["t"]) + equation_params[name]["kappa"] = om_p * units["t"] if verbose and MPI.COMM_WORLD.Get_rank() == 0: print("\nNORMALIZATION PARAMETERS:") - print("- " + species + ":") - for key, val in equation_params[species].items(): + print("- " + name + ":") + for key, val in equation_params[name].items(): print((key + ":").ljust(25), "{:4.3e}".format(val)) - if params.kinetic is not None: - for species in params["kinetic"]: - Z = params["kinetic"][species]["phys_params"]["Z"] - A = params["kinetic"][species]["phys_params"]["A"] + if self.species.kinetic is not None: + for name, species in self.species.kinetic.all.items(): + Z = species.Z + A = species.A # compute equation parameters - om_p = np.sqrt(units["n"] * (Z * e) ** 2 / (eps0 * A * mH)) - om_c = Z * e * units["B"] / (A * mH) - equation_params[species] = {} - equation_params[species]["alpha"] = om_p / om_c - equation_params[species]["epsilon"] = 1.0 / (om_c * units["t"]) - equation_params[species]["kappa"] = om_p * units["t"] + om_p = np.sqrt(units.n * (Z * e) ** 2 / (eps0 * A * mH)) + om_c = Z * e * units.B / (A * mH) + equation_params[name] = {} + equation_params[name]["alpha"] = om_p / om_c + equation_params[name]["epsilon"] = 1.0 / (om_c * units["t"]) + equation_params[name]["kappa"] = om_p * units["t"] if verbose and MPI.COMM_WORLD.Get_rank() == 0: - if "fluid" not in params: + if self.species.fluid is None: print("\nNORMALIZATION PARAMETERS:") - print("- " + species + ":") - for key, val in equation_params[species].items(): + print("- " + name + ":") + for key, val in equation_params[name].items(): print((key + ":").ljust(25), "{:4.3e}".format(val)) return units, equation_params From 3bcb924213d0493aee0a21b1830f92c14cbd0fa4 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 17 Jul 2025 16:09:26 +0200 Subject: [PATCH 017/292] add `set_phys_params` and `set_equation_params` to SubSpecies --- src/struphy/models/species.py | 66 ++++++++++++++++++++++++++++++++--- 1 file changed, 62 insertions(+), 4 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index ca1cd9faa..a32db13fd 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,9 +1,13 @@ from copy import deepcopy from typing import Callable +import numpy as np +from mpi4py import MPI from struphy.fields_background.base import FluidEquilibrium from struphy.initial.base import InitialCondition from struphy.kinetic_background.base import KineticBackground +from struphy.io.options import Units +from struphy.physics.physics import ConstantsOfNature class Species: @@ -87,17 +91,19 @@ def add_background( self, variable: str, background, + verbose=False, ): assert variable in self.all, f"Variable {variable} is not part of model variables {self.all.keys()}" var = getattr(self, variable) assert isinstance(var, Variable) - var.add_background(background) + var.add_background(background, verbose=verbose) def add_perturbation( self, variable: str, perturbation: Callable, given_in_basis: tuple = None, + verbose=False, ): assert variable in self.all, f"Variable {variable} is not part of model variables {self.all.keys()}" var = getattr(self, variable) @@ -105,10 +111,51 @@ def add_perturbation( var.add_perturbation( perturbation=perturbation, given_in_basis=given_in_basis, + verbose=verbose, ) - def set_options(self, propagator): - pass + @property + def Z(self) -> int: + """Charge number in units of elementary charge.""" + return self._Z + + @property + def A(self) -> int: + """Mass number in units of proton mass.""" + return self._A + + def set_phys_params(self, Z: int = 1, A: int = 1): + """Set charge- and mass number.""" + self._Z = Z + self._A = A + + @property + def equation_params(self) -> dict: + if not hasattr(self, "_equation_params"): + self.set_equation_params() + return self._equation_params + + def set_equation_params(self, units: Units = None, verbose=False): + Z = self.Z + A = self.A + + if units is None: + units = Units() + + con = ConstantsOfNature() + + # compute equation parameters + self._equation_params = {} + om_p = np.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_c = Z * con.e * units.B / (A * con.mH) + self._equation_params["alpha"] = om_p / om_c + self._equation_params["epsilon"] = 1.0 / (om_c * units["t"]) + self._equation_params["kappa"] = om_p * units["t"] + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f'Set normalization parameters for speceis {self.name}:') + for key, val in self.equation_params.items(): + print((key + ":").ljust(25), "{:4.3e}".format(val)) class Variable: @@ -159,17 +206,28 @@ def initial_condition(self): ## methods - def add_background(self, background): + def add_background(self, background, verbose=False,): # assert isinstance(...) self._background += [background] + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f"Variable '{self.name}': added background '{background.__class__.__name__}' with:") + for k, v in background.__dict__.items(): + print(f' {k}: {v}') def add_perturbation( self, perturbation: Callable = None, given_in_basis: tuple = None, + verbose=False, ): # assert isinstance(...) self._perturbation += [(perturbation, given_in_basis)] + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f"Variable '{self.name}': added perturbation '{perturbation.__class__.__name__}',{given_in_basis = }, with:") + for k, v in perturbation.__dict__.items(): + print(f' {k}: {v}') def set_initial_condition(self): self._initial_condition = InitialCondition( From a0fec5a10d727e20f94cf93c5de06d3ef6c57d0f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 17 Jul 2025 16:10:54 +0200 Subject: [PATCH 018/292] remove @staticmethod from Model.bulk_species and Model.velocity_scale --- src/struphy/models/toy.py | 30 +++++++++--------------------- 1 file changed, 9 insertions(+), 21 deletions(-) diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index f92b2d38c..0368f4e7b 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,9 +1,10 @@ +from dataclasses import dataclass import numpy as np from struphy.models.base import StruphyModel from struphy.models.species import Species from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.propagators.base import Propagators +from struphy.propagators.hub import Propagators class Maxwell(StruphyModel): @@ -40,13 +41,13 @@ def init_propagators(self): self._propagators.add(propagators_fields.Maxwell, self.species.em_fields.e_field, self.species.em_fields.b_field,) - - @staticmethod - def bulk_species(): + + @property + def bulk_species(self): return None - @staticmethod - def velocity_scale(): + @property + def velocity_scale(self): return "light" # __em_fields__ = [(v.name, v.space) for k, v in species_static().em_fields.all.items() if k != "_name"] @@ -56,22 +57,9 @@ def velocity_scale(): # __velocity_scale__ = velocity_scale() # __propagators__ = [prop.__name__ for prop in propagators_dct()] - def __init__(self, params=None, comm=None, clone_config=None): + def __init__(self, params=None, comm=None, clone_config=None, verbose=False): # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - # extract necessary parameters - # algo = params.em_fields["options"]["Maxwell"]["algo"] - # solver = params.em_fields["options"]["Maxwell"]["solver"] - - # # set keyword arguments for propagators - # self._kwargs[propagators_fields.Maxwell] = { - # "algo": algo, - # "solver": solver, - # } - - # # Initialize propagators used in splitting substeps - # self.init_propagators() + super().__init__(params, comm=comm, clone_config=clone_config, verbose=verbose) # # Scalar variables to be saved during simulation # self.add_scalar("electric energy") From dad496196a30a8ad64d5c62648e7955d41a53880 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 17 Jul 2025 16:11:37 +0200 Subject: [PATCH 019/292] add inner class Options to Propagator base --- src/struphy/propagators/base.py | 48 ++++++++++++++------------------- 1 file changed, 20 insertions(+), 28 deletions(-) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 83546ecb3..df4bafba5 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -1,8 +1,8 @@ "Propagator base class." from abc import ABCMeta, abstractmethod - import numpy as np +from mpi4py import MPI from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import WeightedMassOperators @@ -159,7 +159,7 @@ def projected_equil(self, projected_equil): @property def time_state(self): """A pointer to the time variable of the dynamics ('t').""" - return self._time_state + return self._time_state def add_time_state(self, time_state): """Add a pointer to the time variable of the dynamics ('t'). @@ -296,31 +296,23 @@ def add_eval_kernel( args_eval, ) ] - - -class Propagators: - """Handels the Propagators of a StruphyModel.""" - def __init__(self): - pass - + @property - def all(self): - return self.__dict__ - - def add(self, prop: Propagator, *vars): - print(f'{prop = }') - print(f'{prop.__name__ = }') - for var in vars: - print(var) - setattr(self, prop.__name__, prop(*vars)) + def opts(self): + if not hasattr(self, "_opts"): + self._opts = self.Options() + return self._opts - def set_options( - self, - name: str, - **opts, - ): - print(f'{self.all = }') - assert name in self.all, f"Propagator {name} is not part of model propagators {self.all.keys()}" - prop = getattr(self, name) - assert isinstance(prop, Propagator) - prop.set_options(**opts) \ No newline at end of file + class Options: + def __init__(self, outer_prop, verbose=False): + self._outer_prop = outer_prop.__class__.__name__ # outer class + self._verbose = verbose + + @property + def all(self): + return self.__dict__ + + def add(self, name: str, opt): + setattr(self, name, opt) + if self._verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f"Propagator '{self._outer_prop}': added option '{name}' with value '{opt}'") \ No newline at end of file From 70f534925cd083071223f2be5e2273975e784e1f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 17 Jul 2025 16:12:41 +0200 Subject: [PATCH 020/292] new folder `physics`; holds the natural constants for the moment --- src/struphy/physics/__init__.py | 0 src/struphy/physics/physics.py | 10 ++++++++++ 2 files changed, 10 insertions(+) create mode 100644 src/struphy/physics/__init__.py create mode 100644 src/struphy/physics/physics.py diff --git a/src/struphy/physics/__init__.py b/src/struphy/physics/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/struphy/physics/physics.py b/src/struphy/physics/physics.py new file mode 100644 index 000000000..ce2c4eced --- /dev/null +++ b/src/struphy/physics/physics.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass + +@dataclass +class ConstantsOfNature: + e = 1.602176634e-19 # elementary charge (C) + mH = 1.67262192369e-27 # proton mass (kg) + mu0 = 1.25663706212e-6 # magnetic constant (N/A^2) + eps0 = 8.8541878128e-12 # vacuum permittivity (F/m) + kB = 1.380649e-23 # Boltzmann constant (J/K) + c = 299792458 # speed of light (m/s) \ No newline at end of file From ece3f906e3d2cb9e94bb976287531ac336b16a7d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 17 Jul 2025 16:14:04 +0200 Subject: [PATCH 021/292] new files `propagators/hub.py` and `propagators/hub.pyi`; the Propagators class was moved there becuae the .pyi file takes precedence over .py for type inference --- src/struphy/propagators/hub.py | 29 +++++++++++++++++++++++++++++ src/struphy/propagators/hub.pyi | 28 ++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) create mode 100644 src/struphy/propagators/hub.py create mode 100644 src/struphy/propagators/hub.pyi diff --git a/src/struphy/propagators/hub.py b/src/struphy/propagators/hub.py new file mode 100644 index 000000000..04c8b921d --- /dev/null +++ b/src/struphy/propagators/hub.py @@ -0,0 +1,29 @@ +from struphy.propagators.base import Propagator + + +class Propagators: + """Handels the Propagators of a StruphyModel.""" + def __init__(self): + pass + + @property + def all(self): + return self.__dict__ + + def add(self, prop: Propagator, *vars): + # print(f'{prop = }') + # print(f'{prop.__name__ = }') + # for var in vars: + # print(var) + setattr(self, prop.__name__, prop(*vars)) + + def set_options( + self, + name: str, + **opts, + ): + print(f'{self.all = }') + assert name in self.all, f"Propagator {name} is not part of model propagators {self.all.keys()}" + prop = getattr(self, name) + assert isinstance(prop, Propagator) + prop.set_options(**opts) \ No newline at end of file diff --git a/src/struphy/propagators/hub.pyi b/src/struphy/propagators/hub.pyi new file mode 100644 index 000000000..6e8f188ff --- /dev/null +++ b/src/struphy/propagators/hub.pyi @@ -0,0 +1,28 @@ +from struphy.propagators import propagators_fields +from struphy.propagators.base import Pro + +class Propagator: + + @property + def vars(self): + ... + + @property + def init_kernels(self): + ... + + @property + def eval_kernels(self): + ... + +class Propagators: + + @property + def all(self) -> dict: ... + + def add(self, prop: Propagator, *vars): ... + + @property + def Maxwell(self) -> propagators_fields.Maxwell: ... + + \ No newline at end of file From 0261be4fe595be51e742cb6ab1473586afa15320 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 17 Jul 2025 16:17:50 +0200 Subject: [PATCH 022/292] adapt Maxwell propagator --- src/struphy/propagators/propagators_fields.py | 30 +++++++++++-------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index f996d7b3c..02070ba09 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -45,6 +45,7 @@ from struphy.propagators.base import Propagator from struphy.models.species import Variable from struphy.linear_algebra.solver import SolverParameters +from struphy.io.options import check_option class Maxwell(Propagator): @@ -67,26 +68,29 @@ def __init__( ): super().__init__(e, b) - OPTS_ALGO = Literal["implicit", *ButcherTableau.available_methods(),] - OPTS_SOLVER_TYPE = Literal[("pcg", "MassMatrixPreconditioner"), ("cg", None),] + OptsAlgo = Literal["implicit", *ButcherTableau.available_methods(),] + OptsSolverType = Literal[("pcg", "MassMatrixPreconditioner"), ("cg", None),] def set_options(self, - algo: OPTS_ALGO = "implicit", - solver_type: OPTS_SOLVER_TYPE = ("pcg", "MassMatrixPreconditioner"), + algo: OptsAlgo = "implicit", # type: ignore + solver_type: OptsSolverType = ("pcg", "MassMatrixPreconditioner"), # type: ignore solver_params: SolverParameters = None, + verbose = False, ): - options = get_args(self.OPTS_ALGO) - assert algo in options, f"'{algo}' is not in {options}" - options = get_args(self.OPTS_SOLVER_TYPE) - assert solver_type in options, f"'{solver_type}' is not in {options}" - - self._algo = algo - self._solver = solver_type + + # checks + check_option(algo, self.OptsAlgo) + check_option(solver_type, self.OptsSolverType) + # defaults if solver_params is None: solver_params = SolverParameters() - self._solver_params = solver_params - + + # create options + self._opts = self.Options(self, verbose=verbose) + self.opts.add("algo", algo) + self.opts.add("solver", solver_type) + self.opts.add("solver_params", solver_params) def allocate(self): solver = {} From bf78073c51d7a73572fd9e249720c4927f588c94 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 17 Jul 2025 17:24:17 +0200 Subject: [PATCH 023/292] Object-oriented brainstorming --- src/struphy/io/inp/params_Maxwell_lw.py | 2 + src/struphy/models/species.py | 12 +++--- src/struphy/models/toy.py | 40 ++++++++++++++++++- src/struphy/propagators/base.py | 13 +++--- src/struphy/propagators/propagators_fields.py | 10 +++-- 5 files changed, 61 insertions(+), 16 deletions(-) diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index fe89741c9..e3310da5a 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -31,6 +31,8 @@ species = model.species propagators = model.propagators +model.species2.em_fields.e_field. +model.fluids.ions.add_variable() # model.fluid.set_phys_params("mhd", options.PhysParams()) # model.kinetic.set_phys_params("mhd", options.PhysParams()) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index a32db13fd..98b630e1a 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -55,11 +55,6 @@ def em_fields(self): class MultiSpecies: """Handles multiple fluid or kinetic species.""" - def __init__( - self, - ): - pass - @property def all(self): return self.__dict__ @@ -157,6 +152,13 @@ def set_equation_params(self, units: Units = None, verbose=False): for key, val in self.equation_params.items(): print((key + ":").ljust(25), "{:4.3e}".format(val)) +class FluidSpecies(SubSpecies): + def __init__(self): + pass + +class KineticSpecies(SubSpecies): + def __init__(self): + pass class Variable: """Single variable (unknown) of a StruphyModel.""" diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 0368f4e7b..09575c341 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,3 +1,5 @@ +from struphy.models.species import SubSpecies, MultiSpecies, Variable, KineticSpecies, FluidSpecies + from dataclasses import dataclass import numpy as np @@ -30,14 +32,48 @@ class Maxwell(StruphyModel): :ref:`Model info `: """ - + def init_species(self): self._species = Species(em_fields=True) self._species.em_fields.add_variable(name="e_field", space="Hcurl") self._species.em_fields.add_variable(name="b_field", space="Hdiv") + + + @dataclass + class MaxwellPropagators(Propagators): + # TODO: Make sure the propagators are in the right order + Maxwell = propagators_fields.Maxwell() + + @dataclass + class EMFields: + e_field = Variable(space="Hcurl") + b_field = Variable(space="Hdiv") + + @dataclass + class Ions(KineticSpecies): + density = Variable(space="H0") + pressure = Variable(space="H0") + + @dataclass + class Electrons(FluidSpecies): + density = Variable(space="H0") + pressure = Variable(space="H0") + + def __init__(self): + + self._em_fields = self.EMFields() + self._ions = self.Ions() + self._electrons = self.Electrons() + + self._propagators = self.MaxwellPropagators() + self._propagators.Maxwell.set_variables( + e = self._em_fields.e_field, + b = self._em_fields.b_field, + ) def init_propagators(self): - self._propagators = Propagators() + # self._propagators = Propagators() + self._propagators.add(propagators_fields.Maxwell, self.species.em_fields.e_field, self.species.em_fields.b_field,) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index df4bafba5..2072a8760 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -31,6 +31,14 @@ def __init__(self, *vars): """ comm = None + + # for iterative particle push + self._init_kernels = [] + self._eval_kernels = [] + + self._rank = comm.Get_rank() if comm is not None else 0 + + def set_variables(self, *vars): for var in vars: assert isinstance(var, Variable) if var.has_particles: @@ -41,11 +49,6 @@ def __init__(self, *vars): pass self._vars = vars - # for iterative particle push - self._init_kernels = [] - self._eval_kernels = [] - - self._rank = comm.Get_rank() if comm is not None else 0 @property def vars(self): diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 02070ba09..2ef90d0f5 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -63,14 +63,16 @@ class Maxwell(Propagator): def __init__( self, - e: Variable, - b: Variable, + # e: Variable, + # b: Variable, ): - super().__init__(e, b) + self.e = None + self.b = None + OptsAlgo = Literal["implicit", *ButcherTableau.available_methods(),] OptsSolverType = Literal[("pcg", "MassMatrixPreconditioner"), ("cg", None),] - + def set_options(self, algo: OptsAlgo = "implicit", # type: ignore solver_type: OptsSolverType = ("pcg", "MassMatrixPreconditioner"), # type: ignore From ccdeaba324d18b1ed313ac89bbd1fd8fa8affa25 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 18 Jul 2025 14:10:56 +0200 Subject: [PATCH 024/292] redo Species structure: FieldSpecies, FluidSpecies, KineticSpecies --- src/struphy/models/species.py | 228 +++++++++------------------------- 1 file changed, 59 insertions(+), 169 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 98b630e1a..da164807e 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,3 +1,5 @@ +from abc import ABCMeta, abstractmethod + from copy import deepcopy from typing import Callable import numpy as np @@ -10,57 +12,71 @@ from struphy.physics.physics import ConstantsOfNature -class Species: - """Handles the three species types in StruphyModel: em_fields, fluid, kinetic.""" - - def __init__( - self, - fluid: tuple = None, - kinetic: tuple = None, - em_fields: bool = True, - ): - # fluid species - if fluid is None: - self._fluid = None - else: - self._fluid = MultiSpecies() - for f in fluid: - self._fluid.add_species(name=f) - - # kinetic species - if kinetic is None: - self._kinetic = None - else: - self._kinetic = MultiSpecies() - for k in kinetic: - self._kinetic.add_species(name=k) - - # electromagnetic fields - if em_fields: - self._em_fields = SubSpecies(name="em_fields") - +class Species(metaclass=ABCMeta): + """Single species of a StruphyModel.""" @property - def fluid(self): - return self._fluid - + def charge_number(self) -> int: + """Charge number in units of elementary charge.""" + return self._charge_number + @property - def kinetic(self): - return self._kinetic + def mass_number(self) -> int: + """Mass number in units of proton mass.""" + return self._mass_number + def set_phys_params(self, charge_number: int = 1, mass_number: int = 1): + """Set charge- and mass number.""" + self._charge_number = charge_number + self._mass_number = mass_number + @property - def em_fields(self): - return self._em_fields + def equation_params(self) -> dict: + if not hasattr(self, "_equation_params"): + self.setup_equation_params() + return self._equation_params + + def setup_equation_params(self, units: Units = None, verbose=False): + """Set the following equation parameters: + + * alpha = plasma-frequenca / cyclotron frequency + * epsilon = 1 / (cyclotron frequency * time unit) + * kappa = plasma frequency * time unit + """ + Z = self.charge_number + A = self.mass_number + + if units is None: + units = Units() + con = ConstantsOfNature() + + # compute equation parameters + self._equation_params = {} + om_p = np.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_c = Z * con.e * units.B / (A * con.mH) + self._equation_params["alpha"] = om_p / om_c + self._equation_params["epsilon"] = 1.0 / (om_c * units["t"]) + self._equation_params["kappa"] = om_p * units["t"] -class MultiSpecies: - """Handles multiple fluid or kinetic species.""" + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f'Set normalization parameters for speceis {self.name}:') + for key, val in self.equation_params.items(): + print((key + ":").ljust(25), "{:4.3e}".format(val)) + + +class FieldSpecies(Species): + """Species without mass and charge (so-called 'fields').""" + pass - @property - def all(self): - return self.__dict__ - def add_species(self, name: str = "mhd"): - setattr(self, name, SubSpecies(name)) +class FluidSpecies(Species): + """Single fluid species in 3d configuration space.""" + pass + + +class KineticSpecies(Species): + """Single kinetic species in 3d + vdim phase space.""" + pass class SubSpecies: @@ -109,130 +125,4 @@ def add_perturbation( verbose=verbose, ) - @property - def Z(self) -> int: - """Charge number in units of elementary charge.""" - return self._Z - - @property - def A(self) -> int: - """Mass number in units of proton mass.""" - return self._A - - def set_phys_params(self, Z: int = 1, A: int = 1): - """Set charge- and mass number.""" - self._Z = Z - self._A = A - @property - def equation_params(self) -> dict: - if not hasattr(self, "_equation_params"): - self.set_equation_params() - return self._equation_params - - def set_equation_params(self, units: Units = None, verbose=False): - Z = self.Z - A = self.A - - if units is None: - units = Units() - - con = ConstantsOfNature() - - # compute equation parameters - self._equation_params = {} - om_p = np.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) - om_c = Z * con.e * units.B / (A * con.mH) - self._equation_params["alpha"] = om_p / om_c - self._equation_params["epsilon"] = 1.0 / (om_c * units["t"]) - self._equation_params["kappa"] = om_p * units["t"] - - if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f'Set normalization parameters for speceis {self.name}:') - for key, val in self.equation_params.items(): - print((key + ":").ljust(25), "{:4.3e}".format(val)) - -class FluidSpecies(SubSpecies): - def __init__(self): - pass - -class KineticSpecies(SubSpecies): - def __init__(self): - pass - -class Variable: - """Single variable (unknown) of a StruphyModel.""" - - def __init__(self, name: str, space: str, save_data: bool = True): - self._name = name - self._space = space - self._save_data = save_data - - self._background = [] - self._perturbation = [] - - self._has_particles = False - if "Particles" in space: - self._has_particles = True - ## attributes - - @property - def name(self): - return self._name - - @property - def space(self): - return self._space - - @property - def save_data(self): - return self._save_data - - @property - def background(self): - return self._background - - @property - def perturbation(self): - return self._perturbation - - @property - def has_particles(self): - return self._has_particles - - @property - def initial_condition(self): - if not hasattr(self, "_initial_condition"): - self.set_initial_condition() - return self._initial_condition - - ## methods - - def add_background(self, background, verbose=False,): - # assert isinstance(...) - self._background += [background] - - if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"Variable '{self.name}': added background '{background.__class__.__name__}' with:") - for k, v in background.__dict__.items(): - print(f' {k}: {v}') - - def add_perturbation( - self, - perturbation: Callable = None, - given_in_basis: tuple = None, - verbose=False, - ): - # assert isinstance(...) - self._perturbation += [(perturbation, given_in_basis)] - - if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"Variable '{self.name}': added perturbation '{perturbation.__class__.__name__}',{given_in_basis = }, with:") - for k, v in perturbation.__dict__.items(): - print(f' {k}: {v}') - - def set_initial_condition(self): - self._initial_condition = InitialCondition( - background=self.background, - perturbation=self.perturbation, - ) From 629cc15fc2c30f0c2299cfe6bf7ea1eae3c6c791 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 18 Jul 2025 14:14:24 +0200 Subject: [PATCH 025/292] added new file models/variables.py --- src/struphy/io/inp/params_Maxwell_lw.py | 20 +- src/struphy/models/base.py | 183 ++++++++---------- src/struphy/models/toy.py | 93 ++++----- src/struphy/models/variables.py | 67 +++++++ src/struphy/propagators/base.py | 7 +- src/struphy/propagators/propagators_fields.py | 2 +- 6 files changed, 198 insertions(+), 174 deletions(-) create mode 100644 src/struphy/models/variables.py diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index e3310da5a..ab2df300c 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -27,29 +27,23 @@ derham = options.DerhamOptions() # light-weight instance of model -model = Model(verbose=verbose) -species = model.species +model = Model(units, domain, equil, verbose=verbose) propagators = model.propagators - -model.species2.em_fields.e_field. -model.fluids.ions.add_variable() # model.fluid.set_phys_params("mhd", options.PhysParams()) # model.kinetic.set_phys_params("mhd", options.PhysParams()) # propagator options -propagators.Maxwell.set_options(verbose=verbose) +propagators.maxwell.set_options(verbose=verbose) # initial conditions for model variables (background + perturbation) -species.em_fields.add_background( - "e_field", +model.em_fields.e_field.add_background( options.FieldsBackground( kind="LogicalConst", values=(0.3, 0.15, None), ), verbose=verbose, ) -species.em_fields.add_perturbation( - "e_field", +model.em_fields.e_field.add_perturbation( options.perturbations.TorusModesCos( ms=[[None], [1, 3], [None]], ), @@ -57,16 +51,14 @@ verbose=verbose, ) -species.em_fields.add_background( - "b_field", +model.em_fields.b_field.add_background( options.FieldsBackground( kind="LogicalConst", values=(0.3, 0.15, None), ), verbose=verbose, ) -species.em_fields.add_perturbation( - "b_field", +model.em_fields.b_field.add_perturbation( options.perturbations.TorusModesCos( ms=[[None], [1, 3], [None]], ), diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 2c9ca183c..d2089eff8 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -25,7 +25,7 @@ from struphy.propagators.hub import Propagators from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml -from struphy.models.species import Species, SubSpecies +from struphy.models.species import Species, FieldSpecies, FluidSpecies, KineticSpecies, SubSpecies from struphy.io.setup import derive_units from struphy.io.options import Units from struphy.geometry.base import Domain @@ -54,15 +54,39 @@ class StruphyModel(metaclass=ABCMeta): in one of the modules ``fluid.py``, ``kinetic.py``, ``hybrid.py`` or ``toy.py``. """ - def __init__( + ## abstract methods + + @abstractmethod + class Propagators: + pass + + @property + @abstractmethod + def bulk_species() -> SubSpecies: + """Bulk species of the plasma. Must be an attribute of species_static().""" + + @property + @abstractmethod + def velocity_scale() -> str: + """Velocity unit scale of the model. + Must be one of "alfvén", "cyclotron", "light" or "thermal".""" + + @abstractmethod + def update_scalar_quantities(self): + """Specify an update rule for each item in ``scalar_quantities`` using :meth:`update_scalar`.""" + + ## setup methods + + def setup( self, units: Units = None, domain: Domain = None, equil: FluidEquilibrium = None, comm: MPI.Intracomm = None, clone_config: CloneConfig = None, - verbose: bool = False, + verbose=False, ): + """Light-weight initialization of StruphyModel. Must be called in each Propagator.__init__().""" # defaults if units is None: @@ -83,30 +107,23 @@ def __init__( else: self._rank_world = self.comm_world.Get_rank() - # initialize static version of species and propagators - self.init_species() - self.init_propagators() - # other light-weight inits self._units = units self.units.derive_units(verbose=verbose) - self.init_equation_params(units=self.units, verbose=verbose) - - self.config_domain_and_equil(domain, equil) - - def init_equation_params(self, units: Units, verbose=False): - # set equation parameters for each species - if self.species.fluid is not None: - for _, species in self.species.fluid.all.items(): - assert isinstance(species, SubSpecies) - species.set_equation_params(units=units, verbose=verbose) - - if self.species.kinetic is not None: - for name, species in self.species.kinetic.all.items(): - assert isinstance(species, SubSpecies) - species.set_equation_params(units=units, verbose=verbose) + self.setup_equation_params(units=self.units, verbose=verbose) + self.setup_domain_and_equil(domain, equil) - def config_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): + def setup_equation_params(self, units: Units, verbose=False): + """Set euqation parameters for each fluid and kinetic species.""" + for _, species in self.fluid_species.items(): + assert isinstance(species, FluidSpecies) + species.setup_equation_params(units=units, verbose=verbose) + + for _, species in self.kinetic_species.items(): + assert isinstance(species, KineticSpecies) + species.setup_equation_params(units=units, verbose=verbose) + + def setup_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): """If a numerical equilibirum is used, the domain is taken from this equilibirum. """ if equil is not None: self._equil = equil @@ -118,6 +135,45 @@ def config_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): else: self._domain = domain self._equil = None + + ## species + + @property + def field_species(self) -> dict: + spec = {} + for k, v in self.__dict__.items(): + if isinstance(v, FieldSpecies): + spec[k] = v + return spec + + @property + def fluid_species(self) -> dict: + spec = {} + for k, v in self.__dict__.items(): + if isinstance(v, FluidSpecies): + spec[k] = v + return spec + + @property + def kinetic_species(self) -> dict: + spec = {} + for k, v in self.__dict__.items(): + if isinstance(v, KineticSpecies): + spec[k] = v + return spec + + @property + def all_species(self): + spec = {} + for k, v in self.field_species.items(): + spec[k] = v + for k, v in self.fluid_species.items(): + spec[k] = v + for k, v in self.kinetic_species.items(): + spec[k] = v + return spec + + ## allocate methods def allocate(self, grid=None, @@ -245,55 +301,27 @@ def allocate(self, Propagator.projected_equil = self.projected_equil # create dummy lists/dicts to be filled by the sub-class - self._propagators = [] + self._kwargs = {} self._scalar_quantities = {} return params - ## abstract methods - @abstractmethod - def init_species() -> Species: - """Static version of the species of a model (no runtime parameters).""" - - @abstractmethod - def init_propagators() -> Propagators: - """Static version of the propagators of a model (no runtime parameters).""" - @property - @abstractmethod - def bulk_species() -> SubSpecies: - """Bulk species of the plasma. Must be an attribute of species_static().""" - - @property - @abstractmethod - def velocity_scale() -> str: - """Velocity unit scale of the model. - Must be one of "alfvén", "cyclotron", "light" or "thermal".""" + @staticmethod def diagnostics_dct(): """Diagnostics dictionary. Model specific variables (FemField) which is going to be saved during the simulation. """ - - @abstractmethod - def update_scalar_quantities(self): - """Specify an update rule for each item in ``scalar_quantities`` using :meth:`update_scalar`.""" - pass ## basic properties - @property def species(self) -> Species: return self._species - @property - def propagators(self) -> Propagators: - """A list of propagator instances for the model.""" - return self._propagators - @property def params(self): """Model parameters from :code:`parameters.yml`.""" @@ -332,21 +360,6 @@ def pointer(self): In case of a fluid species, the keys are like "species_variable".""" return self._pointer - @property - def em_fields(self): - """Dictionary of electromagnetic field/potential variables.""" - return self._em_fields - - @property - def fluid(self): - """Dictionary of fluid species.""" - return self._fluid - - @property - def kinetic(self): - """Dictionary of kinetic species.""" - return self._kinetic - @property def diagnostics(self): """Dictionary of diagnostics.""" @@ -645,40 +658,6 @@ def add_time_state(self, time_state): for prop in self.propagators: prop.add_time_state(time_state) - def init_propagators(self): - """Initialize the propagator objects specified in :attr:`~propagators_cls`.""" - if self.rank_world == 0 and self.verbose: - print("\nPROPAGATORS:") - for (prop, variables), (prop2, kwargs_i) in zip(self.propagators_dct().items(), self.kwargs.items()): - assert prop == prop2, ( - f'Propagators {prop} from "self.propagators_dct()" and {prop2} from "self.kwargs" must be identical !!' - ) - - if kwargs_i is None: - if self.rank_world == 0: - print(f'\n-> Propagator "{prop.__name__}" will not be used.') - continue - else: - if self.rank_world == 0 and self.verbose: - print(f'\n-> Initializing propagator "{prop.__name__}"') - print(f"-> for variables {variables}") - print(f"-> with the following parameters:") - for k, v in kwargs_i.items(): - if isinstance(v, StencilVector): - print(f"{k}: {repr(v)}") - else: - print(f"{k}: {v}") - - prop_instance = prop( - *[self.pointer[var] for var in variables], - **kwargs_i, - ) - assert isinstance(prop_instance, Propagator) - self._propagators += [prop_instance] - - if self.rank_world == 0 and self.verbose: - print("\nInitialization of propagators complete.") - def integrate(self, dt, split_algo="LieTrotter"): """ Advance the model by a time step ``dt`` by sequentially calling its Propagators. diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 09575c341..bfe80e92e 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,12 +1,11 @@ -from struphy.models.species import SubSpecies, MultiSpecies, Variable, KineticSpecies, FluidSpecies from dataclasses import dataclass import numpy as np from struphy.models.base import StruphyModel -from struphy.models.species import Species from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.propagators.hub import Propagators +from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable class Maxwell(StruphyModel): @@ -33,50 +32,54 @@ class Maxwell(StruphyModel): :ref:`Model info `: """ - def init_species(self): - self._species = Species(em_fields=True) - self._species.em_fields.add_variable(name="e_field", space="Hcurl") - self._species.em_fields.add_variable(name="b_field", space="Hdiv") - - @dataclass - class MaxwellPropagators(Propagators): - # TODO: Make sure the propagators are in the right order - Maxwell = propagators_fields.Maxwell() + class Propagators: + maxwell = propagators_fields.Maxwell() @dataclass - class EMFields: - e_field = Variable(space="Hcurl") - b_field = Variable(space="Hdiv") + class EMFields(FieldSpecies): + e_field = FEECVariable(space="Hcurl") + b_field = FEECVariable(space="Hdiv") - @dataclass - class Ions(KineticSpecies): - density = Variable(space="H0") - pressure = Variable(space="H0") + # @dataclass + # class Ions(KineticSpecies): + # distribution_function = PICVariable(type="Particles6D") - @dataclass - class Electrons(FluidSpecies): - density = Variable(space="H0") - pressure = Variable(space="H0") + # @dataclass + # class Electrons(FluidSpecies): + # density = FEECVariable(space="H0") + # pressure = FEECVariable(space="H0") - def __init__(self): - + def __init__(self, units, domain, equil, verbose=False): + # 1. instantiate all variales self._em_fields = self.EMFields() - self._ions = self.Ions() - self._electrons = self.Electrons() + # self._ions = self.Ions() + # self._electrons = self.Electrons() - self._propagators = self.MaxwellPropagators() - self._propagators.Maxwell.set_variables( - e = self._em_fields.e_field, - b = self._em_fields.b_field, + # 2. instantiate all propagators + self._propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.maxwell.set_variables( + self.em_fields.e_field, + self.em_fields.b_field, ) - - def init_propagators(self): - # self._propagators = Propagators() - self._propagators.add(propagators_fields.Maxwell, - self.species.em_fields.e_field, - self.species.em_fields.b_field,) + # setup rest of model + self.setup(units=units, domain=domain, equil=equil, verbose=verbose) + + ## variable and propagator attributes + + @property + def propagators(self) -> Propagators: + """A list of propagator instances for the model.""" + return self._propagators + + @property + def em_fields(self) -> EMFields: + return self._em_fields + + ## abstract attributes @property def bulk_species(self): @@ -86,22 +89,6 @@ def bulk_species(self): def velocity_scale(self): return "light" - # __em_fields__ = [(v.name, v.space) for k, v in species_static().em_fields.all.items() if k != "_name"] - # __fluid_species__ = species()["fluid"] - # __kinetic_species__ = species()["kinetic"] - # __bulk_species__ = bulk_species() - # __velocity_scale__ = velocity_scale() - # __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params=None, comm=None, clone_config=None, verbose=False): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config, verbose=verbose) - - # # Scalar variables to be saved during simulation - # self.add_scalar("electric energy") - # self.add_scalar("magnetic energy") - # self.add_scalar("total energy") - def update_scalar_quantities(self): en_E = 0.5 * self.mass_ops.M1.dot_inner(self.pointer["e_field"], self.pointer["e_field"]) en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py new file mode 100644 index 000000000..2b8f48bc9 --- /dev/null +++ b/src/struphy/models/variables.py @@ -0,0 +1,67 @@ +from abc import ABCMeta, abstractmethod +from mpi4py import MPI + +from struphy.initial.base import InitialCondition + + +class Variable(metaclass=ABCMeta): + """Single variable (unknown) of a Species.""" + + @property + def backgrounds(self): + return self._backgrounds + + @property + def perturbations(self): + return self._perturbations + + def add_background(self, background = None, verbose=False,): + # assert isinstance(...) + if not hasattr(self, "_backgrounds"): + self._backgrounds = [] + self._backgrounds += [background] + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f"Variable '{self.__class__.__name__}': added background '{background.__class__.__name__}' with:") + for k, v in background.__dict__.items(): + print(f' {k}: {v}') + + def add_perturbation( + self, + perturbation = None, + given_in_basis: tuple = None, + verbose=False, + ): + # assert isinstance(...) + if not hasattr(self, "_perturbations"): + self._perturbations = [] + self._perturbations += [(perturbation, given_in_basis)] + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f"Variable '{self.__class__.__name__}': added perturbation '{perturbation.__class__.__name__}',{given_in_basis = }, with:") + for k, v in perturbation.__dict__.items(): + print(f' {k}: {v}') + + def define_initial_condition(self): + self._initial_condition = InitialCondition( + background=self.backgrounds, + perturbation=self.perturbations, + ) + + +class FEECVariable(Variable): + def __init__(self, space: str = "H1"): + assert space in ("H1", "Hcurl", "Hdiv", "L2", "H1vec") + self._space = space + + @property + def space(self): + return self._space + + +class PICVariable(Variable): + pass + + +class SPHVariable(Variable): + pass \ No newline at end of file diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 2072a8760..e2d82bd23 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -8,7 +8,7 @@ from struphy.feec.mass import WeightedMassOperators from struphy.feec.psydac_derham import Derham from struphy.geometry.base import Domain -from struphy.models.species import Variable +from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable class Propagator(metaclass=ABCMeta): @@ -41,15 +41,14 @@ def __init__(self, *vars): def set_variables(self, *vars): for var in vars: assert isinstance(var, Variable) - if var.has_particles: + if isinstance(var, (PICVariable, SPHVariable)): # comm = var.obj.mpi_comm pass - else: + elif isinstance(var, FEECVariable): # comm = var.obj.comm pass self._vars = vars - @property def vars(self): """List of Variables to be updated by the propagator. diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 2ef90d0f5..850b2fc28 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -43,7 +43,7 @@ from struphy.pic.particles import Particles5D, Particles6D from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.models.species import Variable +from struphy.models.variables import Variable from struphy.linear_algebra.solver import SolverParameters from struphy.io.options import check_option From 0fddf109f6f743adabd9a23b92d6e7150b09c05b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 18 Jul 2025 14:16:16 +0200 Subject: [PATCH 026/292] remove propagators/hub --- src/struphy/io/options.py | 3 ++- src/struphy/propagators/hub.py | 29 ----------------------------- src/struphy/propagators/hub.pyi | 28 ---------------------------- 3 files changed, 2 insertions(+), 58 deletions(-) delete mode 100644 src/struphy/propagators/hub.py delete mode 100644 src/struphy/propagators/hub.pyi diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index f2fc61c54..8f844073e 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -173,7 +173,6 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk ## field options PolarRegularity = Literal[-1, 1] -BackgroundOpts = Literal["LogicalConst", "FluidEquilibrium"] @dataclass class DerhamOptions: @@ -192,6 +191,8 @@ def __post_init__(self): check_option(self.polar_ck, PolarRegularity) +BackgroundOpts = Literal["LogicalConst", "FluidEquilibrium"] + @dataclass class FieldsBackground: """... diff --git a/src/struphy/propagators/hub.py b/src/struphy/propagators/hub.py deleted file mode 100644 index 04c8b921d..000000000 --- a/src/struphy/propagators/hub.py +++ /dev/null @@ -1,29 +0,0 @@ -from struphy.propagators.base import Propagator - - -class Propagators: - """Handels the Propagators of a StruphyModel.""" - def __init__(self): - pass - - @property - def all(self): - return self.__dict__ - - def add(self, prop: Propagator, *vars): - # print(f'{prop = }') - # print(f'{prop.__name__ = }') - # for var in vars: - # print(var) - setattr(self, prop.__name__, prop(*vars)) - - def set_options( - self, - name: str, - **opts, - ): - print(f'{self.all = }') - assert name in self.all, f"Propagator {name} is not part of model propagators {self.all.keys()}" - prop = getattr(self, name) - assert isinstance(prop, Propagator) - prop.set_options(**opts) \ No newline at end of file diff --git a/src/struphy/propagators/hub.pyi b/src/struphy/propagators/hub.pyi deleted file mode 100644 index 6e8f188ff..000000000 --- a/src/struphy/propagators/hub.pyi +++ /dev/null @@ -1,28 +0,0 @@ -from struphy.propagators import propagators_fields -from struphy.propagators.base import Pro - -class Propagator: - - @property - def vars(self): - ... - - @property - def init_kernels(self): - ... - - @property - def eval_kernels(self): - ... - -class Propagators: - - @property - def all(self) -> dict: ... - - def add(self, prop: Propagator, *vars): ... - - @property - def Maxwell(self) -> propagators_fields.Maxwell: ... - - \ No newline at end of file From 2b7812bcb761c35889ea4aa0f0bc3b6a8fc17398 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 18 Jul 2025 15:28:38 +0200 Subject: [PATCH 027/292] working on struphy main.py --- src/struphy/io/parameters.py | 78 +++-- src/struphy/io/setup.py | 555 +++++++++-------------------------- src/struphy/main.py | 100 +++++-- src/struphy/models/base.py | 5 +- 4 files changed, 244 insertions(+), 494 deletions(-) diff --git a/src/struphy/io/parameters.py b/src/struphy/io/parameters.py index 6d695a326..9a52e49d3 100644 --- a/src/struphy/io/parameters.py +++ b/src/struphy/io/parameters.py @@ -1,78 +1,70 @@ -from struphy.fields_background.base import FluidEquilibrium -from struphy.geometry.base import Domain -from struphy.io.options import DerhamOptions, Time, Units -from struphy.topology import grids +from __future__ import annotations +from mpi4py import MPI +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from struphy.fields_background.base import FluidEquilibrium + from struphy.geometry.base import Domain + from struphy.io.options import DerhamOptions, Time, Units + from struphy.topology import grids + from struphy.models.base import StruphyModel class StruphyParameters: - """Wrapper around Struphy simulation parameters.""" + """Wrapper around simulation parameters.""" def __init__( self, - model: str = None, + model: StruphyModel = None, + units: Units = None, domain: Domain = None, - grid: grids.TensorProductGrid = None, equil: FluidEquilibrium = None, - units: Units = None, time: Time = None, + grid: grids.TensorProductGrid = None, derham: DerhamOptions = None, - em_fields=None, - fluid=None, - kinetic=None, - diagnostic_fields=None, + verbose: bool = False ): self._model = model + self._units = units self._domain = domain - self._grid = grid self._equil = equil - self._units = units self._time = time + self._grid = grid self._derham = derham - self._em_fields = em_fields - self._fluid = fluid - self._kinetic = kinetic - self._diagnostic_fields = diagnostic_fields + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f"{self.model = }") + print(f"{self.units = }") + print(f"{self.domain = }") + print(f"{self.equil = }") + print(f"{self.time = }") + print(f"{self.grid = }") + print(f"{self.derham = }") @property def model(self): return self._model + @property + def units(self): + return self._units + @property def domain(self): return self._domain - @property - def grid(self): - return self._grid - @property def equil(self): return self._equil - - @property - def units(self): - return self._units - + @property def time(self): return self._time + + @property + def grid(self): + return self._grid @property def derham(self): return self._derham - - @property - def em_fields(self): - return self._em_fields - - @property - def fluid(self): - return self._fluid - - @property - def kinetic(self): - return self._kinetic - - @property - def diagnostic_fields(self): - return self._diagnostic_fields diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index 6af443fa6..ffea7a67f 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -1,7 +1,10 @@ import importlib.util import sys +import glob +import os +import shutil +import yaml -import numpy as np from mpi4py import MPI from struphy.io.options import DerhamOptions @@ -10,189 +13,173 @@ from struphy.utils.utils import dict_to_yaml, read_state -def derive_units( - Z_bulk: int = None, - A_bulk: int = None, - x: float = 1.0, - B: float = 1.0, - n: float = 1.0, - kBT: float = None, - velocity_scale: str = "alfvén", +def setup_folders( + path_out: str, + restart: bool, + verbose: bool = False, ): - """Computes units used in Struphy model's :ref:`normalization`. - - Input units from parameter file: + """ + Setup output folders. + """ + if MPI.COMM_WORLD.Get_rank() == 0: + if verbose: + print("\nPREPARATION AND CLEAN-UP:") - * Length (m) - * Magnetic field (T) - * Number density (10^20 1/m^3) - * Thermal energy (keV), optional + # create output folder if it does not exit + if not os.path.exists(path_out): + os.makedirs(path_out, exist_ok=True) + if verbose: + print("Created folder " + path_out) - Velocity unit is defined here: + # create data folder in output folder if it does not exist + if not os.path.exists(os.path.join(path_out, "data/")): + os.mkdir(os.path.join(path_out, "data/")) + if verbose: + print("Created folder " + os.path.join(path_out, "data/")) + else: + # remove post_processing folder + folder = os.path.join(path_out, "post_processing") + if os.path.exists(folder): + shutil.rmtree(folder) + if verbose: + print("Removed existing folder " + folder) - * Velocity (m/s) + # remove meta file + file = os.path.join(path_out, "meta.txt") + if os.path.exists(file): + os.remove(file) + if verbose: + print("Removed existing file " + file) - Derived units using mass and charge number of bulk species: + # remove profiling file + file = os.path.join(path_out, "profile_tmp") + if os.path.exists(file): + os.remove(file) + if verbose: + print("Removed existing file " + file) - * Time (s) - * Pressure (Pa) - * Mass density (kg/m^3) - * Current density (A/m^2) + # remove .png files (if NOT a restart) + if not restart: + files = glob.glob(os.path.join(path_out, "*.png")) + for n, file in enumerate(files): + os.remove(file) + if verbose and n < 10: # print only ten statements in case of many processes + print("Removed existing file " + file) - Parameters - --------- - Z_bulk : int - Charge number of bulk species. + files = glob.glob(os.path.join(path_out, "data", "*.hdf5")) + for n, file in enumerate(files): + os.remove(file) + if verbose and n < 10: # print only ten statements in case of many processes + print("Removed existing file " + file) - A_bulk : int - Mass number of bulk species. - x : float - Unit of length (in meters). +def setup_parameters( + model_name: str, + parameters: dict | str, + path_out: str, + verbose: bool = False, +): + """ + Prepare simulation parameters from .yml or .py file. - B : float - Unit of magnetic field (in Tesla). + Parameters + ---------- + model_name : str + The name of the model to run. - n : float - Unit of particle number density (in 1e20 per cubic meter). + parameters : dict | str + The simulation parameters. Can either be a dictionary OR a string (path of .yml parameter file) - kBT : float - Unit of internal energy (in keV). Only in effect if the velocity scale is set to 'thermal'. + path_out : str + The output directory. Will create a folder if it does not exist OR cleans the folder for new runs. - velocity_scale : str - Velocity scale to be used ("alfvén", "cyclotron", "light" or "thermal"). + verbose : bool + Show full screen output. Returns ------- - units : dict - The Struphy units defined above and some Physics constants. + params : StruphyParameters + The simulation parameters. + + params_path : str + The absolute path to the loaded parameter file. """ - units = {} - - # physics constants - units["elementary charge"] = 1.602176634e-19 # elementary charge (C) - units["proton mass"] = 1.67262192369e-27 # proton mass (kg) - units["mu0"] = 1.25663706212e-6 # magnetic constant (N/A^2) - units["eps0"] = 8.8541878128e-12 # vacuum permittivity (F/m) - units["kB"] = 1.380649e-23 # Boltzmann constant (J/K) - units["speed of light"] = 299792458 # speed of light (m/s) - - e = units["elementary charge"] - mH = units["proton mass"] - mu0 = units["mu0"] - eps0 = units["eps0"] - kB = units["kB"] - c = units["speed of light"] - - # length (m) - units["x"] = x - # magnetic field (T) - units["B"] = B - # number density (1/m^3) - units["n"] = n * 1e20 - - # velocity (m/s) - if velocity_scale is None: - units["v"] = 1.0 - - elif velocity_scale == "light": - units["v"] = 1.0 * c - - elif velocity_scale == "alfvén": - assert A_bulk is not None, 'Need bulk species to choose velocity scale "alfvén".' - units["v"] = units["B"] / np.sqrt(units["n"] * A_bulk * mH * mu0) - - elif velocity_scale == "cyclotron": - assert Z_bulk is not None, 'Need bulk species to choose velocity scale "cyclotron".' - assert A_bulk is not None, 'Need bulk species to choose velocity scale "cyclotron".' - units["v"] = Z_bulk * e * units["B"] / (A_bulk * mH) * units["x"] - - elif velocity_scale == "thermal": - assert A_bulk is not None, 'Need bulk species to choose velocity scale "thermal".' - assert kBT is not None - units["v"] = np.sqrt(kBT * 1000 * e / (mH * A_bulk)) - - # time (s) - units["t"] = units["x"] / units["v"] - if A_bulk is None: - return units - - # pressure (Pa), equal to B^2/mu0 if velocity_scale='alfvén' - units["p"] = A_bulk * mH * units["n"] * units["v"] ** 2 - - # mass density (kg/m^3) - units["rho"] = A_bulk * mH * units["n"] + # save "parameters" dictionary as .yml file + if isinstance(parameters, dict): + params_path = os.path.join(path_out, "parameters.yml") + if MPI.COMM_WORLD.Get_rank() == 0: + dict_to_yaml(parameters, params_path) + params = parameters - # current density (A/m^2) - units["j"] = e * units["n"] * units["v"] + # OR load parameters if "parameters" is a string (path) + else: + params_path = parameters - return units + if ".yml" in parameters or ".yaml" in parameters: + with open(parameters) as file: + params = yaml.load(file, Loader=yaml.FullLoader) + elif ".py" in parameters: + # print(f'{parameters = }') + # Read struphy state file + # state = read_state() + # i_path = state["i_path"] + # load parameter.py + spec = importlib.util.spec_from_file_location("parameters", parameters) + params_in = importlib.util.module_from_spec(spec) + sys.modules["parameters"] = params_in + spec.loader.exec_module(params_in) + if not hasattr(params_in, "model"): + params_in.model = None -def setup_domain_and_equil(params: StruphyParameters): - """ - Creates the domain object and equilibrium for a given parameter file. + if not hasattr(params_in, "domain"): + params_in.domain = None - Parameters - ---------- - params : dict - The full simulation parameter dictionary. + if not hasattr(params_in, "grid"): + params_in.grid = None - Returns - ------- - domain : Domain - The Struphy domain object for evaluating the mapping F : [0, 1]^3 --> R^3 and the corresponding metric coefficients. + if not hasattr(params_in, "equil"): + params_in.equil = None - equil : FluidEquilibrium - The equilibrium object. - """ + if not hasattr(params_in, "units"): + params_in.units = None - # from struphy.fields_background import equils - from struphy.fields_background.base import ( - NumericalFluidEquilibrium, - NumericalFluidEquilibriumWithB, - NumericalMHDequilibrium, - ) - # from struphy.geometry import domains - - if params.equil is not None: - # for eq_type, eq_params in params["fluid_background"].items(): - # eq_class = getattr(equils, eq_type) - # if eq_type in ("EQDSKequilibrium", "GVECequilibrium", "DESCequilibrium"): - # equil = eq_class(**eq_params, units=units) - # else: - # equil = eq_class(**eq_params) - - equil = params.equil - - # for numerical equilibria, the domain comes from the equilibrium - if isinstance(equil, (NumericalMHDequilibrium, NumericalFluidEquilibrium, NumericalFluidEquilibriumWithB)): - domain = equil.domain - # for all other equilibria, the domain can be chosen idependently - else: - # dom_type = params["geometry"]["type"] - # dom_class = getattr(domains, dom_type) + if not hasattr(params_in, "time"): + params_in.time = None - # if dom_type == "Tokamak": - # domain = dom_class(**params["geometry"][dom_type], equilibrium=equil) - # else: - # domain = dom_class(**params["geometry"][dom_type]) - domain = params.domain + if not hasattr(params_in, "derham"): + params_in.derham = None - # set domain attribute in mhd object - equil.domain = params.domain + params = StruphyParameters( + model=params_in.model, + units=params_in.units, + domain=params_in.domain, + equil=params_in.equil, + time=params_in.time, + grid=params_in.grid, + derham=params_in.derham, + verbose=verbose, + ) - # no equilibrium (just load domain) - else: - # dom_type = params["geometry"]["type"] - # dom_class = getattr(domains, dom_type) - # domain = dom_class(**params["geometry"][dom_type]) + if model_name is None: + assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" + model_name = params.model - domain = params.domain - equil = None + if MPI.COMM_WORLD.Get_rank() == 0: + # copy parameter file to output folder + filename = params_path.split("/")[-1] + ext = filename.split(".")[-1] + if params_path != os.path.join(path_out, "parameters." + ext): + shutil.copy2( + params_path, + os.path.join( + path_out, + "parameters." + ext, + ), + ) - return domain, equil + return params, params_path def setup_derham( @@ -282,268 +269,6 @@ def setup_derham( return derham -def pre_processing( - model_name: str, - parameters: dict | str, - path_out: str, - restart: bool, - max_sim_time: int, - save_step: int, - mpi_rank: int, - mpi_size: int, - num_clones: int, - verbose: bool = False, -): - """ - Prepares simulation parameters, output folder and prints some information of the run to the screen. - - Parameters - ---------- - model_name : str - The name of the model to run. - - parameters : dict | str - The simulation parameters. Can either be a dictionary OR a string (path of .yml parameter file) - - path_out : str - The output directory. Will create a folder if it does not exist OR cleans the folder for new runs. - - restart : bool - Whether to restart a run. - - max_sim_time : int - Maximum run time of simulation in minutes. Will finish the time integration once this limit is reached. - - save_step : int - When to save data output: every time step (save_step=1), every second time step (save_step=2). - - mpi_rank : int - The rank of the calling process. - - mpi_size : int - Total number of MPI processes of the run. - - num_clones: int - Number of domain clones. - - verbose : bool - Show full screen output. - - Returns - ------- - params : dict - The simulation parameters. - """ - - import datetime - import glob - import os - import shutil - import sysconfig - - import yaml - - from struphy.models import fluid, hybrid, kinetic, toy - - # prepare output folder - if mpi_rank == 0: - if verbose: - print("\nPREPARATION AND CLEAN-UP:") - - # create output folder if it does not exit - if not os.path.exists(path_out): - os.makedirs(path_out, exist_ok=True) - if verbose: - print("Created folder " + path_out) - - # create data folder in output folder if it does not exist - if not os.path.exists(os.path.join(path_out, "data/")): - os.mkdir(os.path.join(path_out, "data/")) - if verbose: - print("Created folder " + os.path.join(path_out, "data/")) - - # clean output folder if it already exists - else: - # remove post_processing folder - folder = os.path.join(path_out, "post_processing") - if os.path.exists(folder): - shutil.rmtree(folder) - if verbose: - print("Removed existing folder " + folder) - - # remove meta file - file = os.path.join(path_out, "meta.txt") - if os.path.exists(file): - os.remove(file) - if verbose: - print("Removed existing file " + file) - - # remove profiling file - file = os.path.join(path_out, "profile_tmp") - if os.path.exists(file): - os.remove(file) - if verbose: - print("Removed existing file " + file) - - # remove .png files (if NOT a restart) - if not restart: - files = glob.glob(os.path.join(path_out, "*.png")) - for n, file in enumerate(files): - os.remove(file) - if verbose and n < 10: # print only ten statements in case of many processes - print("Removed existing file " + file) - - files = glob.glob(os.path.join(path_out, "data", "*.hdf5")) - for n, file in enumerate(files): - os.remove(file) - if verbose and n < 10: # print only ten statements in case of many processes - print("Removed existing file " + file) - - # save "parameters" dictionary as .yml file - if isinstance(parameters, dict): - parameters_path = os.path.join(path_out, "parameters.yml") - - # write parameters to file and save it in output folder - if mpi_rank == 0: - dict_to_yaml(parameters, parameters_path) - - params = parameters - - # OR load parameters if "parameters" is a string (path) - else: - parameters_path = parameters - - if ".yml" in parameters or ".yaml" in parameters: - with open(parameters) as file: - params = yaml.load(file, Loader=yaml.FullLoader) - elif ".py" in parameters: - # print(f'{parameters = }') - # Read struphy state file - state = read_state() - i_path = state["i_path"] - # load parameter.py - spec = importlib.util.spec_from_file_location("parameters", parameters) - params_in = importlib.util.module_from_spec(spec) - sys.modules["parameters"] = params_in - spec.loader.exec_module(params_in) - - if not hasattr(params_in, "model"): - params_in.model = None - - if not hasattr(params_in, "domain"): - params_in.domain = None - - if not hasattr(params_in, "grid"): - params_in.grid = None - - if not hasattr(params_in, "equil"): - params_in.equil = None - - if not hasattr(params_in, "units"): - params_in.units = None - - if not hasattr(params_in, "time"): - params_in.time = None - - if not hasattr(params_in, "derham"): - params_in.derham = None - - if not hasattr(params_in, "em_fields"): - params_in.em_fields = None - - if not hasattr(params_in, "fluid"): - params_in.fluid = None - - if not hasattr(params_in, "kinetic"): - params_in.kinetic = None - - if not hasattr(params_in, "diagnostic_fields"): - params_in.diagnostic_fields = None - - print(f"{params_in = }") - print(f"{params_in.model = }") - print(f"{params_in.domain = }") - print(f"{params_in.grid = }") - print(f"{params_in.equil = }") - print(f"{params_in.units = }") - print(f"{params_in.time = }") - print(f"{params_in.derham = }") - print(f"{params_in.em_fields = }") - print(f"{params_in.fluid = }") - print(f"{params_in.kinetic = }") - print(f"{params_in.diagnostic_fields = }") - - params = StruphyParameters( - model=params_in.model, - domain=params_in.domain, - grid=params_in.grid, - equil=params_in.equil, - units=params_in.units, - time=params_in.time, - derham=params_in.derham, - em_fields=params_in.em_fields, - fluid=params_in.fluid, - kinetic=params_in.kinetic, - diagnostic_fields=params_in.diagnostic_fields, - ) - - if model_name is None: - assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" - model_name = params.model - - if mpi_rank == 0: - # copy parameter file to output folder - if parameters_path != os.path.join(path_out, "parameters.yml"): - shutil.copy2( - parameters_path, - os.path.join( - path_out, - "parameters.yml", - ), - ) - - # print simulation info - print("\nMETADATA:") - print("platform:".ljust(25), sysconfig.get_platform()) - print("python version:".ljust(25), sysconfig.get_python_version()) - print("model:".ljust(25), model_name) - print("MPI processes:".ljust(25), mpi_size) - print("number of domain clones:".ljust(25), num_clones) - print("parameter file:".ljust(25), parameters_path) - print("output folder:".ljust(25), path_out) - print("restart:".ljust(25), restart) - print("max wall-clock [min]:".ljust(25), max_sim_time) - print("save interval [steps]:".ljust(25), save_step) - - # write meta data to output folder - with open(path_out + "/meta.txt", "w") as f: - f.write( - "date of simulation: ".ljust( - 30, - ) - + str(datetime.datetime.now()) - + "\n", - ) - f.write("platform: ".ljust(30) + sysconfig.get_platform() + "\n") - f.write( - "python version: ".ljust( - 30, - ) - + sysconfig.get_python_version() - + "\n", - ) - f.write("model_name: ".ljust(30) + model_name + "\n") - f.write("processes: ".ljust(30) + str(mpi_size) + "\n") - f.write("output folder:".ljust(30) + path_out + "\n") - f.write("restart:".ljust(30) + str(restart) + "\n") - f.write( - "max wall-clock time [min]:".ljust(30) + str(max_sim_time) + "\n", - ) - f.write("save interval (steps):".ljust(30) + str(save_step) + "\n") - - return params - - def descend_options_dict( d: dict, out: list | dict, diff --git a/src/struphy/main.py b/src/struphy/main.py index a5a2cf3b0..896f43fb9 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -1,4 +1,22 @@ from typing import Optional +import os +import sysconfig +import time +import datetime +import glob +import numpy as np +from mpi4py import MPI +from pyevtk.hl import gridToVTK + +from struphy.feec.psydac_derham import SplineFunction +from struphy.fields_background.base import FluidEquilibriumWithB +from struphy.io.output_handling import DataContainer +from struphy.io.setup import setup_folders, setup_parameters +from struphy.models import fluid, hybrid, kinetic, toy +from struphy.models.base import StruphyModel +from struphy.profiling.profiling import ProfileManager +from struphy.utils.clone_config import CloneConfig +from struphy.pic.base import Particles def main( @@ -50,25 +68,6 @@ def main( Number of domain clones (default=1) """ - import os - import time - - import numpy as np - from mpi4py import MPI - from pyevtk.hl import gridToVTK - - from struphy.feec.psydac_derham import SplineFunction - from struphy.fields_background.base import FluidEquilibriumWithB - from struphy.io.output_handling import DataContainer - from struphy.io.setup import pre_processing - from struphy.models import fluid, hybrid, kinetic, toy - from struphy.models.base import StruphyModel - from struphy.profiling.profiling import ProfileManager - from struphy.utils.clone_config import CloneConfig - - if sort_step: - from struphy.pic.base import Particles - comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() @@ -82,18 +81,55 @@ def main( start_simulation = time.time() # loading of simulation parameters, creating output folder and printing information to screen - params = pre_processing( - model_name=model_name, - parameters=parameters, - path_out=path_out, - restart=restart, - max_sim_time=runtime, - save_step=save_step, - mpi_rank=rank, - mpi_size=size, - num_clones=num_clones, - verbose=verbose, - ) + setup_folders(path_out=path_out, + restart=restart, + verbose=False,) + + params, params_path = setup_parameters(model_name=model_name, + parameters=parameters, + path_out=path_out, + verbose=True,) + + # print simulation info + print("\nMETADATA:") + print("platform:".ljust(25), sysconfig.get_platform()) + print("python version:".ljust(25), sysconfig.get_python_version()) + print("model:".ljust(25), model_name) + print("parameter file:".ljust(25), params_path) + print("output folder:".ljust(25), path_out) + print("MPI processes:".ljust(25), size) + print("number of domain clones:".ljust(25), num_clones) + print("restart:".ljust(25), restart) + print("max wall-clock [min]:".ljust(25), runtime) + print("save interval [steps]:".ljust(25), save_step) + + exit() + + # write meta data to output folder + with open(path_out + "/meta.txt", "w") as f: + f.write( + "date of simulation: ".ljust( + 30, + ) + + str(datetime.datetime.now()) + + "\n", + ) + f.write("platform: ".ljust(30) + sysconfig.get_platform() + "\n") + f.write( + "python version: ".ljust( + 30, + ) + + sysconfig.get_python_version() + + "\n", + ) + f.write("model_name: ".ljust(30) + model_name + "\n") + f.write("processes: ".ljust(30) + str(mpi_size) + "\n") + f.write("output folder:".ljust(30) + path_out + "\n") + f.write("restart:".ljust(30) + str(restart) + "\n") + f.write( + "max wall-clock time [min]:".ljust(30) + str(max_sim_time) + "\n", + ) + f.write("save interval (steps):".ljust(30) + str(save_step) + "\n") if model_name is None: assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" @@ -358,7 +394,7 @@ def main( "--input", type=str, metavar="FILE", - help="absolute path of parameter file (.yml)", + help="absolute path of parameter file", ) # output (absolute path) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index d2089eff8..d52f0e9fc 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -19,14 +19,11 @@ ProjectedMHDequilibrium, ) from struphy.io.parameters import StruphyParameters -from struphy.io.setup import setup_derham, setup_domain_and_equil +from struphy.io.setup import setup_derham from struphy.profiling.profiling import ProfileManager -from struphy.propagators.base import Propagator -from struphy.propagators.hub import Propagators from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml from struphy.models.species import Species, FieldSpecies, FluidSpecies, KineticSpecies, SubSpecies -from struphy.io.setup import derive_units from struphy.io.options import Units from struphy.geometry.base import Domain from struphy.geometry.domains import Cuboid From b27e87907de70d9830e03a850d1bcee9e2f0990b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:00:10 +0200 Subject: [PATCH 028/292] Added DiagnosticSpecies, added attribute variables to Species, added attribute species to each variable --- src/struphy/models/species.py | 68 +++++++++++------------------------ 1 file changed, 21 insertions(+), 47 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index da164807e..94253b581 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,5 +1,5 @@ from abc import ABCMeta, abstractmethod - +from dataclasses import dataclass from copy import deepcopy from typing import Callable import numpy as np @@ -10,10 +10,27 @@ from struphy.kinetic_background.base import KineticBackground from struphy.io.options import Units from struphy.physics.physics import ConstantsOfNature +from struphy.models.variables import Variable class Species(metaclass=ABCMeta): """Single species of a StruphyModel.""" + + # set species attribute for each variable + def __post_init__(self): + for k, v in self.__dict__.items(): + if isinstance(v, Variable): + v._species = self.__class__.__name__ + + @property + def variables(self): + if not hasattr(self, "_variables"): + self._variables = {} + for k, v in self.__dict__.items(): + if isinstance(v, Variable): + self._variables[k] = v + return self._variables + @property def charge_number(self) -> int: """Charge number in units of elementary charge.""" @@ -66,7 +83,6 @@ def setup_equation_params(self, units: Units = None, verbose=False): class FieldSpecies(Species): """Species without mass and charge (so-called 'fields').""" - pass class FluidSpecies(Species): @@ -79,50 +95,8 @@ class KineticSpecies(Species): pass -class SubSpecies: - """Handles the three species types in StruphyModel: em_fields, fluid, kinetic.""" - - def __init__(self, name: str = "mhd"): - self._name = name - - @property - def name(self): - return self._name - - @property - def all(self): - out = deepcopy(self.__dict__) - out.pop("_name") - return out - - def add_variable(self, name: str = "velocity", space: str = "Hdiv"): - setattr(self, name, Variable(name, space)) - - def add_background( - self, - variable: str, - background, - verbose=False, - ): - assert variable in self.all, f"Variable {variable} is not part of model variables {self.all.keys()}" - var = getattr(self, variable) - assert isinstance(var, Variable) - var.add_background(background, verbose=verbose) - - def add_perturbation( - self, - variable: str, - perturbation: Callable, - given_in_basis: tuple = None, - verbose=False, - ): - assert variable in self.all, f"Variable {variable} is not part of model variables {self.all.keys()}" - var = getattr(self, variable) - assert isinstance(var, Variable) - var.add_perturbation( - perturbation=perturbation, - given_in_basis=given_in_basis, - verbose=verbose, - ) +class DiagnosticSpecies(Species): + """Diagnostic species (fields) without mass and charge.""" + pass From 6d6a56a33fd7ebf6aff4caf1551324ab0805b639 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:01:42 +0200 Subject: [PATCH 029/292] added metohd allocate to FEECVariable --- src/struphy/models/variables.py | 81 ++++++++++++++++++++++++++------- 1 file changed, 65 insertions(+), 16 deletions(-) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 2b8f48bc9..cbbbf7a1f 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -2,11 +2,20 @@ from mpi4py import MPI from struphy.initial.base import InitialCondition +from struphy.feec.psydac_derham import Derham, SplineFunction +from struphy.io.options import FieldsBackground +from struphy.initial.perturbations import Perturbation +from struphy.geometry.base import Domain +from struphy.fields_background.base import FluidEquilibrium class Variable(metaclass=ABCMeta): """Single variable (unknown) of a Species.""" + @abstractmethod + def allocate(self): + """Alocate object and memory for variable.""" + @property def backgrounds(self): return self._backgrounds @@ -14,31 +23,49 @@ def backgrounds(self): @property def perturbations(self): return self._perturbations + + @property + def save_data(self): + """Store variable data during simulation (default=True).""" + if not hasattr(self, "_save_data"): + self._save_data = True + return self._save_data + + @save_data.setter + def save_data(self, new): + assert isinstance(new, bool) + self._save_data = new + + @property + def species(self): + if not hasattr(self, "_species"): + self._species = None + return self._species - def add_background(self, background = None, verbose=False,): - # assert isinstance(...) + def add_background(self, background, verbose=False): + """Type inference of added background done in sub class.""" if not hasattr(self, "_backgrounds"): - self._backgrounds = [] - self._backgrounds += [background] + self._backgrounds = background + else: + if not isinstance(self.backgrounds, list): + self._backgrounds = [self.backgrounds] + self._backgrounds += [background] if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"Variable '{self.__class__.__name__}': added background '{background.__class__.__name__}' with:") + print(f"\nVariable '{self.__name__}' of species '{self.species}' - added background '{background.__class__.__name__}' with:") for k, v in background.__dict__.items(): print(f' {k}: {v}') - def add_perturbation( - self, - perturbation = None, - given_in_basis: tuple = None, - verbose=False, - ): - # assert isinstance(...) + def add_perturbation(self, perturbation: Perturbation, verbose=False): if not hasattr(self, "_perturbations"): - self._perturbations = [] - self._perturbations += [(perturbation, given_in_basis)] + self._perturbations = perturbation + else: + if not isinstance(self.perturbations, list): + self._perturbations = [self.perturbations] + self._perturbations += [perturbation] if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"Variable '{self.__class__.__name__}': added perturbation '{perturbation.__class__.__name__}',{given_in_basis = }, with:") + print(f"\nVariable '{self.__name__}' of species '{self.species}' - added perturbation '{perturbation.__class__.__name__}' with:") for k, v in perturbation.__dict__.items(): print(f' {k}: {v}') @@ -50,14 +77,36 @@ def define_initial_condition(self): class FEECVariable(Variable): - def __init__(self, space: str = "H1"): + def __init__(self, name: str = "a_feec_var", space: str = "H1"): assert space in ("H1", "Hcurl", "Hdiv", "L2", "H1vec") + self._name = name self._space = space + @property + def __name__(self): + return self._name + @property def space(self): return self._space + @property + def spline(self) -> SplineFunction: + return self._spline + + def add_background(self, background: FieldsBackground, verbose=False): + super().add_background(background, verbose=verbose) + + def allocate(self, derham: Derham, domain: Domain, equil: FluidEquilibrium,): + self._spline = derham.create_spline_function( + name=self.__name__, + space_id=self.space, + backgrounds=self.backgrounds, + perturbations=self.perturbations, + domain=domain, + equil=equil, + ) + class PICVariable(Variable): pass From 5d320ecf01f3dd604d2a960dc4b7320f1c61b95f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:04:45 +0200 Subject: [PATCH 030/292] - new base class `Perturbation` in initial/base.py - new class `Noise` in perturbations.py - added arguments `given_in_basis` and `comp` to each perturbation --- src/struphy/initial/base.py | 46 +++ src/struphy/initial/perturbations.py | 485 ++++++++++++++------------- 2 files changed, 295 insertions(+), 236 deletions(-) diff --git a/src/struphy/initial/base.py b/src/struphy/initial/base.py index 0899ee8b6..3a5efefa6 100644 --- a/src/struphy/initial/base.py +++ b/src/struphy/initial/base.py @@ -1,8 +1,54 @@ from typing import Callable +from abc import ABCMeta, abstractmethod from struphy.fields_background.base import FluidEquilibrium from struphy.io.options import FieldsBackground from struphy.kinetic_background.base import KineticBackground +from struphy.io.options import GivenInBasis, check_option + + +class Perturbation(metaclass=ABCMeta): + """Base class for perturbations that can be chosen as initial conditions.""" + + @abstractmethod + def __call__(self, eta1, eta2, eta3, flat_eval=False): + pass + + def prepare_eval_pts(self): + # TODO: we could prepare the arguments via a method in this base class (flat_eval, sparse meshgrid, etc.). + pass + + @property + def given_in_basis(self) -> str: + r"""In which basis the perturbation is represented, must be set in child class (use the setter below). + + Either + * '0', '1', '2' or '3' for a p-form basis + * 'v' for a vector-field basis + * 'physical' when defined on the physical (mapped) domain + * 'physical_at_eta' when given the physical components evaluated on the logical domain, u(F(eta)) + * 'norm' when given in the normalized co-variant basis (:math:`\delta_i / |\delta_i|`) + """ + return self._given_in_basis + + @given_in_basis.setter + def given_in_basis(self, new: str): + check_option(new, GivenInBasis) + self._given_in_basis = new + + @property + def comp(self) -> int: + """Which component of vector is perturbed (=0 for scalar-valued functions). + Can be set in child class (use the setter below).""" + if not hasattr(self, "_comp"): + self._comp = 0 + return self._comp + + @comp.setter + def comp(self, new: int): + assert new in (0, 1, 2) + self._comp = new + class InitialCondition: diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 7ce0a4c26..cc1999a45 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -1,12 +1,39 @@ #!/usr/bin/env python3 -"Analytical perturbations (modes)." +"Analytical perturbations." import numpy as np import scipy import scipy.special +from dataclasses import dataclass +from struphy.initial.base import Perturbation +from struphy.io.options import NoiseDirections, GivenInBasis, check_option -class ModesSin: +@dataclass +class Noise(Perturbation): + """White noise for FEEC coefficients. + + Parameters + ---------- + direction: str + The direction(s) of variation of the noise: 'e1', 'e2', 'e3', 'e1e2', etc. + + amp: float + Noise amplitude. + + seed: int + Seed for the random number generator. + """ + direction: NoiseDirections = "e3" + amp: float = 0.0001 + seed: int = None + comp: int = 0 + + def __post_init__(self,): + check_option(self.direction, NoiseDirections) + + +class ModesSin(Perturbation): r"""Sinusoidal function in 3D. .. math:: @@ -25,7 +52,7 @@ class ModesSin: \end{aligned} \right. - Can be used in logical space, where :math:`x \to \eta_1,\, y\to \eta_2,\, z \to \eta_3` + Can be used in logical space (use 'given_in_basis'), where :math:`x \to \eta_1,\, y\to \eta_2,\, z \to \eta_3` and :math:`L_x=L_y=L_z=1.0` (default). Parameters @@ -55,30 +82,12 @@ class ModesSin: Lx, Ly, Lz : float Domain lengths. - - Note - ---- - Example of use in a ``.yml`` parameter file:: - - perturbations : - type : ModesSin - ModesSin : - comps : - scalar_name : '0' # choices: null, 'physical', '0', '3' - vector_name : [null , 'v', '2'] # choices: null, 'physical', '1', '2', 'v', 'norm' - ls : - scalar_name: [1, 3] # two x-modes for scalar variable - vector_name: [null, [0, 1], [4]] # two x-modes for 2nd comp. and one x-mode for third component of vector-valued variable - theta : - scalar_name: [0, 3.1415] - vector_name: [null, [0, 0], [1.5708]] - pfuns : - vector_name: [null, ['localize'], ['Id']] - pfuns_params - vector_name: [null, ['0.1'], [0.]] - Lx : 7.853981633974483 - Ly : 1. - Lz : 1. + + given_in_basis : str + In which basis the perturbation is represented (see base class). + + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ def __init__( @@ -93,6 +102,8 @@ def __init__( Lx=1.0, Ly=1.0, Lz=1.0, + given_in_basis: GivenInBasis = "0", + comp: int = 0, ): if ls is not None: n_modes = len(ls) @@ -142,15 +153,6 @@ def __init__( else: assert len(pfuns_params) == n_modes - self._ls = ls - self._ms = ms - self._ns = ns - self._amps = amps - self._Lx = Lx - self._Ly = Ly - self._Lz = Lz - self._theta = theta - self._pfuns = [] for pfun, params in zip(pfuns, pfuns_params): if pfun == "Id": @@ -161,6 +163,19 @@ def __init__( ] else: raise ValueError(f"Profile function {pfun} is not defined..") + + self._ls = ls + self._ms = ms + self._ns = ns + self._amps = amps + self._Lx = Lx + self._Ly = Ly + self._Lz = Lz + self._theta = theta + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp def __call__(self, x, y, z): val = 0.0 @@ -180,14 +195,14 @@ def __call__(self, x, y, z): return val -class ModesCos: +class ModesCos(Perturbation): r"""Cosinusoidal function in 3D. .. math:: u(x, y, z) = \sum_{s} A_s \cos \left(l_s \frac{2\pi}{L_x} x + m_s \frac{2\pi}{L_y} y + n_s \frac{2\pi}{L_z} z \right) \,. - Can be used in logical space, where :math:`x \to \eta_1,\, y\to \eta_2,\, z \to \eta_3` + Can be used in logical space (use 'given_in_basis'), where :math:`x \to \eta_1,\, y\to \eta_2,\, z \to \eta_3` and :math:`L_x=L_y=L_z=1.0` (default). Parameters @@ -207,25 +222,17 @@ class ModesCos: Lx, Ly, Lz : float Domain lengths. - Note - ---- - Example of use in a ``.yml`` parameter file:: - - perturbations : - type : ModesCos - ModesCos : - comps : - scalar_name : '0' # choices: null, 'physical', '0', '3' - vector_name : [null , 'v', '2'] # choices: null, 'physical', '1', '2', 'v', 'norm' - ls : - scalar_name: [1, 3] # two x-modes for scalar variable - vector_name: [null, [0, 1], [4]] # two x-modes for 2nd comp. and one x-mode for third component of vector-valued variable - Lx : 7.853981633974483 - Ly : 1. - Lz : 1. + given_in_basis : str + In which basis the perturbation is represented (see base class). + + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, ls=None, ms=None, ns=None, amps=(1e-4,), Lx=1.0, Ly=1.0, Lz=1.0): + def __init__(self, ls=None, ms=None, ns=None, amps=(1e-4,), Lx=1.0, Ly=1.0, Lz=1.0, + given_in_basis: GivenInBasis = "0", + comp: int = 0,): + if ls is not None: n_modes = len(ls) elif ms is not None: @@ -263,6 +270,10 @@ def __init__(self, ls=None, ms=None, ns=None, amps=(1e-4,), Lx=1.0, Ly=1.0, Lz=1 self._Lx = Lx self._Ly = Ly self._Lz = Lz + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp def __call__(self, x, y, z): val = 0.0 @@ -275,7 +286,7 @@ def __call__(self, x, y, z): return val -class CoaxialWaveguideElectric_r: +class CoaxialWaveguideElectric_r(Perturbation): r"""Initializes function for Coaxial Waveguide electric field in radial direction. Solutions taken from TUM master thesis of Alicia Robles Pérez: @@ -297,6 +308,10 @@ def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28): self._r2 = a2 self._a = a self._b = b + + # use the setters + self.given_in_basis = "norm" + self.comp = 0 def __call__(self, eta1, eta2, eta3): val = 0.0 @@ -312,7 +327,7 @@ def __call__(self, eta1, eta2, eta3): return val -class CoaxialWaveguideElectric_theta: +class CoaxialWaveguideElectric_theta(Perturbation): r""" Initializes funtion for Coaxial Waveguide electric field in the azimuthal direction. @@ -335,6 +350,10 @@ def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28): self._r2 = a2 self._a = a self._b = b + + # use the setters + self.given_in_basis = "norm" + self.comp = 1 def __call__(self, eta1, eta2, eta3): val = 0.0 @@ -348,7 +367,7 @@ def __call__(self, eta1, eta2, eta3): return val -class CoaxialWaveguideMagnetic: +class CoaxialWaveguideMagnetic(Perturbation): r"""Initializes funtion for Coaxial Waveguide magnetic field in $z$-direction. Solutions taken from TUM master thesis of Alicia Robles Pérez: @@ -370,6 +389,10 @@ def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28): self._r2 = a2 self._a = a self._b = b + + # use the setters + self.given_in_basis = "norm" + self.comp = 2 def __call__(self, eta1, eta2, eta3): val = 0.0 @@ -383,7 +406,7 @@ def __call__(self, eta1, eta2, eta3): return val -class TorusModesSin: +class TorusModesSin(Perturbation): r"""Sinusoidal function in the periodic coordinates of a Torus. .. math:: @@ -404,7 +427,7 @@ class TorusModesSin: \end{aligned} \right. - Can only be defined in logical coordinates. + Can ony be used in logical space (use 'given_in_basis'). Parameters ---------- @@ -423,40 +446,19 @@ class TorusModesSin: pfun_params : tuple | list Provides :math:`[r_0, \sigma]` parameters for each "exp" profile fucntion, and l_s for "sin" and "cos". - Note - ---- - In the parameter .yml, use the following template in the section ``fluid/``:: - - perturbations : - type : TorusModesSin - TorusModesSin : - comps : - n3 : null # choices: null, 'physical', '0', '3' - u2 : ['physical', 'v', '2'] # choices: null, 'physical', '1', '2', 'v', 'norm' - p3 : '0' # choices: null, 'physical', '0', '3' - ms : - n3: null # poloidal mode numbers - u2: [[0], [0], [0]] # poloidal mode numbers - p3: [0] # poloidal mode numbers - ns : - n3: null # toroidal mode numbers - u2: [[1], [1], [1]] # toroidal mode numbers - p3: [1] # toroidal mode numbers - amps : - n3: null # amplitudes of each mode - u2: [[0.001], [0.001], [0.001]] # amplitudes of each mode - p3: [0.01] # amplitudes of each mode - pfuns : - n3: null # profile function in eta1-direction ('sin' or 'cos' or 'exp' or 'd_exp') - u2: [['sin'], ['sin'], ['exp']] # profile function in eta1-direction ('sin' or 'cos' or 'exp' or 'd_exp') - p3: [0.01] # profile function in eta1-direction ('sin' or 'cos' or 'exp' or 'd_exp') - pfun_params : - n3: null # Provides [r_0, sigma] parameters for each "exp" and "d_exp" profile fucntion, and l_s for "sin" and "cos" - u2: [2, null, [[0.5, 1.]]] # Provides [r_0, sigma] parameters for each "exp" and "d_exp" profile fucntion, and l_s for "sin" and "cos" - p3: [0.01] # Provides [r_0, sigma] parameters for each "exp" and "d_exp" profile fucntion, and l_s for "sin" and "cos" + given_in_basis : str + In which basis the perturbation is represented (see base class). + + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, ms=None, ns=None, amps=(1e-4,), pfuns=("sin",), pfun_params=None): + def __init__(self, ms=None, ns=None, amps=(1e-4,), pfuns=("sin",), pfun_params=None, + given_in_basis: GivenInBasis = "0", + comp: int = 0,): + + assert "physical" not in given_in_basis + if ms is not None: n_modes = len(ms) elif ns is not None: @@ -511,6 +513,10 @@ def __init__(self, ms=None, ns=None, amps=(1e-4,), pfuns=("sin",), pfun_params=N ] else: raise ValueError(f"Profile function {pfun} is not defined..") + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp def __call__(self, eta1, eta2, eta3): val = 0.0 @@ -526,7 +532,7 @@ def __call__(self, eta1, eta2, eta3): return val -class TorusModesCos: +class TorusModesCos(Perturbation): r"""Cosinusoidal function in the periodic coordinates of a Torus. .. math:: @@ -547,7 +553,7 @@ class TorusModesCos: \end{aligned} \right. - Can only be defined in logical coordinates. + Can only be used in logical space (use 'given_in_basis'). Parameters ---------- @@ -566,40 +572,19 @@ class TorusModesCos: pfun_params : tuple | list Provides :math:`[r_0, \sigma]` parameters for each "exp" profile fucntion, and l_s for "sin" and "cos". - Note - ---- - In the parameter .yml, use the following template in the section ``fluid/``:: - - perturbations : - type : TorusModesCos - TorusModesCos : - comps : - n3 : null # choices: null, 'physical', '0', '3' - u2 : ['physical', 'v', '2'] # choices: null, 'physical', '1', '2', 'v', 'norm' - p3 : H1 # choices: null, 'physical', '0', '3' - ms : - n3: null # poloidal mode numbers - u2: [[0], [0], [0]] # poloidal mode numbers - p3: [0] # poloidal mode numbers - ns : - n3: null # toroidal mode numbers - u2: [[1], [1], [1]] # toroidal mode numbers - p3: [1] # toroidal mode numbers - amps : - n3: null # amplitudes of each mode - u2: [[0.001], [0.001], [0.001]] # amplitudes of each mode - p3: [0.01] # amplitudes of each mode - pfuns : - n3: null # profile function in eta1-direction ('sin' or 'cos' or 'exp' or 'd_exp') - u2: [['sin'], ['sin'], ['exp']] # profile function in eta1-direction ('sin' or 'cos' or 'exp' or 'd_exp') - p3: [0.01] # profile function in eta1-direction ('sin' or 'cos' or 'exp' or 'd_exp') - pfun_params : - n3: null # Provides [r_0, sigma] parameters for each "exp" and "d_exp" profile fucntion, and l_s for "sin" and "cos". - u2: [2, null, [[0.5, 1.]]] # Provides [r_0, sigma] parameters for each "exp" and "d_exp" profile fucntion, and l_s for "sin" and "cos". - p3: [0.01] # Provides [r_0, sigma] parameters for each "exp" and "d_exp" profile fucntion, and l_s for "sin" and "cos". + given_in_basis : str + In which basis the perturbation is represented (see base class). + + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, ms=None, ns=None, amps=(1e-4,), pfuns=("sin",), pfun_params=None): + def __init__(self, ms=None, ns=None, amps=(1e-4,), pfuns=("sin",), pfun_params=None, + given_in_basis: GivenInBasis = "0", + comp: int = 0,): + + assert "physical" not in given_in_basis + if ms is not None: n_modes = len(ms) elif ns is not None: @@ -658,10 +643,15 @@ def __init__(self, ms=None, ns=None, amps=(1e-4,), pfuns=("sin",), pfun_params=N raise ValueError( 'Profile function must be "sin" or "cos" or "exp".', ) + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp def __call__(self, eta1, eta2, eta3): val = 0.0 for mi, ni, pfun, amp in zip(self._ms, self._ns, self._pfuns, self._amps): + val += ( amp * pfun(eta1) @@ -673,7 +663,7 @@ def __call__(self, eta1, eta2, eta3): return val -class Shear_x: +class Shear_x(Perturbation): r"""Double shear layer in eta1 (-1 in outer regions, 1 in inner regions). .. math:: @@ -689,25 +679,25 @@ class Shear_x: delta : float Characteristic size of the shear layer - - Note - ---- - In the parameter .yml, use the following in the section ``fluid/``:: - - perturbations : - type : Shear_x - Shear_x : - comps : - rho3 : null # choices: null, 'physical', '0', '3' - uv : ['physical', 'v', '2'] # choices: null, 'physical', '1', '2', 'v', 'norm' - s3 : H1 # choices: null, 'physical', '0', '3' - amp : 0.001 # amplitudes of each mode - delta : 0.03333 # characteristic size of the shear layer + + given_in_basis : str + In which basis the perturbation is represented (see base class). + + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, amp=1e-4, delta=1 / 15): + def __init__(self, amp=1e-4, delta=1 / 15, given_in_basis: GivenInBasis = "0", + comp: int = 0,): + + assert "physical" not in given_in_basis, f'Perturbation {self.__name__} can only be used in logical space.' + self._amp = amp self._delta = delta + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp def __call__(self, e1, e2, e3): val = self._amp * (-np.tanh((e1 - 0.75) / self._delta) + np.tanh((e1 - 0.25) / self._delta) - 1) @@ -715,7 +705,7 @@ def __call__(self, e1, e2, e3): return val -class Shear_y: +class Shear_y(Perturbation): r"""Double shear layer in eta2 (-1 in outer regions, 1 in inner regions). .. math:: @@ -731,25 +721,25 @@ class Shear_y: delta : float Characteristic size of the shear layer - - Note - ---- - In the parameter .yml, use the following in the section ``fluid/``:: - - perturbations : - type : Shear_y - Shear_y : - comps : - rho3 : null # choices: null, 'physical', '0', '3' - uv : ['physical', 'v', '2'] # choices: null, 'physical', '1', '2', 'v', 'norm' - s3 : H1 # choices: null, 'physical', '0', '3' - amp : 0.001 # amplitudes of each mode - delta : 0.03333 # characteristic size of the shear layer + + given_in_basis : str + In which basis the perturbation is represented (see base class). + + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, amp=1e-4, delta=1 / 15): + def __init__(self, amp=1e-4, delta=1 / 15, given_in_basis: GivenInBasis = "0", + comp: int = 0,): + + assert "physical" not in given_in_basis, f'Perturbation {self.__name__} can only be used in logical space.' + self._amp = amp self._delta = delta + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp def __call__(self, e1, e2, e3): val = self._amp * (-np.tanh((e2 - 0.75) / self._delta) + np.tanh((e2 - 0.25) / self._delta) - 1) @@ -757,7 +747,7 @@ def __call__(self, e1, e2, e3): return val -class Shear_z: +class Shear_z(Perturbation): r"""Double shear layer in eta3 (-1 in outer regions, 1 in inner regions). .. math:: @@ -774,24 +764,24 @@ class Shear_z: delta : float Characteristic size of the shear layer - Note - ---- - In the parameter .yml, use the following in the section ``fluid/``:: - - perturbations : - type : Shear_y - Shear_y : - comps : - rho3 : null # choices: null, 'physical', '0', '3' - uv : ['physical', 'v', '2'] # choices: null, 'physical', '1', '2', 'v', 'norm' - s3 : H1 # choices: null, 'physical', '0', '3' - amp : 0.001 # amplitudes of each mode - delta : 0.03333 # characteristic size of the shear layer + given_in_basis : str + In which basis the perturbation is represented (see base class). + + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, amp=1e-4, delta=1 / 15): + def __init__(self, amp=1e-4, delta=1 / 15, given_in_basis: GivenInBasis = "0", + comp: int = 0,): + + assert "physical" not in given_in_basis, f'Perturbation {self.__name__} can only be used in logical space.' + self._amp = amp self._delta = delta + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp def __call__(self, e1, e2, e3): val = self._amp * (-np.tanh((e3 - 0.75) / self._delta) + np.tanh((e3 - 0.25) / self._delta) - 1) @@ -799,7 +789,7 @@ def __call__(self, e1, e2, e3): return val -class Erf_z: +class Erf_z(Perturbation): r"""Shear layer in eta3 (-1 in lower regions, 1 in upper regions). .. math:: @@ -815,25 +805,25 @@ class Erf_z: delta : float Characteristic size of the shear layer - - Note - ---- - In the parameter .yml, use the following in the section ``fluid/``:: - - perturbations : - type : Erf_z - Erf_z : - comps : - b2 : ['2', null, null] # choices: null, 'physical', '0', '3' - amp : - b2 : [0.001] # amplitudes of each mode - delta : - b2 : [0.02] # characteristic size of the shear layer + + given_in_basis : str + In which basis the perturbation is represented (see base class). + + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, amp=1e-4, delta=1 / 15): + def __init__(self, amp=1e-4, delta=1 / 15, given_in_basis: GivenInBasis = "0", + comp: int = 0,): + + assert "physical" not in given_in_basis, f'Perturbation {self.__name__} can only be used in logical space.' + self._amp = amp self._delta = delta + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp def __call__(self, e1, e2, e3): from scipy.special import erf @@ -843,7 +833,7 @@ def __call__(self, e1, e2, e3): return val -class RestelliAnalyticSolutionVelocity: +class RestelliAnalyticSolutionVelocity(Perturbation): r"""Analytic solution :math:`u=u_e` of the system: .. math:: @@ -872,8 +862,6 @@ class RestelliAnalyticSolutionVelocity: Parameters ---------- - comp : string - Which component of the solution ('0', '1' or '2'). a : float Minor radius of torus (default: 1.). R0 : float @@ -886,6 +874,8 @@ class RestelliAnalyticSolutionVelocity: (default: 0.1) beta : float (default: 1.0) + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) References ---------- @@ -893,14 +883,17 @@ class RestelliAnalyticSolutionVelocity: in plasma physics, Journal of Computational Physics 2018. """ - def __init__(self, comp="0", a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0): - self._comp = comp + def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0, comp: int = 0,): self._a = a self._R0 = R0 self._B0 = B0 self._Bp = Bp self._alpha = alpha self._beta = beta + + # use the setters + self.given_in_basis = "physical" + self.comp = comp # equilibrium ion velocity def __call__(self, x, y, z): @@ -917,22 +910,22 @@ def __call__(self, x, y, z): ) * (-(R - self._R0)) uphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp - if self._comp == "0": + if self.comp == 0: # ux = np.cos(phi) * uR - R * np.sin(phi) * uphi ux = np.cos(phi) * uR - np.sin(phi) * uphi return ux - elif self._comp == "1": + elif self.comp == 1: # uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi uy = np.sin(phi) * uR + np.cos(phi) * uphi return uy - elif self._comp == "2": + elif self.comp == 2: uz = uZ return uz else: - raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") + raise ValueError(f"Invalid component '{self._comp}'. Must be 0, 1, or 2.") -class RestelliAnalyticSolutionVelocity_2: +class RestelliAnalyticSolutionVelocity_2(Perturbation): r"""Analytic solution :math:`u=u_e` of the system: .. math:: @@ -961,8 +954,6 @@ class RestelliAnalyticSolutionVelocity_2: Parameters ---------- - comp : string - Which component of the solution ('0', '1' or '2'). a : float Minor radius of torus (default: 1.). R0 : float @@ -975,6 +966,8 @@ class RestelliAnalyticSolutionVelocity_2: (default: 0.1) beta : float (default: 1.0) + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) References ---------- @@ -982,14 +975,17 @@ class RestelliAnalyticSolutionVelocity_2: in plasma physics, Journal of Computational Physics 2018. """ - def __init__(self, comp="0", a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0): - self._comp = comp + def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0, comp: int = 0,): self._a = a self._R0 = R0 self._B0 = B0 self._Bp = Bp self._alpha = alpha self._beta = beta + + # use the setter + self.given_in_basis = "physical" + self.comp = comp # equilibrium ion velocity def __call__(self, x, y, z): @@ -1006,22 +1002,22 @@ def __call__(self, x, y, z): ) * (-(R - self._R0)) uphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp - if self._comp == "0": + if self.comp == 0: # ux = np.cos(phi) * uR - R * np.sin(phi) * uphi ux = np.cos(phi) * uR - np.sin(phi) * uphi return ux - elif self._comp == "1": + elif self.comp == 1: # uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi uy = np.sin(phi) * uR + np.cos(phi) * uphi return uy - elif self._comp == "2": + elif self.comp == 2: uz = uZ return uz else: raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") -class RestelliAnalyticSolutionVelocity_3: +class RestelliAnalyticSolutionVelocity_3(Perturbation): r"""Analytic solution :math:`u=u_e` of the system: .. math:: @@ -1050,8 +1046,6 @@ class RestelliAnalyticSolutionVelocity_3: Parameters ---------- - comp : string - Which component of the solution ('0', '1' or '2'). a : float Minor radius of torus (default: 1.). R0 : float @@ -1064,6 +1058,8 @@ class RestelliAnalyticSolutionVelocity_3: (default: 0.1) beta : float (default: 1.0) + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) References ---------- @@ -1071,14 +1067,17 @@ class RestelliAnalyticSolutionVelocity_3: in plasma physics, Journal of Computational Physics 2018. """ - def __init__(self, comp="0", a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0): - self._comp = comp + def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0, comp: int = 0,): self._a = a self._R0 = R0 self._B0 = B0 self._Bp = Bp self._alpha = alpha self._beta = beta + + # use the setters + self.given_in_basis = "physical" + self.comp = comp # equilibrium ion velocity def __call__(self, x, y, z): @@ -1095,22 +1094,22 @@ def __call__(self, x, y, z): ) * (-(R - self._R0)) uphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp - if self._comp == "0": + if self.comp == 0: # ux = np.cos(phi) * uR - R * np.sin(phi) * uphi ux = np.cos(phi) * uR - np.sin(phi) * uphi return ux - elif self._comp == "1": + elif self.comp == 1: # uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi uy = np.sin(phi) * uR + np.cos(phi) * uphi return uy - elif self._comp == "2": + elif self.comp == 2: uz = uZ return uz else: raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") -class RestelliAnalyticSolutionPotential: +class RestelliAnalyticSolutionPotential(Perturbation): r"""Analytic solution :math:`\phi` of the system: .. math:: @@ -1165,6 +1164,9 @@ def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0): self._Bp = Bp self._alpha = alpha self._beta = beta + + # use the setter + self.given_in_basis = "physical" # equilibrium potential def __call__(self, x, y, z): @@ -1175,7 +1177,7 @@ def __call__(self, x, y, z): return pp -class ManufacturedSolutionVelocity: +class ManufacturedSolutionVelocity(Perturbation): r"""Analytic solutions :math:`u` and :math:`u_e` of the system: .. math:: @@ -1207,13 +1209,18 @@ class ManufacturedSolutionVelocity: Defines the manufactured solution to be selected ('1D' or '2D'). b0 : float Magnetic field (default: 1.0). + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, species="Ions", comp="0", dimension="1D", b0=1.0): + def __init__(self, species="Ions", dimension="1D", b0=1.0, comp: int = 0,): self._b = b0 self._species = species - self._comp = comp self._dimension = dimension + + # use the setters + self.given_in_basis = "physical" + self.comp = comp # equilibrium ion velocity def __call__(self, x, y, z): @@ -1234,11 +1241,11 @@ def __call__(self, x, y, z): """z component""" uz = 0.0 * x - if self._comp == "0": + if self.comp == 0: return ux - elif self._comp == "1": + elif self.comp == 1: return uy - elif self._comp == "2": + elif self.comp == 2: return uz else: raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") @@ -1260,11 +1267,11 @@ def __call__(self, x, y, z): """z component""" uz = 0.0 * x - if self._comp == "0": + if self.comp == 0: return ux - if self._comp == "1": + if self.comp == 1: return uy - if self._comp == "2": + if self.comp == 2: return uz else: raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") @@ -1273,7 +1280,7 @@ def __call__(self, x, y, z): raise ValueError(f"Invalid species '{self._species}'. Must be 'Ions' or 'Electrons'.") -class ManufacturedSolutionPotential: +class ManufacturedSolutionPotential(Perturbation): r"""Analytic solution :math:`\phi` of the system: .. math:: @@ -1312,6 +1319,9 @@ class ManufacturedSolutionPotential: def __init__(self, dimension="1D", b0=1.0): self._ab = b0 self._dimension = dimension + + # use the setter + self.given_in_basis = "physical" # equilibrium ion velocity def __call__(self, x, y, z): @@ -1324,7 +1334,7 @@ def __call__(self, x, y, z): return phi -class ManufacturedSolutionVelocity_2: +class ManufacturedSolutionVelocity_2(Perturbation): r"""Analytic solutions :math:`u` and :math:`u_e` of the system: .. math:: @@ -1350,19 +1360,22 @@ class ManufacturedSolutionVelocity_2: ---------- species : string 'Ions' or 'Electrons'. - comp : string - Which component of the solution ('0', '1' or '2'). dimension: string Defines the manufactured solution to be selected ('1D' or '2D'). b0 : float Magnetic field (default: 1.0). + comp : int + Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, species="Ions", comp="0", dimension="1D", b0=1.0): + def __init__(self, species="Ions", dimension="1D", b0=1.0, comp: int = 0,): self._b = b0 self._species = species - self._comp = comp self._dimension = dimension + + # use the setters + self.given_in_basis = "physical" + self.comp = comp # equilibrium ion velocity def __call__(self, x, y, z): @@ -1383,11 +1396,11 @@ def __call__(self, x, y, z): """z component""" uz = 0.0 * x - if self._comp == "0": + if self.comp == 0: return ux - elif self._comp == "1": + elif self.comp == 1: return uy - elif self._comp == "2": + elif self.comp == 2: return uz else: raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") @@ -1409,11 +1422,11 @@ def __call__(self, x, y, z): """z component""" uz = 0.0 * x - if self._comp == "0": + if self.comp == 0: return ux - if self._comp == "1": + if self.comp == 1: return uy - if self._comp == "2": + if self.comp == 2: return uz else: raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") From bf45fe7be3e9cade79ea06d09d7d880b5c1b1bd0 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:08:06 +0200 Subject: [PATCH 031/292] perform initialize_coeffs in __init__ of SplineFunction when background or perturbation is given --- src/struphy/feec/psydac_derham.py | 332 +++++++++++++++++------------- 1 file changed, 193 insertions(+), 139 deletions(-) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index ebf9f05d0..ee634bfb4 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -31,6 +31,10 @@ from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator +from struphy.io.options import FieldsBackground, NoiseDirections, GivenInBasis +from struphy.initial.perturbations import Noise +from struphy.initial.base import Perturbation +from struphy.fields_background.base import FluidEquilibrium class Derham: @@ -871,8 +875,10 @@ def create_spline_function( name: str, space_id: str, coeffs: StencilVector | BlockVector = None, - bckgr_params: dict = None, - pert_params: dict = None, + backgrounds: FieldsBackground | list = None, + perturbations: Perturbation | list = None, + domain: Domain = None, + equil: FluidEquilibrium = None, ): """Creat a callable spline function. @@ -887,19 +893,27 @@ def create_spline_function( coeffs : StencilVector | BlockVector The spline coefficients. - bckgr_params : dict - Field's background parameters. - - pert_params : dict - Field's perturbation parameters for initial condition. + backgrounds : FieldsBackground | list + For the initial condition. + + perturbations : Perturbation | list + For the initial condition. + + domain : Domain + Mapping for pullback/transform of initial condition. + + equil : FLuidEquilibrium + Fluid background used for inital condition. """ return SplineFunction( name, space_id, self, coeffs, - bckgr_params=bckgr_params, - pert_params=pert_params, + backgrounds=backgrounds, + perturbations=perturbations, + domain=domain, + equil=equil, ) def prepare_eval_tp_fixed(self, grids_1d): @@ -1392,11 +1406,14 @@ class SplineFunction: coeffs : StencilVector | BlockVector The spline coefficients (optional). - bckgr_params : dict - Field's background parameters. + backgrounds : FieldsBackground | list + For the initial condition. - pert_params : dict - Field's perturbation parameters for initial condition. + perturbations : Perturbation | list + For the initial condition. + + domain : Domain + Mapping for pullback/transform of initial condition. """ def __init__( @@ -1405,14 +1422,18 @@ def __init__( space_id: str, derham: Derham, coeffs: StencilVector | BlockVector = None, - bckgr_params: dict = None, - pert_params: dict = None, + backgrounds: FieldsBackground | list = None, + perturbations: Perturbation | list = None, + domain: Domain = None, + equil: FluidEquilibrium = None, ): self._name = name self._space_id = space_id self._derham = derham - self._bckgr_params = bckgr_params - self._pert_params = pert_params + self._backgrounds = backgrounds + self._perturbations = perturbations + self._domain = domain + self._equil = equil # initialize field in memory (FEM space, vector and tensor product (stencil) vector) self._space_key = derham.space_to_form[space_id] @@ -1451,6 +1472,11 @@ def __init__( ) else: self._nbasis = [tuple([space.nbasis for space in vec_space.spaces]) for vec_space in self.fem_space.spaces] + + print(f"\nAllocated SplineFuntion '{self.name}' in space '{self.space_id}'.") + + if self.backgrounds is not None or self.perturbations is not None: + self.initialize_coeffs(domain=self.domain, equil=self.equil) @property def name(self): @@ -1471,6 +1497,16 @@ def space_key(self): def derham(self): """3d Derham complex struphy.feec.psydac_derham.Derham.""" return self._derham + + @property + def domain(self): + """Mapping for pullback/transform of initial condition.""" + return self._domain + + @property + def equil(self): + """Fluid equilibirum used for initial condition.""" + return self._equil @property def space(self): @@ -1584,14 +1620,14 @@ def vector_stencil(self): return self._vector_stencil @property - def bckgr_params(self): - """Field's background parameters.""" - return self._bckgr_params + def backgrounds(self) -> FieldsBackground | list: + """For the initial condition.""" + return self._backgrounds @property - def pert_params(self): - """Field's perturbation parameters for initial condition.""" - return self._pert_params + def perturbations(self) -> Perturbation | list: + """For the initial condition.""" + return self._perturbations ############### ### Methods ### @@ -1613,173 +1649,191 @@ def extract_coeffs(self, update_ghost_regions=True): def initialize_coeffs( self, *, - bckgr_params=None, - pert_params=None, - domain=None, - bckgr_obj=None, - species=None, + backgrounds: FieldsBackground | list = None, + perturbations: Perturbation | list = None, + domain: Domain = None, + equil: FluidEquilibrium = None, ): """ - Sets the initial conditions for self.vector. - - Parameters - ---------- - bckgr_params : dict - Field's background parameters. - - pert_params : dict - Field's perturbation parameters for initial condition. - - domain : struphy.geometry.domains - Domain object for metric coefficients, only needed for transform of analytical perturbations. - - bckgr_obj: FluidEquilibrium - Fields background object. - - species : string - Species name (e.g. "mhd") the field belongs to. + Set the initial conditions for self.vector. """ # set background paramters - if bckgr_params is not None: - if self._bckgr_params is not None: - print(f"Attention: overwriting background parameters for {self.name}") - self._bckgr_params = bckgr_params + if backgrounds is not None: + # if self.backgrounds is not None: + # print(f"Attention: overwriting backgrounds for {self.name}") + self._backgrounds = backgrounds # set perturbation paramters - if pert_params is not None: - if self._pert_params is not None: - print(f"Attention: overwriting perturbation parameters for {self.name}") - self._pert_params = pert_params - + if perturbations is not None: + # if self.perturbations is not None: + # print(f"Attention: overwriting perturbation parameters for {self.name}") + self._perturbations = perturbations + + # set domain + if domain is not None: + # if self.domain is not None: + # print(f"Attention: overwriting domain for {self.name}") + self._domain = domain + + if isinstance(self.backgrounds, FieldsBackground): + self._backgrounds = [self.backgrounds] + + if isinstance(self.perturbations, Perturbation): + self._perturbations = [self.perturbations] + + # start from zero coeffs self._vector *= 0.0 if MPI.COMM_WORLD.Get_rank() == 0: print(f"Initializing {self.name} ...") - # add background to initial vector - if self.bckgr_params is not None: - for _type in self.bckgr_params: - _params = self.bckgr_params[_type].copy() - + # add backgrounds to initial vector + if self.backgrounds is not None: + for fb in self.backgrounds: + assert isinstance(fb, FieldsBackground) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"Adding background {fb} ...") + # special case of const - if "LogicalConst" in _type: - _val = _params["values"] + if fb.type == "LogicalConst": + vals = fb.values if self.space_id in {"H1", "L2"}: - assert isinstance(_val, float) or isinstance(_val, int) + assert isinstance(vals, (float, int)) def f_tmp(e1, e2, e3): - return _val + 0.0 * e1 + return vals + 0.0 * e1 fun = f_tmp else: - assert isinstance(_val, list) - assert len(_val) == 3 + assert isinstance(vals, (list, tuple)) + assert len(vals) == 3 fun = [] - for i, _v in enumerate(_val): - assert isinstance(_v, float) or isinstance(_v, int) or _v is None + for i, _v in enumerate(vals): + assert isinstance(_v, (float, int)) or _v is None - if _val[0] is not None: - fun += [lambda e1, e2, e3: _val[0] + 0.0 * e1] + if vals[0] is not None: + fun += [lambda e1, e2, e3: vals[0] + 0.0 * e1] else: fun += [lambda e1, e2, e3: 0.0 * e1] - if _val[1] is not None: - fun += [lambda e1, e2, e3: _val[1] + 0.0 * e1] + if vals[1] is not None: + fun += [lambda e1, e2, e3: vals[1] + 0.0 * e1] else: fun += [lambda e1, e2, e3: 0.0 * e1] - if _val[2] is not None: - fun += [lambda e1, e2, e3: _val[2] + 0.0 * e1] + if vals[2] is not None: + fun += [lambda e1, e2, e3: vals[2] + 0.0 * e1] else: fun += [lambda e1, e2, e3: 0.0 * e1] else: - assert bckgr_obj is not None - _var = _params["variable"] - assert _var in dir(MHDequilibrium), f"{_var = } is not an attribute of any fields background." + assert equil is not None + var = fb.variable + assert var in dir(MHDequilibrium), f"{var = } is not an attribute of any fields background." if self.space_id in {"H1", "L2"}: - fun = getattr(bckgr_obj, _var) + fun = getattr(equil, var) else: - assert (_var + "_1") in dir(MHDequilibrium), ( - f"{(_var + '_1') = } is not an attribute of any fields background." + assert (var + "_1") in dir(MHDequilibrium), ( + f"{(var + '_1') = } is not an attribute of any fields background." ) fun = [ - getattr(bckgr_obj, _var + "_1"), - getattr(bckgr_obj, _var + "_2"), - getattr(bckgr_obj, _var + "_3"), + getattr(equil, var + "_1"), + getattr(equil, var + "_2"), + getattr(equil, var + "_3"), ] - # peform projection + # perform projection self.vector += self.derham.P[self.space_key](fun) # add perturbations to coefficient vector - if self.pert_params is not None: - for _type in self.pert_params: + if self.perturbations is not None: + for ptb in self.perturbations: if MPI.COMM_WORLD.Get_rank() == 0: - print(f"Adding perturbation {_type} ...") - - _params = self.pert_params[_type].copy() + print(f"Adding perturbation {ptb} ...") # special case of white noise in logical space for different components - if "noise" in _type: - # component(s) to perturb - if isinstance(_params["comps"], bool): - comps = [_params["comps"]] - else: - comps = _params["comps"] - _params.pop("comps") - + if isinstance(ptb, Noise): # set white noise FE coefficients if self.space_id in {"H1", "L2"}: - if comps[0]: - self._add_noise(**_params) + assert isinstance(ptb.given_in_basis, str) + self._add_noise(direction=ptb.direction, + amp=ptb.amp, + seed=ptb.seed,) elif self.space_id in {"Hcurl", "Hdiv", "H1vec"}: - for n, comp in enumerate(comps): - if comp: - self._add_noise(**_params, n=n) - - # given function class - elif _type in dir(perturbations): - fun = transform_perturbation(_type, _params, self.space_key, domain) + assert isinstance(ptb.given_in_basis, list) + assert len(ptb.given_in_basis) == 3 + for n, comp in enumerate(ptb.given_in_basis): + if comp is not None: + self._add_noise(direction=ptb.direction, + amp=ptb.amp, + seed=ptb.seed, + n=n) + # perturbation class + elif isinstance(ptb, Perturbation): + if self.space_id in {"H1", "L2"}: + fun = TransformedPformComponent( + ptb, + ptb.given_in_basis, + self.space_key, + domain=domain, + ) + elif self.space_id in {"Hcurl", "Hdiv", "H1vec"}: + fun_vec = [None]*3 + fun_vec[ptb.comp] = ptb + + # pullback callable for each component + fun = [] + for comp in range(3): + fun += [ + TransformedPformComponent( + fun_vec, + ptb.given_in_basis, + self.space_key, + comp=comp, + domain=domain, + ), + ] # peform projection self.vector += self.derham.P[self.space_key](fun) + # TODO: re-add Eigfun and InitFromOutput in new framework + # loading of MHD eigenfunction (legacy code, might not be up to date) - elif "EigFun" in _type: - print("Warning: Eigfun is not regularly tested ...") - from struphy.initial import eigenfunctions - - # select class - funs = getattr(eigenfunctions, _type)( - self.derham, - **_params, - ) - - # select eigenvector and set coefficients - if hasattr(funs, self.name): - eig_vec = getattr(funs, self.name) - - self.vector += eig_vec - - # initialize from existing output file - elif "InitFromOutput" in _type: - # select class - o_data = getattr(utilities, _type)( - self.derham, - self.name, - species, - **_params, - ) - - if isinstance(self.vector, StencilVector): - self.vector._data[:] += o_data.vector - - else: - for n in range(3): - self.vector[n]._data[:] += o_data.vector[n] + # elif "EigFun" in _type: + # print("Warning: Eigfun is not regularly tested ...") + # from struphy.initial import eigenfunctions + + # # select class + # funs = getattr(eigenfunctions, _type)( + # self.derham, + # **_params, + # ) + + # # select eigenvector and set coefficients + # if hasattr(funs, self.name): + # eig_vec = getattr(funs, self.name) + + # self.vector += eig_vec + + # # initialize from existing output file + # elif "InitFromOutput" in _type: + # # select class + # o_data = getattr(utilities, _type)( + # self.derham, + # self.name, + # species, + # **_params, + # ) + + # if isinstance(self.vector, StencilVector): + # self.vector._data[:] += o_data.vector + + # else: + # for n in range(3): + # self.vector[n]._data[:] += o_data.vector[n] # apply boundary operator (in-place) self.derham.boundary_ops[self.space_key].dot( @@ -2172,7 +2226,7 @@ def _flag_pts_not_on_proc(self, *etas): E2[~E2_on_proc] = -1.0 E3[~E3_on_proc] = -1.0 - def _add_noise(self, direction="e3", amp=0.0001, seed=None, n=None): + def _add_noise(self, direction: NoiseDirections = "e3", amp: float = 0.0001, seed: int = None, n: int = None,): """Add noise to a vector component where init_comps==True, otherwise leave at zero. Parameters From 60e0d39d583133c3c036706b138ab25bbc229455 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:11:34 +0200 Subject: [PATCH 032/292] use given_in_basis in TransformedPformComponent; experiment with TYPE_CHECKING to mitigate circualr imports (but not needed anymore here) --- src/struphy/geometry/utilities.py | 50 ++++++++++++++++++------------- 1 file changed, 30 insertions(+), 20 deletions(-) diff --git a/src/struphy/geometry/utilities.py b/src/struphy/geometry/utilities.py index 70b8fc361..f321a1c8b 100644 --- a/src/struphy/geometry/utilities.py +++ b/src/struphy/geometry/utilities.py @@ -1,15 +1,22 @@ +# from __future__ import annotations "Domain-related utility functions." import numpy as np +# from typing import TYPE_CHECKING from scipy.optimize import newton, root, root_scalar from scipy.sparse import csc_matrix from scipy.sparse.linalg import splu +from typing import Callable from struphy.bsplines import bsplines as bsp from struphy.geometry.base import PoloidalSplineTorus from struphy.geometry.utilities_kernels import weighted_arc_lengths_flux_surface from struphy.linear_algebra.linalg_kron import kron_lusolve_2d +# if TYPE_CHECKING: +from struphy.geometry.base import Domain +from struphy.io.options import GivenInBasis + def field_line_tracing( psi, @@ -333,10 +340,10 @@ class TransformedPformComponent: Parameters ---------- - fun : list - Callable function components. Has to be length three for 1-, 2-forms and vector fields, length one otherwise. + fun : Callable | list + Callable function (components). Has to be length three for vector-valued funnctions,. - fun_basis : str + given_in_basis : GivenInBasis In which basis fun is represented: either a p-form, then '0' or '3' for scalar and 'v', '1' or '2' for vector-valued, @@ -346,35 +353,38 @@ class TransformedPformComponent: out_form : str The p-form representation of the output: '0', '1', '2' '3' or 'v'. - + comp : int - Which component of the transformed p-form is returned, 0, 1, or 2 (only needed for vector-valued fun). + Which component of the vector-valued function to return (=0 for scalars). domain: struphy.geometry.domains All things mapping. If None, the input fun is just evaluated and not transformed at __call__. - - Returns - ------- - out : array[float] - The values of the component comp of fun transformed from fun_basis to out_form. """ - def __init__(self, fun: list, fun_basis: str, out_form: str, comp=0, domain=None): - assert len(fun) == 1 or len(fun) == 3 + def __init__(self, + fun: Callable | list, + given_in_basis: GivenInBasis, + out_form: str, + comp: int = 0, + domain: Domain=None, + ): + + if isinstance(fun, list): + assert len(fun) == 1 or len(fun) == 3 + else: + fun = [fun] self._fun = [] for f in fun: if f is None: - def f_zero(x, y, z): return 0 * x - self._fun += [f_zero] else: assert callable(f) self._fun += [f] - self._fun_basis = fun_basis + self._given_in_basis = given_in_basis self._out_form = out_form self._comp = comp self._domain = domain @@ -391,19 +401,19 @@ def f_zero(x, y, z): def __call__(self, eta1, eta2, eta3): """ - Evaluate the component of the transformed p-form specified in self._comp. + Evaluate the component of the transformed p-form specified 'comp'. Depending on the dimension of eta1 either point-wise, tensor-product, slice plane or general (see :ref:`struphy.geometry.base.prepare_arg`). """ - if self._fun_basis == self._out_form or self._domain is None: + if self._given_in_basis == self._out_form or self._domain is None: if self._is_scalar: out = self._fun(eta1, eta2, eta3) else: out = self._fun[self._comp](eta1, eta2, eta3) - elif self._fun_basis == "physical": + elif self._given_in_basis == "physical": if self._is_scalar: out = self._domain.pull( self._fun, @@ -421,7 +431,7 @@ def __call__(self, eta1, eta2, eta3): kind=self._out_form, )[self._comp] - elif self._fun_basis == "physical_at_eta": + elif self._given_in_basis == "physical_at_eta": if self._is_scalar: out = self._domain.pull( self._fun, @@ -442,7 +452,7 @@ def __call__(self, eta1, eta2, eta3): )[self._comp] else: - dict_tran = self._fun_basis + "_to_" + self._out_form + dict_tran = self._given_in_basis + "_to_" + self._out_form if self._is_scalar: out = self._domain.transform( From 88d500d45bbba7501b3b7315db2d148139ac0fff Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:12:36 +0200 Subject: [PATCH 033/292] put Literal options on top of options.py module --- src/struphy/io/options.py | 55 ++++++++++++++++++++++++--------------- 1 file changed, 34 insertions(+), 21 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 8f844073e..cd957f93e 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -3,34 +3,54 @@ import numpy as np # needed for import in StruphyParameters -from struphy.fields_background import equils -from struphy.geometry import domains -from struphy.initial import perturbations -from struphy.kinetic_background import maxwellians -from struphy.topology import grids from struphy.physics.physics import ConstantsOfNature def check_option(opt, options): + """Check if opt is contained in options; if opt is a list, checks for each element.""" opts = get_args(options) - assert opt in opts, f"Option '{opt}' is not in {opts}." + if not isinstance(opt, list): + opt = [opt] + for o in opt: + assert o in opts, f"Option '{o}' is not in {opts}." -## generic options +## Literal options + +# time SplitAlgos = Literal["LieTrotter", "Strang"] +# derham +PolarRegularity = Literal[-1, 1] + +# fields background +BackgroundTypes = Literal["LogicalConst", "FluidEquilibrium"] + +# perturbations +NoiseDirections = Literal["e1", "e2", "e3", "e1e2", "e1e3", "e2e3", "e1e2e3"] +GivenInBasis = Literal['0', '1', '2', '3', 'v', 'physical', 'physical_at_eta', 'norm', None] + + +## Option classes + @dataclass class Time: """... Parameters ---------- - x : float - Unit of length in m. + dt : float + Time step. + + Tend : float + End time. + + split_algo : SplitAlgos + Splitting algorithm (the order of the propagators is defined in the model). """ - dt: float = 1.0 - Tend: float = 1.0 + dt: float = 0.01 + Tend: float = 0.03 split_algo: SplitAlgos = "LieTrotter" def __post_init__(self): @@ -160,6 +180,7 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk # print to screen if verbose and MPI.COMM_WORLD.Get_rank() == 0: units_used = (" m", " T", " m⁻³", "keV", " m/s", " s", " bar", " kg/m³", " A/m²",) + print("") for (k, v), u in zip(self.__dict__.items(), units_used): if v is None: print(f"Unit of {k[1:]} not specified.") @@ -170,10 +191,6 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk ) -## field options - -PolarRegularity = Literal[-1, 1] - @dataclass class DerhamOptions: """... @@ -191,8 +208,6 @@ def __post_init__(self): check_option(self.polar_ck, PolarRegularity) -BackgroundOpts = Literal["LogicalConst", "FluidEquilibrium"] - @dataclass class FieldsBackground: """... @@ -203,12 +218,10 @@ class FieldsBackground: Unit of length in m. """ - kind: str = "LogicalConst" + type: BackgroundTypes = "LogicalConst" values: tuple = (1.5,) variable: str = None def __post_init__(self): - check_option(self.kind, BackgroundOpts) - + check_option(self.type, BackgroundTypes) -## kinetic options From 469a56ba5c4ab02357bbf0a2ddca8612fc7cdd5c Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:14:08 +0200 Subject: [PATCH 034/292] do not allow dict input in setup parameters anymore (only file paths, but can still be .yml) --- src/struphy/io/setup.py | 137 ++++++++++++++++++---------------------- 1 file changed, 60 insertions(+), 77 deletions(-) diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index ffea7a67f..6f1e75b5b 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -11,6 +11,7 @@ from struphy.io.parameters import StruphyParameters from struphy.topology.grids import TensorProductGrid from struphy.utils.utils import dict_to_yaml, read_state +from struphy.geometry.base import Domain def setup_folders( @@ -75,20 +76,20 @@ def setup_folders( def setup_parameters( model_name: str, - parameters: dict | str, + params_path: str, path_out: str, verbose: bool = False, ): """ - Prepare simulation parameters from .yml or .py file. + Prepare simulation parameters from .yml or .py file and save to output folder. Parameters ---------- model_name : str The name of the model to run. - parameters : dict | str - The simulation parameters. Can either be a dictionary OR a string (path of .yml parameter file) + params_path : str + Path to .py parameter file. path_out : str The output directory. Will create a folder if it does not exist OR cleans the folder for new runs. @@ -100,67 +101,53 @@ def setup_parameters( ------- params : StruphyParameters The simulation parameters. - - params_path : str - The absolute path to the loaded parameter file. """ - # save "parameters" dictionary as .yml file - if isinstance(parameters, dict): - params_path = os.path.join(path_out, "parameters.yml") - if MPI.COMM_WORLD.Get_rank() == 0: - dict_to_yaml(parameters, params_path) - params = parameters - - # OR load parameters if "parameters" is a string (path) - else: - params_path = parameters - - if ".yml" in parameters or ".yaml" in parameters: - with open(parameters) as file: - params = yaml.load(file, Loader=yaml.FullLoader) - elif ".py" in parameters: - # print(f'{parameters = }') - # Read struphy state file - # state = read_state() - # i_path = state["i_path"] - # load parameter.py - spec = importlib.util.spec_from_file_location("parameters", parameters) - params_in = importlib.util.module_from_spec(spec) - sys.modules["parameters"] = params_in - spec.loader.exec_module(params_in) - - if not hasattr(params_in, "model"): - params_in.model = None - - if not hasattr(params_in, "domain"): - params_in.domain = None - - if not hasattr(params_in, "grid"): - params_in.grid = None - - if not hasattr(params_in, "equil"): - params_in.equil = None - - if not hasattr(params_in, "units"): - params_in.units = None - - if not hasattr(params_in, "time"): - params_in.time = None - - if not hasattr(params_in, "derham"): - params_in.derham = None - - params = StruphyParameters( - model=params_in.model, - units=params_in.units, - domain=params_in.domain, - equil=params_in.equil, - time=params_in.time, - grid=params_in.grid, - derham=params_in.derham, - verbose=verbose, - ) + if ".yml" in params_path or ".yaml" in params_path: + with open(params_path) as file: + params = yaml.load(file, Loader=yaml.FullLoader) + elif ".py" in params_path: + # print(f'{params_path = }') + # Read struphy state file + # state = read_state() + # i_path = state["i_path"] + # load parameter.py + spec = importlib.util.spec_from_file_location("parameters", params_path) + params_in = importlib.util.module_from_spec(spec) + sys.modules["parameters"] = params_in + spec.loader.exec_module(params_in) + + if not hasattr(params_in, "model"): + params_in.model = None + + if not hasattr(params_in, "domain"): + params_in.domain = None + + if not hasattr(params_in, "grid"): + params_in.grid = None + + if not hasattr(params_in, "equil"): + params_in.equil = None + + if not hasattr(params_in, "units"): + params_in.units = None + + if not hasattr(params_in, "time"): + params_in.time = None + + if not hasattr(params_in, "derham"): + params_in.derham = None + + params = StruphyParameters( + model=params_in.model, + units=params_in.units, + domain=params_in.domain, + equil=params_in.equil, + time=params_in.time, + grid=params_in.grid, + derham=params_in.derham, + verbose=verbose, + ) if model_name is None: assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" @@ -170,24 +157,20 @@ def setup_parameters( # copy parameter file to output folder filename = params_path.split("/")[-1] ext = filename.split(".")[-1] - if params_path != os.path.join(path_out, "parameters." + ext): - shutil.copy2( - params_path, - os.path.join( - path_out, - "parameters." + ext, - ), - ) + shutil.copy2( + params_path, + os.path.join(path_out, "parameters." + ext), + ) - return params, params_path + return params def setup_derham( grid: TensorProductGrid, options: DerhamOptions, - comm=None, - domain=None, - mpi_dims_mask=None, + comm: MPI.Intracomm = None, + domain: Domain = None, + mpi_dims_mask: tuple | list = None, verbose=False, ): """ @@ -201,10 +184,10 @@ def setup_derham( comm: Intracomm MPI communicator (sub_comm if clones are used). - domain : struphy.geometry.base.Domain, optional + domain : Domain, optional The Struphy domain object for evaluating the mapping F : [0, 1]^3 --> R^3 and the corresponding metric coefficients. - mpi_dims_mask: list of bool + mpi_dims_mask: list | tuple[bool] True if the dimension is to be used in the domain decomposition (=default for each dimension). If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. From ae34596aab07fd1f3f5de83d037818f1fcc0508d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:15:08 +0200 Subject: [PATCH 035/292] make ButcherTableau a dataclass --- src/struphy/ode/utils.py | 45 +++++++++++++++++++++------------------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/src/struphy/ode/utils.py b/src/struphy/ode/utils.py index c3fd06bd5..9d8d48b77 100644 --- a/src/struphy/ode/utils.py +++ b/src/struphy/ode/utils.py @@ -1,6 +1,19 @@ import numpy as np +from dataclasses import dataclass +from typing import Literal, get_args +OptsButcher = Literal[ + "rk4", + "forward_euler", + "heun2", + "rk2", + "heun3", + "3/8 rule", + ] + + +@dataclass class ButcherTableau: r""" Butcher tableau for explicit s-stage Runge-Kutta methods. @@ -13,50 +26,40 @@ class ButcherTableau: Parameters ---------- - algo : str + algo : OptsButcher Name of the RK method. """ + + algo: OptsButcher = "rk4" - @staticmethod - def available_methods() -> list: - meth_avail = [ - "rk4", - "forward_euler", - "heun2", - "rk2", - "heun3", - "3/8 rule", - ] - return meth_avail - - def __init__(self, algo: str = "rk4"): + def __post_init__(self): # choose algorithm - if algo == "forward_euler": + if self.algo == "forward_euler": a = () b = (1.0,) c = (0.0,) conv_rate = 1 - elif algo == "heun2": + elif self.algo == "heun2": a = ((1.0,),) b = (1 / 2, 1 / 2) c = (0.0, 1.0) conv_rate = 2 - elif algo == "rk2": + elif self.algo == "rk2": a = ((1 / 2,),) b = (0.0, 1.0) c = (0.0, 1 / 2) conv_rate = 2 - elif algo == "heun3": + elif self.algo == "heun3": a = ((1 / 3,), (0.0, 2 / 3)) b = (1 / 4, 0.0, 3 / 4) c = (0.0, 1 / 3, 2 / 3) conv_rate = 3 - elif algo == "rk4": + elif self.algo == "rk4": a = ((1 / 2,), (0.0, 1 / 2), (0.0, 0.0, 1.0)) b = (1 / 6, 1 / 3, 1 / 3, 1 / 6) c = (0.0, 1 / 2, 1 / 2, 1.0) conv_rate = 4 - elif algo == "3/8 rule": + elif self.algo == "3/8 rule": a = ((1 / 3,), (-1 / 3, 1.0), (1.0, -1.0, 1.0)) b = (1 / 8, 3 / 8, 3 / 8, 1 / 8) c = (0.0, 1 / 3, 2 / 3, 1.0) @@ -78,7 +81,7 @@ def __init__(self, algo: str = "rk4"): self._conv_rate = conv_rate - __available_methods__ = available_methods() + __available_methods__ = get_args(OptsButcher) @property def a(self): From 7dfecfbfe9eeacc341ba98b91b9765bbada77947 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:16:50 +0200 Subject: [PATCH 036/292] Propagator base class: add options attribute and setter; add Options class --- src/struphy/propagators/base.py | 34 +++++++++++++++++---------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index e2d82bd23..7e756cc9a 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -300,21 +300,23 @@ def add_eval_kernel( ] @property - def opts(self): - if not hasattr(self, "_opts"): - self._opts = self.Options() - return self._opts - + def options(self): + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + self._options = new + class Options: - def __init__(self, outer_prop, verbose=False): - self._outer_prop = outer_prop.__class__.__name__ # outer class - self._verbose = verbose - + def __init__(self, prop, **kwargs): + self._kwargs = kwargs + print(f"\nInstance of propagator '{prop.__class__.__name__}' with:") + for k, v in kwargs.items(): + print(f' {k}: {v}') + @property - def all(self): - return self.__dict__ - - def add(self, name: str, opt): - setattr(self, name, opt) - if self._verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"Propagator '{self._outer_prop}': added option '{name}' with value '{opt}'") \ No newline at end of file + def kwargs(self): + return self._kwargs \ No newline at end of file From 1da7b0f6dfcd17090f90dd6d5ca4c9dfde3908f0 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:17:38 +0200 Subject: [PATCH 037/292] New design of Maxwell propagator: - make it a dataclass - new method set_options - use Literals - todo: complete allocate method --- src/struphy/propagators/propagators_fields.py | 59 +++++++++++-------- 1 file changed, 34 insertions(+), 25 deletions(-) diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 850b2fc28..3985e8d13 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -2,6 +2,7 @@ from collections.abc import Callable from copy import deepcopy +from dataclasses import dataclass from typing import Literal, get_args import numpy as np @@ -36,7 +37,7 @@ from struphy.linear_algebra.saddle_point import SaddlePointSolver from struphy.linear_algebra.schur_solver import SchurSolver from struphy.ode.solvers import ODEsolverFEEC -from struphy.ode.utils import ButcherTableau +from struphy.ode.utils import ButcherTableau, OptsButcher from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.pic.accumulation.particles_to_grid import Accumulator, AccumulatorVector from struphy.pic.base import Particles @@ -46,8 +47,10 @@ from struphy.models.variables import Variable from struphy.linear_algebra.solver import SolverParameters from struphy.io.options import check_option +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable +@dataclass class Maxwell(Propagator): r""":ref:`FEEC ` discretization of the following equations: find :math:`\mathbf E \in H(\textnormal{curl})` and :math:`\mathbf B \in H(\textnormal{div})` such that @@ -60,39 +63,45 @@ class Maxwell(Propagator): :ref:`time_discret`: Crank-Nicolson (implicit mid-point). System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`. """ - - def __init__( - self, - # e: Variable, - # b: Variable, - ): - self.e = None - self.b = None - - - OptsAlgo = Literal["implicit", *ButcherTableau.available_methods(),] - OptsSolverType = Literal[("pcg", "MassMatrixPreconditioner"), ("cg", None),] + e: FEECVariable = None + b: FEECVariable = None + + OptsAlgo = Literal["implicit", "explicit"] + OptsSolver = Literal["pcg", "cg"] + OptsPrecond = Literal["MassMatrixPreconditioner", None] def set_options(self, - algo: OptsAlgo = "implicit", # type: ignore - solver_type: OptsSolverType = ("pcg", "MassMatrixPreconditioner"), # type: ignore + algo: OptsAlgo = "implicit", + solver: OptsSolver = "pcg", + precond: OptsPrecond = "MassMatrixPreconditioner", solver_params: SolverParameters = None, - verbose = False, + butcher: ButcherTableau = None, ): # checks check_option(algo, self.OptsAlgo) - check_option(solver_type, self.OptsSolverType) + check_option(solver, self.OptsSolver) + check_option(precond, self.OptsPrecond) # defaults - if solver_params is None: - solver_params = SolverParameters() + if algo == "implicit": + butcher = None + if solver_params is None: + solver_params = SolverParameters() + elif algo == "explicit": + solver = None + precond = None + solver_params = None + if butcher is None: + butcher = ButcherTableau() - # create options - self._opts = self.Options(self, verbose=verbose) - self.opts.add("algo", algo) - self.opts.add("solver", solver_type) - self.opts.add("solver_params", solver_params) + # use setter for options + self.options = self.Options(self, + algo=algo, + solver=solver, + precond=precond, + solver_params=solver_params, + butcher=butcher,) def allocate(self): solver = {} @@ -7201,7 +7210,7 @@ def options(default=False): dct["c_fun"] = ["const"] dct["kappa"] = 1.0 dct["nu"] = 0.01 - dct["algo"] = ButcherTableau.available_methods() + dct["algo"] = get_args(OptsButcher) dct["M0_solver"] = { "type": [ ("pcg", "MassMatrixPreconditioner"), From cd812855d9657e2ae79fd6a6722d82b55395e198 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:23:05 +0200 Subject: [PATCH 038/292] StruphyModel base class: new methods allocate_feec, set_propagators and allocate_variables --- src/struphy/models/base.py | 1198 ++++++++++++++++-------------------- 1 file changed, 526 insertions(+), 672 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index d52f0e9fc..a4b05aa87 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -3,6 +3,9 @@ from abc import ABCMeta, abstractmethod from functools import reduce from typing import Callable +import os +import yaml +import struphy import numpy as np import yaml @@ -18,16 +21,22 @@ ProjectedFluidEquilibriumWithB, ProjectedMHDequilibrium, ) -from struphy.io.parameters import StruphyParameters -from struphy.io.setup import setup_derham +from struphy.io.setup import setup_derham, descend_options_dict from struphy.profiling.profiling import ProfileManager from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml -from struphy.models.species import Species, FieldSpecies, FluidSpecies, KineticSpecies, SubSpecies +from struphy.models.species import Species, FieldSpecies, FluidSpecies, KineticSpecies, DiagnosticSpecies +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.io.options import Units +from struphy.io.options import Time, DerhamOptions +from struphy.topology.grids import TensorProductGrid from struphy.geometry.base import Domain from struphy.geometry.domains import Cuboid from struphy.fields_background.equils import HomogenSlab +from struphy.pic import particles +from struphy.pic.base import Particles +from struphy.propagators.base import Propagator +from struphy.kinetic_background import maxwellians class StruphyModel(metaclass=ABCMeta): @@ -59,7 +68,7 @@ class Propagators: @property @abstractmethod - def bulk_species() -> SubSpecies: + def bulk_species() -> Species: """Bulk species of the plasma. Must be an attribute of species_static().""" @property @@ -132,158 +141,126 @@ def setup_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): else: self._domain = domain self._equil = None + + if MPI.COMM_WORLD.Get_rank() == 0 and self.verbose: + print("\nDOMAIN:") + print(f"type:".ljust(25), self.domain.__class__.__name__) + for key, val in self.domain.params_map.items(): + if key not in {"cx", "cy", "cz"}: + print((key + ":").ljust(25), val) + + print("\nFLUID BACKGROUND:") + if self.equil is not None: + print("type:".ljust(25), self.equil.__class__.__name__) + for key, val in self.equil.params.items(): + print((key + ":").ljust(25), val) + else: + print("None.") ## species @property def field_species(self) -> dict: - spec = {} - for k, v in self.__dict__.items(): - if isinstance(v, FieldSpecies): - spec[k] = v - return spec + if not hasattr(self, "_field_species"): + self._field_species = {} + for k, v in self.__dict__.items(): + if isinstance(v, FieldSpecies): + self._field_species[k] = v + return self._field_species @property def fluid_species(self) -> dict: - spec = {} - for k, v in self.__dict__.items(): - if isinstance(v, FluidSpecies): - spec[k] = v - return spec + if not hasattr(self, "_fluid_species"): + self._fluid_species = {} + for k, v in self.__dict__.items(): + if isinstance(v, FluidSpecies): + self._fluid_species[k] = v + return self._fluid_species @property def kinetic_species(self) -> dict: - spec = {} - for k, v in self.__dict__.items(): - if isinstance(v, KineticSpecies): - spec[k] = v - return spec + if not hasattr(self, "_kinetic_species"): + self._kinetic_species = {} + for k, v in self.__dict__.items(): + if isinstance(v, KineticSpecies): + self._kinetic_species[k] = v + return self._kinetic_species @property - def all_species(self): - spec = {} - for k, v in self.field_species.items(): - spec[k] = v - for k, v in self.fluid_species.items(): - spec[k] = v - for k, v in self.kinetic_species.items(): - spec[k] = v - return spec + def diagnostic_species(self) -> dict: + if not hasattr(self, "_diagnostic_species"): + self._diagnostic_species = {} + for k, v in self.__dict__.items(): + if isinstance(v, DiagnosticSpecies): + self._diagnostic_species[k] = v + return self._diagnostic_species + + @property + def species(self): + if not hasattr(self, "_species"): + self._species = {} + for k, v in self.field_species.items(): + self._species[k] = v + for k, v in self.fluid_species.items(): + self._species[k] = v + for k, v in self.kinetic_species.items(): + self._species[k] = v + return self._species ## allocate methods - def allocate(self, - grid=None, - derham_params=None, - time=None, + def allocate_feec(self, + grid: TensorProductGrid, + derham_params: DerhamOptions=None, + comm: MPI.Intracomm = None, + clone_config: CloneConfig = None, ): - - self.init_derham(grid, derham_params) - - self.allocate_species() - self.allocate_propagators() - - - - - - - if comm.Get_rank() == 0 and self.verbose: - print("\nTIME:") - print( - f"time step:".ljust(25), - "{0} ({1:4.2e} s)".format( - params.time.dt, - params.time.dt * self.units["t"], - ), - ) - print( - f"final time:".ljust(25), - "{0} ({1:4.2e} s)".format( - params.time.Tend, - params.time.Tend * self.units["t"], - ), - ) - print(f"splitting algo:".ljust(25), params.time.split_algo) - - print("\nDOMAIN:") - print(f"type:".ljust(25), self.domain.__class__.__name__) - for key, val in self.domain.params_map.items(): - if key not in {"cx", "cy", "cz"}: - print((key + ":").ljust(25), val) - - print("\nFLUID BACKGROUND:") - if params.equil is not None: - print("type:".ljust(25), self.equil.__class__.__name__) - for key, val in self.equil.params.items(): - print((key + ":").ljust(25), val) - else: - print("None.") # create discrete derham sequence - if params.grid is not None: - dims_mask = params.grid.mpi_dims_mask - if dims_mask is None: - dims_mask = [True] * 3 + dims_mask = grid.mpi_dims_mask + if dims_mask is None: + dims_mask = [True] * 3 - if clone_config is None: - derham_comm = self.comm_world - else: - derham_comm = clone_config.sub_comm - - self._derham = setup_derham( - params.grid, - params.derham, - comm=derham_comm, - domain=self.domain, - mpi_dims_mask=dims_mask, - verbose=self.verbose, - ) + if clone_config is None: + derham_comm = self.comm_world else: - self._derham = None - print("\nDERHAM:\nMeshless simulation - no Derham complex set up.") - - self._projected_equil = None - self._mass_ops = None - if self.derham is not None: - # create projected equilibrium - if isinstance(self.equil, MHDequilibrium): - self._projected_equil = ProjectedMHDequilibrium( - self.equil, - self.derham, - ) - elif isinstance(self.equil, FluidEquilibriumWithB): - self._projected_equil = ProjectedFluidEquilibriumWithB( - self.equil, - self.derham, - ) - elif isinstance(self.equil, FluidEquilibrium): - self._projected_equil = ProjectedFluidEquilibrium( - self.equil, - self.derham, - ) - - # create weighted mass operators - self._mass_ops = WeightedMassOperators( + derham_comm = clone_config.sub_comm + + self._derham = setup_derham( + grid, + derham_params, + comm=derham_comm, + domain=self.domain, + mpi_dims_mask=dims_mask, + verbose=self.verbose, + ) + + # create weighted mass operators + self._mass_ops = WeightedMassOperators( + self.derham, + self.domain, + verbose=self.verbose, + eq_mhd=self.equil, + ) + + # create projected equilibrium + if isinstance(self.equil, MHDequilibrium): + self._projected_equil = ProjectedMHDequilibrium( + self.equil, self.derham, - self.domain, - verbose=self.verbose, - eq_mhd=self.equil, ) - - # allocate memory for variables - self._pointer = {} - self._allocate_variables() - - # store plasma parameters - if self.rank_world == 0: - self._pparams = self._compute_plasma_params(verbose=self.verbose) - else: - self._pparams = self._compute_plasma_params(verbose=False) - - # if self.rank_world == 0: - # self._show_chosen_options() - + elif isinstance(self.equil, FluidEquilibriumWithB): + self._projected_equil = ProjectedFluidEquilibriumWithB( + self.equil, + self.derham, + ) + elif isinstance(self.equil, FluidEquilibrium): + self._projected_equil = ProjectedFluidEquilibrium( + self.equil, + self.derham, + ) + + def set_propagators(self): # set propagators base class attributes (then available to all propagators) Propagator.derham = self.derham Propagator.domain = self.domain @@ -297,17 +274,6 @@ def allocate(self, ) Propagator.projected_equil = self.projected_equil - # create dummy lists/dicts to be filled by the sub-class - - self._kwargs = {} - self._scalar_quantities = {} - - return params - - - - - @staticmethod def diagnostics_dct(): """Diagnostics dictionary. @@ -315,9 +281,6 @@ def diagnostics_dct(): """ ## basic properties - @property - def species(self) -> Species: - return self._species @property def params(self): @@ -392,6 +355,14 @@ def mass_ops(self): """WeighteMassOperators object, see :ref:`mass_ops`.""" return self._mass_ops + @property + def propagators(self): + """Return object that has model propagators in its __dict__.""" + if not hasattr (self, "_propagators"): + self._propagators = None + raise NameError("Propagators object must be named 'self._propagators = ...' in model __init__().") + return self._propagators + @property def prop_fields(self): """Module :mod:`struphy.propagators.propagators_fields`.""" @@ -545,6 +516,9 @@ def add_scalar(self, name, species=None, compute=None, summands=None): str, ), "species must be a string when compute is 'from_particles'" + if not hasattr(self, "_scalar_quantities"): + self._scalar_quantities = {} + self._scalar_quantities[name] = { "value": np.empty(1, dtype=float), "species": species, @@ -652,8 +626,160 @@ def add_time_state(self, time_state): """ assert time_state.size == 1 self._time_state = time_state - for prop in self.propagators: - prop.add_time_state(time_state) + for _, prop in self.propagators.__dict__.items(): + if isinstance(prop, Propagator): + prop.add_time_state(time_state) + + def allocate_variables(self): + """ + Allocate memory for model variables and set initial conditions. + """ + + # pointer directly to data structs of Variables + self._pointer = {} + + # allocate memory for FE coeffs of electromagnetic fields/potentials + if self.field_species: + for _, spec in self.field_species.items(): + assert isinstance(spec, FieldSpecies) + for k, v in spec.variables.items(): + assert isinstance(v, FEECVariable) + v.allocate(derham=self.derham, domain=self.domain, equil=self.equil,) + self._pointer[k] = v.spline.vector + + # allocate memory for FE coeffs of fluid variables + if self.fluid_species: + for species, dct in self.fluid.items(): + for variable, subdct in dct.items(): + if "params" in variable: + continue + else: + subdct["obj"] = self.derham.create_spline_function( + variable, + subdct["space"], + bckgr_params=subdct.get("background"), + pert_params=subdct.get("perturbation"), + ) + + self._pointer[species + "_" + variable] = subdct["obj"].vector + + # marker arrays and plasma parameters of kinetic species + if self.kinetic_species: + for species, val in self.kinetic.items(): + assert any([key in val["params"]["markers"] for key in ["Np", "ppc", "ppb"]]) + + bckgr_params = val["params"].get("background", None) + pert_params = val["params"].get("perturbation", None) + boxes_per_dim = val["params"].get("boxes_per_dim", None) + mpi_dims_mask = val["params"].get("dims_mask", None) + weights_params = val["params"].get("weights", None) + + if self.derham is None: + domain_decomp = None + else: + domain_array = self.derham.domain_array + nprocs = self.derham.domain_decomposition.nprocs + domain_decomp = (domain_array, nprocs) + + kinetic_class = getattr(particles, val["space"]) + + val["obj"] = kinetic_class( + comm_world=self.comm_world, + clone_config=self.clone_config, + **val["params"]["markers"], + weights_params=weights_params, + domain_decomp=domain_decomp, + mpi_dims_mask=mpi_dims_mask, + boxes_per_dim=boxes_per_dim, + name=species, + equation_params=self.equation_params[species], + domain=self.domain, + equil=self.equil, + projected_equil=self.projected_equil, + bckgr_params=bckgr_params, + pert_params=pert_params, + ) + + obj = val["obj"] + assert isinstance(obj, Particles) + + self._pointer[species] = obj + + # for storing markers + val["kinetic_data"] = {} + + # for storing the distribution function + if "f" in val["params"]["save_data"]: + slices = val["params"]["save_data"]["f"]["slices"] + n_bins = val["params"]["save_data"]["f"]["n_bins"] + ranges = val["params"]["save_data"]["f"]["ranges"] + + val["kinetic_data"]["f"] = {} + val["kinetic_data"]["df"] = {} + val["bin_edges"] = {} + if len(slices) > 0: + for i, sli in enumerate(slices): + assert ((len(sli) - 2) / 3).is_integer() + assert len(slices[i].split("_")) == len(ranges[i]) == len(n_bins[i]), ( + f"Number of slices names ({len(slices[i].split('_'))}), number of bins ({len(n_bins[i])}), and number of ranges ({len(ranges[i])}) are inconsistent with each other!\n\n" + ) + val["bin_edges"][sli] = [] + dims = (len(sli) - 2) // 3 + 1 + for j in range(dims): + val["bin_edges"][sli] += [ + np.linspace( + ranges[i][j][0], + ranges[i][j][1], + n_bins[i][j] + 1, + ), + ] + val["kinetic_data"]["f"][sli] = np.zeros( + n_bins[i], + dtype=float, + ) + val["kinetic_data"]["df"][sli] = np.zeros( + n_bins[i], + dtype=float, + ) + + # for storing an sph evaluation of the density n + if "n_sph" in val["params"]["save_data"]: + plot_pts = val["params"]["save_data"]["n_sph"]["plot_pts"] + + val["kinetic_data"]["n_sph"] = [] + val["plot_pts"] = [] + for i, pts in enumerate(plot_pts): + assert len(pts) == 3 + eta1 = np.linspace(0.0, 1.0, pts[0]) + eta2 = np.linspace(0.0, 1.0, pts[1]) + eta3 = np.linspace(0.0, 1.0, pts[2]) + ee1, ee2, ee3 = np.meshgrid( + eta1, + eta2, + eta3, + indexing="ij", + ) + val["plot_pts"] += [(ee1, ee2, ee3)] + val["kinetic_data"]["n_sph"] += [np.zeros(ee1.shape, dtype=float)] + + # other data (wave-particle power exchange, etc.) + # TODO + + # TODO: allocate memory for FE coeffs of diagnostics + # if self.params.diagnostic_fields is not None: + # for key, val in self.diagnostics.items(): + # if "params" in key: + # continue + # else: + # val["obj"] = self.derham.create_spline_function( + # key, + # val["space"], + # bckgr_params=None, + # pert_params=None, + # ) + + # self._pointer[key] = val["obj"].vector + def integrate(self, dt, split_algo="LieTrotter"): """ @@ -668,10 +794,13 @@ def integrate(self, dt, split_algo="LieTrotter"): Splitting algorithm. Currently available: "LieTrotter" and "Strang". """ + if not hasattr(self, "_prop_list"): + self._prop_list = list(self.propagators.__dict__.values()) + # first order in time if split_algo == "LieTrotter": - for propagator in self.propagators: - prop_name = type(propagator).__name__ + for propagator in self._prop_list: + prop_name = propagator.__class__.__name__ with ProfileManager.profile_region(prop_name): propagator(dt) @@ -680,17 +809,17 @@ def integrate(self, dt, split_algo="LieTrotter"): elif split_algo == "Strang": assert len(self.propagators) > 1 - for propagator in self.propagators[:-1]: + for propagator in self._prop_list[:-1]: prop_name = type(propagator).__name__ with ProfileManager.profile_region(prop_name): propagator(dt / 2) - propagator = self.propagators[-1] + propagator = self._prop_list[-1] prop_name = type(propagator).__name__ with ProfileManager.profile_region(prop_name): propagator(dt) - for propagator in self.propagators[:-1][::-1]: + for propagator in self._prop_list[:-1][::-1]: prop_name = type(propagator).__name__ with ProfileManager.profile_region(prop_name): propagator(dt / 2) @@ -707,7 +836,7 @@ def update_markers_to_be_saved(self): from struphy.pic.base import Particles - for val in self.kinetic.values(): + for _, val in self.kinetic_species.items(): obj = val["obj"] assert isinstance(obj, Particles) @@ -750,7 +879,7 @@ def update_distr_functions(self): dim_to_int = {"e1": 0, "e2": 1, "e3": 2, "v1": 3, "v2": 4, "v3": 5} - for val in self.kinetic.values(): + for _, val in self.kinetic_species.items(): obj = val["obj"] assert isinstance(obj, Particles) @@ -805,177 +934,171 @@ def print_scalar_quantities(self): sq_str += key + ": {:14.11f}".format(val[0]) + " " print(sq_str) - def initialize_from_params(self): - """ - Set initial conditions for FE coefficients (electromagnetic and fluid) - and markers according to parameter file. - """ - - from struphy.feec.psydac_derham import Derham - from struphy.pic.base import Particles - - if self.rank_world == 0 and self.verbose: - print("\nINITIAL CONDITIONS:") - - # initialize em fields - if len(self.em_fields) > 0: - with ProfileManager.profile_region("initialize_em_fields"): - for key, val in self.em_fields.items(): - if "params" in key: - continue - else: - obj = val["obj"] - assert isinstance(obj, SplineFunction) - - obj.initialize_coeffs( - domain=self.domain, - bckgr_obj=self.equil, - ) - - if self.rank_world == 0 and self.verbose: - print(f'\nEM field "{key}" was initialized with:') - - _params = self.em_fields["params"] - - if "background" in _params: - if key in _params["background"]: - bckgr_types = _params["background"][key] - if bckgr_types is None: - pass - else: - print("background:") - for _type, _bp in bckgr_types.items(): - print(" " * 4 + _type, ":") - for _pname, _pval in _bp.items(): - print((" " * 8 + _pname + ":").ljust(25), _pval) - else: - print("No background.") - else: - print("No background.") - - if "perturbation" in _params: - if key in _params["perturbation"]: - pert_types = _params["perturbation"][key] - if pert_types is None: - pass - else: - print("perturbation:") - for _type, _pp in pert_types.items(): - print(" " * 4 + _type, ":") - for _pname, _pval in _pp.items(): - print((" " * 8 + _pname + ":").ljust(25), _pval) - else: - print("No perturbation.") - else: - print("No perturbation.") - - if len(self.fluid) > 0: - with ProfileManager.profile_region("initialize_fluids"): - for species, val in self.fluid.items(): - for variable, subval in val.items(): - if "params" in variable: - continue - else: - obj = subval["obj"] - assert isinstance(obj, SplineFunction) - obj.initialize_coeffs( - domain=self.domain, - bckgr_obj=self.equil, - species=species, - ) - - if self.rank_world == 0 and self.verbose: - print( - f'\nFluid species "{species}" was initialized with:', - ) - - _params = val["params"] - - if "background" in _params: - for variable in val: - if "params" in variable: - continue - if variable in _params["background"]: - bckgr_types = _params["background"][variable] - if bckgr_types is None: - pass - else: - print(f"{variable} background:") - for _type, _bp in bckgr_types.items(): - print(" " * 4 + _type, ":") - for _pname, _pval in _bp.items(): - print((" " * 8 + _pname + ":").ljust(25), _pval) - else: - print(f"{variable}: no background.") - else: - print("No background.") - - if "perturbation" in _params: - for variable in val: - if "params" in variable: - continue - if variable in _params["perturbation"]: - pert_types = _params["perturbation"][variable] - if pert_types is None: - pass - else: - print(f"{variable} perturbation:") - for _type, _pp in pert_types.items(): - print(" " * 4 + _type, ":") - for _pname, _pval in _pp.items(): - print((" " * 8 + _pname + ":").ljust(25), _pval) - else: - print(f"{variable}: no perturbation.") - else: - print("No perturbation.") - - # initialize particles - if len(self.kinetic) > 0: - with ProfileManager.profile_region("initialize_particles"): - for species, val in self.kinetic.items(): - obj = val["obj"] - assert isinstance(obj, Particles) - - if self.rank_world == 0 and self.verbose: - _params = val["params"] - assert "background" in _params, "Kinetic species must have background." - - bckgr_types = _params["background"] - print( - f'\nKinetic species "{species}" was initialized with:', - ) - for _type, _bp in bckgr_types.items(): - print(_type, ":") - for _pname, _pval in _bp.items(): - print((" " * 4 + _pname + ":").ljust(25), _pval) - - if "perturbation" in _params: - for variable, pert_types in _params["perturbation"].items(): - if pert_types is None: - pass - else: - print(f"{variable} perturbation:") - for _type, _pp in pert_types.items(): - print(" " * 4 + _type, ":") - for _pname, _pval in _pp.items(): - print((" " * 8 + _pname + ":").ljust(25), _pval) - else: - print("No perturbation.") - - obj.draw_markers(sort=True, verbose=self.verbose) - obj.mpi_sort_markers(do_test=True) - - if not val["params"]["markers"]["loading"] == "restart": - if obj.coords == "vpara_mu": - obj.save_magnetic_moment() - - if val["space"] != "ParticlesSPH" and obj.f0.coords == "constants_of_motion": - obj.save_constants_of_motion() - - obj.initialize_weights( - reject_weights=obj.weights_params["reject_weights"], - threshold=obj.weights_params["threshold"], - ) + # def initialize_from_params(self): + # """ + # Set initial conditions for FE coefficients (electromagnetic and fluid) + # and markers according to parameter file. + # """ + + # # initialize em fields + # if self.field_species: + # with ProfileManager.profile_region("initialize_em_fields"): + # for key, val in self.em_fields.items(): + # if "params" in key: + # continue + # else: + # obj = val["obj"] + # assert isinstance(obj, SplineFunction) + + # obj.initialize_coeffs( + # domain=self.domain, + # bckgr_obj=self.equil, + # ) + + # if self.rank_world == 0 and self.verbose: + # print(f'\nEM field "{key}" was initialized with:') + + # _params = self.em_fields["params"] + + # if "background" in _params: + # if key in _params["background"]: + # bckgr_types = _params["background"][key] + # if bckgr_types is None: + # pass + # else: + # print("background:") + # for _type, _bp in bckgr_types.items(): + # print(" " * 4 + _type, ":") + # for _pname, _pval in _bp.items(): + # print((" " * 8 + _pname + ":").ljust(25), _pval) + # else: + # print("No background.") + # else: + # print("No background.") + + # if "perturbation" in _params: + # if key in _params["perturbation"]: + # pert_types = _params["perturbation"][key] + # if pert_types is None: + # pass + # else: + # print("perturbation:") + # for _type, _pp in pert_types.items(): + # print(" " * 4 + _type, ":") + # for _pname, _pval in _pp.items(): + # print((" " * 8 + _pname + ":").ljust(25), _pval) + # else: + # print("No perturbation.") + # else: + # print("No perturbation.") + + # if len(self.fluid) > 0: + # with ProfileManager.profile_region("initialize_fluids"): + # for species, val in self.fluid.items(): + # for variable, subval in val.items(): + # if "params" in variable: + # continue + # else: + # obj = subval["obj"] + # assert isinstance(obj, SplineFunction) + # obj.initialize_coeffs( + # domain=self.domain, + # bckgr_obj=self.equil, + # species=species, + # ) + + # if self.rank_world == 0 and self.verbose: + # print( + # f'\nFluid species "{species}" was initialized with:', + # ) + + # _params = val["params"] + + # if "background" in _params: + # for variable in val: + # if "params" in variable: + # continue + # if variable in _params["background"]: + # bckgr_types = _params["background"][variable] + # if bckgr_types is None: + # pass + # else: + # print(f"{variable} background:") + # for _type, _bp in bckgr_types.items(): + # print(" " * 4 + _type, ":") + # for _pname, _pval in _bp.items(): + # print((" " * 8 + _pname + ":").ljust(25), _pval) + # else: + # print(f"{variable}: no background.") + # else: + # print("No background.") + + # if "perturbation" in _params: + # for variable in val: + # if "params" in variable: + # continue + # if variable in _params["perturbation"]: + # pert_types = _params["perturbation"][variable] + # if pert_types is None: + # pass + # else: + # print(f"{variable} perturbation:") + # for _type, _pp in pert_types.items(): + # print(" " * 4 + _type, ":") + # for _pname, _pval in _pp.items(): + # print((" " * 8 + _pname + ":").ljust(25), _pval) + # else: + # print(f"{variable}: no perturbation.") + # else: + # print("No perturbation.") + + # # initialize particles + # if len(self.kinetic) > 0: + # with ProfileManager.profile_region("initialize_particles"): + # for species, val in self.kinetic.items(): + # obj = val["obj"] + # assert isinstance(obj, Particles) + + # if self.rank_world == 0 and self.verbose: + # _params = val["params"] + # assert "background" in _params, "Kinetic species must have background." + + # bckgr_types = _params["background"] + # print( + # f'\nKinetic species "{species}" was initialized with:', + # ) + # for _type, _bp in bckgr_types.items(): + # print(_type, ":") + # for _pname, _pval in _bp.items(): + # print((" " * 4 + _pname + ":").ljust(25), _pval) + + # if "perturbation" in _params: + # for variable, pert_types in _params["perturbation"].items(): + # if pert_types is None: + # pass + # else: + # print(f"{variable} perturbation:") + # for _type, _pp in pert_types.items(): + # print(" " * 4 + _type, ":") + # for _pname, _pval in _pp.items(): + # print((" " * 8 + _pname + ":").ljust(25), _pval) + # else: + # print("No perturbation.") + + # obj.draw_markers(sort=True, verbose=self.verbose) + # obj.mpi_sort_markers(do_test=True) + + # if not val["params"]["markers"]["loading"] == "restart": + # if obj.coords == "vpara_mu": + # obj.save_magnetic_moment() + + # if val["space"] != "ParticlesSPH" and obj.f0.coords == "constants_of_motion": + # obj.save_constants_of_motion() + + # obj.initialize_weights( + # reject_weights=obj.weights_params["reject_weights"], + # threshold=obj.weights_params["threshold"], + # ) def initialize_from_restart(self, data): """ @@ -1069,106 +1192,60 @@ def initialize_data_output(self, data, size): else: pass - # save electromagentic fields/potentials data in group 'feec/' - for key, val in self.em_fields.items(): - if "params" in key: - continue - else: - obj = val["obj"] - assert isinstance(obj, SplineFunction) + # save feec data in group 'feec/' + feec_species = self.field_species | self.fluid_species | self.diagnostic_species + for species, val in feec_species.items(): + + assert isinstance(val, FieldSpecies) + + species_path = "feec/" + species + "_" + species_path_restart = "restart/" + species + "_" + + for variable, subval in val.variables.items(): + assert isinstance(subval, FEECVariable) + spline = subval.spline # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! - obj.extract_coeffs(update_ghost_regions=False) + spline.extract_coeffs(update_ghost_regions=False) # save numpy array to be updated each time step. - if val["save_data"]: - key_field = "feec/" + key + if subval.save_data: + key_field = species_path + variable - if isinstance(obj.vector_stencil, StencilVector): + if isinstance(spline.vector_stencil, StencilVector): data.add_data( - {key_field: obj.vector_stencil._data}, + {key_field: spline.vector_stencil._data}, ) else: for n in range(3): key_component = key_field + "/" + str(n + 1) data.add_data( - {key_component: obj.vector_stencil[n]._data}, + {key_component: spline.vector_stencil[n]._data}, ) # save field meta data - data.file[key_field].attrs["space_id"] = obj.space_id - data.file[key_field].attrs["starts"] = obj.starts - data.file[key_field].attrs["ends"] = obj.ends - data.file[key_field].attrs["pads"] = obj.pads + data.file[key_field].attrs["space_id"] = spline.space_id + data.file[key_field].attrs["starts"] = spline.starts + data.file[key_field].attrs["ends"] = spline.ends + data.file[key_field].attrs["pads"] = spline.pads # save numpy array to be updated only at the end of the simulation for restart. - key_field_restart = "restart/" + key + key_field_restart = species_path_restart + variable - if isinstance(obj.vector_stencil, StencilVector): + if isinstance(spline.vector_stencil, StencilVector): data.add_data( - {key_field_restart: obj.vector_stencil._data}, + {key_field_restart: spline.vector_stencil._data}, ) else: for n in range(3): key_component_restart = key_field_restart + "/" + str(n + 1) data.add_data( - {key_component_restart: obj.vector_stencil[n]._data}, - ) - - # save fluid data in group 'feec/' - for species, val in self.fluid.items(): - species_path = "feec/" + species + "_" - species_path_restart = "restart/" + species + "_" - - for variable, subval in val.items(): - if "params" in variable: - continue - else: - obj = subval["obj"] - assert isinstance(obj, SplineFunction) - - # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! - obj.extract_coeffs(update_ghost_regions=False) - - # save numpy array to be updated each time step. - if subval["save_data"]: - key_field = species_path + variable - - if isinstance(obj.vector_stencil, StencilVector): - data.add_data( - {key_field: obj.vector_stencil._data}, - ) - - else: - for n in range(3): - key_component = key_field + "/" + str(n + 1) - data.add_data( - {key_component: obj.vector_stencil[n]._data}, - ) - - # save field meta data - data.file[key_field].attrs["space_id"] = obj.space_id - data.file[key_field].attrs["starts"] = obj.starts - data.file[key_field].attrs["ends"] = obj.ends - data.file[key_field].attrs["pads"] = obj.pads - - # save numpy array to be updated only at the end of the simulation for restart. - key_field_restart = species_path_restart + variable - - if isinstance(obj.vector_stencil, StencilVector): - data.add_data( - {key_field_restart: obj.vector_stencil._data}, + {key_component_restart: spline.vector_stencil[n]._data}, ) - else: - for n in range(3): - key_component_restart = key_field_restart + "/" + str(n + 1) - data.add_data( - {key_component_restart: obj.vector_stencil[n]._data}, - ) # save kinetic data in group 'kinetic/' - for key, val in self.kinetic.items(): + for species, val in self.kinetic_species.items(): obj = val["obj"] assert isinstance(obj, Particles) @@ -1207,53 +1284,6 @@ def initialize_data_output(self, data, size): else: data.add_data({key_dat: val1}) - # save diagnostics data in group 'feec/' - for key, val in self.diagnostics.items(): - if "params" in key: - continue - else: - obj = val["obj"] - assert isinstance(obj, SplineFunction) - - # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! - obj.extract_coeffs(update_ghost_regions=False) - - # save numpy array to be updated each time step. - if val["save_data"]: - key_field = "feec/" + key - - if isinstance(obj.vector_stencil, StencilVector): - data.add_data( - {key_field: obj.vector_stencil._data}, - ) - - else: - for n in range(3): - key_component = key_field + "/" + str(n + 1) - data.add_data( - {key_component: obj.vector_stencil[n]._data}, - ) - - # save field meta data - data.file[key_field].attrs["space_id"] = obj.space_id - data.file[key_field].attrs["starts"] = obj.starts - data.file[key_field].attrs["ends"] = obj.ends - data.file[key_field].attrs["pads"] = obj.pads - - # save numpy array to be updated only at the end of the simulation for restart. - key_field_restart = "restart/" + key - - if isinstance(obj.vector_stencil, StencilVector): - data.add_data( - {key_field_restart: obj.vector_stencil._data}, - ) - else: - for n in range(3): - key_component_restart = key_field_restart + "/" + str(n + 1) - data.add_data( - {key_component_restart: obj.vector_stencil[n]._data}, - ) - # keys to be saved at each time step and only at end (restart) save_keys_all = [] save_keys_end = [] @@ -1518,13 +1548,6 @@ def generate_default_parameter_file( ------- The default parameter dictionary.""" - import os - - import yaml - - import struphy - from struphy.io.setup import descend_options_dict - libpath = struphy.__path__[0] # load a standard parameter file @@ -1792,165 +1815,7 @@ def _init_variable_dicts(self): else: self._diagnostics[var_name]["save_data"] = True - def _allocate_variables(self): - """ - Allocate memory for model variables. - Creates FEM fields for em-fields and fluid variables and a particle class for kinetic species. - """ - - from struphy.feec.psydac_derham import Derham - from struphy.pic import particles - from struphy.pic.base import Particles - - # allocate memory for FE coeffs of electromagnetic fields/potentials - if self.params.em_fields is not None: - for variable, dct in self.em_fields.items(): - if "params" in variable: - continue - else: - dct["obj"] = self.derham.create_spline_function( - variable, - dct["space"], - bckgr_params=dct.get("background"), - pert_params=dct.get("perturbation"), - ) - - self._pointer[variable] = dct["obj"].vector - - # allocate memory for FE coeffs of fluid variables - if self.params.fluid is not None: - for species, dct in self.fluid.items(): - for variable, subdct in dct.items(): - if "params" in variable: - continue - else: - subdct["obj"] = self.derham.create_spline_function( - variable, - subdct["space"], - bckgr_params=subdct.get("background"), - pert_params=subdct.get("perturbation"), - ) - - self._pointer[species + "_" + variable] = subdct["obj"].vector - - # marker arrays and plasma parameters of kinetic species - if self.params.kinetic is not None: - for species, val in self.kinetic.items(): - assert any([key in val["params"]["markers"] for key in ["Np", "ppc", "ppb"]]) - - bckgr_params = val["params"].get("background", None) - pert_params = val["params"].get("perturbation", None) - boxes_per_dim = val["params"].get("boxes_per_dim", None) - mpi_dims_mask = val["params"].get("dims_mask", None) - weights_params = val["params"].get("weights", None) - - if self.derham is None: - domain_decomp = None - else: - domain_array = self.derham.domain_array - nprocs = self.derham.domain_decomposition.nprocs - domain_decomp = (domain_array, nprocs) - - kinetic_class = getattr(particles, val["space"]) - - val["obj"] = kinetic_class( - comm_world=self.comm_world, - clone_config=self.clone_config, - **val["params"]["markers"], - weights_params=weights_params, - domain_decomp=domain_decomp, - mpi_dims_mask=mpi_dims_mask, - boxes_per_dim=boxes_per_dim, - name=species, - equation_params=self.equation_params[species], - domain=self.domain, - equil=self.equil, - projected_equil=self.projected_equil, - bckgr_params=bckgr_params, - pert_params=pert_params, - ) - - obj = val["obj"] - assert isinstance(obj, Particles) - - self._pointer[species] = obj - - # for storing markers - val["kinetic_data"] = {} - - # for storing the distribution function - if "f" in val["params"]["save_data"]: - slices = val["params"]["save_data"]["f"]["slices"] - n_bins = val["params"]["save_data"]["f"]["n_bins"] - ranges = val["params"]["save_data"]["f"]["ranges"] - - val["kinetic_data"]["f"] = {} - val["kinetic_data"]["df"] = {} - val["bin_edges"] = {} - if len(slices) > 0: - for i, sli in enumerate(slices): - assert ((len(sli) - 2) / 3).is_integer() - assert len(slices[i].split("_")) == len(ranges[i]) == len(n_bins[i]), ( - f"Number of slices names ({len(slices[i].split('_'))}), number of bins ({len(n_bins[i])}), and number of ranges ({len(ranges[i])}) are inconsistent with each other!\n\n" - ) - val["bin_edges"][sli] = [] - dims = (len(sli) - 2) // 3 + 1 - for j in range(dims): - val["bin_edges"][sli] += [ - np.linspace( - ranges[i][j][0], - ranges[i][j][1], - n_bins[i][j] + 1, - ), - ] - val["kinetic_data"]["f"][sli] = np.zeros( - n_bins[i], - dtype=float, - ) - val["kinetic_data"]["df"][sli] = np.zeros( - n_bins[i], - dtype=float, - ) - - # for storing an sph evaluation of the density n - if "n_sph" in val["params"]["save_data"]: - plot_pts = val["params"]["save_data"]["n_sph"]["plot_pts"] - - val["kinetic_data"]["n_sph"] = [] - val["plot_pts"] = [] - for i, pts in enumerate(plot_pts): - assert len(pts) == 3 - eta1 = np.linspace(0.0, 1.0, pts[0]) - eta2 = np.linspace(0.0, 1.0, pts[1]) - eta3 = np.linspace(0.0, 1.0, pts[2]) - ee1, ee2, ee3 = np.meshgrid( - eta1, - eta2, - eta3, - indexing="ij", - ) - val["plot_pts"] += [(ee1, ee2, ee3)] - val["kinetic_data"]["n_sph"] += [np.zeros(ee1.shape, dtype=float)] - - # other data (wave-particle power exchange, etc.) - # TODO - - # allocate memory for FE coeffs of diagnostics - if self.params.diagnostic_fields is not None: - for key, val in self.diagnostics.items(): - if "params" in key: - continue - else: - val["obj"] = self.derham.create_spline_function( - key, - val["space"], - bckgr_params=None, - pert_params=None, - ) - - self._pointer[key] = val["obj"].vector - - def _compute_plasma_params(self, verbose=True): + def compute_plasma_params(self, verbose=True): """ Compute and print volume averaged plasma parameters for each species of the model. @@ -1976,18 +1841,9 @@ def _compute_plasma_params(self, verbose=True): - rho/L - alpha = Omega_p/Omega_c - epsilon = 1/(t*Omega_c) - - Returns - ------- - pparams : dict - Plasma parameters for each species. """ - from struphy.fields_background import equils - from struphy.fields_background.base import FluidEquilibriumWithB - from struphy.kinetic_background import maxwellians - - pparams = {} + #TODO: needs re-factoring - create PlasmaParameters class instead of self._pparams dict # physics constants e = 1.602176634e-19 # elementary charge (C) @@ -1997,7 +1853,7 @@ def _compute_plasma_params(self, verbose=True): kB = 1.380649e-23 # Boltzmann constant (J*K^-1) # exit when there is not any plasma species - if len(self.fluid) == 0 and len(self.kinetic) == 0: + if len(self.fluid_species) == 0 and len(self.kinetic_species) == 0: return # compute model units @@ -2058,7 +1914,7 @@ def _compute_plasma_params(self, verbose=True): magnetic_field = np.nan # print("\n+++++++ WARNING +++++++ magnetic field is zero - set to nan !!") - if verbose: + if verbose and self.rank_world == 0: print("\nPLASMA PARAMETERS:") print( f"Plasma volume:".ljust(25), @@ -2082,19 +1938,19 @@ def _compute_plasma_params(self, verbose=True): ) # species dependent parameters - pparams = {} + self._pparams = {} if len(self.fluid) > 0: for species, val in self.fluid.items(): - pparams[species] = {} + self._pparams[species] = {} # type - pparams[species]["type"] = "fluid" + self._pparams[species]["type"] = "fluid" # mass (kg) - pparams[species]["mass"] = val["params"]["phys_params"]["A"] * m_p + self._pparams[species]["mass"] = val["params"]["phys_params"]["A"] * m_p # charge (C) - pparams[species]["charge"] = val["params"]["phys_params"]["Z"] * e + self._pparams[species]["charge"] = val["params"]["phys_params"]["Z"] * e # density (m⁻³) - pparams[species]["density"] = ( + self._pparams[species]["density"] = ( np.mean( self.equil.n0( eta1, @@ -2108,7 +1964,7 @@ def _compute_plasma_params(self, verbose=True): * units["n"] ) # pressure (bar) - pparams[species]["pressure"] = ( + self._pparams[species]["pressure"] = ( np.mean( self.equil.p0( eta1, @@ -2123,7 +1979,7 @@ def _compute_plasma_params(self, verbose=True): * 1e-5 ) # thermal energy (keV) - pparams[species]["kBT"] = pparams[species]["pressure"] * 1e5 / pparams[species]["density"] / e * 1e-3 + self._pparams[species]["kBT"] = self._pparams[species]["pressure"] * 1e5 / self._pparams[species]["density"] / e * 1e-3 if len(self.kinetic) > 0: eta1mg, eta2mg, eta3mg = np.meshgrid( @@ -2134,13 +1990,13 @@ def _compute_plasma_params(self, verbose=True): ) for species, val in self.kinetic.items(): - pparams[species] = {} + self._pparams[species] = {} # type - pparams[species]["type"] = "kinetic" + self._pparams[species]["type"] = "kinetic" # mass (kg) - pparams[species]["mass"] = val["params"]["phys_params"]["A"] * m_p + self._pparams[species]["mass"] = val["params"]["phys_params"]["A"] * m_p # charge (C) - pparams[species]["charge"] = val["params"]["phys_params"]["Z"] * e + self._pparams[species]["charge"] = val["params"]["phys_params"]["Z"] * e # create temp kinetic object for (default) parameter extraction tmp_bckgr = val["params"]["background"] @@ -2171,25 +2027,25 @@ def _compute_plasma_params(self, verbose=True): psi = self.equil.psi_r(r) # density (m⁻³) - pparams[species]["density"] = ( + self._pparams[species]["density"] = ( np.mean(tmp.n(psi) * np.abs(det_tmp)) * units["x"] ** 3 / plasma_volume * units["n"] ) # thermal speed (m/s) - pparams[species]["v_th"] = ( + self._pparams[species]["v_th"] = ( np.mean(tmp.vth(psi) * np.abs(det_tmp)) * units["x"] ** 3 / plasma_volume * units["v"] ) # thermal energy (keV) - pparams[species]["kBT"] = pparams[species]["mass"] * pparams[species]["v_th"] ** 2 / e * 1e-3 + self._pparams[species]["kBT"] = self._pparams[species]["mass"] * self._pparams[species]["v_th"] ** 2 / e * 1e-3 # pressure (bar) - pparams[species]["pressure"] = ( - pparams[species]["kBT"] * e * 1e3 * pparams[species]["density"] * 1e-5 + self._pparams[species]["pressure"] = ( + self._pparams[species]["kBT"] * e * 1e3 * self._pparams[species]["density"] * 1e-5 ) else: # density (m⁻³) - # pparams[species]['density'] = np.mean(tmp.n( + # self._pparams[species]['density'] = np.mean(tmp.n( # eta1mg, eta2mg, eta3mg) * np.abs(det_tmp)) * units['x']**3 / plasma_volume * units['n'] - pparams[species]["density"] = 99.0 + self._pparams[species]["density"] = 99.0 # thermal speeds (m/s) vth = [] # vths = tmp.vth(eta1mg, eta2mg, eta3mg) @@ -2200,55 +2056,55 @@ def _compute_plasma_params(self, verbose=True): ] thermal_speed = 0.0 for dir in range(val["obj"].vdim): - # pparams[species]['vth' + str(dir + 1)] = np.mean(vth[dir]) - pparams[species]["vth" + str(dir + 1)] = 99.0 - thermal_speed += pparams[species]["vth" + str(dir + 1)] + # self._pparams[species]['vth' + str(dir + 1)] = np.mean(vth[dir]) + self._pparams[species]["vth" + str(dir + 1)] = 99.0 + thermal_speed += self._pparams[species]["vth" + str(dir + 1)] # TODO: here it is assumed that background density parameter is called "n", # and that background thermal speeds are called "vthn"; make this a convention? - # pparams[species]['v_th'] = thermal_speed / \ + # self._pparams[species]['v_th'] = thermal_speed / \ # val['obj'].vdim - pparams[species]["v_th"] = 99.0 + self._pparams[species]["v_th"] = 99.0 # thermal energy (keV) - # pparams[species]['kBT'] = pparams[species]['mass'] * \ - # pparams[species]['v_th']**2 / e * 1e-3 - pparams[species]["kBT"] = 99.0 + # self._pparams[species]['kBT'] = self._pparams[species]['mass'] * \ + # self._pparams[species]['v_th']**2 / e * 1e-3 + self._pparams[species]["kBT"] = 99.0 # pressure (bar) - # pparams[species]['pressure'] = pparams[species]['kBT'] * \ - # e * 1e3 * pparams[species]['density'] * 1e-5 - pparams[species]["pressure"] = 99.0 + # self._pparams[species]['pressure'] = self._pparams[species]['kBT'] * \ + # e * 1e3 * self._pparams[species]['density'] * 1e-5 + self._pparams[species]["pressure"] = 99.0 - for species in pparams: + for species in self._pparams: # alfvén speed (m/s) - pparams[species]["v_A"] = magnetic_field / np.sqrt( - mu0 * pparams[species]["mass"] * pparams[species]["density"], + self._pparams[species]["v_A"] = magnetic_field / np.sqrt( + mu0 * self._pparams[species]["mass"] * self._pparams[species]["density"], ) # thermal speed (m/s) - pparams[species]["v_th"] = np.sqrt( - pparams[species]["kBT"] * 1e3 * e / pparams[species]["mass"], + self._pparams[species]["v_th"] = np.sqrt( + self._pparams[species]["kBT"] * 1e3 * e / self._pparams[species]["mass"], ) # thermal frequency (Mrad/s) - pparams[species]["Omega_th"] = pparams[species]["v_th"] / transit_length * 1e-6 + self._pparams[species]["Omega_th"] = self._pparams[species]["v_th"] / transit_length * 1e-6 # cyclotron frequency (Mrad/s) - pparams[species]["Omega_c"] = pparams[species]["charge"] * magnetic_field / pparams[species]["mass"] * 1e-6 + self._pparams[species]["Omega_c"] = self._pparams[species]["charge"] * magnetic_field / self._pparams[species]["mass"] * 1e-6 # plasma frequency (Mrad/s) - pparams[species]["Omega_p"] = ( + self._pparams[species]["Omega_p"] = ( np.sqrt( - pparams[species]["density"] * (pparams[species]["charge"]) ** 2 / eps0 / pparams[species]["mass"], + self._pparams[species]["density"] * (self._pparams[species]["charge"]) ** 2 / eps0 / self._pparams[species]["mass"], ) * 1e-6 ) # alfvén frequency (Mrad/s) - pparams[species]["Omega_A"] = pparams[species]["v_A"] / transit_length * 1e-6 + self._pparams[species]["Omega_A"] = self._pparams[species]["v_A"] / transit_length * 1e-6 # Larmor radius (m) - pparams[species]["rho_th"] = pparams[species]["v_th"] / (pparams[species]["Omega_c"] * 1e6) + self._pparams[species]["rho_th"] = self._pparams[species]["v_th"] / (self._pparams[species]["Omega_c"] * 1e6) # MHD length scale (m) - pparams[species]["v_A/Omega_c"] = pparams[species]["v_A"] / (np.abs(pparams[species]["Omega_c"]) * 1e6) + self._pparams[species]["v_A/Omega_c"] = self._pparams[species]["v_A"] / (np.abs(self._pparams[species]["Omega_c"]) * 1e6) # dim-less ratios - pparams[species]["rho_th/L"] = pparams[species]["rho_th"] / transit_length + self._pparams[species]["rho_th/L"] = self._pparams[species]["rho_th"] / transit_length - if verbose: + if verbose and self.rank_world == 0: print("\nSPECIES PARAMETERS:") - for species, ch in pparams.items(): + for species, ch in self._pparams.items(): print(f"\nname:".ljust(26), species) print(f"type:".ljust(25), ch["type"]) ch.pop("type") @@ -2262,8 +2118,6 @@ def _compute_plasma_params(self, verbose=True): units_affix[kinds], ) - return pparams - class MyDumper(yaml.SafeDumper): # HACK: insert blank lines between top-level objects From 72528d1226c67a4adca08ebba288f0e44f613d6e Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:24:21 +0200 Subject: [PATCH 039/292] added type annotation to FEECVariable declaration for better inference in parameter file --- src/struphy/models/toy.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index bfe80e92e..554a2090f 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -3,9 +3,10 @@ import numpy as np from struphy.models.base import StruphyModel +from struphy.propagators.base import Propagator from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies -from struphy.models.variables import FEECVariable, PICVariable, SPHVariable +from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable class Maxwell(StruphyModel): @@ -38,8 +39,8 @@ class Propagators: @dataclass class EMFields(FieldSpecies): - e_field = FEECVariable(space="Hcurl") - b_field = FEECVariable(space="Hdiv") + e_field: FEECVariable = FEECVariable(name="e_field", space="Hcurl") + b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") # @dataclass # class Ions(KineticSpecies): @@ -65,8 +66,13 @@ def __init__(self, units, domain, equil, verbose=False): self.em_fields.b_field, ) - # setup rest of model + # light-weight setup of model self.setup(units=units, domain=domain, equil=equil, verbose=verbose) + + # define scalars for update_scalar_quantities + self.add_scalar("electric energy") + self.add_scalar("magnetic energy") + self.add_scalar("total energy") ## variable and propagator attributes From 6aa641b2414ae49dddcac4b67c61bc3ec66b71d0 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:25:20 +0200 Subject: [PATCH 040/292] use new methods of model in main.py --- src/struphy/main.py | 184 ++++++++++++++++++-------------------------- 1 file changed, 76 insertions(+), 108 deletions(-) diff --git a/src/struphy/main.py b/src/struphy/main.py index 896f43fb9..e09ff1cb7 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -16,12 +16,15 @@ from struphy.models.base import StruphyModel from struphy.profiling.profiling import ProfileManager from struphy.utils.clone_config import CloneConfig +from struphy.utils.utils import dict_to_yaml from struphy.pic.base import Particles +from struphy.models.species import FieldSpecies +from struphy.models.variables import FEECVariable def main( model_name: Optional[str], - parameters: dict | str, + params_path: Optional[str], path_out: str, *, restart: bool = False, @@ -40,8 +43,8 @@ def main( model_name : str The name of the model to run. Type "struphy run --help" in your terminal to see a list of available models. - parameters : dict | str - The simulation parameters. Can either be a dictionary OR a string (path of .yml parameter file) + params_path : str + Path to .py parameter file. path_out : str The output directory. Will create a folder if it does not exist OR cleans the folder for new runs. @@ -71,75 +74,43 @@ def main( comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() + + # collect meta-data + meta = {} + meta["platform"] = sysconfig.get_platform() + meta["python version"] = sysconfig.get_python_version() + meta["model name"] = model_name + meta["parameter file"] = params_path + meta["output folder"] = path_out + meta["MPI processes"] = size + meta["number of domain clones"] = num_clones + meta["restart"] = restart + meta["max wall-clock [min]"] = runtime + meta["save interval [steps]"] = save_step + + # print meta-data on screen + print("\nMETADATA:") + for k, v in meta.items(): + print(f'{k}:'.ljust(25), v) + + # store meta-data in output folder + dict_to_yaml(meta, os.path.join(path_out, "meta.yml")) - if rank == 0: - print("") - comm.Barrier() - - # synchronize MPI processes to set same start time of simulation for all processes - comm.Barrier() + # start timing start_simulation = time.time() - # loading of simulation parameters, creating output folder and printing information to screen + # creating output folder, loading parameters, extract light-weight model instance setup_folders(path_out=path_out, restart=restart, - verbose=False,) + verbose=verbose,) - params, params_path = setup_parameters(model_name=model_name, - parameters=parameters, - path_out=path_out, - verbose=True,) - - # print simulation info - print("\nMETADATA:") - print("platform:".ljust(25), sysconfig.get_platform()) - print("python version:".ljust(25), sysconfig.get_python_version()) - print("model:".ljust(25), model_name) - print("parameter file:".ljust(25), params_path) - print("output folder:".ljust(25), path_out) - print("MPI processes:".ljust(25), size) - print("number of domain clones:".ljust(25), num_clones) - print("restart:".ljust(25), restart) - print("max wall-clock [min]:".ljust(25), runtime) - print("save interval [steps]:".ljust(25), save_step) - - exit() + params = setup_parameters(model_name=model_name, + params_path=params_path, + path_out=path_out, + verbose=verbose,) + model = params.model - # write meta data to output folder - with open(path_out + "/meta.txt", "w") as f: - f.write( - "date of simulation: ".ljust( - 30, - ) - + str(datetime.datetime.now()) - + "\n", - ) - f.write("platform: ".ljust(30) + sysconfig.get_platform() + "\n") - f.write( - "python version: ".ljust( - 30, - ) - + sysconfig.get_python_version() - + "\n", - ) - f.write("model_name: ".ljust(30) + model_name + "\n") - f.write("processes: ".ljust(30) + str(mpi_size) + "\n") - f.write("output folder:".ljust(30) + path_out + "\n") - f.write("restart:".ljust(30) + str(restart) + "\n") - f.write( - "max wall-clock time [min]:".ljust(30) + str(max_sim_time) + "\n", - ) - f.write("save interval (steps):".ljust(30) + str(save_step) + "\n") - - if model_name is None: - assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" - model_name = params.model - - if rank < 32: - print(f"Rank {rank}: calling struphy/main.py for model {model_name} ...") - if size > 32 and rank == 32: - print(f"Ranks > 31: calling struphy/main.py for model {model_name} ...") - + # config clones if comm is None: clone_config = None else: @@ -152,23 +123,36 @@ def main( # between the clones : inter_comm clone_config = CloneConfig(comm=comm, params=params, num_clones=num_clones) clone_config.print_clone_config() - if "kinetic" in params: + if model.kinetic_species: clone_config.print_particle_config() + comm.Barrier() + + # allocate derham-related objects + if params.grid is not None: + model.allocate_feec(params.grid, params.derham) + else: + model._derham = None + model._mass_ops = None + model._projected_equil = None + print("\nDERHAM:\nMeshless simulation - no Derham complex set up.") + + # allocate variables + model.allocate_variables() + + # pass info to propagators + model.set_propagators() + + # plasma parameters + model.compute_plasma_params(verbose=verbose) - # instantiate Struphy model (will allocate model objects and associated memory) - StruphyModel.verbose = verbose - - objs = [fluid, kinetic, hybrid, toy] - for obj in objs: - try: - model_class = getattr(obj, model_name) - except AttributeError: - pass - - with ProfileManager.profile_region("model_class_setup"): - model = model_class(params=params, comm=comm, clone_config=clone_config) + if model_name is None: + assert model is not None, "If model is not specified, then model: MODEL must be specified in the params!" + model_name = model.__class__.__name__ - assert isinstance(model, StruphyModel) + if rank < 32: + print(f"Rank {rank}: calling struphy/main.py for model {model_name} ...") + if size > 32 and rank == 32: + print(f"Ranks > 31: calling struphy/main.py for model {model_name} ...") # store geometry vtk if rank == 0: @@ -217,12 +201,7 @@ def main( split_algo = params.time.split_algo # set initial conditions for all variables - if not restart: - model.initialize_from_params() - - total_steps = str(int(round(Tend / dt))) - - else: + if restart: model.initialize_from_restart(data) time_state["value"][0] = data.file["restart/time/value"][-1] @@ -230,6 +209,8 @@ def main( time_state["index"][0] = data.file["restart/time/index"][-1] total_steps = str(int(round((Tend - time_state["value"][0]) / dt))) + else: + total_steps = str(int(round(Tend / dt))) # compute initial scalars and kinetic data, pass time state to all propagators model.update_scalar_quantities() @@ -291,7 +272,7 @@ def main( # update time and index (round time to 10 decimals for a clean time grid!) time_state["value"][0] = round(time_state["value"][0] + dt, 10) - time_state["value_sec"][0] = round(time_state["value_sec"][0] + dt * model.units["t"], 10) + time_state["value_sec"][0] = round(time_state["value_sec"][0] + dt * model.units.t, 10) time_state["index"][0] += 1 run_time_now = (time.time() - start_simulation) / 60 @@ -303,28 +284,15 @@ def main( model.update_markers_to_be_saved() model.update_distr_functions() - # extract FEM coefficients - for key, val in model.em_fields.items(): - if "params" not in key: - field = val["obj"] - assert isinstance(field, SplineFunction) - # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! - field.extract_coeffs(update_ghost_regions=False) - - for _, val in model.fluid.items(): - for variable, subval in val.items(): - if "params" not in variable: - field = subval["obj"] - assert isinstance(field, SplineFunction) - # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! - field.extract_coeffs(update_ghost_regions=False) - - for key, val in model.diagnostics.items(): - if "params" not in key: - field = val["obj"] - assert isinstance(field, SplineFunction) + # extract FEEC coefficients + feec_species = model.field_species | model.fluid_species | model.diagnostic_species + for species, val in feec_species.items(): + assert isinstance(val, FieldSpecies) + for variable, subval in val.variables.items(): + assert isinstance(subval, FEECVariable) + spline = subval.spline # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! - field.extract_coeffs(update_ghost_regions=False) + spline.extract_coeffs(update_ghost_regions=False) # save data (everything but restart data) data.save_data(keys=save_keys_all) @@ -336,7 +304,7 @@ def main( message = "time step: " + step + "/" + str(total_steps) message += " | " + "time: {0:10.5f}/{1:10.5f}".format(time_state["value"][0], Tend) message += " | " + "phys. time [s]: {0:12.10f}/{1:12.10f}".format( - time_state["value_sec"][0], Tend * model.units["t"] + time_state["value_sec"][0], Tend * model.units.t ) message += " | " + "wall clock [s]: {0:8.4f} | last step duration [s]: {1:8.4f}".format( run_time_now * 60, t1 - t0 From a3b96017dfae063c63a04f80324ac043b2a26520 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:26:16 +0200 Subject: [PATCH 041/292] have a working light-weight .py parameter file for Maxwell --- src/struphy/io/inp/params_Maxwell_lw.py | 51 +++++++++++++++---------- src/struphy/io/parameters.py | 4 +- src/struphy/linear_algebra/solver.py | 3 +- src/struphy/utils/clone_config.py | 8 ++-- src/struphy/utils/utils.py | 3 +- 5 files changed, 40 insertions(+), 29 deletions(-) diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index ab2df300c..cce261d32 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -1,67 +1,76 @@ -from struphy.io import options +from struphy.io.options import Units, Time, DerhamOptions, FieldsBackground +from struphy.fields_background import equils +from struphy.geometry import domains +from struphy.initial import perturbations +from struphy.kinetic_background import maxwellians +from struphy.topology import grids # import model from struphy.models.toy import Maxwell as Model -verbose = True +verbose = False # units -units = options.Units(x=2.0,) +units = Units() # geometry -domain = options.domains.Cuboid() +domain = domains.Cuboid() # fluid equilibrium (can be used as part of initial conditions) -equil = options.equils.HomogenSlab() +equil = equils.HomogenSlab() # time -time = options.Time() +time = Time() # grid -grid = options.grids.TensorProductGrid( +grid = grids.TensorProductGrid( Nel=(12, 14, 1), p=(2, 3, 1), spl_kind=(False, True, True), ) # derham options -derham = options.DerhamOptions() +derham = DerhamOptions() # light-weight instance of model model = Model(units, domain, equil, verbose=verbose) -propagators = model.propagators # model.fluid.set_phys_params("mhd", options.PhysParams()) # model.kinetic.set_phys_params("mhd", options.PhysParams()) # propagator options -propagators.maxwell.set_options(verbose=verbose) +model.propagators.maxwell.set_options(algo="explicit") # initial conditions for model variables (background + perturbation) model.em_fields.e_field.add_background( - options.FieldsBackground( - kind="LogicalConst", + FieldsBackground( + type="LogicalConst", values=(0.3, 0.15, None), ), verbose=verbose, ) model.em_fields.e_field.add_perturbation( - options.perturbations.TorusModesCos( - ms=[[None], [1, 3], [None]], + perturbations.TorusModesCos( + ms=[1, 3], + given_in_basis="v", + comp=1, ), - given_in_basis=(None, "v", None), verbose=verbose, ) model.em_fields.b_field.add_background( - options.FieldsBackground( - kind="LogicalConst", + FieldsBackground( + type="LogicalConst", values=(0.3, 0.15, None), ), verbose=verbose, ) model.em_fields.b_field.add_perturbation( - options.perturbations.TorusModesCos( - ms=[[None], [1, 3], [None]], + perturbations.TorusModesCos( + ms=[1, 3], + given_in_basis="v", + comp=1, ), - given_in_basis=(None, "v", None), verbose=verbose, -) \ No newline at end of file +) + +# exclude variable from saving +model.em_fields.e_field.save_data = False \ No newline at end of file diff --git a/src/struphy/io/parameters.py b/src/struphy/io/parameters.py index 9a52e49d3..0a4741dc2 100644 --- a/src/struphy/io/parameters.py +++ b/src/struphy/io/parameters.py @@ -33,13 +33,13 @@ def __init__( self._derham = derham if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"{self.model = }") + print(f"\n{self.model = }") print(f"{self.units = }") print(f"{self.domain = }") print(f"{self.equil = }") print(f"{self.time = }") print(f"{self.grid = }") - print(f"{self.derham = }") + print(f"{self.derham = }\n") @property def model(self): diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py index be337f01e..a47f1d592 100644 --- a/src/struphy/linear_algebra/solver.py +++ b/src/struphy/linear_algebra/solver.py @@ -2,8 +2,7 @@ @dataclass class SolverParameters: - """ - Parameters for psydac solvers.""" + """Parameters for psydac solvers.""" tol: float = 1e-8 maxiter: int = 3000 diff --git a/src/struphy/utils/clone_config.py b/src/struphy/utils/clone_config.py index b4894e2cf..046844df5 100644 --- a/src/struphy/utils/clone_config.py +++ b/src/struphy/utils/clone_config.py @@ -1,6 +1,8 @@ import numpy as np from mpi4py import MPI +from struphy.io.parameters import StruphyParameters + class CloneConfig: """ @@ -17,7 +19,7 @@ class CloneConfig: def __init__( self, comm: MPI.Intracomm, - params=None, + params: StruphyParameters, num_clones=1, ): """ @@ -26,8 +28,8 @@ def __init__( Parameters: comm : (MPI.Intracomm) The MPI communicator covering all processes. - params : dict, optional - Dictionary containing simulation parameters. + params : StruphyParameters + Struphy simulation parameters. num_clones : int, optional The number of clones to create. The total number of MPI ranks must be divisible by this number. """ diff --git a/src/struphy/utils/utils.py b/src/struphy/utils/utils.py index 46be2d25b..f52fb3f85 100644 --- a/src/struphy/utils/utils.py +++ b/src/struphy/utils/utils.py @@ -73,7 +73,8 @@ def print_all_attr(obj): print(k.ljust(26), v) -def dict_to_yaml(dictionary, output): +def dict_to_yaml(dictionary: dict, output: str): + """Write dictionary to file and save in output.""" with open(output, "w") as file: yaml.dump( dictionary, From e13bcaf806d93413e393550d477f7dd8e6cfdf54 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 08:58:38 +0200 Subject: [PATCH 042/292] work on allocate of Maxwell propagator --- src/struphy/propagators/propagators_fields.py | 52 +++++++++---------- 1 file changed, 25 insertions(+), 27 deletions(-) diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 3985e8d13..a8e809178 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -104,22 +104,20 @@ def set_options(self, butcher=butcher,) def allocate(self): - solver = {} - # obtain needed matrices M1 = self.mass_ops.M1 M2 = self.mass_ops.M2 curl = self.derham.curl # Preconditioner for M1 + ... - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) - pc = pc_class(self.mass_ops.M1) + pc_class = getattr(preconditioner, self.options.precond) + pc = pc_class(M1) - if self._algo == "implicit": - self._info = solver["info"] + if self.options.algo == "implicit": + self._info = self.options.solver_params.info # Define block matrix [[A B], [C I]] (without time step size dt in the diagonals) _A = M1 @@ -133,11 +131,11 @@ def allocate(self): self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], + self.options.solver, pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, ) # pre-allocate arrays @@ -148,43 +146,43 @@ def allocate(self): # define vector field M1_inv = inverse( M1, - solver["type"][0], + self.options.solver, pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, ) weak_curl = M1_inv @ curl.T @ M2 # allocate output of vector field - out1 = e.space.zeros() - out2 = b.space.zeros() + out1 = self.e.space.zeros() + out2 = self.b.space.zeros() - def f1(t, y1, y2, out=out1): + def f1(t, y1, y2, out: BlockVector = out1): weak_curl.dot(y2, out=out) out.update_ghost_regions() return out - def f2(t, y1, y2, out=out2): + def f2(t, y1, y2, out: BlockVector = out2): curl.dot(y1, out=out) out *= -1.0 out.update_ghost_regions() return out - vector_field = {e: f1, b: f2} - self._ode_solver = ODEsolverFEEC(vector_field, algo=algo) + vector_field = {self.e: f1, self.b: f2} + self._ode_solver = ODEsolverFEEC(vector_field, algo=self.options.butcher) # allocate place-holder vectors to avoid temporary array allocations in __call__ - self._e_tmp1 = e.space.zeros() - self._e_tmp2 = e.space.zeros() - self._b_tmp1 = b.space.zeros() + self._e_tmp1 = self.e.space.zeros() + self._e_tmp2 = self.e.space.zeros() + self._b_tmp1 = self.b.space.zeros() def __call__(self, dt): # current variables - en = self.feec_vars[0] - bn = self.feec_vars[1] + en = self.e + bn = self.b - if self._algo == "implicit": + if self.options.algo == "implicit": # solve for new e coeffs self._B.dot(bn, out=self._byn) From d2b2f56cbcd9ee3cfe90dd348529a49db7d93e36 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 29 Jul 2025 19:45:39 +0200 Subject: [PATCH 043/292] added allocate_propagators and prop_list to StruphyModel --- src/struphy/io/inp/params_Maxwell_lw.py | 2 +- src/struphy/main.py | 4 ++- src/struphy/models/base.py | 26 +++++++++++++------ src/struphy/models/toy.py | 11 ++++---- src/struphy/propagators/propagators_fields.py | 1 - 5 files changed, 27 insertions(+), 17 deletions(-) diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index cce261d32..09e074c33 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -7,7 +7,7 @@ # import model from struphy.models.toy import Maxwell as Model -verbose = False +verbose = True # units units = Units() diff --git a/src/struphy/main.py b/src/struphy/main.py index e09ff1cb7..16f8cd31b 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -140,7 +140,7 @@ def main( model.allocate_variables() # pass info to propagators - model.set_propagators() + model.allocate_propagators() # plasma parameters model.compute_plasma_params(verbose=verbose) @@ -150,6 +150,8 @@ def main( model_name = model.__class__.__name__ if rank < 32: + if rank == 0: + print("") print(f"Rank {rank}: calling struphy/main.py for model {model_name} ...") if size > 32 and rank == 32: print(f"Ranks > 31: calling struphy/main.py for model {model_name} ...") diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index a4b05aa87..5e47b15c4 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -260,7 +260,7 @@ def allocate_feec(self, self.derham, ) - def set_propagators(self): + def allocate_propagators(self): # set propagators base class attributes (then available to all propagators) Propagator.derham = self.derham Propagator.domain = self.domain @@ -273,6 +273,12 @@ def set_propagators(self): eq_mhd=self.equil, ) Propagator.projected_equil = self.projected_equil + + for prop in self.prop_list: + assert isinstance(prop, Propagator) + prop.allocate() + if self.verbose and self.rank_world == 0: + print(f"\nAllocated propagator {prop.__class__.__name__}.") @staticmethod def diagnostics_dct(): @@ -362,6 +368,13 @@ def propagators(self): self._propagators = None raise NameError("Propagators object must be named 'self._propagators = ...' in model __init__().") return self._propagators + + @property + def prop_list(self): + """List of Propagator objects.""" + if not hasattr(self, "_prop_list"): + self._prop_list = list(self.propagators.__dict__.values()) + return self._prop_list @property def prop_fields(self): @@ -794,12 +807,9 @@ def integrate(self, dt, split_algo="LieTrotter"): Splitting algorithm. Currently available: "LieTrotter" and "Strang". """ - if not hasattr(self, "_prop_list"): - self._prop_list = list(self.propagators.__dict__.values()) - # first order in time if split_algo == "LieTrotter": - for propagator in self._prop_list: + for propagator in self.prop_list: prop_name = propagator.__class__.__name__ with ProfileManager.profile_region(prop_name): @@ -809,17 +819,17 @@ def integrate(self, dt, split_algo="LieTrotter"): elif split_algo == "Strang": assert len(self.propagators) > 1 - for propagator in self._prop_list[:-1]: + for propagator in self.prop_list[:-1]: prop_name = type(propagator).__name__ with ProfileManager.profile_region(prop_name): propagator(dt / 2) - propagator = self._prop_list[-1] + propagator = self.prop_list[-1] prop_name = type(propagator).__name__ with ProfileManager.profile_region(prop_name): propagator(dt) - for propagator in self._prop_list[:-1][::-1]: + for propagator in self.prop_list[:-1][::-1]: prop_name = type(propagator).__name__ with ProfileManager.profile_region(prop_name): propagator(dt / 2) diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 554a2090f..e8113ef91 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,5 +1,4 @@ -from dataclasses import dataclass import numpy as np from struphy.models.base import StruphyModel @@ -33,14 +32,14 @@ class Maxwell(StruphyModel): :ref:`Model info `: """ - @dataclass class Propagators: - maxwell = propagators_fields.Maxwell() + def __init__(self): + self.maxwell = propagators_fields.Maxwell() - @dataclass class EMFields(FieldSpecies): - e_field: FEECVariable = FEECVariable(name="e_field", space="Hcurl") - b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") + def __init__(self): + self.e_field = FEECVariable(name="e_field", space="Hcurl") + self.b_field = FEECVariable(name="b_field", space="Hdiv") # @dataclass # class Ions(KineticSpecies): diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index a8e809178..112292a08 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -50,7 +50,6 @@ from struphy.models.variables import FEECVariable, PICVariable, SPHVariable -@dataclass class Maxwell(Propagator): r""":ref:`FEEC ` discretization of the following equations: find :math:`\mathbf E \in H(\textnormal{curl})` and :math:`\mathbf B \in H(\textnormal{div})` such that From ff2c5f13517ed336260d4b506e67731953c5ffa8 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 30 Jul 2025 09:55:26 +0200 Subject: [PATCH 044/292] improved Propagator.set_variables() method; make Maxwell toy model work --- src/struphy/io/options.py | 4 + src/struphy/models/toy.py | 4 +- src/struphy/ode/solvers.py | 10 +- src/struphy/ode/utils.py | 2 +- src/struphy/propagators/base.py | 137 ++++++++---------- src/struphy/propagators/propagators_fields.py | 57 ++++---- 6 files changed, 95 insertions(+), 119 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index cd957f93e..70eb555a9 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -30,6 +30,10 @@ def check_option(opt, options): NoiseDirections = Literal["e1", "e2", "e3", "e1e2", "e1e3", "e2e3", "e1e2e3"] GivenInBasis = Literal['0', '1', '2', '3', 'v', 'physical', 'physical_at_eta', 'norm', None] +# solvers +OptsSymmSolver = Literal["pcg", "cg"] +OptsMassPrecond = Literal["MassMatrixPreconditioner", None] + ## Option classes diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index e8113ef91..913b9d253 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -61,8 +61,8 @@ def __init__(self, units, domain, equil, verbose=False): # 3. assign variables to propagators self.propagators.maxwell.set_variables( - self.em_fields.e_field, - self.em_fields.b_field, + e = self.em_fields.e_field, + b = self.em_fields.b_field, ) # light-weight setup of model diff --git a/src/struphy/ode/solvers.py b/src/struphy/ode/solvers.py index 4706b6b26..12900a878 100644 --- a/src/struphy/ode/solvers.py +++ b/src/struphy/ode/solvers.py @@ -31,10 +31,10 @@ class ODEsolverFEEC: def __init__( self, vector_field: dict, - algo: str = "rk4", + butcher: ButcherTableau = ButcherTableau(), ): # get algorithm - self._butcher = ButcherTableau(algo=algo) + self._butcher = butcher # check arguments and allocate k for each stage self._k = {} @@ -51,7 +51,6 @@ def __init__( self._k[vec] += [vec.space.zeros()] self._vector_field = vector_field - self._algo = algo # collect unknows in list self._y = list(self.vector_field.keys()) @@ -96,11 +95,6 @@ def vector_field(self): values are callables representing the respective component of the vector field.""" return self._vector_field - @property - def algo(self): - """See :class:`~struphy.ode.utils.ButcherTableau` for available algorithms.""" - return self._algo - @property def y(self): """List of variables to be updated.""" diff --git a/src/struphy/ode/utils.py b/src/struphy/ode/utils.py index 9d8d48b77..14af9b848 100644 --- a/src/struphy/ode/utils.py +++ b/src/struphy/ode/utils.py @@ -65,7 +65,7 @@ def __post_init__(self): c = (0.0, 1 / 3, 2 / 3, 1.0) conv_rate = 4 else: - raise NotImplementedError("Chosen algorithm is not implemented.") + raise NotImplementedError(f"Chosen algorithm {self.algo} is not implemented.") self._b = np.array(b) self._c = np.array(c) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 7e756cc9a..0f68300ef 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -9,6 +9,8 @@ from struphy.feec.psydac_derham import Derham from struphy.geometry.base import Domain from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable +from psydac.linalg.stencil import StencilVector +from psydac.linalg.block import BlockVector class Propagator(metaclass=ABCMeta): @@ -21,33 +23,67 @@ class Propagator(metaclass=ABCMeta): Only propagators that update both a FEEC and a PIC species go into ``propagators_coupling.py``. """ - def __init__(self, *vars): - """Create an instance of a Propagator. + @abstractmethod + def set_options(self, **kwargs): + """Set the dynamical options of the propagator (kwargs).""" + + @abstractmethod + def allocate(): + """Allocate all data/objects of the instance.""" + + @abstractmethod + def __call__(self, dt): + """Update variables from t -> t + dt. + Use ``Propagators.feec_vars_update`` to write to FEEC variables to ``Propagator.feec_vars``. Parameters ---------- - vars : Variable - Variables to be updated. + dt : float + Time step size. """ - comm = None - - # for iterative particle push - self._init_kernels = [] - self._eval_kernels = [] - - self._rank = comm.Get_rank() if comm is not None else 0 - - def set_variables(self, *vars): - for var in vars: - assert isinstance(var, Variable) - if isinstance(var, (PICVariable, SPHVariable)): + def set_variables(self, **vars): + """Update variables in __dict__ with user-defined instance variables (allocated).""" + assert len(vars) == len(self.__dict__), f"Variables must be passed in the following order: {self.__dict__}, but is {vars}." + for ((k, v), (kp, vp)) in zip(vars.items(), self.__dict__.items()): + assert k == kp, f"Variable name '{k}' not equal to '{kp}'; variables must be passed in the order {self.__dict__}." + assert isinstance(v, Variable) + if isinstance(v, (PICVariable, SPHVariable)): # comm = var.obj.mpi_comm pass - elif isinstance(var, FEECVariable): + elif isinstance(v, FEECVariable): # comm = var.obj.comm pass - self._vars = vars + self.__dict__.update(vars) + + def update_feec_variables(self, **new_coeffs): + r"""Return max_diff = max(abs(new - old)) for each new_coeffs, + update feec coefficients and update ghost regions. + + Returns + ------- + diffs : dict + max_diff for all feec variables. + """ + diffs = {} + assert len(new_coeffs) == len(self.__dict__), f"Coefficients must be passed in the following order: {self.__dict__}, but is {new_coeffs}." + for ((k, new), (kp, vp)) in zip(new_coeffs.items(), self.__dict__.items()): + assert k == kp, f"Variable name '{k}' not equal to '{kp}'; variables must be passed in the order {self.__dict__}." + assert isinstance(new, (StencilVector, BlockVector)) + assert isinstance(vp, FEECVariable) + old = vp.spline.vector + assert new.space == old.space + + # calculate maximum of difference abs(new - old) + diffs[k] = np.max(np.abs(new.toarray() - old.toarray())) + + # copy new coeffs into old + new.copy(out=old) + + # important: sync processes! + old.update_ghost_regions() + + return diffs @property def vars(self): @@ -77,27 +113,6 @@ def rank(self): """MPI rank, is 0 if no communicator.""" return self._rank - @abstractmethod - def set_options(self, **opts): - """Set the dynamical options of the propagator (kwargs). - """ - - @abstractmethod - def allocate(): - """Allocate all data/objects for an instance. - """ - - @abstractmethod - def __call__(self, dt): - """Update variables from t -> t + dt. - Use ``Propagators.feec_vars_update`` to write to FEEC variables to ``Propagator.feec_vars``. - - Parameters - ---------- - dt : float - Time step size. - """ - @property def derham(self): """Derham spaces and projectors.""" @@ -174,40 +189,6 @@ def add_time_state(self, time_state): assert time_state.size == 1 self._time_state = time_state - def feec_vars_update(self, *variables_new): - r"""Return :math:`\textrm{max}_i |x_i(t + \Delta t) - x_i(t)|` for each unknown in list, - update :method:`~struphy.propagators.base.Propagator.feec_vars` - and update ghost regions. - - Parameters - ---------- - variables_new : list[StencilVector | BlockVector] - Same sequence as in :method:`~struphy.propagators.base.Propagator.feec_vars` - but with the updated variables, - i.e. for feec_vars = [e, b] we must have variables_new = [e_updated, b_updated]. - - Returns - ------- - diffs : list - A list [max(abs(self.feec_vars - variables_new)), ...] for all variables in self.feec_vars and variables_new. - """ - - diffs = [] - - for i, new in enumerate(variables_new): - assert type(new) is type(self.feec_vars[i]) - - # calculate maximum of difference abs(old - new) - diffs += [np.max(np.abs(self.feec_vars[i].toarray() - new.toarray()))] - - # copy new variables into self.feec_vars - new.copy(out=self.feec_vars[i]) - - # important: sync processes! - self.feec_vars[i].update_ghost_regions() - - return diffs - def add_init_kernel( self, kernel, @@ -312,11 +293,11 @@ def options(self, new): class Options: def __init__(self, prop, **kwargs): - self._kwargs = kwargs + self.__dict__.update(kwargs) print(f"\nInstance of propagator '{prop.__class__.__name__}' with:") - for k, v in kwargs.items(): + for k, v in self.__dict__.items(): print(f' {k}: {v}') - @property - def kwargs(self): - return self._kwargs \ No newline at end of file + # @property + # def kwargs(self): + # return self._kwargs \ No newline at end of file diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 112292a08..fa9ab2080 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -46,10 +46,11 @@ from struphy.propagators.base import Propagator from struphy.models.variables import Variable from struphy.linear_algebra.solver import SolverParameters -from struphy.io.options import check_option +from struphy.io.options import check_option, OptsSymmSolver, OptsMassPrecond from struphy.models.variables import FEECVariable, PICVariable, SPHVariable +@dataclass class Maxwell(Propagator): r""":ref:`FEEC ` discretization of the following equations: find :math:`\mathbf E \in H(\textnormal{curl})` and :math:`\mathbf B \in H(\textnormal{div})` such that @@ -62,37 +63,33 @@ class Maxwell(Propagator): :ref:`time_discret`: Crank-Nicolson (implicit mid-point). System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`. """ + # variables to be updated e: FEECVariable = None b: FEECVariable = None + # propagator specific options OptsAlgo = Literal["implicit", "explicit"] - OptsSolver = Literal["pcg", "cg"] - OptsPrecond = Literal["MassMatrixPreconditioner", None] + # abstract methods def set_options(self, algo: OptsAlgo = "implicit", - solver: OptsSolver = "pcg", - precond: OptsPrecond = "MassMatrixPreconditioner", + solver: OptsSymmSolver = "pcg", + precond: OptsMassPrecond = "MassMatrixPreconditioner", solver_params: SolverParameters = None, butcher: ButcherTableau = None, ): # checks check_option(algo, self.OptsAlgo) - check_option(solver, self.OptsSolver) - check_option(precond, self.OptsPrecond) + check_option(solver, OptsSymmSolver) + check_option(precond, OptsMassPrecond) # defaults - if algo == "implicit": - butcher = None - if solver_params is None: - solver_params = SolverParameters() - elif algo == "explicit": - solver = None - precond = None - solver_params = None - if butcher is None: - butcher = ButcherTableau() + if solver_params is None: + solver_params = SolverParameters() + + if algo == "explicit" and butcher is None: + butcher = ButcherTableau() # use setter for options self.options = self.Options(self, @@ -154,8 +151,8 @@ def allocate(self): weak_curl = M1_inv @ curl.T @ M2 # allocate output of vector field - out1 = self.e.space.zeros() - out2 = self.b.space.zeros() + out1 = self.e.spline.vector.space.zeros() + out2 = self.b.spline.vector.space.zeros() def f1(t, y1, y2, out: BlockVector = out1): weak_curl.dot(y2, out=out) @@ -168,18 +165,18 @@ def f2(t, y1, y2, out: BlockVector = out2): out.update_ghost_regions() return out - vector_field = {self.e: f1, self.b: f2} - self._ode_solver = ODEsolverFEEC(vector_field, algo=self.options.butcher) + vector_field = {self.e.spline.vector: f1, self.b.spline.vector: f2} + self._ode_solver = ODEsolverFEEC(vector_field, butcher=self.options.butcher) # allocate place-holder vectors to avoid temporary array allocations in __call__ - self._e_tmp1 = self.e.space.zeros() - self._e_tmp2 = self.e.space.zeros() - self._b_tmp1 = self.b.space.zeros() + self._e_tmp1 = self.e.spline.vector.space.zeros() + self._e_tmp2 = self.e.spline.vector.space.zeros() + self._b_tmp1 = self.b.spline.vector.space.zeros() def __call__(self, dt): - # current variables - en = self.e - bn = self.b + # current FE coeffs + en = self.e.spline.vector + bn = self.b.spline.vector if self.options.algo == "implicit": # solve for new e coeffs @@ -195,7 +192,7 @@ def __call__(self, dt): bn1 += bn # write new coeffs into self.feec_vars - max_de, max_db = self.feec_vars_update(en1, bn1) + diffs = self.update_feec_variables(e=en1, b=bn1) else: self._ode_solver(0.0, dt) @@ -203,8 +200,8 @@ def __call__(self, dt): if self._algo == "implicit": print("Status for Maxwell:", info["success"]) print("Iterations for Maxwell:", info["niter"]) - print("Maxdiff e1 for Maxwell:", max_de) - print("Maxdiff b2 for Maxwell:", max_db) + print("Maxdiff e for Maxwell:", diffs["e"]) + print("Maxdiff b for Maxwell:", diffs["b"]) print() From bd7c452bea10db4a71be480ff01e09b3d6334465 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 30 Jul 2025 12:56:52 +0200 Subject: [PATCH 045/292] Make coaxial test for Maxwell run --- src/struphy/models/toy.py | 12 +++++++++--- src/struphy/models/variables.py | 8 ++++++-- .../post_processing/orbits/orbits_tools.py | 1 - src/struphy/propagators/base.py | 15 ++++++++------- 4 files changed, 23 insertions(+), 13 deletions(-) diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 913b9d253..1aec439d7 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,5 +1,6 @@ import numpy as np +from dataclasses import dataclass from struphy.models.base import StruphyModel from struphy.propagators.base import Propagator @@ -36,10 +37,15 @@ class Propagators: def __init__(self): self.maxwell = propagators_fields.Maxwell() + @dataclass class EMFields(FieldSpecies): - def __init__(self): - self.e_field = FEECVariable(name="e_field", space="Hcurl") - self.b_field = FEECVariable(name="b_field", space="Hdiv") + e_field: FEECVariable = FEECVariable(name="e_field", space="Hcurl") + b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") + + # class EMFields(FieldSpecies): + # def __init__(self): + # self.e_field = FEECVariable(name="e_field", space="Hcurl") + # self.b_field = FEECVariable(name="b_field", space="Hdiv") # @dataclass # class Ions(KineticSpecies): diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index cbbbf7a1f..ba46081d3 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -18,10 +18,14 @@ def allocate(self): @property def backgrounds(self): + if not hasattr(self, "_backgrounds"): + self._backgrounds = None return self._backgrounds @property def perturbations(self): + if not hasattr(self, "_perturbations"): + self._perturbations = None return self._perturbations @property @@ -44,7 +48,7 @@ def species(self): def add_background(self, background, verbose=False): """Type inference of added background done in sub class.""" - if not hasattr(self, "_backgrounds"): + if not hasattr(self, "_backgrounds") or self.backgrounds is None: self._backgrounds = background else: if not isinstance(self.backgrounds, list): @@ -57,7 +61,7 @@ def add_background(self, background, verbose=False): print(f' {k}: {v}') def add_perturbation(self, perturbation: Perturbation, verbose=False): - if not hasattr(self, "_perturbations"): + if not hasattr(self, "_perturbations") or self.perturbations is None: self._perturbations = perturbation else: if not isinstance(self.perturbations, list): diff --git a/src/struphy/post_processing/orbits/orbits_tools.py b/src/struphy/post_processing/orbits/orbits_tools.py index 909b10942..cce0d85b9 100644 --- a/src/struphy/post_processing/orbits/orbits_tools.py +++ b/src/struphy/post_processing/orbits/orbits_tools.py @@ -6,7 +6,6 @@ import yaml from tqdm import tqdm -from struphy.io.setup import setup_domain_and_equil from struphy.post_processing.orbits.orbits_kernels import calculate_guiding_center_from_6d diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 0f68300ef..85ce70a2f 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -43,7 +43,7 @@ def __call__(self, dt): """ def set_variables(self, **vars): - """Update variables in __dict__ with user-defined instance variables (allocated).""" + """Update variables in __dict__ and set self.variables with user-defined instance variables (allocated).""" assert len(vars) == len(self.__dict__), f"Variables must be passed in the following order: {self.__dict__}, but is {vars}." for ((k, v), (kp, vp)) in zip(vars.items(), self.__dict__.items()): assert k == kp, f"Variable name '{k}' not equal to '{kp}'; variables must be passed in the order {self.__dict__}." @@ -55,6 +55,7 @@ def set_variables(self, **vars): # comm = var.obj.comm pass self.__dict__.update(vars) + self._variables = vars def update_feec_variables(self, **new_coeffs): r"""Return max_diff = max(abs(new - old)) for each new_coeffs, @@ -66,9 +67,9 @@ def update_feec_variables(self, **new_coeffs): max_diff for all feec variables. """ diffs = {} - assert len(new_coeffs) == len(self.__dict__), f"Coefficients must be passed in the following order: {self.__dict__}, but is {new_coeffs}." - for ((k, new), (kp, vp)) in zip(new_coeffs.items(), self.__dict__.items()): - assert k == kp, f"Variable name '{k}' not equal to '{kp}'; variables must be passed in the order {self.__dict__}." + assert len(new_coeffs) == len(self.variables), f"Coefficients must be passed in the following order: {self.variables}, but is {new_coeffs}." + for ((k, new), (kp, vp)) in zip(new_coeffs.items(), self.variables.items()): + assert k == kp, f"Variable name '{k}' not equal to '{kp}'; variables must be passed in the order {self.variables}." assert isinstance(new, (StencilVector, BlockVector)) assert isinstance(vp, FEECVariable) old = vp.spline.vector @@ -86,10 +87,10 @@ def update_feec_variables(self, **new_coeffs): return diffs @property - def vars(self): - """List of Variables to be updated by the propagator. + def variables(self): + """Dict of Variables to be updated by the propagator. """ - return self._vars + return self._variables @property def init_kernels(self): From 6a5e4539a5a2420b5cbc76695a934307eb77f7b1 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 30 Jul 2025 15:26:44 +0200 Subject: [PATCH 046/292] add function import_parameters_py; make pproc work with new params --- src/struphy/feec/psydac_derham.py | 8 +- src/struphy/io/setup.py | 20 ++-- src/struphy/main.py | 18 ++-- src/struphy/models/base.py | 4 + src/struphy/models/toy.py | 15 +-- .../post_processing/post_processing_tools.py | 47 ++++----- src/struphy/post_processing/pproc_struphy.py | 98 +++++++------------ 7 files changed, 91 insertions(+), 119 deletions(-) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index ee634bfb4..9451d7fe3 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -879,6 +879,7 @@ def create_spline_function( perturbations: Perturbation | list = None, domain: Domain = None, equil: FluidEquilibrium = None, + verbose: bool = True, ): """Creat a callable spline function. @@ -914,6 +915,7 @@ def create_spline_function( perturbations=perturbations, domain=domain, equil=equil, + verbose=verbose, ) def prepare_eval_tp_fixed(self, grids_1d): @@ -1426,6 +1428,7 @@ def __init__( perturbations: Perturbation | list = None, domain: Domain = None, equil: FluidEquilibrium = None, + verbose: bool = True, ): self._name = name self._space_id = space_id @@ -1472,8 +1475,9 @@ def __init__( ) else: self._nbasis = [tuple([space.nbasis for space in vec_space.spaces]) for vec_space in self.fem_space.spaces] - - print(f"\nAllocated SplineFuntion '{self.name}' in space '{self.space_id}'.") + + if verbose: + print(f"\nAllocated SplineFuntion '{self.name}' in space '{self.space_id}'.") if self.backgrounds is not None or self.perturbations is not None: self.initialize_coeffs(domain=self.domain, equil=self.equil) diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index 6f1e75b5b..fa4d6ec83 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -14,6 +14,16 @@ from struphy.geometry.base import Domain +def import_parameters_py(params_path: str): + """Import a .py parameter file under the module name 'parameters' and return it.""" + assert ".py" in params_path + spec = importlib.util.spec_from_file_location("parameters", params_path) + params_in = importlib.util.module_from_spec(spec) + sys.modules["parameters"] = params_in + spec.loader.exec_module(params_in) + return params_in + + def setup_folders( path_out: str, restart: bool, @@ -75,7 +85,6 @@ def setup_folders( def setup_parameters( - model_name: str, params_path: str, path_out: str, verbose: bool = False, @@ -112,10 +121,7 @@ def setup_parameters( # state = read_state() # i_path = state["i_path"] # load parameter.py - spec = importlib.util.spec_from_file_location("parameters", params_path) - params_in = importlib.util.module_from_spec(spec) - sys.modules["parameters"] = params_in - spec.loader.exec_module(params_in) + params_in = import_parameters_py(params_path) if not hasattr(params_in, "model"): params_in.model = None @@ -149,10 +155,6 @@ def setup_parameters( verbose=verbose, ) - if model_name is None: - assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" - model_name = params.model - if MPI.COMM_WORLD.Get_rank() == 0: # copy parameter file to output folder filename = params_path.split("/")[-1] diff --git a/src/struphy/main.py b/src/struphy/main.py index 16f8cd31b..814046f3a 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -74,6 +74,17 @@ def main( comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() + + start_simulation = time.time() + + # read in paramaters + params = setup_parameters(params_path=params_path, + path_out=path_out, + verbose=verbose,) + + if model_name is None: + assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" + model_name = params.model.__class__.__name__ # collect meta-data meta = {} @@ -96,18 +107,13 @@ def main( # store meta-data in output folder dict_to_yaml(meta, os.path.join(path_out, "meta.yml")) - # start timing - start_simulation = time.time() # creating output folder, loading parameters, extract light-weight model instance setup_folders(path_out=path_out, restart=restart, verbose=verbose,) - params = setup_parameters(model_name=model_name, - params_path=params_path, - path_out=path_out, - verbose=verbose,) + # get model instance model = params.model # config clones diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 5e47b15c4..c62805095 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -369,6 +369,10 @@ def propagators(self): raise NameError("Propagators object must be named 'self._propagators = ...' in model __init__().") return self._propagators + @propagators.setter + def propagators(self, new): + self._propagators = new + @property def prop_list(self): """List of Propagator objects.""" diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 1aec439d7..0dbad68f7 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -58,12 +58,12 @@ class EMFields(FieldSpecies): def __init__(self, units, domain, equil, verbose=False): # 1. instantiate all variales - self._em_fields = self.EMFields() + self.em_fields = self.EMFields() # self._ions = self.Ions() # self._electrons = self.Electrons() # 2. instantiate all propagators - self._propagators = self.Propagators() + self.propagators = self.Propagators() # 3. assign variables to propagators self.propagators.maxwell.set_variables( @@ -79,17 +79,6 @@ def __init__(self, units, domain, equil, verbose=False): self.add_scalar("magnetic energy") self.add_scalar("total energy") - ## variable and propagator attributes - - @property - def propagators(self) -> Propagators: - """A list of propagator instances for the model.""" - return self._propagators - - @property - def em_fields(self) -> EMFields: - return self._em_fields - ## abstract attributes @property diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 127cda4f6..0ad776b4e 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -8,9 +8,10 @@ from tqdm import tqdm from struphy.feec.psydac_derham import Derham -from struphy.io.setup import setup_domain_and_equil from struphy.kinetic_background import maxwellians from struphy.models import fluid, hybrid, kinetic, toy +from struphy.io.setup import import_parameters_py +from struphy.models.base import setup_derham def create_femfields( @@ -35,38 +36,28 @@ def create_femfields( space_ids : dict The space IDs of the fields (H1, Hcurl, Hdiv, L2 or H1vec). space_ids[name] contains the space ID of the field with the name "name". - - model : str - From which model in struphy/models the data has been obtained. """ - # get model name and # of MPI processes from meta.txt file - with open(os.path.join(path, "meta.txt"), "r") as f: - lines = f.readlines() - - model = lines[3].split()[-1] - nproc = lines[4].split()[-1] + with open(os.path.join(path, "meta.yml"), "r") as f: + meta = yaml.load(f, Loader=yaml.FullLoader) + nproc = meta["MPI processes"] - # create Derham sequence from grid parameters - with open(os.path.join(path, "parameters.yml"), "r") as f: - params = yaml.load(f, Loader=yaml.FullLoader) + params_in = import_parameters_py(os.path.join(path, "parameters.py")) - derham = Derham( - params["grid"]["Nel"], - params["grid"]["p"], - params["grid"]["spl_kind"], - ) + derham = setup_derham(params_in.grid, + params_in.derham, + comm=None, + domain=params_in.domain, + mpi_dims_mask=params_in.grid.mpi_dims_mask, + ) # get fields names, space IDs and time grid from 0-th rank hdf5 file file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") - space_ids = {} - for field_name, dset in file["feec"].items(): space_ids[field_name] = dset.attrs["space_id"] t_grid = file["time/value"][::step].copy() - file.close() # create one FemField for each snapshot @@ -74,9 +65,10 @@ def create_femfields( for t in t_grid: fields[t] = {} for field_name, ID in space_ids.items(): - fields[t][field_name] = derham.create_spline_function(field_name, ID) + fields[t][field_name] = derham.create_spline_function(field_name, ID, verbose=False,) # get hdf5 data + print("") for rank in range(int(nproc)): # open hdf5 file file = h5py.File( @@ -142,7 +134,7 @@ def create_femfields( print("Creation of Struphy Fields done.") - return fields, space_ids, model + return fields, space_ids def eval_femfields( @@ -191,17 +183,16 @@ def eval_femfields( assert isinstance(fields, dict) assert isinstance(space_ids, dict) - # domain object according to parameter file and grids - with open(os.path.join(path, "parameters.yml"), "r") as f: - params = yaml.load(f, Loader=yaml.FullLoader) + params_in = import_parameters_py(os.path.join(path, "parameters.py")) - domain = setup_domain_and_equil(params)[0] + # get domain + domain = params_in.domain # create logical and physical grids assert isinstance(celldivide, list) assert len(celldivide) == 3 - Nel = params["grid"]["Nel"] + Nel = params_in.grid.Nel grids_log = [np.linspace(0.0, 1.0, Nel_i * n_i + 1) for Nel_i, n_i in zip(Nel, celldivide)] grids_phy = [ diff --git a/src/struphy/post_processing/pproc_struphy.py b/src/struphy/post_processing/pproc_struphy.py index 7342ce223..e90bd03d1 100644 --- a/src/struphy/post_processing/pproc_struphy.py +++ b/src/struphy/post_processing/pproc_struphy.py @@ -1,3 +1,16 @@ +import os +import pickle +import shutil + +import h5py +import numpy as np +import yaml + +import struphy.post_processing.orbits.orbits_tools as orbits_pproc +import struphy.post_processing.post_processing_tools as pproc +from struphy.io.setup import import_parameters_py + + def main( path: str, *, @@ -37,19 +50,6 @@ def main( time_trace : bool whether to plot the time trace of each measured region """ - - import os - import pickle - import shutil - - import h5py - import numpy as np - import yaml - - import struphy.post_processing.orbits.orbits_tools as orbits_pproc - import struphy.post_processing.post_processing_tools as pproc - from struphy.models import fluid, hybrid, kinetic, toy - print("") # create post-processing folder @@ -75,23 +75,6 @@ def main( # save time grid at which post-processing data is created np.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) - # load parameters.yml - with open(os.path.join(path, "parameters.yml"), "r") as f: - params = yaml.load(f, Loader=yaml.FullLoader) - - # get model class from meta.txt file - with open(os.path.join(path, "meta.txt"), "r") as f: - lines = f.readlines() - model_name = lines[3].split()[-1] - - objs = [fluid, kinetic, hybrid, toy] - - for obj in objs: - try: - model_class = getattr(obj, model_name) - except AttributeError: - pass - if "feec" in file.keys(): exist_fields = True else: @@ -99,34 +82,27 @@ def main( if "kinetic" in file.keys(): exist_kinetic = {"markers": False, "f": False, "n_sph": False} - - kinetic_species = [] - kinetic_kinds = [] - for name in file["kinetic"].keys(): - kinetic_species += [name] - kinetic_kinds += [model_class.species()["kinetic"][name]] - # check for saved markers if "markers" in file["kinetic"][name]: exist_kinetic["markers"] = True - # check for saved distribution function if "f" in file["kinetic"][name]: exist_kinetic["f"] = True - # check for saved sph density if "n_sph" in file["kinetic"][name]: exist_kinetic["n_sph"] = True - else: exist_kinetic = None - + file.close() + + # import parameters + params_in = import_parameters_py(os.path.join(path, "parameters.py")) # field post-processing if exist_fields: - fields, space_ids, _ = pproc.create_femfields(path, step=step) + fields, space_ids = pproc.create_femfields(path, step=step) point_data, grids_log, grids_phy = pproc.eval_femfields( path, fields, space_ids, celldivide=[celldivide, celldivide, celldivide] @@ -149,25 +125,25 @@ def main( # save data dicts for each field for name, val in point_data.items(): aux = name.split("_") - # is em field - if len(aux) == 1 or "field" in name: - subfolder = "em_fields" - new_name = name - try: - os.mkdir(os.path.join(path_fields, subfolder)) - except: - pass - - # is fluid species - else: - subfolder = aux[0] - for au in aux[1:-1]: - subfolder += "_" + au - new_name = aux[-1] - try: - os.mkdir(os.path.join(path_fields, subfolder)) - except: - pass + # # is em field + # if len(aux) == 1 or "field" in name: + # subfolder = "em_fields" + # new_name = name + # try: + # os.mkdir(os.path.join(path_fields, subfolder)) + # except: + # pass + + # # is fluid species + # else: + subfolder = aux[0] + for au in aux[1:-1]: + subfolder += "_" + au + new_name = aux[-1] + try: + os.mkdir(os.path.join(path_fields, subfolder)) + except: + pass print(f"{name = }") print(f"{subfolder = }") From 6721f7666a90fee7b6716fb84753f4d8ab5be883 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 30 Jul 2025 20:47:07 +0200 Subject: [PATCH 047/292] added params file Maxwell_coaxial.py for verification; made pproc work for feec fields in new species-variables framework --- .../io/inp/verification/Maxwell_coaxial.py | 59 ++++ src/struphy/models/base.py | 12 +- .../post_processing/post_processing_tools.py | 316 +++++++++--------- src/struphy/post_processing/pproc_struphy.py | 51 +-- 4 files changed, 242 insertions(+), 196 deletions(-) create mode 100644 src/struphy/io/inp/verification/Maxwell_coaxial.py diff --git a/src/struphy/io/inp/verification/Maxwell_coaxial.py b/src/struphy/io/inp/verification/Maxwell_coaxial.py new file mode 100644 index 000000000..9bf9b836c --- /dev/null +++ b/src/struphy/io/inp/verification/Maxwell_coaxial.py @@ -0,0 +1,59 @@ +from struphy.io.options import Units, Time, DerhamOptions, FieldsBackground +from struphy.fields_background import equils +from struphy.geometry import domains +from struphy.initial import perturbations +from struphy.kinetic_background import maxwellians +from struphy.topology import grids + +# import model +from struphy.models.toy import Maxwell as Model +verbose = True + +# units +units = Units() + +# geometry +a1=2.326744 +a2=3.686839 +domain = domains.HollowCylinder(a1=a1, a2=a2, Lz=2.0) + +# fluid equilibrium (can be used as part of initial conditions) +equil = equils.HomogenSlab() + +# time +time = Time(dt=0.05, Tend=10.0) + +# grid +grid = grids.TensorProductGrid( + Nel=(32, 64, 1), + p=(3, 3, 1), + spl_kind=(False, True, True), + dirichlet_bc=[[True, True], [False, False], [False, False]] +) + +# derham options +derham = DerhamOptions() + +# light-weight instance of model +model = Model(units, domain, equil, verbose=verbose) +# model.fluid.set_phys_params("mhd", options.PhysParams()) +# model.kinetic.set_phys_params("mhd", options.PhysParams()) + +# propagator options +model.propagators.maxwell.set_options(algo="implicit") + +# initial conditions for model variables (background + perturbation) +model.em_fields.e_field.add_perturbation( + perturbations.CoaxialWaveguideElectric_r(m=3, a1=a1, a2=a2), + verbose=verbose, +) + +model.em_fields.e_field.add_perturbation( + perturbations.CoaxialWaveguideElectric_theta(m=3, a1=a1, a2=a2), + verbose=verbose, +) + +model.em_fields.b_field.add_perturbation( + perturbations.CoaxialWaveguideMagnetic(m=3, a1=a1, a2=a2), + verbose=verbose, +) \ No newline at end of file diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index c62805095..21723de60 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1212,8 +1212,8 @@ def initialize_data_output(self, data, size): assert isinstance(val, FieldSpecies) - species_path = "feec/" + species + "_" - species_path_restart = "restart/" + species + "_" + species_path = os.path.join("feec", species) + species_path_restart = os.path.join("restart", species) for variable, subval in val.variables.items(): assert isinstance(subval, FEECVariable) @@ -1224,7 +1224,7 @@ def initialize_data_output(self, data, size): # save numpy array to be updated each time step. if subval.save_data: - key_field = species_path + variable + key_field = os.path.join(species_path, variable) if isinstance(spline.vector_stencil, StencilVector): data.add_data( @@ -1233,7 +1233,7 @@ def initialize_data_output(self, data, size): else: for n in range(3): - key_component = key_field + "/" + str(n + 1) + key_component = os.path.join(key_field, str(n + 1)) data.add_data( {key_component: spline.vector_stencil[n]._data}, ) @@ -1245,7 +1245,7 @@ def initialize_data_output(self, data, size): data.file[key_field].attrs["pads"] = spline.pads # save numpy array to be updated only at the end of the simulation for restart. - key_field_restart = species_path_restart + variable + key_field_restart = os.path.join(species_path_restart, variable) if isinstance(spline.vector_stencil, StencilVector): data.add_data( @@ -1253,7 +1253,7 @@ def initialize_data_output(self, data, size): ) else: for n in range(3): - key_component_restart = key_field_restart + "/" + str(n + 1) + key_component_restart = os.path.join(key_field_restart, str(n + 1)) data.add_data( {key_component_restart: spline.vector_stencil[n]._data}, ) diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 0ad776b4e..d2fafc79f 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -12,10 +12,12 @@ from struphy.models import fluid, hybrid, kinetic, toy from struphy.io.setup import import_parameters_py from struphy.models.base import setup_derham +from struphy.feec.psydac_derham import SplineFunction def create_femfields( path: str, + params_in, *, step: int = 1, ): @@ -25,6 +27,9 @@ def create_femfields( ---------- path : str Absolute path of simulation output folder. + + params_in : parameter module + Imported parameter .py module. step : int Whether to create FEM fields at every time step (step=1, default), every second time step (step=2), etc. @@ -33,17 +38,15 @@ def create_femfields( ------- fields : dict Nested dictionary holding :class:`~struphy.feec.psydac_derham.SplineFunction`: fields[t][name] contains the Field with the name "name" in the hdf5 file at time t. - - space_ids : dict - The space IDs of the fields (H1, Hcurl, Hdiv, L2 or H1vec). space_ids[name] contains the space ID of the field with the name "name". + + t_grid : np.ndarray + Time grid. """ with open(os.path.join(path, "meta.yml"), "r") as f: meta = yaml.load(f, Loader=yaml.FullLoader) nproc = meta["MPI processes"] - params_in = import_parameters_py(os.path.join(path, "parameters.py")) - derham = setup_derham(params_in.grid, params_in.derham, comm=None, @@ -54,8 +57,13 @@ def create_femfields( # get fields names, space IDs and time grid from 0-th rank hdf5 file file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") space_ids = {} - for field_name, dset in file["feec"].items(): - space_ids[field_name] = dset.attrs["space_id"] + print(f"\nReading hdf5 data of following species:") + for species, dset in file["feec"].items(): + space_ids[species] = {} + print(f"{species}:") + for var, ddset in dset.items(): + space_ids[species][var] = ddset.attrs["space_id"] + print(f" {var}:", ddset) t_grid = file["time/value"][::step].copy() file.close() @@ -64,8 +72,10 @@ def create_femfields( fields = {} for t in t_grid: fields[t] = {} - for field_name, ID in space_ids.items(): - fields[t][field_name] = derham.create_spline_function(field_name, ID, verbose=False,) + for species, vars in space_ids.items(): + fields[t][species] = {} + for var, id in vars.items(): + fields[t][species][var] = derham.create_spline_function(var, id, verbose=False,) # get hdf5 data print("") @@ -80,67 +90,66 @@ def create_femfields( "r", ) - for field_name, dset in tqdm(file["feec"].items()): - # get global start indices, end indices and pads - gl_s = dset.attrs["starts"] - gl_e = dset.attrs["ends"] - pads = dset.attrs["pads"] - - assert gl_s.shape == (3,) or gl_s.shape == (3, 3) - assert gl_e.shape == (3,) or gl_e.shape == (3, 3) - assert pads.shape == (3,) or pads.shape == (3, 3) - - # loop over time - for n, t in enumerate(t_grid): - # scalar field - if gl_s.shape == (3,): - s1, s2, s3 = gl_s - e1, e2, e3 = gl_e - p1, p2, p3 = pads - - data = dset[n * step, p1:-p1, p2:-p2, p3:-p3].copy() - - fields[t][field_name].vector[ - s1 : e1 + 1, - s2 : e2 + 1, - s3 : e3 + 1, - ] = data - # update after each data addition, can be made more efficient - fields[t][field_name].vector.update_ghost_regions() - - # vector-valued field - else: - for comp in range(3): - s1, s2, s3 = gl_s[comp] - e1, e2, e3 = gl_e[comp] - p1, p2, p3 = pads[comp] - - data = dset[str(comp + 1)][ - n * step, - p1:-p1, - p2:-p2, - p3:-p3, - ].copy() - - fields[t][field_name].vector[comp][ + for species, dset in file["feec"].items(): + for var, ddset in tqdm(dset.items()): + # get global start indices, end indices and pads + gl_s = ddset.attrs["starts"] + gl_e = ddset.attrs["ends"] + pads = ddset.attrs["pads"] + + assert gl_s.shape == (3,) or gl_s.shape == (3, 3) + assert gl_e.shape == (3,) or gl_e.shape == (3, 3) + assert pads.shape == (3,) or pads.shape == (3, 3) + + # loop over time + for n, t in enumerate(t_grid): + # scalar field + if gl_s.shape == (3,): + s1, s2, s3 = gl_s + e1, e2, e3 = gl_e + p1, p2, p3 = pads + + data = ddset[n * step, p1:-p1, p2:-p2, p3:-p3].copy() + + fields[t][species][var].vector[ s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, ] = data - # update after each data addition, can be made more efficient - fields[t][field_name].vector.update_ghost_regions() + # update after each data addition, can be made more efficient + fields[t][species][var].vector.update_ghost_regions() + # vector-valued field + else: + for comp in range(3): + s1, s2, s3 = gl_s[comp] + e1, e2, e3 = gl_e[comp] + p1, p2, p3 = pads[comp] + + data = ddset[str(comp + 1)][ + n * step, + p1:-p1, + p2:-p2, + p3:-p3, + ].copy() + + fields[t][species][var].vector[comp][ + s1 : e1 + 1, + s2 : e2 + 1, + s3 : e3 + 1, + ] = data + # update after each data addition, can be made more efficient + fields[t][species][var].vector.update_ghost_regions() file.close() print("Creation of Struphy Fields done.") - return fields, space_ids + return fields, t_grid def eval_femfields( - path: str, + params_in, fields: dict, - space_ids: dict, *, celldivide: list = [1, 1, 1], physical: bool = False, @@ -149,15 +158,12 @@ def eval_femfields( Parameters ---------- - path : str - Absolute path of simulation output folder. + params_in : parameter module + Imported parameter .py module. fields : dict Obtained from struphy.diagnostics.post_processing.create_femfields. - space_ids : dict - Obtained from struphy.diagnostics.post_processing.create_femfields. - celldivide : list of ints Grid refinement in each eta direction. @@ -180,15 +186,11 @@ def eval_femfields( Mapped (physical) grids obtained by domain(*grids_log). """ - assert isinstance(fields, dict) - assert isinstance(space_ids, dict) - - params_in = import_parameters_py(os.path.join(path, "parameters.py")) - # get domain domain = params_in.domain # create logical and physical grids + assert isinstance(fields, dict) assert isinstance(celldivide, list) assert len(celldivide) == 3 @@ -203,86 +205,87 @@ def eval_femfields( # evaluate fields at evaluation grid and push-forward point_data = {} - - # one dict for each field - for name in space_ids: - point_data[name] = {} - - # time loop + for species, vars in fields[list(fields.keys())[0]].items(): + point_data[species] = {} + for name, field in vars.items(): + point_data[species][name] = {} + print("Evaluating fields ...") for t in tqdm(fields): - # field loop - for name, field in fields[t].items(): - # space ID - space_id = space_ids[name] - - # field evaluation - temp_val = field(*grids_log) - - point_data[name][t] = [] - - # scalar spaces - if isinstance(temp_val, np.ndarray): - if physical: - # push-forward - if space_id == "H1": - point_data[name][t].append( - domain.push( - temp_val, - *grids_log, - kind="0", - ), - ) - elif space_id == "L2": - point_data[name][t].append( - domain.push( - temp_val, - *grids_log, - kind="3", - ), - ) + for species, vars in fields[t].items(): + for name, field in vars.items(): + + assert isinstance(field, SplineFunction) + space_id = field.space_id - else: - point_data[name][t].append(temp_val) + # field evaluation + temp_val = field(*grids_log) + + point_data[species][name][t] = [] - # vector-valued spaces - else: - for j in range(3): + # scalar spaces + if isinstance(temp_val, np.ndarray): if physical: # push-forward - if space_id == "Hcurl": - point_data[name][t].append( + if space_id == "H1": + point_data[species][name][t].append( domain.push( temp_val, *grids_log, - kind="1", - )[j], + kind="0", + ), ) - elif space_id == "Hdiv": - point_data[name][t].append( + elif space_id == "L2": + point_data[species][name][t].append( domain.push( temp_val, *grids_log, - kind="2", - )[j], - ) - elif space_id == "H1vec": - point_data[name][t].append( - domain.push( - temp_val, - *grids_log, - kind="v", - )[j], + kind="3", + ), ) else: - point_data[name][t].append(temp_val[j]) + point_data[species][name][t].append(temp_val) + # vector-valued spaces + else: + for j in range(3): + if physical: + # push-forward + if space_id == "Hcurl": + point_data[species][name][t].append( + domain.push( + temp_val, + *grids_log, + kind="1", + )[j], + ) + elif space_id == "Hdiv": + point_data[species][name][t].append( + domain.push( + temp_val, + *grids_log, + kind="2", + )[j], + ) + elif space_id == "H1vec": + point_data[species][name][t].append( + domain.push( + temp_val, + *grids_log, + kind="v", + )[j], + ) + + else: + point_data[species][name][t].append(temp_val[j]) + return point_data, grids_log, grids_phy def create_vtk( path: str, + t_grid: np.ndarray, grids_phy: list, point_data: dict, *, @@ -295,6 +298,9 @@ def create_vtk( path : str Absolute path of where to store the .vts files. Will then be in path/vtk/step_.vts. + t_grid : np.ndarray + Time grid. + grids_phy : 3-list Mapped (physical) grids obtained from struphy.diagnostics.post_processing.eval_femfields. @@ -306,46 +312,46 @@ def create_vtk( """ from pyevtk.hl import gridToVTK - - # directory for vtk files - path_vtk = os.path.join(path, "vtk" + physical * "_phy") - - try: - os.mkdir(path_vtk) - except: - shutil.rmtree(path_vtk) - os.mkdir(path_vtk) - - # field names - names = list(point_data.keys()) + + for species, vars in point_data.items(): + species_path = os.path.join(path, species, "vtk" + physical * "_phy") + try: + os.mkdir(species_path) + except: + shutil.rmtree(species_path) + os.mkdir(species_path) # time loop - tgrid = list(point_data[names[0]].keys()) - - nt = len(tgrid) - 1 + nt = len(t_grid) - 1 log_nt = int(np.log10(nt)) + 1 print("Creating vtk ...") - for n, t in enumerate(tqdm(tgrid)): + for n, t in enumerate(tqdm(t_grid)): point_data_n = {} - for name in names: - points_list = point_data[name][t] + for species, vars in point_data.items(): + species_path = os.path.join(path, species, "vtk" + physical * "_phy") + point_data_n[species] = {} + for name, data in vars.items(): + points_list = data[t] - # scalar - if len(points_list) == 1: - point_data_n[name] = points_list[0] + # scalar + if len(points_list) == 1: + point_data_n[species][name] = points_list[0] - # vector - else: - for j in range(3): - point_data_n[name + f"_{j + 1}"] = points_list[j] + # vectorpoint_data[name] + else: + for j in range(3): + point_data_n[species][name + f"_{j + 1}"] = points_list[j] - gridToVTK( - os.path.join(path_vtk, "step_{0:0{1}d}".format(n, log_nt)), - *grids_phy, - pointData=point_data_n, - ) + gridToVTK( + os.path.join(species_path, "step_{0:0{1}d}".format(n, log_nt)), + *grids_phy, + pointData=point_data_n[species], + ) + + if n==0: + print(f" saving .vtk files to {species_path}") def post_process_markers(path_in, path_out, species, kind, step=1): diff --git a/src/struphy/post_processing/pproc_struphy.py b/src/struphy/post_processing/pproc_struphy.py index e90bd03d1..9524abe60 100644 --- a/src/struphy/post_processing/pproc_struphy.py +++ b/src/struphy/post_processing/pproc_struphy.py @@ -102,15 +102,15 @@ def main( # field post-processing if exist_fields: - fields, space_ids = pproc.create_femfields(path, step=step) + fields, t_grid = pproc.create_femfields(path, params_in, step=step) point_data, grids_log, grids_phy = pproc.eval_femfields( - path, fields, space_ids, celldivide=[celldivide, celldivide, celldivide] + params_in, fields, celldivide=[celldivide, celldivide, celldivide] ) if physical: point_data_phy, grids_log, grids_phy = pproc.eval_femfields( - path, fields, space_ids, celldivide=[celldivide, celldivide, celldivide], physical=True + params_in, fields, celldivide=[celldivide, celldivide, celldivide], physical=True ) # directory for field data @@ -123,38 +123,19 @@ def main( os.mkdir(path_fields) # save data dicts for each field - for name, val in point_data.items(): - aux = name.split("_") - # # is em field - # if len(aux) == 1 or "field" in name: - # subfolder = "em_fields" - # new_name = name - # try: - # os.mkdir(os.path.join(path_fields, subfolder)) - # except: - # pass - - # # is fluid species - # else: - subfolder = aux[0] - for au in aux[1:-1]: - subfolder += "_" + au - new_name = aux[-1] - try: - os.mkdir(os.path.join(path_fields, subfolder)) - except: - pass + for species, vars in point_data.items(): + for name, val in vars.items(): + try: + os.mkdir(os.path.join(path_fields, species)) + except: + pass - print(f"{name = }") - print(f"{subfolder = }") - print(f"{new_name = }") + with open(os.path.join(path_fields, species, name + "_log.bin"), "wb") as handle: + pickle.dump(val, handle, protocol=pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_fields, subfolder, new_name + "_log.bin"), "wb") as handle: - pickle.dump(val, handle, protocol=pickle.HIGHEST_PROTOCOL) - - if physical: - with open(os.path.join(path_fields, subfolder, new_name + "_phy.bin"), "wb") as handle: - pickle.dump(point_data_phy[name], handle, protocol=pickle.HIGHEST_PROTOCOL) + if physical: + with open(os.path.join(path_fields, species, name + "_phy.bin"), "wb") as handle: + pickle.dump(point_data_phy[species][name], handle, protocol=pickle.HIGHEST_PROTOCOL) # save grids with open(os.path.join(path_fields, "grids_log.bin"), "wb") as handle: @@ -165,9 +146,9 @@ def main( # create vtk files if not no_vtk: - pproc.create_vtk(path_fields, grids_phy, point_data) + pproc.create_vtk(path_fields, t_grid, grids_phy, point_data) if physical: - pproc.create_vtk(path_fields, grids_phy, point_data_phy, physical=True) + pproc.create_vtk(path_fields, t_grid, grids_phy, point_data_phy, physical=True) # kinetic post-processing if exist_kinetic is not None: From 11408b32b20a8a4510c636bd414a0c8c5ad07759 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 31 Jul 2025 10:48:09 +0200 Subject: [PATCH 048/292] started fixing some unit tests --- src/struphy/initial/perturbations.py | 3 + .../initial/tests/test_init_perturbations.py | 125 +++++++++--------- .../tests/test_saddle_point_propagator.py | 124 ++++++++++------- src/struphy/models/tests/verification.py | 7 +- src/struphy/models/toy.py | 18 +-- src/struphy/models/variables.py | 2 +- src/struphy/ode/tests/test_ode_feec.py | 4 +- src/struphy/propagators/propagators_fields.py | 6 + 8 files changed, 150 insertions(+), 139 deletions(-) diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index cc1999a45..706ca6494 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -31,6 +31,9 @@ class Noise(Perturbation): def __post_init__(self,): check_option(self.direction, NoiseDirections) + + def __call__(self): + pass class ModesSin(Perturbation): diff --git a/src/struphy/initial/tests/test_init_perturbations.py b/src/struphy/initial/tests/test_init_perturbations.py index 7ff965254..ee7628519 100644 --- a/src/struphy/initial/tests/test_init_perturbations.py +++ b/src/struphy/initial/tests/test_init_perturbations.py @@ -1,5 +1,5 @@ import inspect - +from copy import deepcopy import pytest @@ -27,6 +27,8 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False from struphy.feec.psydac_derham import Derham from struphy.geometry import domains from struphy.initial import perturbations + from struphy.initial.base import Perturbation + from struphy.models.variables import FEECVariable comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -87,63 +89,49 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False ): continue - # functions to compare to + # instance of perturbation if "Torus" in key: - fun = val(**kwargs, pfuns=pfuns) + perturbation = val(**kwargs, pfuns=pfuns) else: - fun = val(**kwargs, ls=ls) + perturbation = val(**kwargs, ls=ls) if isinstance(domain, domains.Cuboid) or isinstance(domain, domains.Colella): - fun_xyz = val(**kwargs, ls=ls, Lx=Lx, Ly=Ly, Lz=Lz) + perturbation_xyz = val(**kwargs, ls=ls, Lx=Lx, Ly=Ly, Lz=Lz) + assert isinstance(perturbation, Perturbation) # single component is initialized - for space, name in derham.space_to_form.items(): + for space, form in derham.space_to_form.items(): if do_plot: - plt.figure(key + "_" + name + "-form_e1e2 " + mapping[0], figsize=(24, 16)) - plt.figure(key + "_" + name + "-form_e1e3 " + mapping[0], figsize=(24, 16)) + plt.figure(key + "_" + form + "-form_e1e2 " + mapping[0], figsize=(24, 16)) + plt.figure(key + "_" + form + "-form_e1e3 " + mapping[0], figsize=(24, 16)) - if name in ("0", "3"): + if form in ("0", "3"): for n, fun_form in enumerate(form_scalar): - params = {key: {"given_in_basis": fun_form}} - - if "Modes" in key: - params[key]["ls"] = ls - params[key]["ms"] = kwargs["ms"] - params[key]["ns"] = kwargs["ns"] - params[key]["amps"] = kwargs["amps"] - if fun_form == "physical": - params[key]["Lx"] = Lx - params[key]["Ly"] = Ly - params[key]["Lz"] = Lz - else: - raise ValueError(f'Perturbation {key} not implemented, only "Modes" are testes.') + # use the setter + perturbation.given_in_basis = fun_form - if "Torus" in key: - params[key].pop("ls") - if fun_form == "physical": - continue - params[key]["pfuns"] = pfuns + var = FEECVariable(name=form, space=space) + var.add_perturbation(perturbation) + var.allocate(derham, domain) + field = var.spline - field = derham.create_spline_function(name, space, pert_params=params) - field.initialize_coeffs(domain=domain) - - field_vals_xyz = domain.push(field, e1, e2, e3, kind=name) + field_vals_xyz = domain.push(field, e1, e2, e3, kind=form) x, y, z = domain(e1, e2, e3) r = np.sqrt(x**2 + y**2) if fun_form == "physical": - fun_vals_xyz = fun_xyz(x, y, z) + fun_vals_xyz = perturbation_xyz(x, y, z) elif fun_form == "physical_at_eta": - fun_vals_xyz = fun(eee1, eee2, eee3) + fun_vals_xyz = perturbation(eee1, eee2, eee3) else: - fun_vals_xyz = domain.push(fun, eee1, eee2, eee3, kind=fun_form) + fun_vals_xyz = domain.push(perturbation, eee1, eee2, eee3, kind=fun_form) error = np.max(np.abs(field_vals_xyz - fun_vals_xyz)) / np.max(np.abs(fun_vals_xyz)) - print(f"{rank=}, {key=}, {name=}, {fun_form=}, {error=}") + print(f"{rank=}, {key=}, {form=}, {fun_form=}, {error=}") assert error < 0.02 if do_plot: - plt.figure(key + "_" + name + "-form_e1e2 " + mapping[0]) + plt.figure(key + "_" + form + "-form_e1e2 " + mapping[0]) plt.subplot(2, 4, n + 1) if isinstance(domain, domains.HollowTorus): plt.contourf(r[:, :, 0], z[:, :, 0], field_vals_xyz[:, :, 0]) @@ -172,7 +160,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False ax = plt.gca() ax.set_aspect("equal", adjustable="box") - plt.figure(key + "_" + name + "-form_e1e3 " + mapping[0]) + plt.figure(key + "_" + form + "-form_e1e3 " + mapping[0]) plt.subplot(2, 4, n + 1) if isinstance(domain, domains.HollowTorus): plt.contourf(x[:, 0, :], y[:, 0, :], field_vals_xyz[:, 0, :]) @@ -203,6 +191,10 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False else: for n, fun_form in enumerate(form_vector): + perturbation_0 = perturbation + perturbation_1 = deepcopy(perturbation) + perturbation_2 = deepcopy(perturbation) + params = { key: { "given_in_basis": [fun_form] * 3, @@ -216,24 +208,25 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False else: raise ValueError(f'Perturbation {key} not implemented, only "Modes" are testes.') - if "Torus" in key: - # params[key].pop('ls') - if fun_form == "physical": - continue - params[key]["pfuns"] = [pfuns] * 3 - else: - params[key]["ls"] = [ls] * 3 - if fun_form == "physical": - params[key]["Lx"] = Lx - params[key]["Ly"] = Ly - params[key]["Lz"] = Lz - if isinstance(domain, domains.HollowTorus): - continue - - field = derham.create_spline_function(name, space, pert_params=params) - field.initialize_coeffs(domain=domain) - - f1_xyz, f2_xyz, f3_xyz = domain.push(field, e1, e2, e3, kind=name) + if "Torus" not in key and isinstance(domain, domains.HollowTorus): + continue + + # use the setters + perturbation_0.given_in_basis = fun_form + perturbation_0.comp = 0 + perturbation_1.given_in_basis = fun_form + perturbation_1.comp = 1 + perturbation_2.given_in_basis = fun_form + perturbation_2.comp = 2 + + var = FEECVariable(name=form, space=space) + var.add_perturbation(perturbation_0) + var.add_perturbation(perturbation_1) + var.add_perturbation(perturbation_2) + var.allocate(derham, domain) + field = var.spline + + f1_xyz, f2_xyz, f3_xyz = domain.push(field, e1, e2, e3, kind=form) f_xyz = [f1_xyz, f2_xyz, f3_xyz] x, y, z = domain(e1, e2, e3) @@ -241,20 +234,20 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False # exact values if fun_form == "physical": - fun1_xyz = fun_xyz(x, y, z) - fun2_xyz = fun_xyz(x, y, z) - fun3_xyz = fun_xyz(x, y, z) + fun1_xyz = perturbation_xyz(x, y, z) + fun2_xyz = perturbation_xyz(x, y, z) + fun3_xyz = perturbation_xyz(x, y, z) elif fun_form == "physical_at_eta": - fun1_xyz = fun(eee1, eee2, eee3) - fun2_xyz = fun(eee1, eee2, eee3) - fun3_xyz = fun(eee1, eee2, eee3) + fun1_xyz = perturbation(eee1, eee2, eee3) + fun2_xyz = perturbation(eee1, eee2, eee3) + fun3_xyz = perturbation(eee1, eee2, eee3) elif fun_form == "norm": tmp1, tmp2, tmp3 = domain.transform( - [fun, fun, fun], eee1, eee2, eee3, kind=fun_form + "_to_v" + [perturbation, perturbation, perturbation], eee1, eee2, eee3, kind=fun_form + "_to_v" ) fun1_xyz, fun2_xyz, fun3_xyz = domain.push([tmp1, tmp2, tmp3], eee1, eee2, eee3, kind="v") else: - fun1_xyz, fun2_xyz, fun3_xyz = domain.push([fun, fun, fun], eee1, eee2, eee3, kind=fun_form) + fun1_xyz, fun2_xyz, fun3_xyz = domain.push([perturbation, perturbation, perturbation], eee1, eee2, eee3, kind=fun_form) fun_xyz_vec = [fun1_xyz, fun2_xyz, fun3_xyz] @@ -262,13 +255,13 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False for fi, funi in zip(f_xyz, fun_xyz_vec): error += np.max(np.abs(fi - funi)) / np.max(np.abs(funi)) error /= 3.0 - print(f"{rank=}, {key=}, {name=}, {fun_form=}, {error=}") + print(f"{rank=}, {key=}, {form=}, {fun_form=}, {error=}") assert error < 0.02 if do_plot: rn = len(form_vector) for c, (fi, f) in enumerate(zip(f_xyz, fun_xyz_vec)): - plt.figure(key + "_" + name + "-form_e1e2 " + mapping[0]) + plt.figure(key + "_" + form + "-form_e1e2 " + mapping[0]) plt.subplot(3, rn, rn * c + n + 1) if isinstance(domain, domains.HollowTorus): plt.contourf(r[:, :, 0], z[:, :, 0], fi[:, :, 0]) @@ -285,7 +278,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False ax = plt.gca() ax.set_aspect("equal", adjustable="box") - plt.figure(key + "_" + name + "-form_e1e3 " + mapping[0]) + plt.figure(key + "_" + form + "-form_e1e3 " + mapping[0]) plt.subplot(3, rn, rn * c + n + 1) if isinstance(domain, domains.HollowTorus): plt.contourf(x[:, 0, :], y[:, 0, :], fi[:, 0, :]) diff --git a/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py b/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py index 3fed692d7..3561c49cc 100644 --- a/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py +++ b/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py @@ -21,6 +21,8 @@ def test_propagator1D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): from struphy.fields_background.equils import HomogenSlab from struphy.geometry import domains from struphy.propagators.propagators_fields import TwoFluidQuasiNeutralFull + from struphy.models.variables import FEECVariable + from struphy.initial import perturbations mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -57,42 +59,57 @@ def test_propagator1D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): bas_ops = BasisProjectionOperators(derham, domain, eq_mhd=eq_mhd) # Manufactured solutions - uvec = derham.create_spline_function("u", "Hdiv") - u_evec = derham.create_spline_function("u_e", "Hdiv") - potentialvec = derham.create_spline_function("potential", "L2") - uinitial = derham.create_spline_function("u", "Hdiv") - - pp_u = { - "ManufacturedSolutionVelocity": { - "given_in_basis": ["physical", None, None], - "species": "Ions", - "comp": "0", - "dimension": "1D", - } - } - pp_ue = { - "ManufacturedSolutionVelocity": { - "given_in_basis": ["physical", None, None], - "species": "Electrons", - "comp": "0", - "dimension": "1D", - } - } - pp_potential = { - "ManufacturedSolutionPotential": { - "given_in_basis": "physical", - "dimension": "1D", - } - } - - uvec.initialize_coeffs(domain=domain, pert_params=pp_u) - u_evec.initialize_coeffs(domain=domain, pert_params=pp_ue) - potentialvec.initialize_coeffs(domain=domain, pert_params=pp_potential) + uvec = FEECVariable("u", "Hdiv") + u_evec = FEECVariable("u_e", "Hdiv") + potentialvec = FEECVariable("potential", "L2") + uinitial = FEECVariable("u", "Hdiv") + + pp_u = perturbations.ManufacturedSolutionVelocity() + pp_ue = perturbations.ManufacturedSolutionVelocity(species="Electrons") + pp_potential = perturbations.ManufacturedSolutionPotential() + + # pp_u = { + # "ManufacturedSolutionVelocity": { + # "given_in_basis": ["physical", None, None], + # "species": "Ions", + # "comp": "0", + # "dimension": "1D", + # } + # } + # pp_ue = { + # "ManufacturedSolutionVelocity": { + # "given_in_basis": ["physical", None, None], + # "species": "Electrons", + # "comp": "0", + # "dimension": "1D", + # } + # } + # pp_potential = { + # "ManufacturedSolutionPotential": { + # "given_in_basis": "physical", + # "dimension": "1D", + # } + # } + + uvec.add_perturbation(pp_u) + uvec.allocate(derham, domain, eq_mhd) + + u_evec.add_perturbation(pp_ue) + u_evec.allocate(derham, domain, eq_mhd) + + potentialvec.add_perturbation(pp_potential) + potentialvec.allocate(derham, domain, eq_mhd) + + uinitial.allocate(derham, domain, eq_mhd) + + # uvec.initialize_coeffs(domain=domain, pert_params=pp_u) + # u_evec.initialize_coeffs(domain=domain, pert_params=pp_ue) + # potentialvec.initialize_coeffs(domain=domain, pert_params=pp_potential) # Save manufactured solution to compare it later with the outcome of the propagator - uvec_initial = uvec.vector.copy() - u_evec_initial = u_evec.vector.copy() - potentialvec_initial = potentialvec.vector.copy() + uvec_initial = uvec.spline.vector.copy() + u_evec_initial = u_evec.spline.vector.copy() + potentialvec_initial = potentialvec.spline.vector.copy() solver = {} solver["type"] = ["gmres", None] @@ -109,9 +126,9 @@ def test_propagator1D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): # Starting with initial condition u=0 and ue and phi start with manufactured solution prop = TwoFluidQuasiNeutralFull( - uinitial.vector, - u_evec.vector, - potentialvec.vector, + uinitial.spline.vector, + u_evec.spline.vector, + potentialvec.spline.vector, stab_sigma=epsilon, D1_dt=dt, variant="Uzawa", @@ -250,9 +267,9 @@ def test_propagator2D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): bas_ops = BasisProjectionOperators(derham, domain, eq_mhd=eq_mhd) # Manufactured solutions - uvec = derham.create_spline_function("u", "Hdiv") - u_evec = derham.create_spline_function("u_e", "Hdiv") - potentialvec = derham.create_spline_function("potential", "L2") + uvec = FEECVariable("u", "Hdiv") + u_evec = FEECVariable("u_e", "Hdiv") + potentialvec = FEECVariable("potential", "L2") pp_u = { "ManufacturedSolutionVelocity": { @@ -386,6 +403,15 @@ def test_propagator2D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): if __name__ == "__main__": + test_propagator1D( + [16, 1, 1], + [2, 2, 1], + [True, True, True], + [[False, False], [False, False], [False, False]], + ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}], + 0.001, + 0.01, + ) # test_propagator2D( # [16, 16, 1], # [1, 1, 1], @@ -395,15 +421,15 @@ def test_propagator2D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): # 0.001, # 0.01, # ) - test_propagator2D( - [16, 16, 1], - [2, 2, 1], - [True, True, True], - [[False, False], [False, False], [False, False]], - ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}], - 0.001, - 0.01, - ) + # test_propagator2D( + # [16, 16, 1], + # [2, 2, 1], + # [True, True, True], + # [[False, False], [False, False], [False, False]], + # ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}], + # 0.001, + # 0.01, + # ) # test_propagator2D( # [32, 32, 1], # [2, 2, 1], diff --git a/src/struphy/models/tests/verification.py b/src/struphy/models/tests/verification.py index aa122c5de..f70ef7402 100644 --- a/src/struphy/models/tests/verification.py +++ b/src/struphy/models/tests/verification.py @@ -1,6 +1,7 @@ import os import pickle from pathlib import Path +from mpi4py import MPI import h5py import numpy as np @@ -238,7 +239,6 @@ def IsothermalEulerSPH_soundwave( def Maxwell_coaxial( path_out: str, - rank: int, show_plots: bool = False, ): """Verification test for coaxial cable with Maxwell equations. Comparison w.r.t analytic solution. @@ -251,12 +251,11 @@ def Maxwell_coaxial( path_out : str Simulation output folder (absolute path). - rank : int - MPI rank. - show_plots: bool Whether to show plots.""" + rank = MPI.COMM_WORLD.Get_rank() + if rank == 0: pproc_struphy.main(path_out, physical=True) MPI.COMM_WORLD.Barrier() diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 0dbad68f7..701434c80 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -41,26 +41,10 @@ def __init__(self): class EMFields(FieldSpecies): e_field: FEECVariable = FEECVariable(name="e_field", space="Hcurl") b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") - - # class EMFields(FieldSpecies): - # def __init__(self): - # self.e_field = FEECVariable(name="e_field", space="Hcurl") - # self.b_field = FEECVariable(name="b_field", space="Hdiv") - - # @dataclass - # class Ions(KineticSpecies): - # distribution_function = PICVariable(type="Particles6D") - - # @dataclass - # class Electrons(FluidSpecies): - # density = FEECVariable(space="H0") - # pressure = FEECVariable(space="H0") def __init__(self, units, domain, equil, verbose=False): # 1. instantiate all variales self.em_fields = self.EMFields() - # self._ions = self.Ions() - # self._electrons = self.Electrons() # 2. instantiate all propagators self.propagators = self.Propagators() @@ -79,7 +63,7 @@ def __init__(self, units, domain, equil, verbose=False): self.add_scalar("magnetic energy") self.add_scalar("total energy") - ## abstract attributes + ## abstract methods @property def bulk_species(self): diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index ba46081d3..04db10aae 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -101,7 +101,7 @@ def spline(self) -> SplineFunction: def add_background(self, background: FieldsBackground, verbose=False): super().add_background(background, verbose=verbose) - def allocate(self, derham: Derham, domain: Domain, equil: FluidEquilibrium,): + def allocate(self, derham: Derham, domain: Domain = None, equil: FluidEquilibrium = None,): self._spline = derham.create_spline_function( name=self.__name__, space_id=self.space, diff --git a/src/struphy/ode/tests/test_ode_feec.py b/src/struphy/ode/tests/test_ode_feec.py index 7825d2537..2df6aca95 100644 --- a/src/struphy/ode/tests/test_ode_feec.py +++ b/src/struphy/ode/tests/test_ode_feec.py @@ -1,6 +1,6 @@ import pytest -from struphy.ode.utils import ButcherTableau +from struphy.ode.utils import OptsButcher @pytest.mark.mpi(min_size=2) @@ -14,7 +14,7 @@ ("1", "0", "2"), ], ) -@pytest.mark.parametrize("algo", ButcherTableau.available_methods()) +@pytest.mark.parametrize("algo", OptsButcher) def test_exp_growth(spaces, algo, show_plots=False): """Solve dy/dt = omega*y for different feec variables y and with all available solvers from the ButcherTableau.""" diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index fa9ab2080..45b8fee11 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -7412,6 +7412,12 @@ class TwoFluidQuasiNeutralFull(Propagator): :ref:`time_discret`: fully implicit. """ + def allocate(self): + pass + + def set_options(self, **kwargs): + pass + @staticmethod def options(default=False): dct = {} From 0445f62c1d39e199fed31986466f0c749723c1b3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 31 Jul 2025 15:46:37 +0200 Subject: [PATCH 049/292] - remove StruphyModel.setup() - remove StruphyModel.pointer - start re-factoring LinearMHD --- src/struphy/io/inp/params_Maxwell_lw.py | 20 +- src/struphy/io/options.py | 2 + src/struphy/linear_algebra/schur_solver.py | 17 +- src/struphy/main.py | 48 ++- src/struphy/models/base.py | 81 +----- src/struphy/models/fluid.py | 225 +++++++------- src/struphy/models/toy.py | 20 +- src/struphy/models/variables.py | 6 +- src/struphy/propagators/propagators_fields.py | 275 +++++++++--------- 9 files changed, 330 insertions(+), 364 deletions(-) diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index 09e074c33..bd2689a41 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -7,7 +7,9 @@ # import model from struphy.models.toy import Maxwell as Model -verbose = True + +# light-weight model instance +model = Model() # units units = Units() @@ -31,21 +33,15 @@ # derham options derham = DerhamOptions() -# light-weight instance of model -model = Model(units, domain, equil, verbose=verbose) -# model.fluid.set_phys_params("mhd", options.PhysParams()) -# model.kinetic.set_phys_params("mhd", options.PhysParams()) - # propagator options model.propagators.maxwell.set_options(algo="explicit") -# initial conditions for model variables (background + perturbation) +# initial conditions (background + perturbation) model.em_fields.e_field.add_background( FieldsBackground( type="LogicalConst", values=(0.3, 0.15, None), ), - verbose=verbose, ) model.em_fields.e_field.add_perturbation( perturbations.TorusModesCos( @@ -53,7 +49,6 @@ given_in_basis="v", comp=1, ), - verbose=verbose, ) model.em_fields.b_field.add_background( @@ -61,7 +56,6 @@ type="LogicalConst", values=(0.3, 0.15, None), ), - verbose=verbose, ) model.em_fields.b_field.add_perturbation( perturbations.TorusModesCos( @@ -69,8 +63,8 @@ given_in_basis="v", comp=1, ), - verbose=verbose, ) -# exclude variable from saving -model.em_fields.e_field.save_data = False \ No newline at end of file +# exclude variables from saving +# model.em_fields.e_field.save_data = False +# model.em_fields.b_field.save_data = False \ No newline at end of file diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 70eb555a9..cb9ca39e4 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -22,6 +22,7 @@ def check_option(opt, options): # derham PolarRegularity = Literal[-1, 1] +OptsVecSpace = Literal["Hcurl", "Hdiv", "H1vec"] # fields background BackgroundTypes = Literal["LogicalConst", "FluidEquilibrium"] @@ -32,6 +33,7 @@ def check_option(opt, options): # solvers OptsSymmSolver = Literal["pcg", "cg"] +OptsGenSolver = Literal["pbicgstab", "bicgstab"] OptsMassPrecond = Literal["MassMatrixPreconditioner", None] diff --git a/src/struphy/linear_algebra/schur_solver.py b/src/struphy/linear_algebra/schur_solver.py index c29a50db8..fc154f01c 100644 --- a/src/struphy/linear_algebra/schur_solver.py +++ b/src/struphy/linear_algebra/schur_solver.py @@ -1,6 +1,7 @@ from psydac.linalg.basic import IdentityOperator, LinearOperator, Vector from psydac.linalg.block import BlockLinearOperator, BlockVector from psydac.linalg.solvers import inverse +from struphy.linear_algebra.solver import SolverParameters class SchurSolver: @@ -46,12 +47,18 @@ class SchurSolver: Must correspond to the chosen solver. """ - def __init__(self, A: LinearOperator, BC: LinearOperator, solver_name: str, **solver_params): + def __init__(self, A: LinearOperator, BC: LinearOperator, + solver_name: str, + precond = None, # TODO: add Preconditioner base class + solver_params: SolverParameters = None,): assert isinstance(A, LinearOperator) assert isinstance(BC, LinearOperator) assert A.domain == BC.domain assert A.codomain == BC.codomain + + if solver_params is None: + solver_params = SolverParameters() # linear operators self._A = A @@ -64,10 +71,12 @@ def __init__(self, A: LinearOperator, BC: LinearOperator, solver_name: str, **so # initialize solver with dummy matrix A self._solver_name = solver_name - if solver_params["pc"] is None: - solver_params.pop("pc") + kwargs = solver_params.__dict__ + kwargs.pop("info") + if precond is not None: + kwargs["pc"] = precond - self._solver = inverse(A, solver_name, **solver_params) + self._solver = inverse(A, solver_name, **kwargs) # right-hand side vector (avoids temporary memory allocation!) self._rhs = A.codomain.zeros() diff --git a/src/struphy/main.py b/src/struphy/main.py index 814046f3a..c7fa7e5be 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -77,16 +77,20 @@ def main( start_simulation = time.time() - # read in paramaters + # load paramaters and set defaults params = setup_parameters(params_path=params_path, path_out=path_out, verbose=verbose,) + # check model + model = params.model + assert hasattr(model, "propagators"), "Attribute 'self.propagators' must be set in model __init__!" + if model_name is None: - assert params.model is not None, "If model is not specified, then model: MODEL must be specified in the params!" - model_name = params.model.__class__.__name__ + assert model is not None, "If model is not specified, then model: MODEL must be specified in the params!" + model_name = model.__class__.__name__ - # collect meta-data + # meta-data meta = {} meta["platform"] = sysconfig.get_platform() meta["python version"] = sysconfig.get_python_version() @@ -99,23 +103,17 @@ def main( meta["max wall-clock [min]"] = runtime meta["save interval [steps]"] = save_step - # print meta-data on screen print("\nMETADATA:") for k, v in meta.items(): print(f'{k}:'.ljust(25), v) - # store meta-data in output folder dict_to_yaml(meta, os.path.join(path_out, "meta.yml")) - - # creating output folder, loading parameters, extract light-weight model instance + # creating output folders setup_folders(path_out=path_out, restart=restart, verbose=verbose,) - # get model instance - model = params.model - # config clones if comm is None: clone_config = None @@ -133,6 +131,33 @@ def main( clone_config.print_particle_config() comm.Barrier() + ## configure model instance + + # mpi config + model._comm_world = comm + model._clone_config = clone_config + + if model.comm_world is None: + model._rank_world = 0 + else: + model._rank_world = model.comm_world.Get_rank() + + # units + model._units = params.units + if model.bulk_species is None: + A_bulk = None + Z_bulk = None + else: + A_bulk = model.bulk_species.mass_number + Z_bulk = model.bulk_species.charge_number + model.units.derive_units(velocity_scale=model.velocity_scale, + A_bulk=A_bulk, + Z_bulk=Z_bulk, + verbose=verbose,) + + # domain and fluid bckground + model.setup_domain_and_equil(params.domain, params.equil) + # allocate derham-related objects if params.grid is not None: model.allocate_feec(params.grid, params.derham) @@ -150,6 +175,7 @@ def main( # plasma parameters model.compute_plasma_params(verbose=verbose) + model.setup_equation_params(units=model.units, verbose=verbose) if model_name is None: assert model is not None, "If model is not specified, then model: MODEL must be specified in the params!" diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 21723de60..bfd6c118b 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -83,42 +83,6 @@ def update_scalar_quantities(self): ## setup methods - def setup( - self, - units: Units = None, - domain: Domain = None, - equil: FluidEquilibrium = None, - comm: MPI.Intracomm = None, - clone_config: CloneConfig = None, - verbose=False, - ): - """Light-weight initialization of StruphyModel. Must be called in each Propagator.__init__().""" - - # defaults - if units is None: - units = Units() - - if domain is None: - domain = Cuboid() - - if equil is None: - equil = HomogenSlab() - - # mpi config - self._comm_world = comm - self._clone_config = clone_config - - if self.comm_world is None: - self._rank_world = 0 - else: - self._rank_world = self.comm_world.Get_rank() - - # other light-weight inits - self._units = units - self.units.derive_units(verbose=verbose) - self.setup_equation_params(units=self.units, verbose=verbose) - self.setup_domain_and_equil(domain, equil) - def setup_equation_params(self, units: Units, verbose=False): """Set euqation parameters for each fluid and kinetic species.""" for _, species in self.fluid_species.items(): @@ -318,14 +282,6 @@ def clone_config(self): """Config in case domain clones are used.""" return self._clone_config - @property - def pointer(self): - """Dictionary pointing to the data structures of the species (Stencil/BlockVector or "Particle" class). - - The keys are the keys from the "species" property. - In case of a fluid species, the keys are like "species_variable".""" - return self._pointer - @property def diagnostics(self): """Dictionary of diagnostics.""" @@ -360,18 +316,6 @@ def units(self): def mass_ops(self): """WeighteMassOperators object, see :ref:`mass_ops`.""" return self._mass_ops - - @property - def propagators(self): - """Return object that has model propagators in its __dict__.""" - if not hasattr (self, "_propagators"): - self._propagators = None - raise NameError("Propagators object must be named 'self._propagators = ...' in model __init__().") - return self._propagators - - @propagators.setter - def propagators(self, new): - self._propagators = new @property def prop_list(self): @@ -651,34 +595,21 @@ def allocate_variables(self): """ Allocate memory for model variables and set initial conditions. """ - - # pointer directly to data structs of Variables - self._pointer = {} - # allocate memory for FE coeffs of electromagnetic fields/potentials if self.field_species: - for _, spec in self.field_species.items(): + for species, spec in self.field_species.items(): assert isinstance(spec, FieldSpecies) for k, v in spec.variables.items(): assert isinstance(v, FEECVariable) v.allocate(derham=self.derham, domain=self.domain, equil=self.equil,) - self._pointer[k] = v.spline.vector # allocate memory for FE coeffs of fluid variables if self.fluid_species: - for species, dct in self.fluid.items(): - for variable, subdct in dct.items(): - if "params" in variable: - continue - else: - subdct["obj"] = self.derham.create_spline_function( - variable, - subdct["space"], - bckgr_params=subdct.get("background"), - pert_params=subdct.get("perturbation"), - ) - - self._pointer[species + "_" + variable] = subdct["obj"].vector + for species, spec in self.fluid_species.items(): + assert isinstance(spec, FluidSpecies) + for k, v in spec.variables.items(): + assert isinstance(v, FEECVariable) + v.allocate(derham=self.derham, domain=self.domain, equil=self.equil,) # marker arrays and plasma parameters of kinetic species if self.kinetic_species: diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 765ef7cc0..e5abfaafa 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -1,7 +1,10 @@ import numpy as np +from dataclasses import dataclass from struphy.models.base import StruphyModel from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers +from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies +from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable class LinearMHD(StruphyModel): @@ -35,88 +38,42 @@ class LinearMHD(StruphyModel): :ref:`Model info `: """ - - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["em_fields"]["b_field"] = "Hdiv" - dct["fluid"]["mhd"] = {"density": "L2", "velocity": "Hdiv", "pressure": "L2"} - return dct - - @staticmethod - def bulk_species(): - return "mhd" - - @staticmethod - def velocity_scale(): - return "alfvén" - - @staticmethod - def propagators_dct(): - return { - propagators_fields.ShearAlfven: ["mhd_velocity", "b_field"], - propagators_fields.Magnetosonic: ["mhd_density", "mhd_velocity", "mhd_pressure"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - # add special options - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["fluid", "mhd"], - key="u_space", - option="Hdiv", - dct=dct, - ) - return dct - - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - from struphy.polar.basic import PolarVector - - # extract necessary parameters - u_space = params["fluid"]["mhd"]["options"]["u_space"] - alfven_solver = params["fluid"]["mhd"]["options"]["ShearAlfven"]["solver"] - alfven_algo = params["fluid"]["mhd"]["options"]["ShearAlfven"]["algo"] - sonic_solver = params["fluid"]["mhd"]["options"]["Magnetosonic"]["solver"] - - # project background magnetic field (2-form) and pressure (3-form) - self._b_eq = self.projected_equil.b2 - self._p_eq = self.projected_equil.p3 - self._ones = self._p_eq.space.zeros() - - if isinstance(self._ones, PolarVector): - self._ones.tp[:] = 1.0 - else: - self._ones[:] = 1.0 - - # set keyword arguments for propagators - self._kwargs[propagators_fields.ShearAlfven] = { - "u_space": u_space, - "solver": alfven_solver, - "algo": alfven_algo, - } - - self._kwargs[propagators_fields.Magnetosonic] = { - "b": self.pointer["b_field"], - "u_space": u_space, - "solver": sonic_solver, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation + @dataclass + class EMFields(FieldSpecies): + b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") + + @dataclass + class MHD(FluidSpecies): + density: FEECVariable = FEECVariable(name="density", space="L2") + velocity: FEECVariable = FEECVariable(name="velocity", space="Hdiv") + pressure: FEECVariable = FEECVariable(name="pressure", space="L2") + + class Propagators: + def __init__(self): + self.shear_alf = propagators_fields.ShearAlfven() + self.mag_sonic = propagators_fields.Magnetosonic() + + def __init__(self): + # 1. instantiate all variales + self.em_fields = self.EMFields() + self.mhd = self.MHD() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.shear_alf.set_variables( + u = self.mhd.velocity, + b = self.em_fields.b_field, + ) + + self.propagators.mag_sonic.set_variables( + n = self.mhd.density, + u = self.mhd.velocity, + p = self.mhd.pressure, + ) + + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_p") self.add_scalar("en_B") @@ -124,40 +81,98 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("en_B_eq") self.add_scalar("en_B_tot") self.add_scalar("en_tot") + + ## abstract methods - # vectors for computing scalar quantities - self._tmp_b1 = self.derham.Vh["2"].zeros() - self._tmp_b2 = self.derham.Vh["2"].zeros() + @property + def bulk_species(self): + return self.mhd + + @property + def velocity_scale(self): + return "alfvén" + + # def __init__old(self, params, comm, clone_config=None): + # # initialize base class + # super().__init__(params, comm=comm, clone_config=clone_config) + + # from struphy.polar.basic import PolarVector + + # # extract necessary parameters + # u_space = params["fluid"]["mhd"]["options"]["u_space"] + # alfven_solver = params["fluid"]["mhd"]["options"]["ShearAlfven"]["solver"] + # alfven_algo = params["fluid"]["mhd"]["options"]["ShearAlfven"]["algo"] + # sonic_solver = params["fluid"]["mhd"]["options"]["Magnetosonic"]["solver"] + + # # project background magnetic field (2-form) and pressure (3-form) + # self._b_eq = self.projected_equil.b2 + # self._p_eq = self.projected_equil.p3 + # self._ones = self._p_eq.space.zeros() + + # if isinstance(self._ones, PolarVector): + # self._ones.tp[:] = 1.0 + # else: + # self._ones[:] = 1.0 + + # # set keyword arguments for propagators + # self._kwargs[propagators_fields.ShearAlfven] = { + # "u_space": u_space, + # "solver": alfven_solver, + # "algo": alfven_algo, + # } + + # self._kwargs[propagators_fields.Magnetosonic] = { + # "b": self.pointer["b_field"], + # "u_space": u_space, + # "solver": sonic_solver, + # } + + # # Initialize propagators used in splitting substeps + # self.init_propagators() + + # # Scalar variables to be saved during simulation + # self.add_scalar("en_U") + # self.add_scalar("en_p") + # self.add_scalar("en_B") + # self.add_scalar("en_p_eq") + # self.add_scalar("en_B_eq") + # self.add_scalar("en_B_tot") + # self.add_scalar("en_tot") + + # # vectors for computing scalar quantities + # self._tmp_b1 = self.derham.Vh["2"].zeros() + # self._tmp_b2 = self.derham.Vh["2"].zeros() def update_scalar_quantities(self): - # perturbed fields - en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.pointer["mhd_velocity"], self.pointer["mhd_velocity"]) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) - en_p = self.pointer["mhd_pressure"].inner(self._ones) / (5 / 3 - 1) + pass + # # perturbed fields + # en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.pointer["mhd_velocity"], self.pointer["mhd_velocity"]) + # en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) + # en_p = self.pointer["mhd_pressure"].inner(self._ones) / (5 / 3 - 1) - self.update_scalar("en_U", en_U) - self.update_scalar("en_B", en_B) - self.update_scalar("en_p", en_p) - self.update_scalar("en_tot", en_U + en_B + en_p) + # self.update_scalar("en_U", en_U) + # self.update_scalar("en_B", en_B) + # self.update_scalar("en_p", en_p) + # self.update_scalar("en_tot", en_U + en_B + en_p) - # background fields - self.mass_ops.M2.dot(self._b_eq, apply_bc=False, out=self._tmp_b1) + # # background fields + # self.mass_ops.M2.dot(self._b_eq, apply_bc=False, out=self._tmp_b1) - en_B0 = self._b_eq.inner(self._tmp_b1) / 2 - en_p0 = self._p_eq.inner(self._ones) / (5 / 3 - 1) + # en_B0 = self._b_eq.inner(self._tmp_b1) / 2 + # en_p0 = self._p_eq.inner(self._ones) / (5 / 3 - 1) - self.update_scalar("en_B_eq", en_B0) - self.update_scalar("en_p_eq", en_p0) + # self.update_scalar("en_B_eq", en_B0) + # self.update_scalar("en_p_eq", en_p0) - # total magnetic field - self._b_eq.copy(out=self._tmp_b1) - self._tmp_b1 += self.pointer["b_field"] + # # total magnetic field + # self._b_eq.copy(out=self._tmp_b1) + # self._tmp_b1 += self.pointer["b_field"] - self.mass_ops.M2.dot(self._tmp_b1, apply_bc=False, out=self._tmp_b2) + # self.mass_ops.M2.dot(self._tmp_b1, apply_bc=False, out=self._tmp_b2) - en_Btot = self._tmp_b1.inner(self._tmp_b2) / 2 + # en_Btot = self._tmp_b1.inner(self._tmp_b2) / 2 - self.update_scalar("en_B_tot", en_Btot) + # self.update_scalar("en_B_tot", en_Btot) class LinearExtendedMHDuniform(StruphyModel): diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 701434c80..ca0b85bef 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -32,17 +32,16 @@ class Maxwell(StruphyModel): :ref:`Model info `: """ - - class Propagators: - def __init__(self): - self.maxwell = propagators_fields.Maxwell() - @dataclass class EMFields(FieldSpecies): e_field: FEECVariable = FEECVariable(name="e_field", space="Hcurl") b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") + + class Propagators: + def __init__(self): + self.maxwell = propagators_fields.Maxwell() - def __init__(self, units, domain, equil, verbose=False): + def __init__(self): # 1. instantiate all variales self.em_fields = self.EMFields() @@ -55,16 +54,13 @@ def __init__(self, units, domain, equil, verbose=False): b = self.em_fields.b_field, ) - # light-weight setup of model - self.setup(units=units, domain=domain, equil=equil, verbose=verbose) - # define scalars for update_scalar_quantities self.add_scalar("electric energy") self.add_scalar("magnetic energy") self.add_scalar("total energy") ## abstract methods - + @property def bulk_species(self): return None @@ -74,8 +70,8 @@ def velocity_scale(self): return "light" def update_scalar_quantities(self): - en_E = 0.5 * self.mass_ops.M1.dot_inner(self.pointer["e_field"], self.pointer["e_field"]) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) + en_E = 0.5 * self.mass_ops.M1.dot_inner(self.em_fields.e_field.spline.vector, self.em_fields.e_field.spline.vector) + en_B = 0.5 * self.mass_ops.M2.dot_inner(self.em_fields.b_field.spline.vector, self.em_fields.b_field.spline.vector) self.update_scalar("electric energy", en_E) self.update_scalar("magnetic energy", en_B) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 04db10aae..ddf546038 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -46,7 +46,7 @@ def species(self): self._species = None return self._species - def add_background(self, background, verbose=False): + def add_background(self, background, verbose=True): """Type inference of added background done in sub class.""" if not hasattr(self, "_backgrounds") or self.backgrounds is None: self._backgrounds = background @@ -60,7 +60,7 @@ def add_background(self, background, verbose=False): for k, v in background.__dict__.items(): print(f' {k}: {v}') - def add_perturbation(self, perturbation: Perturbation, verbose=False): + def add_perturbation(self, perturbation: Perturbation, verbose=True): if not hasattr(self, "_perturbations") or self.perturbations is None: self._perturbations = perturbation else: @@ -98,7 +98,7 @@ def space(self): def spline(self) -> SplineFunction: return self._spline - def add_background(self, background: FieldsBackground, verbose=False): + def add_background(self, background: FieldsBackground, verbose=True): super().add_background(background, verbose=verbose) def allocate(self, derham: Derham, domain: Domain = None, equil: FluidEquilibrium = None,): diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 45b8fee11..3fcc50297 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -46,7 +46,7 @@ from struphy.propagators.base import Propagator from struphy.models.variables import Variable from struphy.linear_algebra.solver import SolverParameters -from struphy.io.options import check_option, OptsSymmSolver, OptsMassPrecond +from struphy.io.options import (check_option, OptsSymmSolver, OptsMassPrecond, OptsGenSolver, OptsVecSpace) from struphy.models.variables import FEECVariable, PICVariable, SPHVariable @@ -70,7 +70,7 @@ class Maxwell(Propagator): # propagator specific options OptsAlgo = Literal["implicit", "explicit"] - # abstract methods + ## abstract methods def set_options(self, algo: OptsAlgo = "implicit", solver: OptsSymmSolver = "pcg", @@ -128,10 +128,8 @@ def allocate(self): _A, _BC, self.options.solver, - pc=pc, - tol=self.options.solver_params.tol, - maxiter=self.options.solver_params.maxiter, - verbose=self.options.solver_params.verbose, + precond=pc, + solver_params=self.options.solver_params, ) # pre-allocate arrays @@ -191,13 +189,12 @@ def __call__(self, dt): bn1 *= -dt bn1 += bn - # write new coeffs into self.feec_vars diffs = self.update_feec_variables(e=en1, b=bn1) else: self._ode_solver(0.0, dt) if self._info and MPI.COMM_WORLD.Get_rank() == 0: - if self._algo == "implicit": + if self.options.algo == "implicit": print("Status for Maxwell:", info["success"]) print("Iterations for Maxwell:", info["niter"]) print("Maxdiff e for Maxwell:", diffs["e"]) @@ -422,6 +419,7 @@ def __call__(self, dt): print() +@dataclass class ShearAlfven(Propagator): r""":ref:`FEEC ` discretization of the following equations: find :math:`\mathbf U \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\}` and :math:`\mathbf B \in H(\textnormal{div})` such that @@ -445,44 +443,48 @@ class ShearAlfven(Propagator): where :math:`\alpha \in \{1, 2, v\}` and :math:`\mathbb M^\rho_\alpha` is a weighted mass matrix in :math:`\alpha`-space, the weight being :math:`\rho_0`, the MHD equilibirum density. The solution of the above system is based on the :ref:`Schur complement `. """ + # variables to be updated + u: FEECVariable = None + b: FEECVariable = None + + # propagator specific options + OptsAlgo = Literal["implicit", "explicit"] + + ## abstract methods + def set_options(self, + u_space: OptsVecSpace = "Hdiv", + algo: OptsAlgo = "implicit", + solver: OptsSymmSolver = "pcg", + precond: OptsMassPrecond = "MassMatrixPreconditioner", + solver_params: SolverParameters = None, + butcher: ButcherTableau = None, + ): + + # checks + check_option(u_space, OptsVecSpace) + check_option(algo, self.OptsAlgo) + check_option(solver, OptsSymmSolver) + check_option(precond, OptsMassPrecond) + + # defaults + if solver_params is None: + solver_params = SolverParameters() + + if algo == "explicit" and butcher is None: + butcher = ButcherTableau() + + # use setter for options + self.options = self.Options(self, + u_space=u_space, + algo=algo, + solver=solver, + precond=precond, + solver_params=solver_params, + butcher=butcher,) - @staticmethod - def options(default=False): - dct = {} - dct["algo"] = ["implicit", "rk4", "forward_euler", "heun2", "rk2", "heun3"] - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["turn_off"] = False - - if default: - dct = descend_options_dict(dct, []) - - return dct - - def __init__( - self, - u: BlockVector, - b: BlockVector, - *, - u_space: str, - algo: dict = options(default=True)["algo"], - solver: dict = options(default=True)["solver"], - ): - super().__init__(u, b) - - assert u_space in {"Hcurl", "Hdiv", "H1vec"} - - self._algo = algo - + def allocate(self): + u_space = self.options.u_space + # define block matrix [[A B], [C I]] (without time step size dt in the diagonals) id_M = "M" + self.derham.space_to_form[u_space] + "n" id_T = "T" + self.derham.space_to_form[u_space] @@ -494,14 +496,14 @@ def __init__( curl = self.derham.curl # Preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(getattr(self.mass_ops, id_M)) - if self._algo == "implicit": - self._info = solver["info"] + if self.options.algo == "implicit": + self._info = self.options.solver_params.info self._B = -1 / 2 * _T.T @ curl.T @ _M2 self._C = 1 / 2 * curl @ _T @@ -512,12 +514,9 @@ def __init__( self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, ) # pre-allocate arrays @@ -529,44 +528,44 @@ def __init__( # define vector field A_inv = inverse( _A, - solver["type"][0], + self.options.solver, pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, ) _f1 = A_inv @ _T.T @ curl.T @ _M2 _f2 = curl @ _T # allocate output of vector field - out1 = u.space.zeros() - out2 = b.space.zeros() + out1 = self.u.space.zeros() + out2 = self.b.space.zeros() - def f1(t, y1, y2, out=out1): + def f1(t, y1, y2, out: BlockVector = out1): _f1.dot(y2, out=out) out.update_ghost_regions() return out - def f2(t, y1, y2, out=out2): + def f2(t, y1, y2, out: BlockVector = out2): _f2.dot(y1, out=out) out *= -1.0 out.update_ghost_regions() return out - vector_field = {u: f1, b: f2} - self._ode_solver = ODEsolverFEEC(vector_field, algo=algo) + vector_field = {self.u.spline.vector: f1, self.b.spline.vector: f2} + self._ode_solver = ODEsolverFEEC(vector_field, butcher=self.options.butcher) # allocate dummy vectors to avoid temporary array allocations - self._u_tmp1 = u.space.zeros() - self._u_tmp2 = u.space.zeros() - self._b_tmp1 = b.space.zeros() + self._u_tmp1 = self.u.spline.vector.space.zeros() + self._u_tmp2 = self.u.spline.vector.space.zeros() + self._b_tmp1 = self.b.spline.vector.space.zeros() def __call__(self, dt): - # current variables - un = self.feec_vars[0] - bn = self.feec_vars[1] + # current FE coeffs + un = self.u.spline.vector + bn = self.b.spline.vector - if self._algo == "implicit": + if self.options.algo == "implicit": # solve for new u coeffs byn = self._B.dot(bn, out=self._byn) @@ -579,18 +578,16 @@ def __call__(self, dt): bn1 *= -dt bn1 += bn - # write new coeffs into self.feec_vars - max_du, max_db = self.feec_vars_update(un1, bn1) - + diffs = self.update_feec_variables(u=un1, b=bn1) else: self._ode_solver(0.0, dt) if self._info and MPI.COMM_WORLD.Get_rank() == 0: - if self._algo == "implicit": + if self.options.algo == "implicit": print("Status for ShearAlfven:", info["success"]) print("Iterations for ShearAlfven:", info["niter"]) - print("Maxdiff up for ShearAlfven:", max_du) - print("Maxdiff b2 for ShearAlfven:", max_db) + print("Maxdiff up for ShearAlfven:", diffs["u"]) + print("Maxdiff b2 for ShearAlfven:", diffs["b"]) print() @@ -843,6 +840,7 @@ def __call__(self, dt): print() +@dataclass class Magnetosonic(Propagator): r""" :ref:`FEEC ` discretization of the following equations: @@ -876,43 +874,45 @@ class Magnetosonic(Propagator): \boldsymbol{\rho}^{n+1} = \boldsymbol{\rho}^n - \frac{\Delta t}{2} \mathbb D \mathcal Q^\alpha (\mathbf u^{n+1} + \mathbf u^n) \,. """ + # variables to be updated + n: FEECVariable = None + u: FEECVariable = None + p: FEECVariable = None + + ## abstract methods + def set_options(self, + b_field: FEECVariable = None, + u_space: OptsVecSpace = "Hdiv", + solver: OptsGenSolver = "pbicgstab", + precond: OptsMassPrecond = "MassMatrixPreconditioner", + solver_params: SolverParameters = None, + ): + + # checks + check_option(u_space, OptsVecSpace) + check_option(solver, OptsGenSolver) + check_option(precond, OptsMassPrecond) + + # defaults + if b_field is None: + b_field = FEECVariable(space="Hdiv") + b_field.allocate(self.derham, self.domain) + if solver_params is None: + solver_params = SolverParameters() + + # use setter for options + self.options = self.Options(self, + u_space=u_space, + b_field=b_field, + solver=solver, + precond=precond, + solver_params=solver_params, + ) - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pbicgstab", "MassMatrixPreconditioner"), - ("bicgstab", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["turn_off"] = False - - if default: - dct = descend_options_dict(dct, []) - - return dct - - def __init__( - self, - n: StencilVector, - u: BlockVector, - p: StencilVector, - *, - u_space: str, - b: BlockVector, - solver: dict = options(default=True)["solver"], - ): - super().__init__(n, u, p) - - assert u_space in {"Hcurl", "Hdiv", "H1vec"} + def allocate(self): + u_space = self.options.u_space - self._info = solver["info"] + self._info = self.options.solver_params.info self._bc = self.derham.dirichlet_bc # define block matrix [[A B], [C I]] (without time step size dt in the diagonals) @@ -931,7 +931,8 @@ def __init__( _K = getattr(self.basis_ops, id_K) if id_U is None: - _U, _UT = IdentityOperator(u.space), IdentityOperator(u.space) + _U = IdentityOperator(self.u.spline.vector.space) + _UT = IdentityOperator(self.u.spline.vector.space) else: _U = getattr(self.basis_ops, id_U) _UT = _U.T @@ -942,13 +943,13 @@ def __init__( self._MJ = getattr(self.mass_ops, id_MJ) self._DQ = self.derham.div @ getattr(self.basis_ops, id_Q) - self._b = b + self._b = self.options.b_field.spline.vector # preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(getattr(self.mass_ops, id_Mn)) # instantiate Schur solver (constant in this case) @@ -957,29 +958,26 @@ def __init__( self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, ) # allocate dummy vectors to avoid temporary array allocations - self._u_tmp1 = u.space.zeros() - self._u_tmp2 = u.space.zeros() - self._p_tmp1 = p.space.zeros() - self._n_tmp1 = n.space.zeros() + self._u_tmp1 = self.u.spline.vector.space.zeros() + self._u_tmp2 = self.u.spline.vector.space.zeros() + self._p_tmp1 = self.p.spline.vector.space.zeros() + self._n_tmp1 = self.n.spline.vector.space.zeros() self._b_tmp1 = self._b.space.zeros() self._byn1 = self._B.codomain.zeros() self._byn2 = self._B.codomain.zeros() def __call__(self, dt): - # current variables - nn = self.feec_vars[0] - un = self.feec_vars[1] - pn = self.feec_vars[2] + # current FE coeffs + nn = self.n.spline.vector + un = self.u.spline.vector + pn = self.p.spline.vector # solve for new u coeffs (no tmps created here) byn1 = self._B.dot(pn, out=self._byn1) @@ -1000,19 +998,14 @@ def __call__(self, dt): nn1 *= -dt / 2 nn1 += nn - # write new coeffs into self.feec_vars - max_dn, max_du, max_dp = self.feec_vars_update( - nn1, - un1, - pn1, - ) - + diffs = self.update_feec_variables(n=nn1, u=un1, p=pn1) + if self._info and MPI.COMM_WORLD.Get_rank() == 0: print("Status for Magnetosonic:", info["success"]) print("Iterations for Magnetosonic:", info["niter"]) - print("Maxdiff n3 for Magnetosonic:", max_dn) - print("Maxdiff up for Magnetosonic:", max_du) - print("Maxdiff p3 for Magnetosonic:", max_dp) + print("Maxdiff n3 for Magnetosonic:", diffs["n"]) + print("Maxdiff up for Magnetosonic:", diffs["u"]) + print("Maxdiff p3 for Magnetosonic:", diffs["p"]) print() From 6fdae08ca9f9f9b1d171c956bdedec9297611d6f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 31 Jul 2025 17:07:29 +0200 Subject: [PATCH 050/292] LinearMHD model is running --- src/struphy/io/parameters.py | 2 +- src/struphy/main.py | 4 +- src/struphy/models/base.py | 524 ++++++++++++---------------------- src/struphy/models/species.py | 6 +- 4 files changed, 196 insertions(+), 340 deletions(-) diff --git a/src/struphy/io/parameters.py b/src/struphy/io/parameters.py index 0a4741dc2..c3be1f594 100644 --- a/src/struphy/io/parameters.py +++ b/src/struphy/io/parameters.py @@ -39,7 +39,7 @@ def __init__( print(f"{self.equil = }") print(f"{self.time = }") print(f"{self.grid = }") - print(f"{self.derham = }\n") + print(f"{self.derham = }") @property def model(self): diff --git a/src/struphy/main.py b/src/struphy/main.py index c7fa7e5be..305ec2ab3 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -18,7 +18,7 @@ from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml from struphy.pic.base import Particles -from struphy.models.species import FieldSpecies +from struphy.models.species import Species from struphy.models.variables import FEECVariable @@ -321,7 +321,7 @@ def main( # extract FEEC coefficients feec_species = model.field_species | model.fluid_species | model.diagnostic_species for species, val in feec_species.items(): - assert isinstance(val, FieldSpecies) + assert isinstance(val, Species) for variable, subval in val.variables.items(): assert isinstance(subval, FEECVariable) spline = subval.spline diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index bfd6c118b..4f5f59a01 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -238,11 +238,12 @@ def allocate_propagators(self): ) Propagator.projected_equil = self.projected_equil + assert len(self.prop_list) > 0, "No propagators in this model, check the model class." for prop in self.prop_list: assert isinstance(prop, Propagator) prop.allocate() - if self.verbose and self.rank_world == 0: - print(f"\nAllocated propagator {prop.__class__.__name__}.") + if self.rank_world == 0: + print(f"\nAllocated propagator '{prop.__class__.__name__}'.") @staticmethod def diagnostics_dct(): @@ -308,7 +309,7 @@ def projected_equil(self): return self._projected_equil @property - def units(self): + def units(self) -> Units: """All Struphy units.""" return self._units @@ -1140,8 +1141,7 @@ def initialize_data_output(self, data, size): # save feec data in group 'feec/' feec_species = self.field_species | self.fluid_species | self.diagnostic_species for species, val in feec_species.items(): - - assert isinstance(val, FieldSpecies) + assert isinstance(val, Species) species_path = os.path.join("feec", species) species_path_restart = os.path.join("restart", species) @@ -1245,131 +1245,6 @@ def initialize_data_output(self, data, size): # Class methods : ################### - def model_units( - self, - units: Units, - verbose: bool = False, - comm: MPI.Intracomm = None, - ): - """ - Return model units and print them to screen. - - Returns - ------- - units_basic : dict - Basic units for time, length, mass and magnetic field. - - units_der : dict - Derived units for velocity, pressure, mass density and particle density. - """ - - print(f'{self.species.em_fields.all = }') - - # look for bulk species in fluid OR kinetic parameter dictionaries - Z_bulk = None - A_bulk = None - if self.bulk_species is not None: - Z_bulk = self.bulk_species.Z - A_bulk = self.bulk_species.A - - # compute model units - kBT = units.kBT - - units = derive_units( - Z_bulk=Z_bulk, - A_bulk=A_bulk, - x=units.x, - B=units.B, - n=units.n, - kBT=kBT, - velocity_scale=self.velocity_scale, - ) - - # print to screen - if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print("\nUNITS:") - print( - f"Unit of length:".ljust(25), - "{:4.3e}".format(units["x"]) + " m", - ) - print( - f"Unit of time:".ljust(25), - "{:4.3e}".format(units["t"]) + " s", - ) - print( - f"Unit of velocity:".ljust(25), - "{:4.3e}".format(units["v"]) + " m/s", - ) - print( - f"Unit of magnetic field:".ljust(25), - "{:4.3e}".format(units["B"]) + " T", - ) - - if A_bulk is not None: - print( - f"Unit of particle density:".ljust(25), - "{:4.3e}".format(units["n"]) + " m⁻³", - ) - print( - f"Unit of mass density:".ljust(25), - "{:4.3e}".format(units["rho"]) + " kg/m³", - ) - print( - f"Unit of pressure:".ljust(25), - "{:4.3e}".format(units["p"] * 1e-5) + " bar", - ) - print( - f"Unit of current density:".ljust(25), - "{:4.3e}".format(units["j"]) + " A/m²", - ) - - # compute equation parameters for each species - e = 1.602176634e-19 # elementary charge (C) - mH = 1.67262192369e-27 # proton mass (kg) - eps0 = 8.8541878128e-12 # vacuum permittivity (F/m) - - equation_params = {} - if self.species.fluid is not None: - for name, species in self.species.fluid.all.items(): - Z = species.Z - A = species.A - - # compute equation parameters - om_p = np.sqrt(units.n * (Z * e) ** 2 / (eps0 * A * mH)) - om_c = Z * e * units.B / (A * mH) - equation_params[name] = {} - equation_params[name]["alpha"] = om_p / om_c - equation_params[name]["epsilon"] = 1.0 / (om_c * units["t"]) - equation_params[name]["kappa"] = om_p * units["t"] - - if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print("\nNORMALIZATION PARAMETERS:") - print("- " + name + ":") - for key, val in equation_params[name].items(): - print((key + ":").ljust(25), "{:4.3e}".format(val)) - - if self.species.kinetic is not None: - for name, species in self.species.kinetic.all.items(): - Z = species.Z - A = species.A - - # compute equation parameters - om_p = np.sqrt(units.n * (Z * e) ** 2 / (eps0 * A * mH)) - om_c = Z * e * units.B / (A * mH) - equation_params[name] = {} - equation_params[name]["alpha"] = om_p / om_c - equation_params[name]["epsilon"] = 1.0 / (om_c * units["t"]) - equation_params[name]["kappa"] = om_p * units["t"] - - if verbose and MPI.COMM_WORLD.Get_rank() == 0: - if self.species.fluid is None: - print("\nNORMALIZATION PARAMETERS:") - print("- " + name + ":") - for key, val in equation_params[name].items(): - print((key + ":").ljust(25), "{:4.3e}".format(val)) - - return units, equation_params - @classmethod def show_options(cls): """Print available model options to screen.""" @@ -1788,26 +1663,6 @@ def compute_plasma_params(self, verbose=True): - epsilon = 1/(t*Omega_c) """ - #TODO: needs re-factoring - create PlasmaParameters class instead of self._pparams dict - - # physics constants - e = 1.602176634e-19 # elementary charge (C) - m_p = 1.67262192369e-27 # proton mass (kg) - mu0 = 1.25663706212e-6 # magnetic constant (N*A^-2) - eps0 = 8.8541878128e-12 # vacuum permittivity (F*m^-1) - kB = 1.380649e-23 # Boltzmann constant (J*K^-1) - - # exit when there is not any plasma species - if len(self.fluid_species) == 0 and len(self.kinetic_species) == 0: - return - - # compute model units - units, equation_params = self.model_units( - self.params, - verbose=False, - comm=self.comm_world, - ) - # units affices for printing units_affix = {} units_affix["plasma volume"] = " m³" @@ -1838,12 +1693,13 @@ def compute_plasma_params(self, verbose=True): eta2 = np.linspace(h / 2.0, 1.0 - h / 2.0, 20) eta3 = np.linspace(h / 2.0, 1.0 - h / 2.0, 20) - # global parameters + ## global parameters + # plasma volume (hat x^3) det_tmp = self.domain.jacobian_det(eta1, eta2, eta3) vol1 = np.mean(np.abs(det_tmp)) # plasma volume (m⁻³) - plasma_volume = vol1 * units["x"] ** 3 + plasma_volume = vol1 * self.units.x ** 3 # transit length (m) transit_length = plasma_volume ** (1 / 3) # magnetic field (T) @@ -1851,9 +1707,9 @@ def compute_plasma_params(self, verbose=True): B_tmp = self.equil.absB0(eta1, eta2, eta3) else: B_tmp = np.zeros((eta1.size, eta2.size, eta3.size)) - magnetic_field = np.mean(B_tmp * np.abs(det_tmp)) / vol1 * units["B"] - B_max = np.max(B_tmp) * units["B"] - B_min = np.min(B_tmp) * units["B"] + magnetic_field = np.mean(B_tmp * np.abs(det_tmp)) / vol1 * self.units.B + B_max = np.max(B_tmp) * self.units.B + B_min = np.min(B_tmp) * self.units.B if magnetic_field < 1e-14: magnetic_field = np.nan @@ -1882,186 +1738,186 @@ def compute_plasma_params(self, verbose=True): "{:4.3e}".format(B_min) + units_affix["magnetic field"], ) - # species dependent parameters - self._pparams = {} - - if len(self.fluid) > 0: - for species, val in self.fluid.items(): - self._pparams[species] = {} - # type - self._pparams[species]["type"] = "fluid" - # mass (kg) - self._pparams[species]["mass"] = val["params"]["phys_params"]["A"] * m_p - # charge (C) - self._pparams[species]["charge"] = val["params"]["phys_params"]["Z"] * e - # density (m⁻³) - self._pparams[species]["density"] = ( - np.mean( - self.equil.n0( - eta1, - eta2, - eta3, - ) - * np.abs(det_tmp), - ) - * units["x"] ** 3 - / plasma_volume - * units["n"] - ) - # pressure (bar) - self._pparams[species]["pressure"] = ( - np.mean( - self.equil.p0( - eta1, - eta2, - eta3, - ) - * np.abs(det_tmp), - ) - * units["x"] ** 3 - / plasma_volume - * units["p"] - * 1e-5 - ) - # thermal energy (keV) - self._pparams[species]["kBT"] = self._pparams[species]["pressure"] * 1e5 / self._pparams[species]["density"] / e * 1e-3 - - if len(self.kinetic) > 0: - eta1mg, eta2mg, eta3mg = np.meshgrid( - eta1, - eta2, - eta3, - indexing="ij", - ) - - for species, val in self.kinetic.items(): - self._pparams[species] = {} - # type - self._pparams[species]["type"] = "kinetic" - # mass (kg) - self._pparams[species]["mass"] = val["params"]["phys_params"]["A"] * m_p - # charge (C) - self._pparams[species]["charge"] = val["params"]["phys_params"]["Z"] * e - - # create temp kinetic object for (default) parameter extraction - tmp_bckgr = val["params"]["background"] - - if val["space"] != "ParticlesSPH": - tmp = None - for fi, maxw_params in tmp_bckgr.items(): - if fi[-2] == "_": - fi_type = fi[:-2] - else: - fi_type = fi - - if tmp is None: - tmp = getattr(maxwellians, fi_type)( - maxw_params=maxw_params, - equil=self.equil, - ) - else: - tmp = tmp + getattr(maxwellians, fi_type)( - maxw_params=maxw_params, - equil=self.equil, - ) - - if val["space"] != "ParticlesSPH" and tmp.coords == "constants_of_motion": - # call parameters - a1 = self.domain.params_map["a1"] - r = eta1mg * (1 - a1) + a1 - psi = self.equil.psi_r(r) - - # density (m⁻³) - self._pparams[species]["density"] = ( - np.mean(tmp.n(psi) * np.abs(det_tmp)) * units["x"] ** 3 / plasma_volume * units["n"] - ) - # thermal speed (m/s) - self._pparams[species]["v_th"] = ( - np.mean(tmp.vth(psi) * np.abs(det_tmp)) * units["x"] ** 3 / plasma_volume * units["v"] - ) - # thermal energy (keV) - self._pparams[species]["kBT"] = self._pparams[species]["mass"] * self._pparams[species]["v_th"] ** 2 / e * 1e-3 - # pressure (bar) - self._pparams[species]["pressure"] = ( - self._pparams[species]["kBT"] * e * 1e3 * self._pparams[species]["density"] * 1e-5 - ) - - else: - # density (m⁻³) - # self._pparams[species]['density'] = np.mean(tmp.n( - # eta1mg, eta2mg, eta3mg) * np.abs(det_tmp)) * units['x']**3 / plasma_volume * units['n'] - self._pparams[species]["density"] = 99.0 - # thermal speeds (m/s) - vth = [] - # vths = tmp.vth(eta1mg, eta2mg, eta3mg) - vths = [99.0] - for k in range(len(vths)): - vth += [ - vths[k] * np.abs(det_tmp) * units["x"] ** 3 / plasma_volume * units["v"], - ] - thermal_speed = 0.0 - for dir in range(val["obj"].vdim): - # self._pparams[species]['vth' + str(dir + 1)] = np.mean(vth[dir]) - self._pparams[species]["vth" + str(dir + 1)] = 99.0 - thermal_speed += self._pparams[species]["vth" + str(dir + 1)] - # TODO: here it is assumed that background density parameter is called "n", - # and that background thermal speeds are called "vthn"; make this a convention? - # self._pparams[species]['v_th'] = thermal_speed / \ - # val['obj'].vdim - self._pparams[species]["v_th"] = 99.0 - # thermal energy (keV) - # self._pparams[species]['kBT'] = self._pparams[species]['mass'] * \ - # self._pparams[species]['v_th']**2 / e * 1e-3 - self._pparams[species]["kBT"] = 99.0 - # pressure (bar) - # self._pparams[species]['pressure'] = self._pparams[species]['kBT'] * \ - # e * 1e3 * self._pparams[species]['density'] * 1e-5 - self._pparams[species]["pressure"] = 99.0 - - for species in self._pparams: - # alfvén speed (m/s) - self._pparams[species]["v_A"] = magnetic_field / np.sqrt( - mu0 * self._pparams[species]["mass"] * self._pparams[species]["density"], - ) - # thermal speed (m/s) - self._pparams[species]["v_th"] = np.sqrt( - self._pparams[species]["kBT"] * 1e3 * e / self._pparams[species]["mass"], - ) - # thermal frequency (Mrad/s) - self._pparams[species]["Omega_th"] = self._pparams[species]["v_th"] / transit_length * 1e-6 - # cyclotron frequency (Mrad/s) - self._pparams[species]["Omega_c"] = self._pparams[species]["charge"] * magnetic_field / self._pparams[species]["mass"] * 1e-6 - # plasma frequency (Mrad/s) - self._pparams[species]["Omega_p"] = ( - np.sqrt( - self._pparams[species]["density"] * (self._pparams[species]["charge"]) ** 2 / eps0 / self._pparams[species]["mass"], - ) - * 1e-6 - ) - # alfvén frequency (Mrad/s) - self._pparams[species]["Omega_A"] = self._pparams[species]["v_A"] / transit_length * 1e-6 - # Larmor radius (m) - self._pparams[species]["rho_th"] = self._pparams[species]["v_th"] / (self._pparams[species]["Omega_c"] * 1e6) - # MHD length scale (m) - self._pparams[species]["v_A/Omega_c"] = self._pparams[species]["v_A"] / (np.abs(self._pparams[species]["Omega_c"]) * 1e6) - # dim-less ratios - self._pparams[species]["rho_th/L"] = self._pparams[species]["rho_th"] / transit_length + # # species dependent parameters + # self._pparams = {} + + # if len(self.fluid_species) > 0: + # for species, val in self.fluid_species.items(): + # self._pparams[species] = {} + # # type + # self._pparams[species]["type"] = "fluid" + # # mass (kg) + # self._pparams[species]["mass"] = val["params"]["phys_params"]["A"] * m_p + # # charge (C) + # self._pparams[species]["charge"] = val["params"]["phys_params"]["Z"] * e + # # density (m⁻³) + # self._pparams[species]["density"] = ( + # np.mean( + # self.equil.n0( + # eta1, + # eta2, + # eta3, + # ) + # * np.abs(det_tmp), + # ) + # * self.units.x ** 3 + # / plasma_volume + # * self.units.n + # ) + # # pressure (bar) + # self._pparams[species]["pressure"] = ( + # np.mean( + # self.equil.p0( + # eta1, + # eta2, + # eta3, + # ) + # * np.abs(det_tmp), + # ) + # * self.units.x ** 3 + # / plasma_volume + # * self.units.p + # * 1e-5 + # ) + # # thermal energy (keV) + # self._pparams[species]["kBT"] = self._pparams[species]["pressure"] * 1e5 / self._pparams[species]["density"] / e * 1e-3 + + # if len(self.kinetic) > 0: + # eta1mg, eta2mg, eta3mg = np.meshgrid( + # eta1, + # eta2, + # eta3, + # indexing="ij", + # ) + + # for species, val in self.kinetic.items(): + # self._pparams[species] = {} + # # type + # self._pparams[species]["type"] = "kinetic" + # # mass (kg) + # self._pparams[species]["mass"] = val["params"]["phys_params"]["A"] * m_p + # # charge (C) + # self._pparams[species]["charge"] = val["params"]["phys_params"]["Z"] * e + + # # create temp kinetic object for (default) parameter extraction + # tmp_bckgr = val["params"]["background"] + + # if val["space"] != "ParticlesSPH": + # tmp = None + # for fi, maxw_params in tmp_bckgr.items(): + # if fi[-2] == "_": + # fi_type = fi[:-2] + # else: + # fi_type = fi + + # if tmp is None: + # tmp = getattr(maxwellians, fi_type)( + # maxw_params=maxw_params, + # equil=self.equil, + # ) + # else: + # tmp = tmp + getattr(maxwellians, fi_type)( + # maxw_params=maxw_params, + # equil=self.equil, + # ) + + # if val["space"] != "ParticlesSPH" and tmp.coords == "constants_of_motion": + # # call parameters + # a1 = self.domain.params_map["a1"] + # r = eta1mg * (1 - a1) + a1 + # psi = self.equil.psi_r(r) + + # # density (m⁻³) + # self._pparams[species]["density"] = ( + # np.mean(tmp.n(psi) * np.abs(det_tmp)) * self.units.x ** 3 / plasma_volume * self.units.n + # ) + # # thermal speed (m/s) + # self._pparams[species]["v_th"] = ( + # np.mean(tmp.vth(psi) * np.abs(det_tmp)) * self.units.x ** 3 / plasma_volume * self.units.v + # ) + # # thermal energy (keV) + # self._pparams[species]["kBT"] = self._pparams[species]["mass"] * self._pparams[species]["v_th"] ** 2 / e * 1e-3 + # # pressure (bar) + # self._pparams[species]["pressure"] = ( + # self._pparams[species]["kBT"] * e * 1e3 * self._pparams[species]["density"] * 1e-5 + # ) - if verbose and self.rank_world == 0: - print("\nSPECIES PARAMETERS:") - for species, ch in self._pparams.items(): - print(f"\nname:".ljust(26), species) - print(f"type:".ljust(25), ch["type"]) - ch.pop("type") - print(f"is bulk:".ljust(25), species == self.bulk_species()) - for kinds, vals in ch.items(): - print( - kinds.ljust(25), - "{:+4.3e}".format( - vals, - ), - units_affix[kinds], - ) + # else: + # # density (m⁻³) + # # self._pparams[species]['density'] = np.mean(tmp.n( + # # eta1mg, eta2mg, eta3mg) * np.abs(det_tmp)) * units['x']**3 / plasma_volume * units['n'] + # self._pparams[species]["density"] = 99.0 + # # thermal speeds (m/s) + # vth = [] + # # vths = tmp.vth(eta1mg, eta2mg, eta3mg) + # vths = [99.0] + # for k in range(len(vths)): + # vth += [ + # vths[k] * np.abs(det_tmp) * self.units.x ** 3 / plasma_volume * self.units.v, + # ] + # thermal_speed = 0.0 + # for dir in range(val["obj"].vdim): + # # self._pparams[species]['vth' + str(dir + 1)] = np.mean(vth[dir]) + # self._pparams[species]["vth" + str(dir + 1)] = 99.0 + # thermal_speed += self._pparams[species]["vth" + str(dir + 1)] + # # TODO: here it is assumed that background density parameter is called "n", + # # and that background thermal speeds are called "vthn"; make this a convention? + # # self._pparams[species]['v_th'] = thermal_speed / \ + # # val['obj'].vdim + # self._pparams[species]["v_th"] = 99.0 + # # thermal energy (keV) + # # self._pparams[species]['kBT'] = self._pparams[species]['mass'] * \ + # # self._pparams[species]['v_th']**2 / e * 1e-3 + # self._pparams[species]["kBT"] = 99.0 + # # pressure (bar) + # # self._pparams[species]['pressure'] = self._pparams[species]['kBT'] * \ + # # e * 1e3 * self._pparams[species]['density'] * 1e-5 + # self._pparams[species]["pressure"] = 99.0 + + # for species in self._pparams: + # # alfvén speed (m/s) + # self._pparams[species]["v_A"] = magnetic_field / np.sqrt( + # mu0 * self._pparams[species]["mass"] * self._pparams[species]["density"], + # ) + # # thermal speed (m/s) + # self._pparams[species]["v_th"] = np.sqrt( + # self._pparams[species]["kBT"] * 1e3 * e / self._pparams[species]["mass"], + # ) + # # thermal frequency (Mrad/s) + # self._pparams[species]["Omega_th"] = self._pparams[species]["v_th"] / transit_length * 1e-6 + # # cyclotron frequency (Mrad/s) + # self._pparams[species]["Omega_c"] = self._pparams[species]["charge"] * magnetic_field / self._pparams[species]["mass"] * 1e-6 + # # plasma frequency (Mrad/s) + # self._pparams[species]["Omega_p"] = ( + # np.sqrt( + # self._pparams[species]["density"] * (self._pparams[species]["charge"]) ** 2 / eps0 / self._pparams[species]["mass"], + # ) + # * 1e-6 + # ) + # # alfvén frequency (Mrad/s) + # self._pparams[species]["Omega_A"] = self._pparams[species]["v_A"] / transit_length * 1e-6 + # # Larmor radius (m) + # self._pparams[species]["rho_th"] = self._pparams[species]["v_th"] / (self._pparams[species]["Omega_c"] * 1e6) + # # MHD length scale (m) + # self._pparams[species]["v_A/Omega_c"] = self._pparams[species]["v_A"] / (np.abs(self._pparams[species]["Omega_c"]) * 1e6) + # # dim-less ratios + # self._pparams[species]["rho_th/L"] = self._pparams[species]["rho_th"] / transit_length + + # if verbose and self.rank_world == 0: + # print("\nSPECIES PARAMETERS:") + # for species, ch in self._pparams.items(): + # print(f"\nname:".ljust(26), species) + # print(f"type:".ljust(25), ch["type"]) + # ch.pop("type") + # print(f"is bulk:".ljust(25), species == self.bulk_species()) + # for kinds, vals in ch.items(): + # print( + # kinds.ljust(25), + # "{:+4.3e}".format( + # vals, + # ), + # units_affix[kinds], + # ) class MyDumper(yaml.SafeDumper): diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 94253b581..5dac4caed 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -72,11 +72,11 @@ def setup_equation_params(self, units: Units = None, verbose=False): om_p = np.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) self._equation_params["alpha"] = om_p / om_c - self._equation_params["epsilon"] = 1.0 / (om_c * units["t"]) - self._equation_params["kappa"] = om_p * units["t"] + self._equation_params["epsilon"] = 1.0 / (om_c * units.t) + self._equation_params["kappa"] = om_p * units.t if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f'Set normalization parameters for speceis {self.name}:') + print(f'Set normalization parameters for species {self.__class__.__name__}:') for key, val in self.equation_params.items(): print((key + ":").ljust(25), "{:4.3e}".format(val)) From 7526d9a8a54aa2c61e323f4af5433c8c8b9689c8 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 1 Aug 2025 08:29:12 +0200 Subject: [PATCH 051/292] new abstract method StruphyModel.allocate_helpers() --- .../fields_background/projected_equils.py | 6 +- src/struphy/main.py | 1 + src/struphy/models/base.py | 4 + src/struphy/models/fluid.py | 114 ++++++------------ 4 files changed, 48 insertions(+), 77 deletions(-) diff --git a/src/struphy/fields_background/projected_equils.py b/src/struphy/fields_background/projected_equils.py index afd3d15fa..427398ee8 100644 --- a/src/struphy/fields_background/projected_equils.py +++ b/src/struphy/fields_background/projected_equils.py @@ -4,6 +4,8 @@ FluidEquilibriumWithB, MHDequilibrium, ) +from psydac.linalg.stencil import StencilVector +from psydac.linalg.block import BlockVector class ProjectedFluidEquilibrium: @@ -100,7 +102,7 @@ def absB3(self): return coeffs @property - def p3(self): + def p3(self) -> StencilVector: tmp = self._P3(self.equil.p3) coeffs = self._E3T.dot(tmp) coeffs.update_ghost_regions() @@ -256,7 +258,7 @@ def a1(self): # 2-forms # # ---------# @property - def b2(self): + def b2(self) -> BlockVector: tmp = self._P2([self.equil.b2_1, self.equil.b2_2, self.equil.b2_3]) coeffs = self._E2T.dot(tmp) coeffs.update_ghost_regions() diff --git a/src/struphy/main.py b/src/struphy/main.py index 305ec2ab3..9d053f854 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -169,6 +169,7 @@ def main( # allocate variables model.allocate_variables() + model.allocate_helpers() # pass info to propagators model.allocate_propagators() diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 4f5f59a01..c84a97e6a 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -77,6 +77,10 @@ def velocity_scale() -> str: """Velocity unit scale of the model. Must be one of "alfvén", "cyclotron", "light" or "thermal".""" + @abstractmethod + def allocate_helpers(self): + """Allocate helper arrays that are needed during simulation.""" + @abstractmethod def update_scalar_quantities(self): """Specify an update rule for each item in ``scalar_quantities`` using :meth:`update_scalar`.""" diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index e5abfaafa..9ac27adb9 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -5,6 +5,8 @@ from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable +from struphy.polar.basic import PolarVector +from psydac.linalg.block import BlockVector class LinearMHD(StruphyModel): @@ -38,21 +40,27 @@ class LinearMHD(StruphyModel): :ref:`Model info `: """ + ## species + @dataclass class EMFields(FieldSpecies): b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") - + @dataclass class MHD(FluidSpecies): density: FEECVariable = FEECVariable(name="density", space="L2") velocity: FEECVariable = FEECVariable(name="velocity", space="Hdiv") pressure: FEECVariable = FEECVariable(name="pressure", space="L2") + + ## propagators class Propagators: def __init__(self): self.shear_alf = propagators_fields.ShearAlfven() self.mag_sonic = propagators_fields.Magnetosonic() + ## abstract methods + def __init__(self): # 1. instantiate all variales self.em_fields = self.EMFields() @@ -81,8 +89,6 @@ def __init__(self): self.add_scalar("en_B_eq") self.add_scalar("en_B_tot") self.add_scalar("en_tot") - - ## abstract methods @property def bulk_species(self): @@ -92,87 +98,45 @@ def bulk_species(self): def velocity_scale(self): return "alfvén" - # def __init__old(self, params, comm, clone_config=None): - # # initialize base class - # super().__init__(params, comm=comm, clone_config=clone_config) - - # from struphy.polar.basic import PolarVector - - # # extract necessary parameters - # u_space = params["fluid"]["mhd"]["options"]["u_space"] - # alfven_solver = params["fluid"]["mhd"]["options"]["ShearAlfven"]["solver"] - # alfven_algo = params["fluid"]["mhd"]["options"]["ShearAlfven"]["algo"] - # sonic_solver = params["fluid"]["mhd"]["options"]["Magnetosonic"]["solver"] - - # # project background magnetic field (2-form) and pressure (3-form) - # self._b_eq = self.projected_equil.b2 - # self._p_eq = self.projected_equil.p3 - # self._ones = self._p_eq.space.zeros() - - # if isinstance(self._ones, PolarVector): - # self._ones.tp[:] = 1.0 - # else: - # self._ones[:] = 1.0 - - # # set keyword arguments for propagators - # self._kwargs[propagators_fields.ShearAlfven] = { - # "u_space": u_space, - # "solver": alfven_solver, - # "algo": alfven_algo, - # } - - # self._kwargs[propagators_fields.Magnetosonic] = { - # "b": self.pointer["b_field"], - # "u_space": u_space, - # "solver": sonic_solver, - # } - - # # Initialize propagators used in splitting substeps - # self.init_propagators() - - # # Scalar variables to be saved during simulation - # self.add_scalar("en_U") - # self.add_scalar("en_p") - # self.add_scalar("en_B") - # self.add_scalar("en_p_eq") - # self.add_scalar("en_B_eq") - # self.add_scalar("en_B_tot") - # self.add_scalar("en_tot") - - # # vectors for computing scalar quantities - # self._tmp_b1 = self.derham.Vh["2"].zeros() - # self._tmp_b2 = self.derham.Vh["2"].zeros() + def allocate_helpers(self): + self._ones = self.projected_equil.p3.space.zeros() + if isinstance(self._ones, PolarVector): + self._ones.tp[:] = 1.0 + else: + self._ones[:] = 1.0 + + self._tmp_b1: BlockVector = self.derham.Vh["2"].zeros() # TODO: replace derham.Vh dict by class + self._tmp_b2 = self.derham.Vh["2"].zeros() - def update_scalar_quantities(self): - pass - # # perturbed fields - # en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.pointer["mhd_velocity"], self.pointer["mhd_velocity"]) - # en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) - # en_p = self.pointer["mhd_pressure"].inner(self._ones) / (5 / 3 - 1) + def update_scalar_quantities(self): + # perturbed fields + en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.mhd.velocity.spline.vector, self.mhd.velocity.spline.vector,) + en_B = 0.5 * self.mass_ops.M2.dot_inner(self.em_fields.b_field.spline.vector, self.em_fields.b_field.spline.vector,) + en_p = self.mhd.pressure.spline.vector.inner(self._ones) / (5 / 3 - 1) - # self.update_scalar("en_U", en_U) - # self.update_scalar("en_B", en_B) - # self.update_scalar("en_p", en_p) - # self.update_scalar("en_tot", en_U + en_B + en_p) + self.update_scalar("en_U", en_U) + self.update_scalar("en_B", en_B) + self.update_scalar("en_p", en_p) + self.update_scalar("en_tot", en_U + en_B + en_p) - # # background fields - # self.mass_ops.M2.dot(self._b_eq, apply_bc=False, out=self._tmp_b1) + # background fields + self.mass_ops.M2.dot(self.projected_equil.b2, apply_bc=False, out=self._tmp_b1) - # en_B0 = self._b_eq.inner(self._tmp_b1) / 2 - # en_p0 = self._p_eq.inner(self._ones) / (5 / 3 - 1) + en_B0 = self.projected_equil.b2.inner(self._tmp_b1) / 2 + en_p0 = self.projected_equil.p3.inner(self._ones) / (5 / 3 - 1) - # self.update_scalar("en_B_eq", en_B0) - # self.update_scalar("en_p_eq", en_p0) + self.update_scalar("en_B_eq", en_B0) + self.update_scalar("en_p_eq", en_p0) - # # total magnetic field - # self._b_eq.copy(out=self._tmp_b1) - # self._tmp_b1 += self.pointer["b_field"] + # total magnetic field + self.projected_equil.b2.copy(out=self._tmp_b1) + self._tmp_b1 += self.em_fields.b_field.spline.vector - # self.mass_ops.M2.dot(self._tmp_b1, apply_bc=False, out=self._tmp_b2) + self.mass_ops.M2.dot(self._tmp_b1, apply_bc=False, out=self._tmp_b2) - # en_Btot = self._tmp_b1.inner(self._tmp_b2) / 2 + en_Btot = self._tmp_b1.inner(self._tmp_b2) / 2 - # self.update_scalar("en_B_tot", en_Btot) + self.update_scalar("en_B_tot", en_Btot) class LinearExtendedMHDuniform(StruphyModel): From 0c74a66fb5c1c7e06639fdc5a39025674fe88edc Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 1 Aug 2025 13:54:28 +0200 Subject: [PATCH 052/292] new generate_default_parameter_file for .py parameters --- src/struphy/console/main.py | 2 +- src/struphy/console/params.py | 31 ++- src/struphy/feec/linear_operators.py | 4 +- src/struphy/feec/psydac_derham.py | 2 +- src/struphy/feec/utilities.py | 9 +- src/struphy/geometry/domains.py | 28 +-- src/struphy/initial/perturbations.py | 11 +- src/struphy/io/inp/params_Maxwell.py | 83 +++----- src/struphy/io/inp/params_Maxwell_lw.py | 8 +- src/struphy/io/options.py | 32 ++- src/struphy/io/setup.py | 26 +-- src/struphy/models/base.py | 262 +++++++++--------------- src/struphy/models/toy.py | 11 +- src/struphy/topology/grids.py | 102 +-------- 14 files changed, 228 insertions(+), 383 deletions(-) diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index 51c0e3c59..74294f9e2 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -690,7 +690,7 @@ def add_parser_params(subparsers, list_models, model_message): ) parser_params.add_argument( - "model", + "model_name", type=str, choices=list_models, metavar="MODEL", diff --git a/src/struphy/console/params.py b/src/struphy/console/params.py index dace206ab..bf7548b57 100644 --- a/src/struphy/console/params.py +++ b/src/struphy/console/params.py @@ -1,34 +1,31 @@ -def struphy_params(model, file, yes=False, options=False): +from struphy.models.base import StruphyModel +from struphy.models import fluid, hybrid, kinetic, toy + + +def struphy_params(model_name: str, file, yes=False, options=False): """Create a model's default parameter file and save in current input path. Parameters ---------- - model : str + model_name : str The name of the Struphy model. - yes : bool - If true, say yes on prompt to overwrite .yml FILE - file : str An alternative file name to the default params_.yml. + + yes : bool + If true, say yes on prompt to overwrite .yml FILE show_options : bool Whether to print to screen all possible options for the model. """ - - from struphy.models import fluid, hybrid, kinetic, toy - - # load model class objs = [fluid, kinetic, hybrid, toy] for obj in objs: try: - model_class = getattr(obj, model) + model_class = getattr(obj, model_name) + model: StruphyModel = model_class() except AttributeError: pass - - # print units - if options: - model_class.show_options() - else: - prompt = not yes - params = model_class.generate_default_parameter_file(file=file, prompt=prompt) + + prompt = not yes + model.generate_default_parameter_file(file_name=file, prompt=prompt) diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 9d42fc9c3..2bb14757e 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -294,7 +294,7 @@ class BoundaryOperator(LinOpWithTransp): space_id : str Symbolic space ID of vector_space (H1, Hcurl, Hdiv, L2 or H1vec). - dirichlet_bc : list[list[bool]] + dirichlet_bc : tuple[tuple[bool]] Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). """ @@ -310,7 +310,7 @@ def __init__(self, vector_space, space_id, dirichlet_bc): self._space_id = space_id self._bc = dirichlet_bc - assert isinstance(dirichlet_bc, list) + assert isinstance(dirichlet_bc, tuple) assert len(dirichlet_bc) == 3 # number of non-zero elements in poloidal/toroidal direction diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 9451d7fe3..4bab49a49 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -118,7 +118,7 @@ def __init__( if dirichlet_bc is not None: assert len(dirichlet_bc) == 3 # make sure that boundary conditions are compatible with spline space - assert np.all([bc == [False, False] for i, bc in enumerate(dirichlet_bc) if spl_kind[i]]) + assert np.all([bc == (False, False) for i, bc in enumerate(dirichlet_bc) if spl_kind[i]]) self._dirichlet_bc = dirichlet_bc diff --git a/src/struphy/feec/utilities.py b/src/struphy/feec/utilities.py index 292eb61e6..d0c89108b 100644 --- a/src/struphy/feec/utilities.py +++ b/src/struphy/feec/utilities.py @@ -4,6 +4,7 @@ from psydac.fem.vector import VectorFemSpace from psydac.linalg.block import BlockLinearOperator, BlockVector from psydac.linalg.stencil import StencilMatrix, StencilVector +from psydac.linalg.basic import Vector import struphy.feec.utilities_kernels as kernels from struphy.feec import banded_to_stencil_kernels as bts @@ -283,7 +284,7 @@ def compare_arrays(arr_psy, arr, rank, atol=1e-14, verbose=False): ) -def apply_essential_bc_to_array(space_id, vector, bc): +def apply_essential_bc_to_array(space_id: str, vector: Vector, bc: tuple): """ Sets entries corresponding to boundary B-splines to zero. @@ -292,15 +293,15 @@ def apply_essential_bc_to_array(space_id, vector, bc): space_id : str The name of the continuous functions space the given vector belongs to (H1, Hcurl, Hdiv, L2 or H1vec). - vector : StencilVector | BlockVector + vector : Vector The vector whose boundary values shall be set to zero. - bc : list[list[bool]] + bc : tuple[tuple[bool]] Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). """ assert isinstance(vector, (StencilVector, BlockVector, PolarVector)) - assert isinstance(bc, list) + assert isinstance(bc, tuple) assert len(bc) == 3 if isinstance(vector, PolarVector): diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index 2e142fa1f..8b18a56c1 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -558,22 +558,22 @@ class Cuboid(Domain): r3 : 1. # end of z-interval, r3>l3 """ - def __init__(self, **params): + def __init__(self, l1: float = 0.0, + r1: float = 2.0, + l2: float = 0.0, + r2: float = 3.0, + l3: float = 0.0, + r3: float = 6.0,): + self._kind_map = 10 - # set default parameters and remove wrong/not needed keys - params_default = { - "l1": 0.0, - "r1": 2.0, - "l2": 0.0, - "r2": 3.0, - "l3": 0.0, - "r3": 6.0, - } - - self._params_map, self._params_numpy = Domain.prepare_params_map( - params, - params_default, + self._params_map, self._params_numpy = Domain.prepare_params_map_new( + l1=l1, + r1=r1, + l2=l2, + r2=r2, + l3=l3, + r3=r3, ) # periodicity in eta3-direction and pole at eta1=0 diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 706ca6494..2ea51bee3 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -560,16 +560,16 @@ class TorusModesCos(Perturbation): Parameters ---------- - ms : tuple | list[int] + ms : tuple[int] Poloidal mode numbers. - ns : tuple | list[int] + ns : tuple[int] Toroidal mode numbers. - pfuns : tuple | list[str] + pfuns : tuple[str] "sin" or "cos" or "exp" to define the profile functions. - amps : tuple | list[float] + amps : tuple[float] Amplitudes of each mode (m_i, n_i). pfun_params : tuple | list @@ -582,7 +582,8 @@ class TorusModesCos(Perturbation): Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, ms=None, ns=None, amps=(1e-4,), pfuns=("sin",), pfun_params=None, + def __init__(self, ms: tuple = (2,), ns: tuple = (1,), amps: tuple = (1e-4,), + pfuns: tuple = ("sin",), pfun_params=None, given_in_basis: GivenInBasis = "0", comp: int = 0,): diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index 73cf77da6..921f6311b 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -1,62 +1,41 @@ -from struphy.io import options +from struphy.io.options import Units, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.initial import perturbations +from struphy.topology import grids +from struphy.io.options import DerhamOptions +from struphy.io.options import FieldsBackground -# model -model = "Maxwell" +# import model +from struphy.models.toy import Maxwell as Model + +# light-weight model instance +model = Model() # units -units = options.Units( - x=1.0, - B=1.0, - n=1.0, - kBT=1.0, -) +units = Units() -# time -time = options.Time(split_algo="LieTrotter") +# time stepping +time = Time() # geometry -domain = options.domains.Cuboid() +domain = domains.Cuboid() + +# fluid equilibrium (can be used as part of initial conditions) +equil = equils.HomogenSlab() # grid -grid = options.grids.TensorProductGrid( - Nel=(12, 14, 1), - p=(2, 3, 1), - spl_kind=(False, True, True), -) +grid = grids.TensorProductGrid() # derham options -derham = options.DerhamOptions() - -# fluid equilibrium -equil = options.equils.HomogenSlab() - -# FOR NOW: initial conditions and options -em_fields = {} -em_fields["background"] = {} -em_fields["perturbation"] = {} -em_fields["options"] = {} - -em_fields["background"]["e_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} -em_fields["background"]["b_field"] = {"LogicalConst": {"values": [0.3, 0.15, None]}} - -em_fields["perturbation"]["e_field"] = {} -em_fields["perturbation"]["b_field"] = {} -em_fields["perturbation"]["e_field"]["TorusModesCos"] = { - "given_in_basis": [None, "v", None], - "ms": [[None], [1, 3], [None]], -} -em_fields["perturbation"]["b_field"]["TorusModesCos"] = { - "given_in_basis": [None, "v", None], - "ms": [[None], [1, 3], [None]], -} - -solver = { - "type": ["pcg", "MassMatrixPreconditioner"], - "tol": 1.0e-08, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, -} - -em_fields["options"]["Maxwell"] = {"algo": "implicit", "solver": solver} +derham = DerhamOptions() + +# propagator options +model.propagators.maxwell.set_options() + +# initial conditions (background + perturbation) +model.em_fields.b_field.add_background(FieldsBackground()) +model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos()) + +# optional: exclude variables from saving +# model.em_fields.b_field.save_data = False diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index bd2689a41..5ac1d1ef4 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -24,14 +24,10 @@ time = Time() # grid -grid = grids.TensorProductGrid( - Nel=(12, 14, 1), - p=(2, 3, 1), - spl_kind=(False, True, True), -) +grid = grids.TensorProductGrid(Nel=(12, 14, 1)) # derham options -derham = DerhamOptions() +derham = DerhamOptions(p=(2, 3, 1), spl_kind=(False, True, True)) # propagator options model.propagators.maxwell.set_options(algo="explicit") diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index cb9ca39e4..17b0dee9a 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -199,15 +199,37 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk @dataclass class DerhamOptions: - """... + """Options for the Derham spaces. Parameters ---------- - x : float - Unit of length in m. - """ + p : tuple[int] + Spline degree in each direction. + + spl_kind : tuple[bool] + Kind of spline in each direction (True=periodic, False=clamped). - polar_ck: int = -1 + dirichlet_bc : tuple[tuple[bool]] + Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). + + nquads : tuple[int] + Number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials). + + nq_pr : tuple[int] + Number of Gauss-Legendre quadrature points in each direction for geometric projectors (default = p+1, leads to exact integration of degree 2p+1 polynomials). + + polar_ck : PolarRegularity + Smoothness at a polar singularity at eta_1=0 (default -1 : standard tensor product splines, OR 1 : C1 polar splines) + + local_projectors : bool + Whether to build the local commuting projectors based on quasi-inter-/histopolation. + """ + p: tuple = (1, 1, 1) + spl_kind: tuple = (True, True, True) + dirichlet_bc: tuple = ((False, False), (False, False), (False, False)) + nquads: tuple = None + nq_pr: tuple = None + polar_ck: PolarRegularity = -1 local_projectors: bool = False def __post_init__(self): diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index fa4d6ec83..017931375 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -12,6 +12,8 @@ from struphy.topology.grids import TensorProductGrid from struphy.utils.utils import dict_to_yaml, read_state from struphy.geometry.base import Domain +from struphy.geometry.domains import Cuboid +from struphy.io.options import Units, Time def import_parameters_py(params_path: str): @@ -127,7 +129,7 @@ def setup_parameters( params_in.model = None if not hasattr(params_in, "domain"): - params_in.domain = None + params_in.domain = Cuboid() if not hasattr(params_in, "grid"): params_in.grid = None @@ -136,10 +138,10 @@ def setup_parameters( params_in.equil = None if not hasattr(params_in, "units"): - params_in.units = None + params_in.units = Units() if not hasattr(params_in, "time"): - params_in.time = None + params_in.time = Time() if not hasattr(params_in, "derham"): params_in.derham = None @@ -172,7 +174,6 @@ def setup_derham( options: DerhamOptions, comm: MPI.Intracomm = None, domain: Domain = None, - mpi_dims_mask: tuple | list = None, verbose=False, ): """ @@ -189,10 +190,6 @@ def setup_derham( domain : Domain, optional The Struphy domain object for evaluating the mapping F : [0, 1]^3 --> R^3 and the corresponding metric coefficients. - mpi_dims_mask: list | tuple[bool] - True if the dimension is to be used in the domain decomposition (=default for each dimension). - If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. - verbose : bool Show info on screen. @@ -206,16 +203,19 @@ def setup_derham( # number of grid cells Nel = grid.Nel + # mpi + mpi_dims_mask = grid.mpi_dims_mask + # spline degrees - p = grid.p + p = options.p # spline types (clamped vs. periodic) - spl_kind = grid.spl_kind + spl_kind = options.spl_kind # boundary conditions (Homogeneous Dirichlet or None) - dirichlet_bc = grid.dirichlet_bc + dirichlet_bc = options.dirichlet_bc # Number of quadrature points per histopolation cell - nq_pr = grid.nq_pr + nq_pr = options.nq_pr # Number of quadrature points per grid cell for L^2 - nquads = grid.nquads + nquads = options.nquads # C^k smoothness at eta_1=0 for polar domains polar_ck = options.polar_ck # local commuting projectors diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index c84a97e6a..224fe975e 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -24,7 +24,7 @@ from struphy.io.setup import setup_derham, descend_options_dict from struphy.profiling.profiling import ProfileManager from struphy.utils.clone_config import CloneConfig -from struphy.utils.utils import dict_to_yaml +from struphy.utils.utils import dict_to_yaml, read_state from struphy.models.species import Species, FieldSpecies, FluidSpecies, KineticSpecies, DiagnosticSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.io.options import Units @@ -166,29 +166,19 @@ def diagnostic_species(self) -> dict: @property def species(self): if not hasattr(self, "_species"): - self._species = {} - for k, v in self.field_species.items(): - self._species[k] = v - for k, v in self.fluid_species.items(): - self._species[k] = v - for k, v in self.kinetic_species.items(): - self._species[k] = v + self._species = self.field_species | self.fluid_species | self.kinetic_species return self._species ## allocate methods def allocate_feec(self, grid: TensorProductGrid, - derham_params: DerhamOptions=None, + derham_params: DerhamOptions, comm: MPI.Intracomm = None, clone_config: CloneConfig = None, ): # create discrete derham sequence - dims_mask = grid.mpi_dims_mask - if dims_mask is None: - dims_mask = [True] * 3 - if clone_config is None: derham_comm = self.comm_world else: @@ -199,7 +189,6 @@ def allocate_feec(self, derham_params, comm=derham_comm, domain=self.domain, - mpi_dims_mask=dims_mask, verbose=self.verbose, ) @@ -1345,11 +1334,9 @@ def write_parameters_to_file(cls, parameters=None, file=None, save=True, prompt= else: pass - @classmethod def generate_default_parameter_file( - cls, - file: str = None, - save: bool = True, + self, + file_name: str = None, prompt: bool = True, ): """Generate a parameter file with default options for each species, @@ -1359,156 +1346,105 @@ def generate_default_parameter_file( Parameters ---------- - file : str - Alternative filename to params_.yml. - - save : bool - Whether to save the parameter file in the current input path. + file_name : str + Alternative filename to params_.py. prompt : bool Whether to prompt for overwriting the specified .yml file. + """ - Returns - ------- - The default parameter dictionary.""" - - libpath = struphy.__path__[0] - - # load a standard parameter file - with open(os.path.join(libpath, "io/inp/parameters.yml")) as tmp: - parameters = yaml.load(tmp, Loader=yaml.FullLoader) - - parameters["model"] = cls.__name__ - - # extract default em_fields parameters - bckgr_params_1_em = parameters["em_fields"]["background"]["var_1"] - bckgr_params_2_em = parameters["em_fields"]["background"]["var_2"] - parameters["em_fields"].pop("background") - - pert_params_1_em = parameters["em_fields"]["perturbation"]["var_1"] - pert_params_2_em = parameters["em_fields"]["perturbation"]["var_2"] - parameters["em_fields"].pop("perturbation") - - # extract default fluid parameters - bckgr_params_1_fluid = parameters["fluid"]["species_name"]["background"]["var_1"] - bckgr_params_2_fluid = parameters["fluid"]["species_name"]["background"]["var_2"] - parameters["fluid"]["species_name"].pop("background") - - pert_params_1_fluid = parameters["fluid"]["species_name"]["perturbation"]["var_1"] - pert_params_2_fluid = parameters["fluid"]["species_name"]["perturbation"]["var_2"] - parameters["fluid"]["species_name"].pop("perturbation") - - # standard Maxwellians - parameters["kinetic"]["species_name"].pop("background") - maxw_name = { - "6D": "Maxwellian3D", - "5D": "GyroMaxwellian2D", - "4D": "Maxwellian1D", - "3D": "ColdPlasma", - "PH": "ConstantVelocity", - } - - # init options dicts - d_opts = {"em_fields": [], "fluid": {}, "kinetic": {}} - - # set the correct names in the parameter file - if len(cls.species()["em_fields"]) > 0: - parameters["em_fields"]["background"] = {} - parameters["em_fields"]["perturbation"] = {} - for name, space in cls.species()["em_fields"].items(): - if space in {"H1", "L2"}: - parameters["em_fields"]["background"][name] = bckgr_params_1_em - parameters["em_fields"]["perturbation"][name] = pert_params_1_em - elif space in {"Hcurl", "Hdiv", "H1vec"}: - parameters["em_fields"]["background"][name] = bckgr_params_2_em - parameters["em_fields"]["perturbation"][name] = pert_params_2_em - else: - parameters.pop("em_fields") - - # find out the default em_fields options of the model - if "options" in cls.options()["em_fields"]: - # create the default options parameters - d_default = descend_options_dict( - cls.options()["em_fields"]["options"], - d_opts["em_fields"], - ) - parameters["em_fields"]["options"] = d_default - - # fluid - fluid_params = parameters["fluid"].pop("species_name") - - if len(cls.species()["fluid"]) > 0: - for name, dct in cls.species()["fluid"].items(): - parameters["fluid"][name] = fluid_params - parameters["fluid"][name]["background"] = {} - parameters["fluid"][name]["perturbation"] = {} - - # find out the default fluid options of the model - if name in cls.options()["fluid"]: - d_opts["fluid"][name] = [] - - # create the default options parameters - d_default = descend_options_dict( - cls.options()["fluid"][name]["options"], - d_opts["fluid"][name], - ) - - parameters["fluid"][name]["options"] = d_default - - # set the correct names parameter file - for sub_name, space in dct.items(): - if space in {"H1", "L2"}: - parameters["fluid"][name]["background"][sub_name] = bckgr_params_1_fluid - parameters["fluid"][name]["perturbation"][sub_name] = pert_params_1_fluid - elif space in {"Hcurl", "Hdiv", "H1vec"}: - parameters["fluid"][name]["background"][sub_name] = bckgr_params_2_fluid - parameters["fluid"][name]["perturbation"][sub_name] = pert_params_2_fluid - else: - parameters.pop("fluid") - - # kinetic - kinetic_params = parameters["kinetic"].pop("species_name") - - if len(cls.species()["kinetic"]) > 0: - parameters["kinetic"] = {} - - for name, kind in cls.species()["kinetic"].items(): - parameters["kinetic"][name] = kinetic_params - - # find out the default kinetic options of the model - if name in cls.options()["kinetic"]: - d_opts["kinetic"][name] = [] - - # create the default options parameters - d_default = descend_options_dict( - cls.options()["kinetic"][name]["options"], - d_opts["kinetic"][name], - ) - - parameters["kinetic"][name]["options"] = d_default - - # set the background - dim = kind[-2:] - parameters["kinetic"][name]["background"] = { - maxw_name[dim]: {"n": 0.05}, - } - else: - parameters.pop("kinetic") - - # diagnostics - if cls.diagnostics_dct() is not None: - parameters["diagnostics"] = {} - for name, space in cls.diagnostics_dct().items(): - parameters["diagnostics"][name] = {"save_data": True} - - cls.write_parameters_to_file( - parameters=parameters, - file=file, - save=save, - prompt=prompt, - ) + # Read struphy state file + state = read_state() + i_path = state["i_path"] + assert os.path.exists(i_path), f"The path '{i_path}' does not exist. Set path with `struphy --set-i PATH`" + + if file_name is None: + file_name = os.path.join(i_path, f"params_{self.__class__.__name__}.py") - return parameters + # create new default file + try: + file = open(file_name, "x") + except FileExistsError: + if not prompt: + yn = "Y" + else: + yn = input(f"File {file_name} exists, overwrite (Y/n)? ") + if yn in ("", "Y", "y", "yes", "Yes"): + file = open(file_name, "w") + else: + print("exiting ...") + return + + + # generic options for all models + file.write("from struphy.io.options import Units, Time\n") + file.write("from struphy.geometry import domains\n") + file.write("from struphy.fields_background import equils\n") + file.write("from struphy.initial import perturbations\n") + + has_feec = False + has_pic = False + has_sph = False + for sn, species in self.species.items(): + assert isinstance(species, Species) + for vn, var in species.variables.items(): + if isinstance(var, FEECVariable): + has_feec = True + init_bckgr_feec = f"model.{sn}.{vn}.add_background(FieldsBackground())\n" + init_pert_feec = f"model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos())\n" + exclude_feec = f"# model.{sn}.{vn}.save_data = False\n" + elif isinstance(var, PICVariable): + has_pic = True + elif isinstance(var, SPHVariable): + has_sph = True + + if has_feec: + file.write("from struphy.topology import grids\n") + file.write("from struphy.io.options import DerhamOptions\n") + file.write("from struphy.io.options import FieldsBackground\n") + + if has_pic or has_sph: + file.write("from struphy.kinetic_background import maxwellians\n") + + file.write("\n# import model\n") + file.write(f"from {self.__module__} import {self.__class__.__name__} as Model\n") + + file.write("\n# light-weight model instance\n") + file.write("model = Model()\n") + + file.write("\n# units\n") + file.write("units = Units()\n") + + file.write("\n# time stepping\n") + file.write("time = Time()\n") + + file.write("\n# geometry\n") + file.write("domain = domains.Cuboid()\n") + + file.write("\n# fluid equilibrium (can be used as part of initial conditions)\n") + file.write("equil = equils.HomogenSlab()\n") + + if has_feec: + file.write("\n# grid\n") + file.write("grid = grids.TensorProductGrid()\n") + + file.write("\n# derham options\n") + file.write("derham = DerhamOptions()\n") + + file.write("\n# propagator options\n") + for prop in self.propagators.__dict__: + file.write(f"model.propagators.{prop}.set_options()\n") + + file.write("\n# initial conditions (background + perturbation)\n") + if has_feec: + file.write(init_bckgr_feec) + file.write(init_pert_feec) + + file.write("\n# optional: exclude variables from saving\n") + file.write(exclude_feec) + + print(f"Default parameter file for '{self.__class__.__name__}' has been created.\n\ +You can now launch with 'struphy run {self.__class__.__name__}' or with 'struphy run -i params_{self.__class__.__name__}.py'") ################### # Private methods : diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index ca0b85bef..93b69b271 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -32,15 +32,21 @@ class Maxwell(StruphyModel): :ref:`Model info `: """ + ## species + @dataclass class EMFields(FieldSpecies): e_field: FEECVariable = FEECVariable(name="e_field", space="Hcurl") b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") + ## propagators + class Propagators: def __init__(self): self.maxwell = propagators_fields.Maxwell() + ## abstract methods + def __init__(self): # 1. instantiate all variales self.em_fields = self.EMFields() @@ -58,8 +64,6 @@ def __init__(self): self.add_scalar("electric energy") self.add_scalar("magnetic energy") self.add_scalar("total energy") - - ## abstract methods @property def bulk_species(self): @@ -69,6 +73,9 @@ def bulk_species(self): def velocity_scale(self): return "light" + def allocate_helpers(self): + pass + def update_scalar_quantities(self): en_E = 0.5 * self.mass_ops.M1.dot_inner(self.em_fields.e_field.spline.vector, self.em_fields.e_field.spline.vector) en_B = 0.5 * self.mass_ops.M2.dot_inner(self.em_fields.b_field.spline.vector, self.em_fields.b_field.spline.vector) diff --git a/src/struphy/topology/grids.py b/src/struphy/topology/grids.py index cda18c021..e66dae1ea 100644 --- a/src/struphy/topology/grids.py +++ b/src/struphy/topology/grids.py @@ -1,6 +1,8 @@ import numpy as np +from dataclasses import dataclass +@dataclass class TensorProductGrid: """Grid as a tensor product of 1d grids. @@ -9,105 +11,9 @@ class TensorProductGrid: Nel : tuple[int] Number of elements in each direction. - p : tuple[int] - Spline degree in each direction. - - spl_kind : tuple[bool] - Kind of spline in each direction (True=periodic, False=clamped). - - dirichlet_bc : tuple[tuple[bool]] - Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). - - nquads : tuple[int] - Number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials). - - nq_pr : tuple[int] - Number of Gauss-Legendre quadrature points in each direction for geometric projectors (default = p+1, leads to exact integration of degree 2p+1 polynomials). - mpi_dims_mask: Tuple of bool True if the dimension is to be used in the domain decomposition (=default for each dimension). If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. """ - - def __init__( - self, - Nel: tuple, - p: tuple, - spl_kind: tuple, - *, - dirichlet_bc: tuple = None, - nquads: tuple = None, - nq_pr: tuple = None, - mpi_dims_mask: tuple = None, - ): - assert len(Nel) == len(p) == len(spl_kind) - - self._Nel = Nel - self._p = p - self._spl_kind = spl_kind - - # boundary conditions at eta=0 and eta=1 in each direction (None for periodic, 'd' for homogeneous Dirichlet) - if dirichlet_bc is not None: - assert len(dirichlet_bc) == len(Nel) - # make sure that boundary conditions are compatible with spline space - assert np.all([bc == [False, False] for i, bc in enumerate(dirichlet_bc) if spl_kind[i]]) - - self._dirichlet_bc = dirichlet_bc - - # default p: exact integration of degree 2p+1 polynomials - if nquads is None: - self._nquads = tuple([pi + 1 for pi in p]) - else: - assert len(nquads) == len(Nel) - self._nquads = nquads - - # default p + 1 : exact integration of degree 2p+1 polynomials - if nq_pr is None: - self._nq_pr = tuple([pi + 1 for pi in p]) - else: - assert len(nq_pr) == len(Nel) - self._nq_pr = nq_pr - - # mpi domain decomposition directions - if mpi_dims_mask is None: - self._mpi_dims_mask = (True,) * len(Nel) - else: - assert len(mpi_dims_mask) == len(Nel) - self._mpi_dims_mask = mpi_dims_mask - - @property - def Nel(self): - """Tuple of number of elements (=cells) in each direction.""" - return self._Nel - - @property - def p(self): - """Tuple of B-spline degrees in each direction.""" - return self._p - - @property - def spl_kind(self): - """Tuple of spline type (periodic=True or clamped=False) in each direction.""" - return self._spl_kind - - @property - def dirichlet_bc(self): - """None, or Tuple of boundary conditions in each direction. - Each entry is a list with two entries (left and right boundary), "d" (hom. Dirichlet) or None (periodic). - """ - return self._dirichlet_bc - - @property - def nquads(self): - """Tuple of number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials).""" - return self._nquads - - @property - def nq_pr(self): - """Tuple of number of Gauss-Legendre quadrature points in histopolation (default = p + 1) in each direction.""" - return self._nq_pr - - @property - def mpi_dims_mask(self): - """Tuple of bool; whether to use direction in domain decomposition.""" - return self._mpi_dims_mask + Nel: tuple = (16, 1, 1) + mpi_dims_mask: tuple = (True, True, True) From c3bab8b9728a8f76bdac23f6b6d1a7d5136dbeab Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 4 Aug 2025 09:05:05 +0200 Subject: [PATCH 053/292] include call to struphy main in parameter file; new class MetaOptions --- src/struphy/io/inp/params_Maxwell.py | 27 ++++- src/struphy/io/options.py | 74 +++++++++--- src/struphy/main.py | 165 +++++++++++---------------- src/struphy/models/base.py | 74 ++++++++---- 4 files changed, 201 insertions(+), 139 deletions(-) diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index 921f6311b..5caef842e 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -1,16 +1,19 @@ -from struphy.io.options import Units, Time +from struphy.io.options import Units, Time, MetaOptions from struphy.geometry import domains from struphy.fields_background import equils from struphy.initial import perturbations from struphy.topology import grids from struphy.io.options import DerhamOptions from struphy.io.options import FieldsBackground +from struphy.kinetic_background import maxwellians +from struphy import main -# import model +# import model, set verbosity from struphy.models.toy import Maxwell as Model +verbose = False -# light-weight model instance -model = Model() +# meta options +meta = MetaOptions() # units units = Units() @@ -30,6 +33,9 @@ # derham options derham = DerhamOptions() +# light-weight model instance +model = Model() + # propagator options model.propagators.maxwell.set_options() @@ -39,3 +45,16 @@ # optional: exclude variables from saving # model.em_fields.b_field.save_data = False + +# start run +main.main(model, + params_path=__file__, + units=units, + time_opts=time, + domain=domain, + equil=equil, + grid=grid, + derham=derham, + meta=meta, + verbose=verbose, + ) \ No newline at end of file diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 17b0dee9a..d4f60b268 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -1,20 +1,12 @@ from dataclasses import dataclass from typing import Literal, get_args import numpy as np +import os # needed for import in StruphyParameters from struphy.physics.physics import ConstantsOfNature -def check_option(opt, options): - """Check if opt is contained in options; if opt is a list, checks for each element.""" - opts = get_args(options) - if not isinstance(opt, list): - opt = [opt] - for o in opt: - assert o in opts, f"Option '{o}' is not in {opts}." - - ## Literal options # time @@ -41,7 +33,7 @@ def check_option(opt, options): @dataclass class Time: - """... + """Time stepping options. Parameters ---------- @@ -238,12 +230,18 @@ def __post_init__(self): @dataclass class FieldsBackground: - """... - + """Options for backgrounds in configuration (=position) space. + Parameters ---------- - x : float - Unit of length in m. + type : BackgroundTypes + Type of background. + + values : tuple[float] + Values for LogicalConst on the unit cube. + + variable : str + Name of the function in FluidEquilibrium that should be the background. """ type: BackgroundTypes = "LogicalConst" @@ -252,4 +250,52 @@ class FieldsBackground: def __post_init__(self): check_option(self.type, BackgroundTypes) + + +@dataclass +class MetaOptions: + """Meta options for launching run on current architecture + (these options do not influence the simulation result). + + Parameters + ---------- + out_folders : str + The directory where all sim_folders are stored. + + sim_folder : str + Folder in 'out_folders/' for the current simulation (default='sim_1'). + Will create the folder if it does not exist OR cleans the folder for new runs. + + restart : bool + Whether to restart a run (default=False). + + max_runtime : int, + Maximum run time of simulation in minutes. Will finish the time integration once this limit is reached (default=300). + + save_step : int + When to save data output: every time step (save_step=1), every second time step (save_step=2), etc (default=1). + + sort_step: int, optional + Sort markers in memory every N time steps (default=0, which means markers are sorted only at the start of simulation) + + num_clones: int, optional + Number of domain clones (default=1) + """ + + out_folders: str = os.getcwd() + sim_folder: str = "sim_1" + restart: bool = False + max_runtime: int = 300 + save_step: int = 1 + sort_step: int = 0 + num_clones: int = 1 + + +def check_option(opt, options): + """Check if opt is contained in options; if opt is a list, checks for each element.""" + opts = get_args(options) + if not isinstance(opt, list): + opt = [opt] + for o in opt: + assert o in opts, f"Option '{o}' is not in {opts}." diff --git a/src/struphy/main.py b/src/struphy/main.py index 9d053f854..9890e7662 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -7,6 +7,7 @@ import numpy as np from mpi4py import MPI from pyevtk.hl import gridToVTK +import shutil from struphy.feec.psydac_derham import SplineFunction from struphy.fields_background.base import FluidEquilibriumWithB @@ -20,56 +21,42 @@ from struphy.pic.base import Particles from struphy.models.species import Species from struphy.models.variables import FEECVariable +from struphy.io.options import Units, Time, MetaOptions def main( - model_name: Optional[str], - params_path: Optional[str], - path_out: str, + model: StruphyModel, + params_path: str, *, - restart: bool = False, - runtime: int = 300, - save_step: int = 1, + path_out: str = None, + units: Units = None, + time_opts: Time = None, + domain = None, + equil = None, + grid = None, + derham = None, + meta: MetaOptions = None, verbose: bool = False, - supress_out: bool = False, - sort_step: int = 0, - num_clones: int = 1, ): """ Run a Struphy model. Parameters ---------- - model_name : str - The name of the model to run. Type "struphy run --help" in your terminal to see a list of available models. + model : StruphyModel + The model to run. Check https://struphy.pages.mpcdf.de/struphy/sections/models.html for available models. params_path : str - Path to .py parameter file. - + Absolute path to .py parameter file. + path_out : str - The output directory. Will create a folder if it does not exist OR cleans the folder for new runs. - - restart : bool, optional - Whether to restart a run (default=False). - - runtime : int, optional - Maximum run time of simulation in minutes. Will finish the time integration once this limit is reached (default=300). - - save_step : int, optional - When to save data output: every time step (save_step=1), every second time step (save_step=2), etc (default=1). + The output directory (default is 'sim_1' in cwd). Will create a folder if it does not exist OR cleans the folder for new runs. - verbose : bool - Show full screen output. - - supress_out : bool - Whether to supress screen output during time integration. - - sort_step: int, optional - Sort markers in memory every N time steps (default=0, which means markers are sorted only at the start of simulation) - - num_clones: int, optional - Number of domain clones (default=1) """ + + # check model + assert hasattr(model, "propagators"), "Attribute 'self.propagators' must be set in model __init__!" + model_name = model.__class__.__name__ comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -77,43 +64,48 @@ def main( start_simulation = time.time() - # load paramaters and set defaults - params = setup_parameters(params_path=params_path, - path_out=path_out, - verbose=verbose,) - - # check model - model = params.model - assert hasattr(model, "propagators"), "Attribute 'self.propagators' must be set in model __init__!" + # meta-data + out_folders = meta.out_folders + sim_folder = meta.sim_folder + path_out = os.path.join(out_folders, sim_folder) - if model_name is None: - assert model is not None, "If model is not specified, then model: MODEL must be specified in the params!" - model_name = model.__class__.__name__ + restart = meta.restart + max_runtime = meta.max_runtime + save_step = meta.save_step + sort_step = meta.sort_step + num_clones = meta.num_clones - # meta-data - meta = {} - meta["platform"] = sysconfig.get_platform() - meta["python version"] = sysconfig.get_python_version() - meta["model name"] = model_name - meta["parameter file"] = params_path - meta["output folder"] = path_out - meta["MPI processes"] = size - meta["number of domain clones"] = num_clones - meta["restart"] = restart - meta["max wall-clock [min]"] = runtime - meta["save interval [steps]"] = save_step + meta_dct = {} + meta_dct["platform"] = sysconfig.get_platform() + meta_dct["python version"] = sysconfig.get_python_version() + meta_dct["model name"] = model_name + meta_dct["parameter file"] = params_path + meta_dct["output folder"] = path_out + meta_dct["MPI processes"] = size + meta_dct["number of domain clones"] = num_clones + meta_dct["restart"] = restart + meta_dct["max wall-clock [min]"] = max_runtime + meta_dct["save interval [steps]"] = save_step print("\nMETADATA:") - for k, v in meta.items(): + for k, v in meta_dct.items(): print(f'{k}:'.ljust(25), v) - - dict_to_yaml(meta, os.path.join(path_out, "meta.yml")) - + # creating output folders setup_folders(path_out=path_out, restart=restart, verbose=verbose,) + # save meta-data + dict_to_yaml(meta_dct, os.path.join(path_out, "meta.yml")) + + # save parameter file + if rank == 0: + shutil.copy2( + params_path, + os.path.join(path_out, "parameters.py"), + ) + # config clones if comm is None: clone_config = None @@ -125,25 +117,18 @@ def main( # MPI.COMM_WORLD : comm # within a clone: : sub_comm # between the clones : inter_comm - clone_config = CloneConfig(comm=comm, params=params, num_clones=num_clones) + clone_config = CloneConfig(comm=comm, params=None, num_clones=num_clones) clone_config.print_clone_config() if model.kinetic_species: clone_config.print_particle_config() + + model.clone_config = clone_config comm.Barrier() ## configure model instance - - # mpi config - model._comm_world = comm - model._clone_config = clone_config - - if model.comm_world is None: - model._rank_world = 0 - else: - model._rank_world = model.comm_world.Get_rank() # units - model._units = params.units + model.units = units if model.bulk_species is None: A_bulk = None Z_bulk = None @@ -156,11 +141,11 @@ def main( verbose=verbose,) # domain and fluid bckground - model.setup_domain_and_equil(params.domain, params.equil) + model.setup_domain_and_equil(domain, equil) # allocate derham-related objects - if params.grid is not None: - model.allocate_feec(params.grid, params.derham) + if grid is not None: + model.allocate_feec(grid, derham) else: model._derham = None model._mass_ops = None @@ -177,11 +162,7 @@ def main( # plasma parameters model.compute_plasma_params(verbose=verbose) model.setup_equation_params(units=model.units, verbose=verbose) - - if model_name is None: - assert model is not None, "If model is not specified, then model: MODEL must be specified in the params!" - model_name = model.__class__.__name__ - + if rank < 32: if rank == 0: print("") @@ -231,9 +212,9 @@ def main( data.add_data({key_time_restart: val}) # retrieve time parameters - dt = params.time.dt - Tend = params.time.Tend - split_algo = params.time.split_algo + dt = time_opts.dt + Tend = time_opts.Tend + split_algo = time_opts.split_algo # set initial conditions for all variables if restart: @@ -271,7 +252,7 @@ def main( # stop time loop? break_cond_1 = time_state["value"][0] >= Tend - break_cond_2 = run_time_now > runtime + break_cond_2 = run_time_now > max_runtime if break_cond_1 or break_cond_2: # save restart data (other data already saved below) @@ -292,7 +273,7 @@ def main( if isinstance(val, Particles): val.do_sort() t1 = time.time() - if rank == 0 and not supress_out: + if rank == 0 and verbose: message = "Particles sorted | wall clock [s]: {0:8.4f} | sorting duration [s]: {1:8.4f}".format( run_time_now * 60, t1 - t0 ) @@ -333,7 +314,7 @@ def main( data.save_data(keys=save_keys_all) # print current time and scalar quantities to screen - if rank == 0 and not supress_out: + if rank == 0 and verbose: step = str(time_state["index"][0]).zfill(len(total_steps)) message = "time step: " + step + "/" + str(total_steps) @@ -418,9 +399,9 @@ def main( action="store_true", ) - # runtime + # max_runtime parser.add_argument( - "--runtime", + "--max-runtime", type=int, metavar="N", help="maximum wall-clock time of program in minutes (default=300)", @@ -462,13 +443,6 @@ def main( action="store_true", ) - # supress screen output - parser.add_argument( - "--supress-out", - help="supress screen output during time integration", - action="store_true", - ) - parser.add_argument( "--likwid", help="run with Likwid", @@ -511,7 +485,6 @@ def main( runtime=args.runtime, save_step=args.save_step, verbose=args.verbose, - supress_out=args.supress_out, sort_step=args.sort_step, num_clones=args.nclones, ) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 224fe975e..711b8ab72 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -180,7 +180,7 @@ def allocate_feec(self, # create discrete derham sequence if clone_config is None: - derham_comm = self.comm_world + derham_comm = MPI.COMM_WORLD else: derham_comm = clone_config.sub_comm @@ -261,20 +261,15 @@ def equation_params(self): """Parameters appearing in model equation due to Struphy normalization.""" return self._equation_params - @property - def comm_world(self): - """MPI_COMM_WORLD communicator.""" - return self._comm_world - - @property - def rank_world(self): - """Global rank.""" - return self._rank_world - @property def clone_config(self): """Config in case domain clones are used.""" return self._clone_config + + @clone_config.setter + def clone_config(self, new): + assert isinstance(new, CloneConfig) or new is None + self._clone_config = new @property def diagnostics(self): @@ -305,6 +300,11 @@ def projected_equil(self): def units(self) -> Units: """All Struphy units.""" return self._units + + @units.setter + def units(self, new) : + assert isinstance(new, Units) + self._units = new @property def mass_ops(self): @@ -1377,7 +1377,7 @@ def generate_default_parameter_file( # generic options for all models - file.write("from struphy.io.options import Units, Time\n") + file.write("from struphy.io.options import Units, Time, MetaOptions\n") file.write("from struphy.geometry import domains\n") file.write("from struphy.fields_background import equils\n") file.write("from struphy.initial import perturbations\n") @@ -1398,19 +1398,21 @@ def generate_default_parameter_file( elif isinstance(var, SPHVariable): has_sph = True - if has_feec: - file.write("from struphy.topology import grids\n") - file.write("from struphy.io.options import DerhamOptions\n") - file.write("from struphy.io.options import FieldsBackground\n") + # if has_feec: + file.write("from struphy.topology import grids\n") + file.write("from struphy.io.options import DerhamOptions\n") + file.write("from struphy.io.options import FieldsBackground\n") - if has_pic or has_sph: - file.write("from struphy.kinetic_background import maxwellians\n") + # if has_pic or has_sph: + file.write("from struphy.kinetic_background import maxwellians\n") + file.write("from struphy import main\n") - file.write("\n# import model\n") + file.write("\n# import model, set verbosity\n") file.write(f"from {self.__module__} import {self.__class__.__name__} as Model\n") + file.write("verbose = False\n") - file.write("\n# light-weight model instance\n") - file.write("model = Model()\n") + file.write("\n# meta options\n") + file.write("meta = MetaOptions()\n") file.write("\n# units\n") file.write("units = Units()\n") @@ -1425,11 +1427,20 @@ def generate_default_parameter_file( file.write("equil = equils.HomogenSlab()\n") if has_feec: - file.write("\n# grid\n") - file.write("grid = grids.TensorProductGrid()\n") + grid = "grid = grids.TensorProductGrid()\n" + derham = "derham = DerhamOptions()\n" + else: + grid = "grid = None\n" + derham = "derham = None\n" - file.write("\n# derham options\n") - file.write("derham = DerhamOptions()\n") + file.write("\n# grid\n") + file.write(grid) + + file.write("\n# derham options\n") + file.write(derham) + + file.write("\n# light-weight model instance\n") + file.write("model = Model()\n") file.write("\n# propagator options\n") for prop in self.propagators.__dict__: @@ -1443,6 +1454,19 @@ def generate_default_parameter_file( file.write("\n# optional: exclude variables from saving\n") file.write(exclude_feec) + file.write("\n# start run\n") + file.write("main.main(model, \n\ + params_path=__file__, \n\ + units=units, \n\ + time_opts=time, \n\ + domain=domain, \n\ + equil=equil, \n\ + grid=grid, \n\ + derham=derham, \n\ + meta=meta, \n\ + verbose=verbose, \n\ + )") + print(f"Default parameter file for '{self.__class__.__name__}' has been created.\n\ You can now launch with 'struphy run {self.__class__.__name__}' or with 'struphy run -i params_{self.__class__.__name__}.py'") From 95b6e17721d4a59afd47a70659033feada348616 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 4 Aug 2025 09:22:14 +0200 Subject: [PATCH 054/292] renam main.py -> struphy.py and main() -> run() --- src/struphy/io/inp/params_Maxwell.py | 24 ++++++++-------- src/struphy/models/base.py | 24 ++++++++-------- src/struphy/models/tests/util.py | 12 ++++---- src/struphy/{main.py => struphy.py} | 4 +-- src/struphy/tutorials/tests/test_tutorials.py | 28 +++++++++---------- 5 files changed, 46 insertions(+), 46 deletions(-) rename src/struphy/{main.py => struphy.py} (99%) diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index 5caef842e..f1d968e99 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -6,7 +6,7 @@ from struphy.io.options import DerhamOptions from struphy.io.options import FieldsBackground from struphy.kinetic_background import maxwellians -from struphy import main +from struphy import struphy # import model, set verbosity from struphy.models.toy import Maxwell as Model @@ -47,14 +47,14 @@ # model.em_fields.b_field.save_data = False # start run -main.main(model, - params_path=__file__, - units=units, - time_opts=time, - domain=domain, - equil=equil, - grid=grid, - derham=derham, - meta=meta, - verbose=verbose, - ) \ No newline at end of file +struphy.run(model, + params_path=__file__, + units=units, + time_opts=time, + domain=domain, + equil=equil, + grid=grid, + derham=derham, + meta=meta, + verbose=verbose, + ) \ No newline at end of file diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 711b8ab72..64e2da5e9 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1405,7 +1405,7 @@ def generate_default_parameter_file( # if has_pic or has_sph: file.write("from struphy.kinetic_background import maxwellians\n") - file.write("from struphy import main\n") + file.write("from struphy import struphy\n") file.write("\n# import model, set verbosity\n") file.write(f"from {self.__module__} import {self.__class__.__name__} as Model\n") @@ -1455,17 +1455,17 @@ def generate_default_parameter_file( file.write(exclude_feec) file.write("\n# start run\n") - file.write("main.main(model, \n\ - params_path=__file__, \n\ - units=units, \n\ - time_opts=time, \n\ - domain=domain, \n\ - equil=equil, \n\ - grid=grid, \n\ - derham=derham, \n\ - meta=meta, \n\ - verbose=verbose, \n\ - )") + file.write("struphy.run(model, \n\ + params_path=__file__, \n\ + units=units, \n\ + time_opts=time, \n\ + domain=domain, \n\ + equil=equil, \n\ + grid=grid, \n\ + derham=derham, \n\ + meta=meta, \n\ + verbose=verbose, \n\ + )") print(f"Default parameter file for '{self.__class__.__name__}' has been created.\n\ You can now launch with 'struphy run {self.__class__.__name__}' or with 'struphy run -i params_{self.__class__.__name__}.py'") diff --git a/src/struphy/models/tests/util.py b/src/struphy/models/tests/util.py index 3f6c20c3a..bbd330088 100644 --- a/src/struphy/models/tests/util.py +++ b/src/struphy/models/tests/util.py @@ -9,7 +9,7 @@ import struphy from struphy.console.main import recursive_get_files from struphy.io.setup import descend_options_dict -from struphy.main import main +from struphy.struphy import run from struphy.models.base import StruphyModel libpath = struphy.__path__[0] @@ -181,7 +181,7 @@ def call_test( parameters["time"]["Tend"] = Tend if MPI.COMM_WORLD.Get_rank() == 0: print_test_params(parameters) - main( + run( model_name, parameters, path_out, @@ -196,7 +196,7 @@ def call_test( # run with default if MPI.COMM_WORLD.Get_rank() == 0: print_test_params(parameters) - main( + run( model_name, parameters, path_out, @@ -230,7 +230,7 @@ def call_test( test_list += [opt] if MPI.COMM_WORLD.Get_rank() == 0: print_test_params(parameters) - main( + run( model_name, parameters, path_out, @@ -252,7 +252,7 @@ def call_test( test_list += [opt] if MPI.COMM_WORLD.Get_rank() == 0: print_test_params(parameters) - main( + run( model_name, parameters, path_out, @@ -274,7 +274,7 @@ def call_test( test_list += [opt] if MPI.COMM_WORLD.Get_rank() == 0: print_test_params(parameters) - main( + run( model_name, parameters, path_out, diff --git a/src/struphy/main.py b/src/struphy/struphy.py similarity index 99% rename from src/struphy/main.py rename to src/struphy/struphy.py index 9890e7662..07b5e7ed6 100644 --- a/src/struphy/main.py +++ b/src/struphy/struphy.py @@ -24,7 +24,7 @@ from struphy.io.options import Units, Time, MetaOptions -def main( +def run( model: StruphyModel, params_path: str, *, @@ -477,7 +477,7 @@ def main( pylikwid_markerinit() with ProfileManager.profile_region("main"): # solve the model - main( + run( args.model, args.input, args.output, diff --git a/src/struphy/tutorials/tests/test_tutorials.py b/src/struphy/tutorials/tests/test_tutorials.py index cf1292b1b..cfa726aee 100644 --- a/src/struphy/tutorials/tests/test_tutorials.py +++ b/src/struphy/tutorials/tests/test_tutorials.py @@ -5,7 +5,7 @@ from mpi4py import MPI import struphy -from struphy.main import main +from struphy.struphy import run from struphy.post_processing import pproc_struphy comm = MPI.COMM_WORLD @@ -18,7 +18,7 @@ @pytest.mark.mpi(min_size=2) def test_tutorial_02(): - main( + run( "LinearMHDVlasovCC", os.path.join(i_path, "tutorials", "params_02.yml"), os.path.join(o_path, "tutorial_02"), @@ -28,7 +28,7 @@ def test_tutorial_02(): @pytest.mark.mpi(min_size=2) def test_tutorial_03(): - main( + run( "LinearMHD", os.path.join(i_path, "tutorials", "params_03.yml"), os.path.join(o_path, "tutorial_03"), @@ -42,7 +42,7 @@ def test_tutorial_03(): @pytest.mark.mpi(min_size=2) def test_tutorial_04(fast): - main( + run( "Maxwell", os.path.join(i_path, "tutorials", "params_04a.yml"), os.path.join(o_path, "tutorial_04a"), @@ -53,7 +53,7 @@ def test_tutorial_04(fast): if rank == 0: pproc_struphy.main(os.path.join(o_path, "tutorial_04a")) - main( + run( "LinearMHD", os.path.join(i_path, "tutorials", "params_04b.yml"), os.path.join(o_path, "tutorial_04b"), @@ -65,7 +65,7 @@ def test_tutorial_04(fast): pproc_struphy.main(os.path.join(o_path, "tutorial_04b")) if not fast: - main( + run( "VariationalMHD", os.path.join(i_path, "tutorials", "params_04c.yml"), os.path.join(o_path, "tutorial_04c"), @@ -78,7 +78,7 @@ def test_tutorial_04(fast): def test_tutorial_05(): - main( + run( "Vlasov", os.path.join(i_path, "tutorials", "params_05a.yml"), os.path.join(o_path, "tutorial_05a"), @@ -89,7 +89,7 @@ def test_tutorial_05(): if rank == 0: pproc_struphy.main(os.path.join(o_path, "tutorial_05a")) - main( + run( "Vlasov", os.path.join(i_path, "tutorials", "params_05b.yml"), os.path.join(o_path, "tutorial_05b"), @@ -100,7 +100,7 @@ def test_tutorial_05(): if rank == 0: pproc_struphy.main(os.path.join(o_path, "tutorial_05b")) - main( + run( "GuidingCenter", os.path.join(i_path, "tutorials", "params_05c.yml"), os.path.join(o_path, "tutorial_05c"), @@ -111,7 +111,7 @@ def test_tutorial_05(): if rank == 0: pproc_struphy.main(os.path.join(o_path, "tutorial_05c")) - main( + run( "GuidingCenter", os.path.join(i_path, "tutorials", "params_05d.yml"), os.path.join(o_path, "tutorial_05d"), @@ -122,7 +122,7 @@ def test_tutorial_05(): if rank == 0: pproc_struphy.main(os.path.join(o_path, "tutorial_05d")) - main( + run( "GuidingCenter", os.path.join(i_path, "tutorials", "params_05e.yml"), os.path.join(o_path, "tutorial_05e"), @@ -133,7 +133,7 @@ def test_tutorial_05(): if rank == 0: pproc_struphy.main(os.path.join(o_path, "tutorial_05e")) - main( + run( "GuidingCenter", os.path.join(i_path, "tutorials", "params_05f.yml"), os.path.join(o_path, "tutorial_05f"), @@ -146,7 +146,7 @@ def test_tutorial_05(): def test_tutorial_12(): - main( + run( "Vlasov", os.path.join(i_path, "tutorials", "params_12a.yml"), os.path.join(o_path, "tutorial_12a"), @@ -158,7 +158,7 @@ def test_tutorial_12(): if rank == 0: pproc_struphy.main(os.path.join(o_path, "tutorial_12a")) - main( + run( "GuidingCenter", os.path.join(i_path, "tutorials", "params_12b.yml"), os.path.join(o_path, "tutorial_12b"), From b76d94ea8d7a88ecb3983ed8a548f997e1d8e876 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 4 Aug 2025 09:43:37 +0200 Subject: [PATCH 055/292] rename MetaOptions -> EnvironmentOptions --- src/struphy/io/inp/params_Maxwell.py | 4 +- src/struphy/io/options.py | 4 +- src/struphy/struphy.py | 67 ++++++++++++++-------------- 3 files changed, 38 insertions(+), 37 deletions(-) diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index f1d968e99..624590225 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -1,4 +1,4 @@ -from struphy.io.options import Units, Time, MetaOptions +from struphy.io.options import Units, Time, EnvironmentOptions from struphy.geometry import domains from struphy.fields_background import equils from struphy.initial import perturbations @@ -13,7 +13,7 @@ verbose = False # meta options -meta = MetaOptions() +meta = EnvironmentOptions() # units units = Units() diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index d4f60b268..ad44a4dd6 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -253,8 +253,8 @@ def __post_init__(self): @dataclass -class MetaOptions: - """Meta options for launching run on current architecture +class EnvironmentOptions: + """Environment options for launching run on current architecture (these options do not influence the simulation result). Parameters diff --git a/src/struphy/struphy.py b/src/struphy/struphy.py index 07b5e7ed6..8a94ce0ac 100644 --- a/src/struphy/struphy.py +++ b/src/struphy/struphy.py @@ -21,21 +21,26 @@ from struphy.pic.base import Particles from struphy.models.species import Species from struphy.models.variables import FEECVariable -from struphy.io.options import Units, Time, MetaOptions +from struphy.io.options import Units, Time, EnvironmentOptions +from struphy.geometry.base import Domain +from struphy.geometry.domains import Cuboid +from struphy.fields_background.base import FluidEquilibrium +from struphy.fields_background.equils import HomogenSlab +from struphy.topology.grids import TensorProductGrid +from struphy.io.options import DerhamOptions def run( model: StruphyModel, params_path: str, *, - path_out: str = None, - units: Units = None, - time_opts: Time = None, - domain = None, - equil = None, - grid = None, - derham = None, - meta: MetaOptions = None, + env: EnvironmentOptions = EnvironmentOptions(), + units: Units = Units(), + time_opts: Time = Time(), + domain: Domain = Cuboid(), + equil: FluidEquilibrium = HomogenSlab(), + grid: TensorProductGrid = None, + derham: DerhamOptions = None, verbose: bool = False, ): """ @@ -48,10 +53,6 @@ def run( params_path : str Absolute path to .py parameter file. - - path_out : str - The output directory (default is 'sim_1' in cwd). Will create a folder if it does not exist OR cleans the folder for new runs. - """ # check model @@ -65,30 +66,30 @@ def run( start_simulation = time.time() # meta-data - out_folders = meta.out_folders - sim_folder = meta.sim_folder + out_folders = env.out_folders + sim_folder = env.sim_folder path_out = os.path.join(out_folders, sim_folder) - restart = meta.restart - max_runtime = meta.max_runtime - save_step = meta.save_step - sort_step = meta.sort_step - num_clones = meta.num_clones + restart = env.restart + max_runtime = env.max_runtime + save_step = env.save_step + sort_step = env.sort_step + num_clones = env.num_clones - meta_dct = {} - meta_dct["platform"] = sysconfig.get_platform() - meta_dct["python version"] = sysconfig.get_python_version() - meta_dct["model name"] = model_name - meta_dct["parameter file"] = params_path - meta_dct["output folder"] = path_out - meta_dct["MPI processes"] = size - meta_dct["number of domain clones"] = num_clones - meta_dct["restart"] = restart - meta_dct["max wall-clock [min]"] = max_runtime - meta_dct["save interval [steps]"] = save_step + meta = {} + meta["platform"] = sysconfig.get_platform() + meta["python version"] = sysconfig.get_python_version() + meta["model name"] = model_name + meta["parameter file"] = params_path + meta["output folder"] = path_out + meta["MPI processes"] = size + meta["number of domain clones"] = num_clones + meta["restart"] = restart + meta["max wall-clock [min]"] = max_runtime + meta["save interval [steps]"] = save_step print("\nMETADATA:") - for k, v in meta_dct.items(): + for k, v in meta.items(): print(f'{k}:'.ljust(25), v) # creating output folders @@ -97,7 +98,7 @@ def run( verbose=verbose,) # save meta-data - dict_to_yaml(meta_dct, os.path.join(path_out, "meta.yml")) + dict_to_yaml(meta, os.path.join(path_out, "meta.yml")) # save parameter file if rank == 0: From ad4f3feaee4873cedf0a8fd9c812779cd63d5e27 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 4 Aug 2025 14:02:53 +0200 Subject: [PATCH 056/292] New framework for model tests: - base class in models/tests/bast_test.py - first subclass `LightWaveDispersion` in models/tests/test_toy.py - a test must be allocated and then passed to struphy.run() - new defaults for Cuboid mapping (unit cube) --- src/struphy/feec/psydac_derham.py | 18 ++------ src/struphy/geometry/domains.py | 26 +++-------- src/struphy/initial/perturbations.py | 1 + src/struphy/io/inp/params_Maxwell.py | 64 ++++++++++++++------------- src/struphy/models/base.py | 35 +++++++-------- src/struphy/models/tests/base_test.py | 38 ++++++++++++++++ src/struphy/models/tests/test_toy.py | 34 ++++++++++++++ src/struphy/models/toy.py | 3 +- src/struphy/struphy.py | 34 ++++++++++---- 9 files changed, 161 insertions(+), 92 deletions(-) create mode 100644 src/struphy/models/tests/base_test.py create mode 100644 src/struphy/models/tests/test_toy.py diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 4bab49a49..4dda194fb 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1760,20 +1760,10 @@ def f_tmp(e1, e2, e3): # special case of white noise in logical space for different components if isinstance(ptb, Noise): # set white noise FE coefficients - if self.space_id in {"H1", "L2"}: - assert isinstance(ptb.given_in_basis, str) - self._add_noise(direction=ptb.direction, - amp=ptb.amp, - seed=ptb.seed,) - elif self.space_id in {"Hcurl", "Hdiv", "H1vec"}: - assert isinstance(ptb.given_in_basis, list) - assert len(ptb.given_in_basis) == 3 - for n, comp in enumerate(ptb.given_in_basis): - if comp is not None: - self._add_noise(direction=ptb.direction, - amp=ptb.amp, - seed=ptb.seed, - n=n) + self._add_noise(direction=ptb.direction, + amp=ptb.amp, + seed=ptb.seed, + n=ptb.comp,) # perturbation class elif isinstance(ptb, Perturbation): if self.space_id in {"H1", "L2"}: diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index 8b18a56c1..ea7f543d6 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -533,37 +533,23 @@ class Cuboid(Domain): l1 : float Start of x-interval (default: 0.). r1 : float - End of x-interval, r1>l1 (default: 2.). + End of x-interval, r1>l1 (default: 1.). l2 : float Start of y-interval (default: 0.). r2 : float - End of y-interval, r2>l2 (default: 3.). + End of y-interval, r2>l2 (default: 1.). l3 : float Start of z-interval (default: 0.). r3 : float - End of z-interval, r3>l3 (default: 6.). - - Note - ---- - In the parameter .yml, use the following in the section `geometry`:: - - geometry : - type : Cuboid - Cuboid : - l1 : 0. # start of x-interval - r1 : 2. # end of x-interval, r1>l1 - l2 : 0. # start of y-interval - r2 : 2. # end of y-interval, r2>l2 - l3 : 0. # start of z-interval - r3 : 1. # end of z-interval, r3>l3 + End of z-interval, r3>l3 (default: 1.). """ def __init__(self, l1: float = 0.0, - r1: float = 2.0, + r1: float = 1.0, l2: float = 0.0, - r2: float = 3.0, + r2: float = 1.0, l3: float = 0.0, - r3: float = 6.0,): + r3: float = 1.0,): self._kind_map = 10 diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 2ea51bee3..58931822a 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -28,6 +28,7 @@ class Noise(Perturbation): amp: float = 0.0001 seed: int = None comp: int = 0 + given_in_basis: GivenInBasis = "0" def __post_init__(self,): check_option(self.direction, NoiseDirections) diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index 624590225..b36497884 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -1,60 +1,64 @@ -from struphy.io.options import Units, Time, EnvironmentOptions +from struphy.io.options import EnvironmentOptions, Units, Time from struphy.geometry import domains from struphy.fields_background import equils -from struphy.initial import perturbations from struphy.topology import grids from struphy.io.options import DerhamOptions from struphy.io.options import FieldsBackground +from struphy.initial import perturbations from struphy.kinetic_background import maxwellians from struphy import struphy # import model, set verbosity from struphy.models.toy import Maxwell as Model -verbose = False +verbose = True -# meta options -meta = EnvironmentOptions() +# environment options +env = EnvironmentOptions() -# units -units = Units() +# # units +# units = Units() -# time stepping -time = Time() +# # time stepping +# time = Time() -# geometry -domain = domains.Cuboid() +# # geometry +# domain = domains.Cuboid() -# fluid equilibrium (can be used as part of initial conditions) -equil = equils.HomogenSlab() +# # fluid equilibrium (can be used as part of initial conditions) +# equil = equils.HomogenSlab() -# grid -grid = grids.TensorProductGrid() +# # grid +# grid = grids.TensorProductGrid() -# derham options -derham = DerhamOptions() +# # derham options +# derham_opts = DerhamOptions() # light-weight model instance model = Model() -# propagator options -model.propagators.maxwell.set_options() +# # propagator options +# model.propagators.maxwell.set_options() -# initial conditions (background + perturbation) -model.em_fields.b_field.add_background(FieldsBackground()) -model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos()) +# # initial conditions (background + perturbation) +# model.em_fields.b_field.add_background(FieldsBackground()) +# model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos()) # optional: exclude variables from saving # model.em_fields.b_field.save_data = False +from struphy.models.tests import test_toy +test = test_toy.LightWaveDispersion(model) + # start run struphy.run(model, - params_path=__file__, - units=units, - time_opts=time, - domain=domain, - equil=equil, - grid=grid, - derham=derham, - meta=meta, + __file__, + env=env, + # units=units, + # time_opts=time, + # domain=domain, + # equil=equil, + # grid=grid, + # derham_opts=derham_opts, verbose=verbose, + test=test, ) \ No newline at end of file diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 64e2da5e9..87491b2e2 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -27,8 +27,7 @@ from struphy.utils.utils import dict_to_yaml, read_state from struphy.models.species import Species, FieldSpecies, FluidSpecies, KineticSpecies, DiagnosticSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable -from struphy.io.options import Units -from struphy.io.options import Time, DerhamOptions +from struphy.io.options import Units, Time, DerhamOptions from struphy.topology.grids import TensorProductGrid from struphy.geometry.base import Domain from struphy.geometry.domains import Cuboid @@ -173,7 +172,7 @@ def species(self): def allocate_feec(self, grid: TensorProductGrid, - derham_params: DerhamOptions, + derham_opts: DerhamOptions, comm: MPI.Intracomm = None, clone_config: CloneConfig = None, ): @@ -186,7 +185,7 @@ def allocate_feec(self, self._derham = setup_derham( grid, - derham_params, + derham_opts, comm=derham_comm, domain=self.domain, verbose=self.verbose, @@ -216,6 +215,8 @@ def allocate_feec(self, self.equil, self.derham, ) + else: + self._projected_equil = None def allocate_propagators(self): # set propagators base class attributes (then available to all propagators) @@ -235,7 +236,7 @@ def allocate_propagators(self): for prop in self.prop_list: assert isinstance(prop, Propagator) prop.allocate() - if self.rank_world == 0: + if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nAllocated propagator '{prop.__class__.__name__}'.") @staticmethod @@ -1375,12 +1376,9 @@ def generate_default_parameter_file( print("exiting ...") return - - # generic options for all models - file.write("from struphy.io.options import Units, Time, MetaOptions\n") + file.write("from struphy.io.options import EnvironmentOptions, Units, Time\n") file.write("from struphy.geometry import domains\n") file.write("from struphy.fields_background import equils\n") - file.write("from struphy.initial import perturbations\n") has_feec = False has_pic = False @@ -1398,12 +1396,11 @@ def generate_default_parameter_file( elif isinstance(var, SPHVariable): has_sph = True - # if has_feec: file.write("from struphy.topology import grids\n") file.write("from struphy.io.options import DerhamOptions\n") file.write("from struphy.io.options import FieldsBackground\n") + file.write("from struphy.initial import perturbations\n") - # if has_pic or has_sph: file.write("from struphy.kinetic_background import maxwellians\n") file.write("from struphy import struphy\n") @@ -1411,8 +1408,8 @@ def generate_default_parameter_file( file.write(f"from {self.__module__} import {self.__class__.__name__} as Model\n") file.write("verbose = False\n") - file.write("\n# meta options\n") - file.write("meta = MetaOptions()\n") + file.write("\n# environment options\n") + file.write("env = EnvironmentOptions()\n") file.write("\n# units\n") file.write("units = Units()\n") @@ -1428,10 +1425,10 @@ def generate_default_parameter_file( if has_feec: grid = "grid = grids.TensorProductGrid()\n" - derham = "derham = DerhamOptions()\n" + derham = "derham_opts = DerhamOptions()\n" else: grid = "grid = None\n" - derham = "derham = None\n" + derham = "derham_opts = None\n" file.write("\n# grid\n") file.write(grid) @@ -1456,14 +1453,14 @@ def generate_default_parameter_file( file.write("\n# start run\n") file.write("struphy.run(model, \n\ - params_path=__file__, \n\ + __file__, \n\ + env=env, \n\ units=units, \n\ time_opts=time, \n\ domain=domain, \n\ equil=equil, \n\ grid=grid, \n\ - derham=derham, \n\ - meta=meta, \n\ + derham_opts=derham_opts, \n\ verbose=verbose, \n\ )") @@ -1679,7 +1676,7 @@ def compute_plasma_params(self, verbose=True): magnetic_field = np.nan # print("\n+++++++ WARNING +++++++ magnetic field is zero - set to nan !!") - if verbose and self.rank_world == 0: + if verbose and MPI.COMM_WORLD.Get_rank() == 0: print("\nPLASMA PARAMETERS:") print( f"Plasma volume:".ljust(25), diff --git a/src/struphy/models/tests/base_test.py b/src/struphy/models/tests/base_test.py new file mode 100644 index 000000000..a87396b76 --- /dev/null +++ b/src/struphy/models/tests/base_test.py @@ -0,0 +1,38 @@ +from abc import ABCMeta, abstractmethod +from dataclasses import dataclass + +from struphy.io.options import EnvironmentOptions, Units, Time, DerhamOptions +from struphy.topology.grids import TensorProductGrid +from struphy.geometry.base import Domain +from struphy.geometry.domains import Cuboid +from struphy.fields_background.equils import HomogenSlab +from struphy.pic import particles +from struphy.pic.base import Particles +from struphy.propagators.base import Propagator +from struphy.kinetic_background import maxwellians +from struphy.fields_background.base import FluidEquilibrium +from struphy.models.base import StruphyModel + + +@dataclass +class ModelTest(metaclass=ABCMeta): + """Base class for model verification tests.""" + model: StruphyModel = None + units: Units = None + time_opts: Time = None + domain: Domain = None + equil: FluidEquilibrium = None + grid: TensorProductGrid = None + derham_opts: DerhamOptions = None + + @abstractmethod + def set_propagator_opts(self): + """Set propagator options for test.""" + + @abstractmethod + def set_variables_inits(self): + """Set initial conditions for test.""" + + @abstractmethod + def verification_script(self): + """Contains assert statements (and plots) that verify the test.""" \ No newline at end of file diff --git a/src/struphy/models/tests/test_toy.py b/src/struphy/models/tests/test_toy.py new file mode 100644 index 000000000..9e55b8ffa --- /dev/null +++ b/src/struphy/models/tests/test_toy.py @@ -0,0 +1,34 @@ +from struphy.models.tests.base_test import ModelTest +from dataclasses import dataclass + +from struphy.models.toy import Maxwell +from struphy.io.options import EnvironmentOptions, Units, Time, DerhamOptions +from struphy.geometry.base import Domain +from struphy.fields_background.base import FluidEquilibrium +from struphy.topology.grids import TensorProductGrid +from struphy.geometry.domains import Cuboid +from struphy.fields_background.equils import HomogenSlab +from struphy.initial import perturbations + + +class LightWaveDispersion(ModelTest): + """Simulate light wave dispersion relation in Cartesian coordinates (Cuboid mapping).""" + + def __init__(self, model: Maxwell): + self.model: Maxwell = model + self.units: Units = Units() + self.time_opts: Time = Time(dt=0.05, Tend=100.0 - 1e-6) + self.domain: Domain = Cuboid(r3=20.0) + self.equil: FluidEquilibrium = None + self.grid: TensorProductGrid = TensorProductGrid(Nel=(1, 1, 128)) + self.derham_opts: DerhamOptions = DerhamOptions(p=(1, 1, 3)) + + def set_propagator_opts(self): + self.model.propagators.maxwell.set_options() + + def set_variables_inits(self): + self.model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0)) + self.model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1)) + + def verification_script(self): + assert -1 < 0 \ No newline at end of file diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 93b69b271..c00283e5b 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,6 +1,7 @@ import numpy as np from dataclasses import dataclass +import inspect from struphy.models.base import StruphyModel from struphy.propagators.base import Propagator @@ -83,7 +84,7 @@ def update_scalar_quantities(self): self.update_scalar("electric energy", en_E) self.update_scalar("magnetic energy", en_B) self.update_scalar("total energy", en_E + en_B) - + class Vlasov(StruphyModel): r"""Vlasov equation in static background magnetic field. diff --git a/src/struphy/struphy.py b/src/struphy/struphy.py index 8a94ce0ac..9966e9acd 100644 --- a/src/struphy/struphy.py +++ b/src/struphy/struphy.py @@ -13,7 +13,6 @@ from struphy.fields_background.base import FluidEquilibriumWithB from struphy.io.output_handling import DataContainer from struphy.io.setup import setup_folders, setup_parameters -from struphy.models import fluid, hybrid, kinetic, toy from struphy.models.base import StruphyModel from struphy.profiling.profiling import ProfileManager from struphy.utils.clone_config import CloneConfig @@ -28,6 +27,7 @@ from struphy.fields_background.equils import HomogenSlab from struphy.topology.grids import TensorProductGrid from struphy.io.options import DerhamOptions +from struphy.models.tests.base_test import ModelTest def run( @@ -40,8 +40,9 @@ def run( domain: Domain = Cuboid(), equil: FluidEquilibrium = HomogenSlab(), grid: TensorProductGrid = None, - derham: DerhamOptions = None, + derham_opts: DerhamOptions = None, verbose: bool = False, + test: ModelTest = None, ): """ Run a Struphy model. @@ -55,16 +56,33 @@ def run( Absolute path to .py parameter file. """ - # check model - assert hasattr(model, "propagators"), "Attribute 'self.propagators' must be set in model __init__!" - model_name = model.__class__.__name__ - comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() start_simulation = time.time() + # check for test run + if test is not None: + model = test.model + units = test.units + time_opts = test.time_opts + domain = test.domain + equil = test.equil + grid = test.grid + derham_opts = test.derham_opts + + print(f"\nRunning test '{test.__class__.__name__}':") + for k, v in test.__dict__.items(): + print(f" {k}: {v}") + + test.set_propagator_opts() + test.set_variables_inits() + + # check model + assert hasattr(model, "propagators"), "Attribute 'self.propagators' must be set in model __init__!" + model_name = model.__class__.__name__ + # meta-data out_folders = env.out_folders sim_folder = env.sim_folder @@ -146,12 +164,12 @@ def run( # allocate derham-related objects if grid is not None: - model.allocate_feec(grid, derham) + model.allocate_feec(grid, derham_opts) else: model._derham = None model._mass_ops = None model._projected_equil = None - print("\nDERHAM:\nMeshless simulation - no Derham complex set up.") + print("GRID:\nNo grid specified - no Derham complex set up.") # allocate variables model.allocate_variables() From 78a357853b01b0217309b4291110753ee54f1ca8 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 4 Aug 2025 14:37:31 +0200 Subject: [PATCH 057/292] delete StruphyParameters --- src/struphy/io/inp/params_Maxwell.py | 50 +++++------ src/struphy/io/options.py | 1 - src/struphy/io/parameters.py | 70 --------------- src/struphy/io/setup.py | 85 ------------------- src/struphy/models/base.py | 11 --- src/struphy/models/tests/test_toy.py | 34 -------- .../post_processing/post_processing_tools.py | 9 +- src/struphy/struphy.py | 21 +---- src/struphy/utils/clone_config.py | 4 +- 9 files changed, 30 insertions(+), 255 deletions(-) delete mode 100644 src/struphy/io/parameters.py delete mode 100644 src/struphy/models/tests/test_toy.py diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index b36497884..24d573ac1 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -15,50 +15,46 @@ # environment options env = EnvironmentOptions() -# # units -# units = Units() +# units +units = Units() -# # time stepping -# time = Time() +# time stepping +time = Time() -# # geometry -# domain = domains.Cuboid() +# geometry +domain = domains.Cuboid() -# # fluid equilibrium (can be used as part of initial conditions) -# equil = equils.HomogenSlab() +# fluid equilibrium (can be used as part of initial conditions) +equil = equils.HomogenSlab() -# # grid -# grid = grids.TensorProductGrid() +# grid +grid = grids.TensorProductGrid() -# # derham options -# derham_opts = DerhamOptions() +# derham options +derham_opts = DerhamOptions() # light-weight model instance model = Model() -# # propagator options -# model.propagators.maxwell.set_options() +# propagator options +model.propagators.maxwell.set_options() -# # initial conditions (background + perturbation) -# model.em_fields.b_field.add_background(FieldsBackground()) -# model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos()) +# initial conditions (background + perturbation) +model.em_fields.b_field.add_background(FieldsBackground()) +model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos()) # optional: exclude variables from saving # model.em_fields.b_field.save_data = False -from struphy.models.tests import test_toy -test = test_toy.LightWaveDispersion(model) - # start run struphy.run(model, __file__, env=env, - # units=units, - # time_opts=time, - # domain=domain, - # equil=equil, - # grid=grid, - # derham_opts=derham_opts, + units=units, + time_opts=time, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, verbose=verbose, - test=test, ) \ No newline at end of file diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index ad44a4dd6..cf235d99e 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -3,7 +3,6 @@ import numpy as np import os -# needed for import in StruphyParameters from struphy.physics.physics import ConstantsOfNature diff --git a/src/struphy/io/parameters.py b/src/struphy/io/parameters.py deleted file mode 100644 index c3be1f594..000000000 --- a/src/struphy/io/parameters.py +++ /dev/null @@ -1,70 +0,0 @@ -from __future__ import annotations -from mpi4py import MPI -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from struphy.fields_background.base import FluidEquilibrium - from struphy.geometry.base import Domain - from struphy.io.options import DerhamOptions, Time, Units - from struphy.topology import grids - from struphy.models.base import StruphyModel - - -class StruphyParameters: - """Wrapper around simulation parameters.""" - - def __init__( - self, - model: StruphyModel = None, - units: Units = None, - domain: Domain = None, - equil: FluidEquilibrium = None, - time: Time = None, - grid: grids.TensorProductGrid = None, - derham: DerhamOptions = None, - verbose: bool = False - ): - self._model = model - self._units = units - self._domain = domain - self._equil = equil - self._time = time - self._grid = grid - self._derham = derham - - if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"\n{self.model = }") - print(f"{self.units = }") - print(f"{self.domain = }") - print(f"{self.equil = }") - print(f"{self.time = }") - print(f"{self.grid = }") - print(f"{self.derham = }") - - @property - def model(self): - return self._model - - @property - def units(self): - return self._units - - @property - def domain(self): - return self._domain - - @property - def equil(self): - return self._equil - - @property - def time(self): - return self._time - - @property - def grid(self): - return self._grid - - @property - def derham(self): - return self._derham diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index 017931375..3ad0ccd7c 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -8,9 +8,7 @@ from mpi4py import MPI from struphy.io.options import DerhamOptions -from struphy.io.parameters import StruphyParameters from struphy.topology.grids import TensorProductGrid -from struphy.utils.utils import dict_to_yaml, read_state from struphy.geometry.base import Domain from struphy.geometry.domains import Cuboid from struphy.io.options import Units, Time @@ -86,89 +84,6 @@ def setup_folders( print("Removed existing file " + file) -def setup_parameters( - params_path: str, - path_out: str, - verbose: bool = False, -): - """ - Prepare simulation parameters from .yml or .py file and save to output folder. - - Parameters - ---------- - model_name : str - The name of the model to run. - - params_path : str - Path to .py parameter file. - - path_out : str - The output directory. Will create a folder if it does not exist OR cleans the folder for new runs. - - verbose : bool - Show full screen output. - - Returns - ------- - params : StruphyParameters - The simulation parameters. - """ - - if ".yml" in params_path or ".yaml" in params_path: - with open(params_path) as file: - params = yaml.load(file, Loader=yaml.FullLoader) - elif ".py" in params_path: - # print(f'{params_path = }') - # Read struphy state file - # state = read_state() - # i_path = state["i_path"] - # load parameter.py - params_in = import_parameters_py(params_path) - - if not hasattr(params_in, "model"): - params_in.model = None - - if not hasattr(params_in, "domain"): - params_in.domain = Cuboid() - - if not hasattr(params_in, "grid"): - params_in.grid = None - - if not hasattr(params_in, "equil"): - params_in.equil = None - - if not hasattr(params_in, "units"): - params_in.units = Units() - - if not hasattr(params_in, "time"): - params_in.time = Time() - - if not hasattr(params_in, "derham"): - params_in.derham = None - - params = StruphyParameters( - model=params_in.model, - units=params_in.units, - domain=params_in.domain, - equil=params_in.equil, - time=params_in.time, - grid=params_in.grid, - derham=params_in.derham, - verbose=verbose, - ) - - if MPI.COMM_WORLD.Get_rank() == 0: - # copy parameter file to output folder - filename = params_path.split("/")[-1] - ext = filename.split(".")[-1] - shutil.copy2( - params_path, - os.path.join(path_out, "parameters." + ext), - ) - - return params - - def setup_derham( grid: TensorProductGrid, options: DerhamOptions, diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 87491b2e2..46a527efc 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -42,17 +42,6 @@ class StruphyModel(metaclass=ABCMeta): """ Base class for all Struphy models. - Parameters - ---------- - params : StruphyParameters - Simulation parameters. - - comm : mpi4py.MPI.Intracomm - MPI communicator for parallel runs. - - clone_config: struphy.utils.CloneConfig - Contains the # TODO - Note ---- All Struphy models are subclasses of ``StruphyModel`` and should be added to ``struphy/models/`` diff --git a/src/struphy/models/tests/test_toy.py b/src/struphy/models/tests/test_toy.py deleted file mode 100644 index 9e55b8ffa..000000000 --- a/src/struphy/models/tests/test_toy.py +++ /dev/null @@ -1,34 +0,0 @@ -from struphy.models.tests.base_test import ModelTest -from dataclasses import dataclass - -from struphy.models.toy import Maxwell -from struphy.io.options import EnvironmentOptions, Units, Time, DerhamOptions -from struphy.geometry.base import Domain -from struphy.fields_background.base import FluidEquilibrium -from struphy.topology.grids import TensorProductGrid -from struphy.geometry.domains import Cuboid -from struphy.fields_background.equils import HomogenSlab -from struphy.initial import perturbations - - -class LightWaveDispersion(ModelTest): - """Simulate light wave dispersion relation in Cartesian coordinates (Cuboid mapping).""" - - def __init__(self, model: Maxwell): - self.model: Maxwell = model - self.units: Units = Units() - self.time_opts: Time = Time(dt=0.05, Tend=100.0 - 1e-6) - self.domain: Domain = Cuboid(r3=20.0) - self.equil: FluidEquilibrium = None - self.grid: TensorProductGrid = TensorProductGrid(Nel=(1, 1, 128)) - self.derham_opts: DerhamOptions = DerhamOptions(p=(1, 1, 3)) - - def set_propagator_opts(self): - self.model.propagators.maxwell.set_options() - - def set_variables_inits(self): - self.model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0)) - self.model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1)) - - def verification_script(self): - assert -1 < 0 \ No newline at end of file diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index d2fafc79f..6c15fea80 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -17,7 +17,6 @@ def create_femfields( path: str, - params_in, *, step: int = 1, ): @@ -27,9 +26,6 @@ def create_femfields( ---------- path : str Absolute path of simulation output folder. - - params_in : parameter module - Imported parameter .py module. step : int Whether to create FEM fields at every time step (step=1, default), every second time step (step=2), etc. @@ -46,6 +42,11 @@ def create_femfields( with open(os.path.join(path, "meta.yml"), "r") as f: meta = yaml.load(f, Loader=yaml.FullLoader) nproc = meta["MPI processes"] + + with open(os.path.join(path, "parameters.py"), "r") as params: + print(f"{params = }") + + exit() derham = setup_derham(params_in.grid, params_in.derham, diff --git a/src/struphy/struphy.py b/src/struphy/struphy.py index 9966e9acd..ec48a9f51 100644 --- a/src/struphy/struphy.py +++ b/src/struphy/struphy.py @@ -12,7 +12,7 @@ from struphy.feec.psydac_derham import SplineFunction from struphy.fields_background.base import FluidEquilibriumWithB from struphy.io.output_handling import DataContainer -from struphy.io.setup import setup_folders, setup_parameters +from struphy.io.setup import setup_folders from struphy.models.base import StruphyModel from struphy.profiling.profiling import ProfileManager from struphy.utils.clone_config import CloneConfig @@ -27,7 +27,6 @@ from struphy.fields_background.equils import HomogenSlab from struphy.topology.grids import TensorProductGrid from struphy.io.options import DerhamOptions -from struphy.models.tests.base_test import ModelTest def run( @@ -42,7 +41,6 @@ def run( grid: TensorProductGrid = None, derham_opts: DerhamOptions = None, verbose: bool = False, - test: ModelTest = None, ): """ Run a Struphy model. @@ -61,23 +59,6 @@ def run( size = comm.Get_size() start_simulation = time.time() - - # check for test run - if test is not None: - model = test.model - units = test.units - time_opts = test.time_opts - domain = test.domain - equil = test.equil - grid = test.grid - derham_opts = test.derham_opts - - print(f"\nRunning test '{test.__class__.__name__}':") - for k, v in test.__dict__.items(): - print(f" {k}: {v}") - - test.set_propagator_opts() - test.set_variables_inits() # check model assert hasattr(model, "propagators"), "Attribute 'self.propagators' must be set in model __init__!" diff --git a/src/struphy/utils/clone_config.py b/src/struphy/utils/clone_config.py index 046844df5..081f4d862 100644 --- a/src/struphy/utils/clone_config.py +++ b/src/struphy/utils/clone_config.py @@ -1,8 +1,6 @@ import numpy as np from mpi4py import MPI -from struphy.io.parameters import StruphyParameters - class CloneConfig: """ @@ -19,7 +17,7 @@ class CloneConfig: def __init__( self, comm: MPI.Intracomm, - params: StruphyParameters, + params: None, num_clones=1, ): """ From c138e26331f1f28f71a765244fc6d7bcb2adb032 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 Aug 2025 08:40:10 +0200 Subject: [PATCH 058/292] allow pickling parameter objects when no params_path is given --- src/struphy/compile_struphy.mk | 2 +- src/struphy/console/main.py | 4 +- src/struphy/io/inp/params_Maxwell.py | 27 ++- src/struphy/io/options.py | 4 + src/struphy/models/base.py | 27 ++- .../post_processing/post_processing_tools.py | 77 +++++- src/struphy/struphy.py | 223 +++++++++++++++++- 7 files changed, 320 insertions(+), 44 deletions(-) diff --git a/src/struphy/compile_struphy.mk b/src/struphy/compile_struphy.mk index 62b310eab..2aaec6cc6 100644 --- a/src/struphy/compile_struphy.mk +++ b/src/struphy/compile_struphy.mk @@ -5,7 +5,7 @@ PYTHON := python3 SO_EXT := $(shell $(PYTHON) -c "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX'))") LIBDIR := $(shell $(PYTHON) -c "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))") -struphy_path := $(shell $(PYTHON) -c "import struphy as _; print(_.__path__[0])") +struphy_path := $(shell $(PYTHON) -c "import struphy; print(struphy.__path__[0])") # Arguments to this script are: STRUPHY_SOURCES := $(sources) diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index 74294f9e2..b8f64a3ac 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -15,10 +15,10 @@ import yaml # struphy path -import struphy +import struphy as _ from struphy.utils import utils -libpath = struphy.__path__[0] +libpath = _.__path__[0] __version__ = importlib.metadata.version("struphy") # version message diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index 24d573ac1..b07376994 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -10,7 +10,7 @@ # import model, set verbosity from struphy.models.toy import Maxwell as Model -verbose = True +verbose = False # environment options env = EnvironmentOptions() @@ -46,15 +46,16 @@ # optional: exclude variables from saving # model.em_fields.b_field.save_data = False -# start run -struphy.run(model, - __file__, - env=env, - units=units, - time_opts=time, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=verbose, - ) \ No newline at end of file +if __name__ == "__main__": + # start run + struphy.run(model, + __file__, + env=env, + units=units, + time_opts=time, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) \ No newline at end of file diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index cf235d99e..6be51e8d9 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -283,12 +283,16 @@ class EnvironmentOptions: out_folders: str = os.getcwd() sim_folder: str = "sim_1" + path_out: str = None restart: bool = False max_runtime: int = 300 save_step: int = 1 sort_step: int = 0 num_clones: int = 1 + def __post_init__(self): + self.path_out: str = os.path.join(self.out_folders, self.sim_folder) + def check_option(opt, options): """Check if opt is contained in options; if opt is a list, checks for each element.""" diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 46a527efc..85abdf41b 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1404,7 +1404,7 @@ def generate_default_parameter_file( file.write("units = Units()\n") file.write("\n# time stepping\n") - file.write("time = Time()\n") + file.write("time_opts = Time()\n") file.write("\n# geometry\n") file.write("domain = domains.Cuboid()\n") @@ -1440,18 +1440,19 @@ def generate_default_parameter_file( file.write("\n# optional: exclude variables from saving\n") file.write(exclude_feec) - file.write("\n# start run\n") - file.write("struphy.run(model, \n\ - __file__, \n\ - env=env, \n\ - units=units, \n\ - time_opts=time, \n\ - domain=domain, \n\ - equil=equil, \n\ - grid=grid, \n\ - derham_opts=derham_opts, \n\ - verbose=verbose, \n\ - )") + file.write('\nif __name__ == "__main__":\n') + file.write(" # start run\n") + file.write(" struphy.run(model, \n\ + params_path=__file__, \n\ + env=env, \n\ + units=units, \n\ + time_opts=time_opts, \n\ + domain=domain, \n\ + equil=equil, \n\ + grid=grid, \n\ + derham_opts=derham_opts, \n\ + verbose=verbose, \n\ + )") print(f"Default parameter file for '{self.__class__.__name__}' has been created.\n\ You can now launch with 'struphy run {self.__class__.__name__}' or with 'struphy run -i params_{self.__class__.__name__}.py'") diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 6c15fea80..ecccfbbe1 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -1,11 +1,11 @@ import os import shutil - import h5py import matplotlib.pyplot as plt import numpy as np import yaml from tqdm import tqdm +import pickle from struphy.feec.psydac_derham import Derham from struphy.kinetic_background import maxwellians @@ -13,6 +13,73 @@ from struphy.io.setup import import_parameters_py from struphy.models.base import setup_derham from struphy.feec.psydac_derham import SplineFunction +from struphy.io.options import EnvironmentOptions, Units, Time +from struphy.topology.grids import TensorProductGrid + + +class ParamsIn: + """Holds the input parameters of a Struphy simulation as attributes.""" + def __init__(self, + env: EnvironmentOptions = None, + units: Units = None, + time_opts: Time = None, + equil = None, + grid: TensorProductGrid = None, + derham_opts = None): + self.env = env + self.units = units + self.time_opts = time_opts + self.equil = equil + self.grid = grid + self.derham_opts = derham_opts + + +def get_params_of_run(path: str) -> ParamsIn: + """Retrieve parameters of finished Struphy run. + + Parameters + ---------- + path : str + Absolute path of simulation output folder. + """ + + params_path = os.path.join(path, "parameters.py") + bin_path = os.path.join(path, "env.bin") + + if os.path.exists(params_path): + params_in = import_parameters_py(params_path) + env = params_in.env + units = params_in.units + time_opts = params_in.time_opts + domain = params_in.domain + equil = params_in.equil + grid = params_in.grid + derham_opts = params_in.derham_opts + + elif os.path.exists(bin_path): + with open(os.path.join(path, "env.bin"), "rb") as f: + env = pickle.load(f) + with open(os.path.join(path, "units.bin"), "rb") as f: + units = pickle.load(f) + with open(os.path.join(path, "time_opts.bin"), "rb") as f: + time_opts = pickle.load(f) + with open(os.path.join(path, "equil.bin"), "rb") as f: + equil = pickle.load(f) + with open(os.path.join(path, "grid.bin"), "rb") as f: + grid = pickle.load(f) + with open(os.path.join(path, "derham_opts.bin"), "rb") as f: + derham_opts = pickle.load(f) + + else: + raise FileNotFoundError(f"Neither of the paths {params_path} or {bin_path} exists.") + + return ParamsIn(env=env, + units=units, + time_opts=time_opts, + equil=equil, + grid=grid, + derham_opts=derham_opts, + ) def create_femfields( @@ -43,13 +110,11 @@ def create_femfields( meta = yaml.load(f, Loader=yaml.FullLoader) nproc = meta["MPI processes"] - with open(os.path.join(path, "parameters.py"), "r") as params: - print(f"{params = }") - - exit() + # import parameters + params_in = get_params_of_run(path) derham = setup_derham(params_in.grid, - params_in.derham, + params_in.derham_opts, comm=None, domain=params_in.domain, mpi_dims_mask=params_in.grid.mpi_dims_mask, diff --git a/src/struphy/struphy.py b/src/struphy/struphy.py index ec48a9f51..be956c93b 100644 --- a/src/struphy/struphy.py +++ b/src/struphy/struphy.py @@ -8,6 +8,8 @@ from mpi4py import MPI from pyevtk.hl import gridToVTK import shutil +import pickle +import h5py from struphy.feec.psydac_derham import SplineFunction from struphy.fields_background.base import FluidEquilibriumWithB @@ -27,12 +29,13 @@ from struphy.fields_background.equils import HomogenSlab from struphy.topology.grids import TensorProductGrid from struphy.io.options import DerhamOptions +from struphy.post_processing.post_processing_tools import create_femfields, eval_femfields def run( model: StruphyModel, - params_path: str, *, + params_path: str = None, env: EnvironmentOptions = EnvironmentOptions(), units: Units = Units(), time_opts: Time = Time(), @@ -65,10 +68,7 @@ def run( model_name = model.__class__.__name__ # meta-data - out_folders = env.out_folders - sim_folder = env.sim_folder - path_out = os.path.join(out_folders, sim_folder) - + path_out = env.path_out restart = env.restart max_runtime = env.max_runtime save_step = env.save_step @@ -101,10 +101,29 @@ def run( # save parameter file if rank == 0: - shutil.copy2( - params_path, - os.path.join(path_out, "parameters.py"), - ) + # save python param file + if params_path is not None: + assert params_path[-3:] == ".py" + shutil.copy2( + params_path, + os.path.join(path_out, "parameters.py"), + ) + # pickle struphy objects + else: + with open(os.path.join(path_out, "env.bin"), 'wb') as f: + pickle.dump(env, f, pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_out, "units.bin"), 'wb') as f: + pickle.dump(units, f, pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_out, "time_opts.bin"), 'wb') as f: + pickle.dump(time_opts, f, pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_out, "domain.bin"), 'wb') as f: + pickle.dump(domain, f, pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_out, "equil.bin"), 'wb') as f: + pickle.dump(equil, f, pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_out, "grid.bin"), 'wb') as f: + pickle.dump(grid, f, pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_out, "derham_opts.bin"), 'wb') as f: + pickle.dump(derham_opts, f, pickle.HIGHEST_PROTOCOL) # config clones if comm is None: @@ -343,6 +362,192 @@ def run( clone_config.free() +def pproc( + path: str, + *, + step: int = 1, + celldivide: int = 1, + physical: bool = False, + guiding_center: bool = False, + classify: bool = False, + no_vtk: bool = False, + time_trace: bool = False, +): + """Post-processing finished Struphy runs. + + Parameters + ---------- + path : str + Absolute path of simulation output folder to post-process. + + step : int + Whether to do post-processing at every time step (step=1, default), every second time step (step=2), etc. + + celldivide : int + Grid refinement in evaluation of FEM fields. E.g. celldivide=2 evaluates two points per grid cell. + + physical : bool + Wether to do post-processing into push-forwarded physical (xyz) components of fields. + + guiding_center : bool + Compute guiding-center coordinates (only from Particles6D). + + classify : bool + Classify guiding-center trajectories (passing, trapped or lost). + + no_vtk : bool + whether vtk files creation should be skipped + + time_trace : bool + whether to plot the time trace of each measured region + """ + print("") + + # create post-processing folder + path_pproc = os.path.join(path, "post_processing") + + try: + os.mkdir(path_pproc) + except: + shutil.rmtree(path_pproc) + os.mkdir(path_pproc) + + if time_trace: + from struphy.post_processing.likwid.plot_time_traces import plot_gantt_chart, plot_time_vs_duration + + path_time_trace = os.path.join(path, "profiling_time_trace.pkl") + plot_time_vs_duration(path_time_trace, output_path=path_pproc) + plot_gantt_chart(path_time_trace, output_path=path_pproc) + return + + # check for fields and kinetic data in hdf5 file that need post processing + file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") + + # save time grid at which post-processing data is created + np.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) + + if "feec" in file.keys(): + exist_fields = True + else: + exist_fields = False + + if "kinetic" in file.keys(): + exist_kinetic = {"markers": False, "f": False, "n_sph": False} + for name in file["kinetic"].keys(): + # check for saved markers + if "markers" in file["kinetic"][name]: + exist_kinetic["markers"] = True + # check for saved distribution function + if "f" in file["kinetic"][name]: + exist_kinetic["f"] = True + # check for saved sph density + if "n_sph" in file["kinetic"][name]: + exist_kinetic["n_sph"] = True + else: + exist_kinetic = None + + file.close() + + # field post-processing + if exist_fields: + fields, t_grid = create_femfields(path, step=step) + print("exiting after fem_fields ...") + exit() + + point_data, grids_log, grids_phy = eval_femfields( + path, fields, celldivide=[celldivide, celldivide, celldivide] + ) + + if physical: + point_data_phy, grids_log, grids_phy = eval_femfields( + path, fields, celldivide=[celldivide, celldivide, celldivide], physical=True + ) + + # directory for field data + path_fields = os.path.join(path_pproc, "fields_data") + + try: + os.mkdir(path_fields) + except: + shutil.rmtree(path_fields) + os.mkdir(path_fields) + + # save data dicts for each field + for species, vars in point_data.items(): + for name, val in vars.items(): + try: + os.mkdir(os.path.join(path_fields, species)) + except: + pass + + with open(os.path.join(path_fields, species, name + "_log.bin"), "wb") as handle: + pickle.dump(val, handle, protocol=pickle.HIGHEST_PROTOCOL) + + if physical: + with open(os.path.join(path_fields, species, name + "_phy.bin"), "wb") as handle: + pickle.dump(point_data_phy[species][name], handle, protocol=pickle.HIGHEST_PROTOCOL) + + # save grids + with open(os.path.join(path_fields, "grids_log.bin"), "wb") as handle: + pickle.dump(grids_log, handle, protocol=pickle.HIGHEST_PROTOCOL) + + with open(os.path.join(path_fields, "grids_phy.bin"), "wb") as handle: + pickle.dump(grids_phy, handle, protocol=pickle.HIGHEST_PROTOCOL) + + # create vtk files + if not no_vtk: + pproc.create_vtk(path_fields, t_grid, grids_phy, point_data) + if physical: + pproc.create_vtk(path_fields, t_grid, grids_phy, point_data_phy, physical=True) + + # kinetic post-processing + if exist_kinetic is not None: + # directory for kinetic data + path_kinetics = os.path.join(path_pproc, "kinetic_data") + + try: + os.mkdir(path_kinetics) + except: + shutil.rmtree(path_kinetics) + os.mkdir(path_kinetics) + + # kinetic post-processing for each species + for n, species in enumerate(kinetic_species): + # directory for each species + path_kinetics_species = os.path.join(path_kinetics, species) + + try: + os.mkdir(path_kinetics_species) + except: + shutil.rmtree(path_kinetics_species) + os.mkdir(path_kinetics_species) + + # markers + if exist_kinetic["markers"]: + pproc.post_process_markers(path, path_kinetics_species, species, kinetic_kinds[n], step) + + if guiding_center: + assert kinetic_kinds[n] == "Particles6D" + orbits_pproc.post_process_orbit_guiding_center(path, path_kinetics_species, species) + + if classify: + orbits_pproc.post_process_orbit_classification(path_kinetics_species, species) + + # distribution function + if exist_kinetic["f"]: + if kinetic_kinds[n] == "DeltaFParticles6D": + compute_bckgr = True + else: + compute_bckgr = False + + pproc.post_process_f(path, path_kinetics_species, species, step, compute_bckgr=compute_bckgr) + + # sph density + if exist_kinetic["n_sph"]: + pproc.post_process_n_sph(path, path_kinetics_species, species, step, compute_bckgr=compute_bckgr) + + + if __name__ == "__main__": import argparse import os From 419be1d77da5cb8bc553eba907c77eb5954a13a5 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 Aug 2025 09:11:50 +0200 Subject: [PATCH 059/292] pickle of compiled domain doesnt work --- src/struphy/{struphy.py => main.py} | 0 src/struphy/models/base.py | 4 ++-- src/struphy/post_processing/post_processing_tools.py | 5 ++++- 3 files changed, 6 insertions(+), 3 deletions(-) rename src/struphy/{struphy.py => main.py} (100%) diff --git a/src/struphy/struphy.py b/src/struphy/main.py similarity index 100% rename from src/struphy/struphy.py rename to src/struphy/main.py diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 85abdf41b..3ed9d4d8c 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1391,7 +1391,7 @@ def generate_default_parameter_file( file.write("from struphy.initial import perturbations\n") file.write("from struphy.kinetic_background import maxwellians\n") - file.write("from struphy import struphy\n") + file.write("from struphy import main\n") file.write("\n# import model, set verbosity\n") file.write(f"from {self.__module__} import {self.__class__.__name__} as Model\n") @@ -1442,7 +1442,7 @@ def generate_default_parameter_file( file.write('\nif __name__ == "__main__":\n') file.write(" # start run\n") - file.write(" struphy.run(model, \n\ + file.write(" main.run(model, \n\ params_path=__file__, \n\ env=env, \n\ units=units, \n\ diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index ecccfbbe1..f401151d5 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -23,12 +23,14 @@ def __init__(self, env: EnvironmentOptions = None, units: Units = None, time_opts: Time = None, + domain = None, equil = None, grid: TensorProductGrid = None, derham_opts = None): self.env = env self.units = units self.time_opts = time_opts + self.domain = domain self.equil = equil self.grid = grid self.derham_opts = derham_opts @@ -63,6 +65,8 @@ def get_params_of_run(path: str) -> ParamsIn: units = pickle.load(f) with open(os.path.join(path, "time_opts.bin"), "rb") as f: time_opts = pickle.load(f) + with open(os.path.join(path, "domain.bin"), "rb") as f: + domain = pickle.load(f) with open(os.path.join(path, "equil.bin"), "rb") as f: equil = pickle.load(f) with open(os.path.join(path, "grid.bin"), "rb") as f: @@ -117,7 +121,6 @@ def create_femfields( params_in.derham_opts, comm=None, domain=params_in.domain, - mpi_dims_mask=params_in.grid.mpi_dims_mask, ) # get fields names, space IDs and time grid from 0-th rank hdf5 file From 2198811128e7742cca90065e720c42eb715ee83d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 5 Aug 2025 10:10:22 +0200 Subject: [PATCH 060/292] implement work-around for pickling domain --- src/struphy/geometry/base.py | 4 +- src/struphy/main.py | 13 ++- src/struphy/models/tests/test_maxwell.py | 79 +++++++++++++++++++ .../post_processing/post_processing_tools.py | 25 ++++-- 4 files changed, 106 insertions(+), 15 deletions(-) create mode 100644 src/struphy/models/tests/test_maxwell.py diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index f8630eefc..0b344b687 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -156,13 +156,13 @@ def kind_map(self): @property @abstractmethod - def params_map(self): + def params_map(self) -> dict: """Mapping parameters as dictionary.""" pass @property @abstractmethod - def params_numpy(self): + def params_numpy(self) -> np.ndarray: """Mapping parameters as numpy array (can be empty).""" pass diff --git a/src/struphy/main.py b/src/struphy/main.py index be956c93b..099530edf 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -29,7 +29,7 @@ from struphy.fields_background.equils import HomogenSlab from struphy.topology.grids import TensorProductGrid from struphy.io.options import DerhamOptions -from struphy.post_processing.post_processing_tools import create_femfields, eval_femfields +from struphy.post_processing.post_processing_tools import create_femfields, eval_femfields, create_vtk def run( @@ -117,7 +117,9 @@ def run( with open(os.path.join(path_out, "time_opts.bin"), 'wb') as f: pickle.dump(time_opts, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "domain.bin"), 'wb') as f: - pickle.dump(domain, f, pickle.HIGHEST_PROTOCOL) + # WORKAROUND: cannot pickle pyccelized classes at the moment + tmp_dct = {"name": domain.__class__.__name__, "params_map": domain.params_map} + pickle.dump(tmp_dct, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "equil.bin"), 'wb') as f: pickle.dump(equil, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "grid.bin"), 'wb') as f: @@ -401,7 +403,6 @@ def pproc( time_trace : bool whether to plot the time trace of each measured region """ - print("") # create post-processing folder path_pproc = os.path.join(path, "post_processing") @@ -451,8 +452,6 @@ def pproc( # field post-processing if exist_fields: fields, t_grid = create_femfields(path, step=step) - print("exiting after fem_fields ...") - exit() point_data, grids_log, grids_phy = eval_femfields( path, fields, celldivide=[celldivide, celldivide, celldivide] @@ -496,9 +495,9 @@ def pproc( # create vtk files if not no_vtk: - pproc.create_vtk(path_fields, t_grid, grids_phy, point_data) + create_vtk(path_fields, t_grid, grids_phy, point_data) if physical: - pproc.create_vtk(path_fields, t_grid, grids_phy, point_data_phy, physical=True) + create_vtk(path_fields, t_grid, grids_phy, point_data_phy, physical=True) # kinetic post-processing if exist_kinetic is not None: diff --git a/src/struphy/models/tests/test_maxwell.py b/src/struphy/models/tests/test_maxwell.py new file mode 100644 index 000000000..acbade824 --- /dev/null +++ b/src/struphy/models/tests/test_maxwell.py @@ -0,0 +1,79 @@ +from mpi4py import MPI +import os + +from struphy.io.options import EnvironmentOptions, Units, 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.io.options import FieldsBackground +from struphy.initial import perturbations +from struphy.kinetic_background import maxwellians +from struphy import main +from struphy.post_processing.post_processing_tools import create_femfields + +test_folder = os.path.join(os.getcwd(), "verification_tests") + + +def test_light_wave_1d(): + # import model, set verbosity + from struphy.models.toy import Maxwell as Model + verbose = True + + # environment options + out_folders = os.path.join(test_folder, "maxwell") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="light_wave_1d") + + # units + units = Units() + + # time stepping + time_opts = Time(dt=0.05, Tend=1.0 - 1e-6) + + # geometry + domain = domains.Cuboid(r3=20.0) + + # fluid equilibrium (can be used as part of initial conditions) + equil = None + + # grid + grid = grids.TensorProductGrid(Nel=(1, 1, 128)) + + # derham options + derham_opts = DerhamOptions(p=(1, 1, 3)) + + # light-weight model instance + model = Model() + + # propagator options + model.propagators.maxwell.set_options() + + # initial conditions (background + perturbation) + model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0)) + model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1)) + + # optional: exclude variables from saving + # model.em_fields.b_field.save_data = False + + # start run + main.run(model, + params_path=None, + env=env, + units=units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) + + # post processing + if MPI.COMM_WORLD.Get_rank() == 0: + main.pproc(env.path_out) + + # diagnostics + + +if __name__ == "__main__": + test_light_wave_1d() \ No newline at end of file diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index f401151d5..aa65d7f22 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -15,6 +15,7 @@ from struphy.feec.psydac_derham import SplineFunction from struphy.io.options import EnvironmentOptions, Units, Time from struphy.topology.grids import TensorProductGrid +from struphy.geometry import domains class ParamsIn: @@ -45,6 +46,8 @@ def get_params_of_run(path: str) -> ParamsIn: Absolute path of simulation output folder. """ + print(f"\nReading in paramters from {path} ... ") + params_path = os.path.join(path, "parameters.py") bin_path = os.path.join(path, "env.bin") @@ -66,7 +69,11 @@ def get_params_of_run(path: str) -> ParamsIn: with open(os.path.join(path, "time_opts.bin"), "rb") as f: time_opts = pickle.load(f) with open(os.path.join(path, "domain.bin"), "rb") as f: - domain = pickle.load(f) + # WORKAROUND: cannot pickle pyccelized classes at the moment + domain_dct = pickle.load(f) + name = domain_dct["name"] + params_map = domain_dct["params_map"] + domain = getattr(domains, name)(**params_map) with open(os.path.join(path, "equil.bin"), "rb") as f: equil = pickle.load(f) with open(os.path.join(path, "grid.bin"), "rb") as f: @@ -77,9 +84,12 @@ def get_params_of_run(path: str) -> ParamsIn: else: raise FileNotFoundError(f"Neither of the paths {params_path} or {bin_path} exists.") + print("done.") + return ParamsIn(env=env, units=units, time_opts=time_opts, + domain=domain, equil=equil, grid=grid, derham_opts=derham_opts, @@ -217,7 +227,7 @@ def create_femfields( def eval_femfields( - params_in, + path: str, fields: dict, *, celldivide: list = [1, 1, 1], @@ -227,8 +237,8 @@ def eval_femfields( Parameters ---------- - params_in : parameter module - Imported parameter .py module. + path : str + Absolute path of simulation output folder. fields : dict Obtained from struphy.diagnostics.post_processing.create_femfields. @@ -255,6 +265,9 @@ def eval_femfields( Mapped (physical) grids obtained by domain(*grids_log). """ + # import parameters + params_in = get_params_of_run(path) + # get domain domain = params_in.domain @@ -279,7 +292,7 @@ def eval_femfields( for name, field in vars.items(): point_data[species][name] = {} - print("Evaluating fields ...") + print("\nEvaluating fields ...") for t in tqdm(fields): for species, vars in fields[t].items(): for name, field in vars.items(): @@ -394,7 +407,7 @@ def create_vtk( nt = len(t_grid) - 1 log_nt = int(np.log10(nt)) + 1 - print("Creating vtk ...") + print("\nCreating vtk ...") for n, t in enumerate(tqdm(t_grid)): point_data_n = {} From 02541efbf7cd38d322db1d63b4360507dd6cef7b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 6 Aug 2025 07:22:48 +0200 Subject: [PATCH 061/292] new function `load_data` in main.py; added fitting to power_spectrum_2d; test_maxwell.py is running --- src/struphy/diagnostics/diagn_tools.py | 98 +++++++++++++++++------- src/struphy/main.py | 88 +++++++++++++++++++-- src/struphy/models/tests/test_maxwell.py | 54 ++++++++----- 3 files changed, 188 insertions(+), 52 deletions(-) diff --git a/src/struphy/diagnostics/diagn_tools.py b/src/struphy/diagnostics/diagn_tools.py index 7a2005482..8ac7d1450 100644 --- a/src/struphy/diagnostics/diagn_tools.py +++ b/src/struphy/diagnostics/diagn_tools.py @@ -14,19 +14,22 @@ def power_spectrum_2d( - values, - name, - code, - grids, - grids_mapped=None, - component=0, - slice_at=(None, 0, 0), - do_plot=False, - disp_name=None, - disp_params={}, - save_plot=False, - save_name=None, - file_format="png", + values: dict, + name: str, + grids: tuple, + grids_mapped: tuple = None, + component: int = 0, + slice_at: tuple = (None, 0, 0), + do_plot: bool = False, + disp_name: str = None, + disp_params: dict = {}, + fit_branches: int = 0, + noise_level: float = 0.1, + extr_order: int = 3, + fit_degree: tuple = (1,), + save_plot: bool = False, + save_name: str = None, + file_format: str = "png", ): """Perform fft in space-time, (t, x) -> (omega, k), where x can be a logical or physical coordinate. Returns values if plot=False. @@ -40,19 +43,16 @@ def power_spectrum_2d( name : str Name of the FemField. - code : str - From which code the data has been obtained. - - grids : 3-list + grids : 3-tuple 1d logical grids in each eta-direction with Nel[i]*npts_per_cell[i] + 1 entries in each direction. - grids_mapped : 3-list + grids_mapped : 3-tuple Mapped grids obtained by domain(). If None, the fft is performed on the logical grids. component : int Which component of a FemField to consider; is 0 for 0-and 3-forms, is in {0, 1, 2} for 1- and 2-forms. - slice_at : 3-list + slice_at : 3-tuple At which indices i, j the 1d slice data (t, eta)_(i, j) should be obtained. One entry must be "None"; this is the direction of the fft. Default: [None, 0, 0] performs the eta1-fft at (eta2[0], eta3[0]). @@ -66,6 +66,20 @@ def power_spectrum_2d( disp_params : dict Parameters needed for analytical dispersion relation, see struphy.dispersion_relations.analytic. + fit_branches: int + How many branches to fit in the dispersion relation. + Default=0 means no fits are made. + + noise_level: float + Sets the threshold above which local maxima in the power spectrum are taken into account. + Computed as threshold = max(spectrum) * noise_level. + + extr_oder: int + Order given to argrelextrema. + + fit_degree: tuple[int] + Degree of fitting polynomial for each branch (fit_branches) of power spectrum. + save_plot : boolean Save figure if True. Then a path has to be given. @@ -84,16 +98,13 @@ def power_spectrum_2d( 1d array of wave vector. dispersion : np.array - 2d array of shape (omega.size, kvce.size) holding the fft. + 2d array of shape (omega.size, kvec.size) holding the fft. """ - print(f"code: {code}") - keys = list(values.keys()) # check uniform grid in time dt = keys[1] - keys[0] - print(f"time step: {dt}") assert np.all([np.abs(y - x - dt) < 1e-12 for x, y in zip(keys[:-1], keys[1:])]) # create 4d np.array with shape (time, eta1, eta2, eta3) @@ -131,13 +142,40 @@ def power_spectrum_2d( Nt = data.shape[0] Nx = grid.size dx = grid[1] - grid[0] - print(f"space step: {dx}") assert np.allclose(grid[1:] - grid[:-1], dx * np.ones_like(grid[:-1])) dispersion = (2.0 / Nt) * (2.0 / Nx) * np.abs(fftn(data))[: Nt // 2, : Nx // 2] kvec = 2 * np.pi * fftfreq(Nx, dx)[: Nx // 2] omega = 2 * np.pi * fftfreq(Nt, dt)[: Nt // 2] + if fit_branches > 0: + assert len(fit_degree) == fit_branches + # determine maxima for each k + siz = kvec.size // 2 # take only first half of k-vector + k_fit = [] + omega_fit = {} + for n in range(fit_branches): + omega_fit[n] = [] + for k, f_of_omega in zip(kvec[:siz], dispersion[:, :siz].T): + threshold = np.max(f_of_omega) * noise_level + extrms = argrelextrema(f_of_omega, np.greater, order=3)[0] + above_noise = np.nonzero(f_of_omega > threshold)[0] + intersec = list(set(extrms) & set(above_noise)) + if not intersec: + continue + assert len(intersec) == fit_branches, f"Number of found branches {len(intersec)} is not {fit_branches = }! \ + Try to lower 'noise_level' or increase 'extr_order'." + # print(f"{intersec = }, {omega[intersec[0]] = }") + k_fit += [k] + for n in range(fit_branches): + omega_fit[n] += [omega[intersec[n]]] + + # fit + coeffs = [] + for m, om in omega_fit.items(): + coeffs += [np.polyfit(k_fit, om, deg=fit_degree[n])] + print(f"\nFitted {coeffs = }") + if do_plot: _, ax = plt.subplots(1, 1, figsize=(10, 10)) colormap = "plasma" @@ -156,10 +194,17 @@ def power_spectrum_2d( mappable=disp_plot, format="%.0e", ) - title = name + " component " + str(component + 1) + " from code: " + code + title = name + ", component " + str(component + 1) ax.set_title(title) ax.set_xlabel("$k$ [a.u.]") ax.set_ylabel(r"$\omega$ [a.u.]") + for cs in coeffs: + def fun(k): + out = k*0.0 + for i, c in enumerate(np.flip(cs)): + out += c * k**i + return out + ax.plot(kvec, fun(kvec), "r:", label="fit") # analytic solution: disp_class = getattr(analytic, disp_name) @@ -190,8 +235,7 @@ def power_spectrum_2d( else: plt.show() - else: - return kvec, omega, dispersion + return omega, kvec, dispersion def plot_scalars( diff --git a/src/struphy/main.py b/src/struphy/main.py index 099530edf..50d946677 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -96,9 +96,6 @@ def run( restart=restart, verbose=verbose,) - # save meta-data - dict_to_yaml(meta, os.path.join(path_out, "meta.yml")) - # save parameter file if rank == 0: # save python param file @@ -353,11 +350,12 @@ def run( # =================================================================== - with open(path_out + "/meta.txt", "a") as f: - # f.write('wall-clock time [min]:'.ljust(30) + str((end_simulation - start_simulation)/60.) + '\n') - f.write(f"{rank} {'wall-clock time[min]: '.ljust(30)}{(end_simulation - start_simulation) / 60}\n") + meta["wall-clock time[min]"] = (end_simulation - start_simulation) / 60 comm.Barrier() + if rank == 0: + # save meta-data + dict_to_yaml(meta, os.path.join(path_out, "meta.yml")) print("Struphy run finished.") if clone_config is not None: @@ -546,6 +544,84 @@ def pproc( pproc.post_process_n_sph(path, path_kinetics_species, species, step, compute_bckgr=compute_bckgr) +class SimData: + """Holds post-processed Struphy data as attributes. + + Parameters + ---------- + path : str + Absolute path of simulation output folder to post-process. + """ + def __init__(self, path: str): + self.path = path + self.feec_species = {} + self.pic_species = {} + self.sph_species = {} + self.arrays = {} + self.grids_log = None + self.grids_phy = None + + @property + def spline_grid_resolution(self): + if self.grids_log is not None: + res = [x.size for x in self.grids_log] + else: + res = None + return res + + +def load_data(path: str) -> SimData: + """Load data generated during post-processing. + + Parameters + ---------- + path : str + Absolute path of simulation output folder to post-process. + """ + + + path_pproc = os.path.join(path, "post_processing") + assert os.path.exists(path_pproc), f"Path {path_pproc} does not exist, run 'pproc' first?" + print("\nLoading post-processed simulation data ...") + + path_fields = os.path.join(path_pproc, "fields_data") + + simdata = SimData(path) + + if os.path.exists(path_fields): + + # grids + with open(os.path.join(path_fields, "grids_log.bin"), "rb") as f: + simdata.grids_log = pickle.load(f) + with open(os.path.join(path_fields, "grids_phy.bin"), "rb") as f: + simdata.grids_phy = pickle.load(f) + + # species folders + species = next(os.walk(path_fields))[1] + for spec in species: + simdata.feec_species[spec] = [] + simdata.arrays[spec] = {} + path_spec = os.path.join(path_fields, spec) + wlk = os.walk(path_spec) + files = next(wlk)[2] + for file in files: + if ".bin" in file: + var = file.split(".")[0] + with open(os.path.join(path_spec, file), "rb") as f: + simdata.feec_species[spec] += [var] + simdata.arrays[spec][var] = pickle.load(f) + + print("Done - the following data has been loaded:") + print(f"{simdata.spline_grid_resolution = }") + print(f"{simdata.feec_species = }") + print(f"{simdata.pic_species = }") + print(f"{simdata.sph_species = }") + + return simdata + + + + if __name__ == "__main__": import argparse diff --git a/src/struphy/models/tests/test_maxwell.py b/src/struphy/models/tests/test_maxwell.py index acbade824..24a5b7ed9 100644 --- a/src/struphy/models/tests/test_maxwell.py +++ b/src/struphy/models/tests/test_maxwell.py @@ -1,5 +1,6 @@ from mpi4py import MPI import os +import numpy as np from struphy.io.options import EnvironmentOptions, Units, Time from struphy.geometry import domains @@ -10,12 +11,12 @@ from struphy.initial import perturbations from struphy.kinetic_background import maxwellians from struphy import main -from struphy.post_processing.post_processing_tools import create_femfields +from struphy.diagnostics.diagn_tools import power_spectrum_2d test_folder = os.path.join(os.getcwd(), "verification_tests") -def test_light_wave_1d(): +def test_light_wave_1d(do_plot: bool = False): # import model, set verbosity from struphy.models.toy import Maxwell as Model verbose = True @@ -28,7 +29,7 @@ def test_light_wave_1d(): units = Units() # time stepping - time_opts = Time(dt=0.05, Tend=1.0 - 1e-6) + time_opts = Time(dt=0.05, Tend=50.0) # geometry domain = domains.Cuboid(r3=20.0) @@ -56,24 +57,39 @@ def test_light_wave_1d(): # model.em_fields.b_field.save_data = False # start run - main.run(model, - params_path=None, - env=env, - units=units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=verbose, - ) + # main.run(model, + # params_path=None, + # env=env, + # units=units, + # time_opts=time_opts, + # domain=domain, + # equil=equil, + # grid=grid, + # derham_opts=derham_opts, + # verbose=verbose, + # ) - # post processing - if MPI.COMM_WORLD.Get_rank() == 0: - main.pproc(env.path_out) + # # post processing + # if MPI.COMM_WORLD.Get_rank() == 0: + # main.pproc(env.path_out) # diagnostics - + if MPI.COMM_WORLD.Get_rank() == 0: + simdata = main.load_data(env.path_out) + + # fft in (t, z) of first component of e_field on physical grid + Ex_of_t = simdata.arrays["em_fields"]["e_field_log"] + power_spectrum_2d(Ex_of_t, + "e_field_log", + grids=simdata.grids_log, + grids_mapped=simdata.grids_phy, + component=0, + slice_at=[0, 0, None], + do_plot=do_plot, + disp_name='Maxwell1D', + fit_branches=True, + noise_level=0.5,) + if __name__ == "__main__": - test_light_wave_1d() \ No newline at end of file + test_light_wave_1d(do_plot=True) \ No newline at end of file From f96f77b1f108afca5aa252d69b4712d54ea93ffe Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 6 Aug 2025 08:15:17 +0200 Subject: [PATCH 062/292] clean up screen output and maxwell test --- src/struphy/diagnostics/diagn_tools.py | 12 ++++-- src/struphy/feec/psydac_derham.py | 2 +- src/struphy/main.py | 9 ++-- src/struphy/models/tests/test_maxwell.py | 43 ++++++++++--------- src/struphy/models/toy.py | 2 +- .../post_processing/post_processing_tools.py | 5 +-- src/struphy/propagators/base.py | 7 +-- 7 files changed, 44 insertions(+), 36 deletions(-) diff --git a/src/struphy/diagnostics/diagn_tools.py b/src/struphy/diagnostics/diagn_tools.py index 8ac7d1450..8e9c620d5 100644 --- a/src/struphy/diagnostics/diagn_tools.py +++ b/src/struphy/diagnostics/diagn_tools.py @@ -25,7 +25,7 @@ def power_spectrum_2d( disp_params: dict = {}, fit_branches: int = 0, noise_level: float = 0.1, - extr_order: int = 3, + extr_order: int = 10, fit_degree: tuple = (1,), save_plot: bool = False, save_name: str = None, @@ -99,6 +99,9 @@ def power_spectrum_2d( dispersion : np.array 2d array of shape (omega.size, kvec.size) holding the fft. + + coeffs : list[list] + List of fitting coefficients (lenght is fit_branches). """ keys = list(values.keys()) @@ -148,6 +151,7 @@ def power_spectrum_2d( kvec = 2 * np.pi * fftfreq(Nx, dx)[: Nx // 2] omega = 2 * np.pi * fftfreq(Nt, dt)[: Nt // 2] + coeffs = None if fit_branches > 0: assert len(fit_degree) == fit_branches # determine maxima for each k @@ -158,14 +162,14 @@ def power_spectrum_2d( omega_fit[n] = [] for k, f_of_omega in zip(kvec[:siz], dispersion[:, :siz].T): threshold = np.max(f_of_omega) * noise_level - extrms = argrelextrema(f_of_omega, np.greater, order=3)[0] + extrms = argrelextrema(f_of_omega, np.greater, order=extr_order)[0] above_noise = np.nonzero(f_of_omega > threshold)[0] intersec = list(set(extrms) & set(above_noise)) if not intersec: continue + print(f"{intersec = }, {omega[intersec[0]] = }") assert len(intersec) == fit_branches, f"Number of found branches {len(intersec)} is not {fit_branches = }! \ Try to lower 'noise_level' or increase 'extr_order'." - # print(f"{intersec = }, {omega[intersec[0]] = }") k_fit += [k] for n in range(fit_branches): omega_fit[n] += [omega[intersec[n]]] @@ -235,7 +239,7 @@ def fun(k): else: plt.show() - return omega, kvec, dispersion + return omega, kvec, dispersion, coeffs def plot_scalars( diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 4dda194fb..3769d3a9a 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1476,7 +1476,7 @@ def __init__( else: self._nbasis = [tuple([space.nbasis for space in vec_space.spaces]) for vec_space in self.fem_space.spaces] - if verbose: + if verbose and MPI.COMM_WORLD.Get_rank() == 0: print(f"\nAllocated SplineFuntion '{self.name}' in space '{self.space_id}'.") if self.backgrounds is not None or self.perturbations is not None: diff --git a/src/struphy/main.py b/src/struphy/main.py index 50d946677..68e9af312 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -87,9 +87,10 @@ def run( meta["max wall-clock [min]"] = max_runtime meta["save interval [steps]"] = save_step - print("\nMETADATA:") - for k, v in meta.items(): - print(f'{k}:'.ljust(25), v) + if rank == 0: + print("\nMETADATA:") + for k, v in meta.items(): + print(f'{k}:'.ljust(25), v) # creating output folders setup_folders(path_out=path_out, @@ -184,7 +185,9 @@ def run( if rank < 32: if rank == 0: print("") + comm.Barrier() print(f"Rank {rank}: calling struphy/main.py for model {model_name} ...") + if size > 32 and rank == 32: print(f"Ranks > 31: calling struphy/main.py for model {model_name} ...") diff --git a/src/struphy/models/tests/test_maxwell.py b/src/struphy/models/tests/test_maxwell.py index 24a5b7ed9..5cdbed6b9 100644 --- a/src/struphy/models/tests/test_maxwell.py +++ b/src/struphy/models/tests/test_maxwell.py @@ -53,25 +53,22 @@ def test_light_wave_1d(do_plot: bool = False): model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0)) model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1)) - # optional: exclude variables from saving - # model.em_fields.b_field.save_data = False - # start run - # main.run(model, - # params_path=None, - # env=env, - # units=units, - # time_opts=time_opts, - # domain=domain, - # equil=equil, - # grid=grid, - # derham_opts=derham_opts, - # verbose=verbose, - # ) + main.run(model, + params_path=None, + env=env, + units=units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) - # # post processing - # if MPI.COMM_WORLD.Get_rank() == 0: - # main.pproc(env.path_out) + # post processing + if MPI.COMM_WORLD.Get_rank() == 0: + main.pproc(env.path_out) # diagnostics if MPI.COMM_WORLD.Get_rank() == 0: @@ -79,7 +76,7 @@ def test_light_wave_1d(do_plot: bool = False): # fft in (t, z) of first component of e_field on physical grid Ex_of_t = simdata.arrays["em_fields"]["e_field_log"] - power_spectrum_2d(Ex_of_t, + _1, _2, _3, coeffs = power_spectrum_2d(Ex_of_t, "e_field_log", grids=simdata.grids_log, grids_mapped=simdata.grids_phy, @@ -87,8 +84,14 @@ def test_light_wave_1d(do_plot: bool = False): slice_at=[0, 0, None], do_plot=do_plot, disp_name='Maxwell1D', - fit_branches=True, - noise_level=0.5,) + fit_branches=1, + noise_level=0.5, + extr_order=10, + fit_degree=(1,), + ) + + # test + assert np.abs(coeffs[0][0] - 1.0) < 0.02 if __name__ == "__main__": diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index c00283e5b..3cfd499bf 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -48,7 +48,7 @@ def __init__(self): ## abstract methods - def __init__(self): + def __init__(self): # 1. instantiate all variales self.em_fields = self.EMFields() diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index aa65d7f22..58dcc3bbb 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -407,7 +407,7 @@ def create_vtk( nt = len(t_grid) - 1 log_nt = int(np.log10(nt)) + 1 - print("\nCreating vtk ...") + print(f"\nCreating vtk in {path} ...") for n, t in enumerate(tqdm(t_grid)): point_data_n = {} @@ -431,9 +431,6 @@ def create_vtk( *grids_phy, pointData=point_data_n[species], ) - - if n==0: - print(f" saving .vtk files to {species_path}") def post_process_markers(path_in, path_out, species, kind, step=1): diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 85ce70a2f..51e8fd25c 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -295,9 +295,10 @@ def options(self, new): class Options: def __init__(self, prop, **kwargs): self.__dict__.update(kwargs) - print(f"\nInstance of propagator '{prop.__class__.__name__}' with:") - for k, v in self.__dict__.items(): - print(f' {k}: {v}') + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nInstance of propagator '{prop.__class__.__name__}' with:") + for k, v in self.__dict__.items(): + print(f' {k}: {v}') # @property # def kwargs(self): From 92616fcf9c82528aef5451a8c905d1482e299dc4 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 6 Aug 2025 08:30:07 +0200 Subject: [PATCH 063/292] nicer screen output, inidicate code sections --- src/struphy/diagnostics/diagn_tools.py | 2 +- src/struphy/main.py | 18 ++++++++++++------ src/struphy/models/toy.py | 7 ++++++- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/src/struphy/diagnostics/diagn_tools.py b/src/struphy/diagnostics/diagn_tools.py index 8e9c620d5..3a1eef1f4 100644 --- a/src/struphy/diagnostics/diagn_tools.py +++ b/src/struphy/diagnostics/diagn_tools.py @@ -167,7 +167,7 @@ def power_spectrum_2d( intersec = list(set(extrms) & set(above_noise)) if not intersec: continue - print(f"{intersec = }, {omega[intersec[0]] = }") + # print(f"{intersec = }, {omega[intersec[0]] = }") assert len(intersec) == fit_branches, f"Number of found branches {len(intersec)} is not {fit_branches = }! \ Try to lower 'noise_level' or increase 'extr_order'." k_fit += [k] diff --git a/src/struphy/main.py b/src/struphy/main.py index 68e9af312..1d81e7c13 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -31,6 +31,10 @@ from struphy.io.options import DerhamOptions from struphy.post_processing.post_processing_tools import create_femfields, eval_femfields, create_vtk +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +size = comm.Get_size() + def run( model: StruphyModel, @@ -56,10 +60,6 @@ def run( params_path : str Absolute path to .py parameter file. """ - - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - size = comm.Get_size() start_simulation = time.time() @@ -67,6 +67,9 @@ def run( assert hasattr(model, "propagators"), "Attribute 'self.propagators' must be set in model __init__!" model_name = model.__class__.__name__ + if rank == 0: + print(f"\n*** Starting run for model '{model_name}':") + # meta-data path_out = env.path_out restart = env.restart @@ -405,6 +408,9 @@ def pproc( whether to plot the time trace of each measured region """ + if rank == 0: + print(f"\n*** Start post-processing of {path}:") + # create post-processing folder path_pproc = os.path.join(path, "post_processing") @@ -585,7 +591,7 @@ def load_data(path: str) -> SimData: path_pproc = os.path.join(path, "post_processing") assert os.path.exists(path_pproc), f"Path {path_pproc} does not exist, run 'pproc' first?" - print("\nLoading post-processed simulation data ...") + print("\n*** Loading post-processed simulation data:") path_fields = os.path.join(path_pproc, "fields_data") @@ -614,7 +620,7 @@ def load_data(path: str) -> SimData: simdata.feec_species[spec] += [var] simdata.arrays[spec][var] = pickle.load(f) - print("Done - the following data has been loaded:") + print("\nThe following data has been loaded:") print(f"{simdata.spline_grid_resolution = }") print(f"{simdata.feec_species = }") print(f"{simdata.pic_species = }") diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 3cfd499bf..e7740b9d5 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,7 +1,7 @@ import numpy as np from dataclasses import dataclass -import inspect +from mpi4py import MPI from struphy.models.base import StruphyModel from struphy.propagators.base import Propagator @@ -9,6 +9,8 @@ from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable +rank = MPI.COMM_WORLD.Get_rank() + class Maxwell(StruphyModel): r"""Maxwell's equations in vacuum. @@ -49,6 +51,9 @@ def __init__(self): ## abstract methods def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + # 1. instantiate all variales self.em_fields = self.EMFields() From 186545c81b57b30ddc93f3b1d71d5d77ca3715f2 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 6 Aug 2025 08:49:47 +0200 Subject: [PATCH 064/292] make test_maxwell a pytest --- src/struphy/main.py | 10 +++++----- src/struphy/models/tests/test_maxwell.py | 16 +++++++++------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/struphy/main.py b/src/struphy/main.py index 1d81e7c13..ae3b34ce8 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -30,10 +30,6 @@ from struphy.topology.grids import TensorProductGrid from struphy.io.options import DerhamOptions from struphy.post_processing.post_processing_tools import create_femfields, eval_femfields, create_vtk - -comm = MPI.COMM_WORLD -rank = comm.Get_rank() -size = comm.Get_size() def run( @@ -61,6 +57,10 @@ def run( Absolute path to .py parameter file. """ + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + start_simulation = time.time() # check model @@ -408,7 +408,7 @@ def pproc( whether to plot the time trace of each measured region """ - if rank == 0: + if MPI.COMM_WORLD.Get_rank() == 0: print(f"\n*** Start post-processing of {path}:") # create post-processing folder diff --git a/src/struphy/models/tests/test_maxwell.py b/src/struphy/models/tests/test_maxwell.py index 5cdbed6b9..30565f32c 100644 --- a/src/struphy/models/tests/test_maxwell.py +++ b/src/struphy/models/tests/test_maxwell.py @@ -1,13 +1,13 @@ from mpi4py import MPI import os import numpy as np +import pytest from struphy.io.options import EnvironmentOptions, Units, 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.io.options import FieldsBackground +from struphy.io.options import DerhamOptions, FieldsBackground from struphy.initial import perturbations from struphy.kinetic_background import maxwellians from struphy import main @@ -16,7 +16,9 @@ test_folder = os.path.join(os.getcwd(), "verification_tests") -def test_light_wave_1d(do_plot: bool = False): +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize('algo', ["implicit", "explicit"]) +def test_light_wave_1d(algo: str, do_plot: bool = False): # import model, set verbosity from struphy.models.toy import Maxwell as Model verbose = True @@ -47,11 +49,11 @@ def test_light_wave_1d(do_plot: bool = False): model = Model() # propagator options - model.propagators.maxwell.set_options() + model.propagators.maxwell.set_options(algo=algo) # initial conditions (background + perturbation) - model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0)) - model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1)) + model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0, seed=123)) + model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1, seed=123)) # start run main.run(model, @@ -95,4 +97,4 @@ def test_light_wave_1d(do_plot: bool = False): if __name__ == "__main__": - test_light_wave_1d(do_plot=True) \ No newline at end of file + test_light_wave_1d(algo="explicit", do_plot=True) \ No newline at end of file From 5d787959bd11407821bc3fc8ebfd150a281f42de Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 6 Aug 2025 10:35:09 +0200 Subject: [PATCH 065/292] add coaxial test to test_maxwell --- src/struphy/main.py | 19 ++- src/struphy/models/tests/test_maxwell.py | 169 ++++++++++++++++++++++- 2 files changed, 181 insertions(+), 7 deletions(-) diff --git a/src/struphy/main.py b/src/struphy/main.py index ae3b34ce8..c9fa9ea0a 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -567,8 +567,9 @@ def __init__(self, path: str): self.pic_species = {} self.sph_species = {} self.arrays = {} - self.grids_log = None - self.grids_phy = None + self.grids_log: list[np.ndarray] = None + self.grids_phy: list[np.ndarray] = None + self.t_grid: np.ndarray = None @property def spline_grid_resolution(self): @@ -577,6 +578,10 @@ def spline_grid_resolution(self): else: res = None return res + + @property + def time_grid_size(self): + return self.t_grid.size def load_data(path: str) -> SimData: @@ -588,15 +593,18 @@ def load_data(path: str) -> SimData: Absolute path of simulation output folder to post-process. """ - path_pproc = os.path.join(path, "post_processing") assert os.path.exists(path_pproc), f"Path {path_pproc} does not exist, run 'pproc' first?" print("\n*** Loading post-processed simulation data:") - path_fields = os.path.join(path_pproc, "fields_data") - simdata = SimData(path) + # load time grid + simdata.t_grid = np.load(os.path.join(path_pproc, "t_grid.npy")) + + # load point data + path_fields = os.path.join(path_pproc, "fields_data") + if os.path.exists(path_fields): # grids @@ -621,6 +629,7 @@ def load_data(path: str) -> SimData: simdata.arrays[spec][var] = pickle.load(f) print("\nThe following data has been loaded:") + print(f"{simdata.time_grid_size = }") print(f"{simdata.spline_grid_resolution = }") print(f"{simdata.feec_species = }") print(f"{simdata.pic_species = }") diff --git a/src/struphy/models/tests/test_maxwell.py b/src/struphy/models/tests/test_maxwell.py index 30565f32c..aab37f21e 100644 --- a/src/struphy/models/tests/test_maxwell.py +++ b/src/struphy/models/tests/test_maxwell.py @@ -2,6 +2,8 @@ import os import numpy as np import pytest +from scipy.special import jv, yn +from matplotlib import pyplot as plt from struphy.io.options import EnvironmentOptions, Units, Time from struphy.geometry import domains @@ -92,9 +94,172 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): fit_degree=(1,), ) - # test + # assert assert np.abs(coeffs[0][0] - 1.0) < 0.02 + + +@pytest.mark.mpi(min_size=1) +def test_coaxial(do_plot: bool = False): + # import model, set verbosity + from struphy.models.toy import Maxwell as Model + verbose = True + + # environment options + out_folders = os.path.join(test_folder, "maxwell") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="coaxial") + + # units + units = Units() + + # time + time_opts = Time(dt=0.05, Tend=10.0) + + # geometry + a1=2.326744 + a2=3.686839 + Lz = 2.0 + domain = domains.HollowCylinder(a1=a1, a2=a2, Lz=Lz) + + # fluid equilibrium (can be used as part of initial conditions) + equil = equils.HomogenSlab() + + # grid + grid = grids.TensorProductGrid(Nel=(32, 64, 1)) + + # derham options + derham_opts = DerhamOptions(p=(3, 3, 1), + spl_kind=(False, True, True), + dirichlet_bc=((True, True), (False, False), (False, False)), + ) + + # light-weight model instance + model = Model() + + # propagator options + model.propagators.maxwell.set_options(algo="implicit") + + # initial conditions (background + perturbation) + m = 3 + model.em_fields.e_field.add_perturbation(perturbations.CoaxialWaveguideElectric_r(m=m, a1=a1, a2=a2)) + model.em_fields.e_field.add_perturbation(perturbations.CoaxialWaveguideElectric_theta(m=m, a1=a1, a2=a2)) + model.em_fields.b_field.add_perturbation(perturbations.CoaxialWaveguideMagnetic(m=m, a1=a1, a2=a2)) + + # start run + main.run(model, + params_path=None, + env=env, + units=units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) + + # post processing + if MPI.COMM_WORLD.Get_rank() == 0: + main.pproc(env.path_out, physical=True) + + # diagnostics + if MPI.COMM_WORLD.Get_rank() == 0: + # get parameters + dt = time_opts.dt + split_algo = time_opts.split_algo + Nel = grid.Nel + modes = m + + # load data + simdata = main.load_data(env.path_out) + + t_grid = simdata.t_grid + grids_phy = simdata.grids_phy + e_field_phy = simdata.arrays["em_fields"]["e_field_phy"] + b_field_phy = simdata.arrays["em_fields"]["b_field_phy"] + + # define analytic solution + def B_z(X, Y, Z, m, t): + """Magnetic field in z direction of coaxial cabel""" + r = (X**2 + Y**2) ** 0.5 + theta = np.arctan2(Y, X) + return (jv(m, r) - 0.28 * yn(m, r)) * np.cos(m * theta - t) + + def E_r(X, Y, Z, m, t): + """Electrical field in radial direction of coaxial cabel""" + r = (X**2 + Y**2) ** 0.5 + theta = np.arctan2(Y, X) + return -m / r * (jv(m, r) - 0.28 * yn(m, r)) * np.cos(m * theta - t) + + def E_theta(X, Y, Z, m, t): + """Electrical field in azimuthal direction of coaxial cabel""" + r = (X**2 + Y**2) ** 0.5 + theta = np.arctan2(Y, X) + return ((m / r * jv(m, r) - jv(m + 1, r)) - 0.28 * (m / r * yn(m, r) - yn(m + 1, r))) * np.sin(m * theta - t) + + def to_E_r(X, Y, E_x, E_y): + r = (X**2 + Y**2) ** 0.5 + theta = np.arctan2(Y, X) + return np.cos(theta) * E_x + np.sin(theta) * E_y + + def to_E_theta(X, Y, E_x, E_y): + r = (X**2 + Y**2) ** 0.5 + theta = np.arctan2(Y, X) + return -np.sin(theta) * E_x + np.cos(theta) * E_y + + # plot + if do_plot: + X = grids_phy[0][:, :, 0] + Y = grids_phy[1][:, :, 0] + + vmin = E_theta(X, Y, grids_phy[0], modes, 0).min() + vmax = E_theta(X, Y, grids_phy[0], modes, 0).max() + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4)) + plot_exac = ax1.contourf( + X, Y, E_theta(X, Y, grids_phy[0], modes, t_grid[-1]), cmap="plasma", levels=100, vmin=vmin, vmax=vmax + ) + ax2.contourf( + X, + Y, + to_E_theta(X, Y, e_field_phy[t_grid[-1]][0][:, :, 0], e_field_phy[t_grid[-1]][1][:, :, 0]), + cmap="plasma", + levels=100, + vmin=vmin, + vmax=vmax, + ) + fig.colorbar(plot_exac, ax=[ax1, ax2], orientation="vertical", shrink=0.9) + ax1.set_xlabel("Exact") + ax2.set_xlabel("Numerical") + fig.suptitle(f"Exact and Simulated $E_\\theta$ Field {dt=}, {split_algo=}, {Nel=}", fontsize=14) + plt.show() + + # assert + Ex_tend = e_field_phy[t_grid[-1]][0][:, :, 0] + Ey_tend = e_field_phy[t_grid[-1]][1][:, :, 0] + Er_exact = E_r(X, Y, grids_phy[0], modes, t_grid[-1]) + Etheta_exact = E_theta(X, Y, grids_phy[0], modes, t_grid[-1]) + Bz_tend = b_field_phy[t_grid[-1]][2][:, :, 0] + Bz_exact = B_z(X, Y, grids_phy[0], modes, t_grid[-1]) + + error_Er = np.max(np.abs((to_E_r(X, Y, Ex_tend, Ey_tend) - Er_exact))) + error_Etheta = np.max(np.abs((to_E_theta(X, Y, Ex_tend, Ey_tend) - Etheta_exact))) + error_Bz = np.max(np.abs((Bz_tend - Bz_exact))) + + rel_err_Er = error_Er / np.max(np.abs(Er_exact)) + rel_err_Etheta = error_Etheta / np.max(np.abs(Etheta_exact)) + rel_err_Bz = error_Bz / np.max(np.abs(Bz_exact)) + + print("") + assert rel_err_Bz < 0.0021, f"Assertion for magnetic field Maxwell failed: {rel_err_Bz = }" + print(f"Assertion for magnetic field Maxwell passed ({rel_err_Bz = }).") + assert rel_err_Etheta < 0.0021, ( + f"Assertion for electric (E_theta) field Maxwell failed: {rel_err_Etheta = }" + ) + print(f"Assertion for electric field Maxwell passed ({rel_err_Etheta = }).") + assert rel_err_Er < 0.0021, f"Assertion for electric (E_r) field Maxwell failed: {rel_err_Er = }" + print(f"Assertion for electric field Maxwell passed ({rel_err_Er = }).") + if __name__ == "__main__": - test_light_wave_1d(algo="explicit", do_plot=True) \ No newline at end of file + # test_light_wave_1d(algo="explicit", do_plot=True) + test_coaxial(do_plot=True) \ No newline at end of file From 66d78b1999fc660a3b6f5b377f90de83c513f5e9 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 6 Aug 2025 11:00:40 +0200 Subject: [PATCH 066/292] make __init__ an abstract method of StruphyModel; require 3 ranks for 1d tests and 4 ranks for 2d tests --- src/struphy/models/base.py | 4 ++++ src/struphy/models/fluid.py | 10 ++++++++-- src/struphy/models/tests/test_maxwell.py | 10 +++++----- src/struphy/models/toy.py | 2 +- 4 files changed, 18 insertions(+), 8 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 3ed9d4d8c..5068d8dfc 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -54,6 +54,10 @@ class StruphyModel(metaclass=ABCMeta): class Propagators: pass + @abstractmethod + def __init__(self): + """Light-weight init of model.""" + @property @abstractmethod def bulk_species() -> Species: diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 9ac27adb9..15526bf44 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -1,5 +1,6 @@ import numpy as np from dataclasses import dataclass +from mpi4py import MPI from struphy.models.base import StruphyModel from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers @@ -8,6 +9,8 @@ from struphy.polar.basic import PolarVector from psydac.linalg.block import BlockVector +rank = MPI.COMM_WORLD.Get_rank() + class LinearMHD(StruphyModel): r"""Linear ideal MHD with zero-flow equilibrium (:math:`\mathbf U_0 = 0`). @@ -62,7 +65,10 @@ def __init__(self): ## abstract methods def __init__(self): - # 1. instantiate all variales + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species, variables self.em_fields = self.EMFields() self.mhd = self.MHD() @@ -106,7 +112,7 @@ def allocate_helpers(self): self._ones[:] = 1.0 self._tmp_b1: BlockVector = self.derham.Vh["2"].zeros() # TODO: replace derham.Vh dict by class - self._tmp_b2 = self.derham.Vh["2"].zeros() + self._tmp_b2: BlockVector = self.derham.Vh["2"].zeros() def update_scalar_quantities(self): # perturbed fields diff --git a/src/struphy/models/tests/test_maxwell.py b/src/struphy/models/tests/test_maxwell.py index aab37f21e..32a9ffec2 100644 --- a/src/struphy/models/tests/test_maxwell.py +++ b/src/struphy/models/tests/test_maxwell.py @@ -18,7 +18,7 @@ test_folder = os.path.join(os.getcwd(), "verification_tests") -@pytest.mark.mpi(min_size=2) +@pytest.mark.mpi(min_size=3) @pytest.mark.parametrize('algo', ["implicit", "explicit"]) def test_light_wave_1d(algo: str, do_plot: bool = False): # import model, set verbosity @@ -98,7 +98,7 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): assert np.abs(coeffs[0][0] - 1.0) < 0.02 -@pytest.mark.mpi(min_size=1) +@pytest.mark.mpi(min_size=4) def test_coaxial(do_plot: bool = False): # import model, set verbosity from struphy.models.toy import Maxwell as Model @@ -177,6 +177,9 @@ def test_coaxial(do_plot: bool = False): e_field_phy = simdata.arrays["em_fields"]["e_field_phy"] b_field_phy = simdata.arrays["em_fields"]["b_field_phy"] + X = grids_phy[0][:, :, 0] + Y = grids_phy[1][:, :, 0] + # define analytic solution def B_z(X, Y, Z, m, t): """Magnetic field in z direction of coaxial cabel""" @@ -208,9 +211,6 @@ def to_E_theta(X, Y, E_x, E_y): # plot if do_plot: - X = grids_phy[0][:, :, 0] - Y = grids_phy[1][:, :, 0] - vmin = E_theta(X, Y, grids_phy[0], modes, 0).min() vmax = E_theta(X, Y, grids_phy[0], modes, 0).max() fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4)) diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index e7740b9d5..043b36bef 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -54,7 +54,7 @@ def __init__(self): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # 1. instantiate all variales + # 1. instantiate all species, variables self.em_fields = self.EMFields() # 2. instantiate all propagators From 8e79ab8c3ccb52d6d1f55921b4312ab71f1c62e4 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 6 Aug 2025 14:05:13 +0200 Subject: [PATCH 067/292] new test test_LinearMHD.py is working --- src/struphy/diagnostics/diagn_tools.py | 30 ++-- src/struphy/models/base.py | 10 ++ src/struphy/models/tests/test_LinearMHD.py | 155 ++++++++++++++++++ .../{test_maxwell.py => test_Maxwell.py} | 13 +- src/struphy/propagators/propagators_fields.py | 4 +- 5 files changed, 192 insertions(+), 20 deletions(-) create mode 100644 src/struphy/models/tests/test_LinearMHD.py rename src/struphy/models/tests/{test_maxwell.py => test_Maxwell.py} (96%) diff --git a/src/struphy/diagnostics/diagn_tools.py b/src/struphy/diagnostics/diagn_tools.py index 3a1eef1f4..841bff70c 100644 --- a/src/struphy/diagnostics/diagn_tools.py +++ b/src/struphy/diagnostics/diagn_tools.py @@ -155,19 +155,23 @@ def power_spectrum_2d( if fit_branches > 0: assert len(fit_degree) == fit_branches # determine maxima for each k - siz = kvec.size // 2 # take only first half of k-vector + k_start = kvec.size // 8 # take only first half of k-vector + k_end = kvec.size // 2 # take only first half of k-vector k_fit = [] omega_fit = {} for n in range(fit_branches): omega_fit[n] = [] - for k, f_of_omega in zip(kvec[:siz], dispersion[:, :siz].T): + for k, f_of_omega in zip(kvec[k_start:k_end], dispersion[:, k_start:k_end].T): threshold = np.max(f_of_omega) * noise_level extrms = argrelextrema(f_of_omega, np.greater, order=extr_order)[0] above_noise = np.nonzero(f_of_omega > threshold)[0] intersec = list(set(extrms) & set(above_noise)) + # intersec = list(set(extrms)) if not intersec: continue - # print(f"{intersec = }, {omega[intersec[0]] = }") + intersec.sort() + # print(f"{intersec = }") + # print(f"{[omega[intersec[n]] for n in range(fit_branches)]}") assert len(intersec) == fit_branches, f"Number of found branches {len(intersec)} is not {fit_branches = }! \ Try to lower 'noise_level' or increase 'extr_order'." k_fit += [k] @@ -177,8 +181,8 @@ def power_spectrum_2d( # fit coeffs = [] for m, om in omega_fit.items(): - coeffs += [np.polyfit(k_fit, om, deg=fit_degree[n])] - print(f"\nFitted {coeffs = }") + coeffs += [np.polyfit(k_fit, om, deg=fit_degree[n])] + print(f"\nFitted {coeffs = }") if do_plot: _, ax = plt.subplots(1, 1, figsize=(10, 10)) @@ -202,13 +206,15 @@ def power_spectrum_2d( ax.set_title(title) ax.set_xlabel("$k$ [a.u.]") ax.set_ylabel(r"$\omega$ [a.u.]") - for cs in coeffs: - def fun(k): - out = k*0.0 - for i, c in enumerate(np.flip(cs)): - out += c * k**i - return out - ax.plot(kvec, fun(kvec), "r:", label="fit") + + if fit_branches > 0: + for n, cs in enumerate(coeffs): + def fun(k): + out = k*0.0 + for i, c in enumerate(np.flip(cs)): + out += c * k**i + return out + ax.plot(kvec, fun(kvec), "r:", label=f"fit_{n + 1}") # analytic solution: disp_class = getattr(analytic, disp_name) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 5068d8dfc..2dd3737cb 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1373,11 +1373,18 @@ def generate_default_parameter_file( file.write("from struphy.geometry import domains\n") file.write("from struphy.fields_background import equils\n") + species_params = "\n# species parameters" + has_plasma = False has_feec = False has_pic = False has_sph = False for sn, species in self.species.items(): assert isinstance(species, Species) + + if isinstance(species, (FluidSpecies, KineticSpecies)): + has_plasma = True + species_params += f"\nmodel.{sn}.set_phys_params()" + for vn, var in species.variables.items(): if isinstance(var, FEECVariable): has_feec = True @@ -1431,6 +1438,9 @@ def generate_default_parameter_file( file.write("\n# light-weight model instance\n") file.write("model = Model()\n") + + if has_plasma: + file.write(species_params) file.write("\n# propagator options\n") for prop in self.propagators.__dict__: diff --git a/src/struphy/models/tests/test_LinearMHD.py b/src/struphy/models/tests/test_LinearMHD.py new file mode 100644 index 000000000..4b28813e6 --- /dev/null +++ b/src/struphy/models/tests/test_LinearMHD.py @@ -0,0 +1,155 @@ +from mpi4py import MPI +import os +import numpy as np +import pytest + +from struphy.io.options import EnvironmentOptions, Units, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.topology import grids +from struphy.io.options import DerhamOptions, FieldsBackground +from struphy.initial import perturbations +from struphy.kinetic_background import maxwellians +from struphy import main +from struphy.diagnostics.diagn_tools import power_spectrum_2d + +test_folder = os.path.join(os.getcwd(), "verification_tests") + + +@pytest.mark.mpi(min_size=3) +@pytest.mark.parametrize('algo', ["implicit", "explicit"]) +def test_slab_waves_1d(algo: str, do_plot: bool = False): + # import model, set verbosity + from struphy.models.fluid import LinearMHD as Model + verbose = True + + # environment options + out_folders = os.path.join(test_folder, "LinearMHD") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="slab_waves_1d") + + # units + units = Units() + + # time stepping + time_opts = Time(dt=0.15, Tend=180.0) + + # geometry + domain = domains.Cuboid(r3=60.0) + + # fluid equilibrium (can be used as part of initial conditions) + B0x = 0.0 + B0y = 1.0 + B0z = 1.0 + beta = 3.0 + n0 = 0.7 + equil = equils.HomogenSlab(B0x=B0x, B0y=B0y, B0z=B0z, beta=beta, n0=n0) + + # grid + grid = grids.TensorProductGrid(Nel=(1, 1, 64)) + + # derham options + derham_opts = DerhamOptions(p=(1, 1, 3)) + + # light-weight model instance + model = Model() + + # species parameters + model.mhd.set_phys_params() + + # propagator options + model.propagators.shear_alf.set_options(algo=algo) + model.propagators.mag_sonic.set_options(b_field=model.em_fields.b_field) + + # initial conditions (background + perturbation) + model.mhd.velocity.add_perturbation(perturbations.Noise(amp=0.1, comp=0, seed=123)) + model.mhd.velocity.add_perturbation(perturbations.Noise(amp=0.1, comp=1, seed=123)) + model.mhd.velocity.add_perturbation(perturbations.Noise(amp=0.1, comp=2, seed=123)) + + # start run + main.run(model, + params_path=None, + env=env, + units=units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) + + # post processing + if MPI.COMM_WORLD.Get_rank() == 0: + main.pproc(env.path_out) + + # diagnostics + if MPI.COMM_WORLD.Get_rank() == 0: + simdata = main.load_data(env.path_out) + + # first fft + u_of_t = simdata.arrays["mhd"]["velocity_log"] + + Bsquare = (B0x**2 + B0y**2 + B0z**2) + p0 = beta * Bsquare / 2 + + disp_params = {'B0x': B0x, + 'B0y': B0y, + 'B0z': B0z, + 'p0': p0, + 'n0': n0, + 'gamma': 5/3} + + _1, _2, _3, coeffs = power_spectrum_2d(u_of_t, + "velocity_log", + grids=simdata.grids_log, + grids_mapped=simdata.grids_phy, + component=0, + slice_at=[0, 0, None], + do_plot=do_plot, + disp_name='MHDhomogenSlab', + disp_params=disp_params, + fit_branches=1, + noise_level=0.5, + extr_order=10, + fit_degree=(1,), + ) + + # assert + vA = np.sqrt(Bsquare / n0) + v_alfven = vA * B0z / np.sqrt(Bsquare) + print(f"{v_alfven = }") + assert np.abs(coeffs[0][0] - v_alfven) < 0.07 + + # second fft + p_of_t = simdata.arrays["mhd"]["pressure_log"] + + _1, _2, _3, coeffs = power_spectrum_2d(p_of_t, + "pressure_log", + grids=simdata.grids_log, + grids_mapped=simdata.grids_phy, + component=0, + slice_at=[0, 0, None], + do_plot=do_plot, + disp_name='MHDhomogenSlab', + disp_params=disp_params, + fit_branches=2, + noise_level=0.4, + extr_order=10, + fit_degree=(1, 1), + ) + + # assert + gamma = 5/3 + cS = np.sqrt(gamma * p0 / n0) + + delta = (4 * B0z ** 2 * cS**2 * vA**2) / ((cS**2 + vA**2) ** 2 * Bsquare) + v_slow = np.sqrt(1 / 2 * (cS**2 + vA**2) * (1 - np.sqrt(1 - delta))) + v_fast = np.sqrt(1 / 2 * (cS**2 + vA**2) * (1 + np.sqrt(1 - delta))) + print(f"{v_slow = }") + print(f"{v_fast = }") + assert np.abs(coeffs[0][0] - v_slow) < 0.05 + assert np.abs(coeffs[1][0] - v_fast) < 0.19 + + +if __name__ == '__main__': + test_slab_waves_1d(algo="implicit", do_plot=True) \ No newline at end of file diff --git a/src/struphy/models/tests/test_maxwell.py b/src/struphy/models/tests/test_Maxwell.py similarity index 96% rename from src/struphy/models/tests/test_maxwell.py rename to src/struphy/models/tests/test_Maxwell.py index 32a9ffec2..8b22e6c4a 100644 --- a/src/struphy/models/tests/test_maxwell.py +++ b/src/struphy/models/tests/test_Maxwell.py @@ -26,7 +26,7 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): verbose = True # environment options - out_folders = os.path.join(test_folder, "maxwell") + out_folders = os.path.join(test_folder, "Maxwell") env = EnvironmentOptions(out_folders=out_folders, sim_folder="light_wave_1d") # units @@ -78,9 +78,9 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): if MPI.COMM_WORLD.Get_rank() == 0: simdata = main.load_data(env.path_out) - # fft in (t, z) of first component of e_field on physical grid - Ex_of_t = simdata.arrays["em_fields"]["e_field_log"] - _1, _2, _3, coeffs = power_spectrum_2d(Ex_of_t, + # fft + E_of_t = simdata.arrays["em_fields"]["e_field_log"] + _1, _2, _3, coeffs = power_spectrum_2d(E_of_t, "e_field_log", grids=simdata.grids_log, grids_mapped=simdata.grids_phy, @@ -95,7 +95,8 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): ) # assert - assert np.abs(coeffs[0][0] - 1.0) < 0.02 + c_light_speed = 1.0 + assert np.abs(coeffs[0][0] - c_light_speed) < 0.02 @pytest.mark.mpi(min_size=4) @@ -105,7 +106,7 @@ def test_coaxial(do_plot: bool = False): verbose = True # environment options - out_folders = os.path.join(test_folder, "maxwell") + out_folders = os.path.join(test_folder, "Maxwell") env = EnvironmentOptions(out_folders=out_folders, sim_folder="coaxial") # units diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 3fcc50297..aca2b66a0 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -538,8 +538,8 @@ def allocate(self): _f2 = curl @ _T # allocate output of vector field - out1 = self.u.space.zeros() - out2 = self.b.space.zeros() + out1 = self.u.spline.vector.space.zeros() + out2 = self.b.spline.vector.space.zeros() def f1(t, y1, y2, out: BlockVector = out1): _f1.dot(y2, out=out) From bfca873cf10af627a1558aec8212056835965880 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 6 Aug 2025 15:37:07 +0200 Subject: [PATCH 068/292] refactor `struphy params`, use getcwd()/params_MODEL.py as default path --- src/struphy/console/main.py | 55 ++++++++++++---------------- src/struphy/console/params.py | 9 ++--- src/struphy/feec/psydac_derham.py | 8 +--- src/struphy/initial/perturbations.py | 2 +- src/struphy/io/inp/params_Maxwell.py | 16 ++++---- src/struphy/io/options.py | 5 ++- src/struphy/main.py | 4 +- src/struphy/models/base.py | 36 +++++++++--------- 8 files changed, 63 insertions(+), 72 deletions(-) diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index b8f64a3ac..ed5bd895c 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -61,15 +61,15 @@ def struphy(): batch_files = get_batch_files(b_path) # Load the models and messages + model_message = "All models are listed on https://struphy.pages.mpcdf.de/struphy/sections/models.html" list_models = [] - model_message = fluid_message = kinetic_message = hybrid_message = toy_message = "" try: with open(os.path.join(libpath, "models", "models_list"), "rb") as fp: list_models = pickle.load(fp) - with open(os.path.join(libpath, "models", "models_message"), "rb") as fp: - model_message, fluid_message, kinetic_message, hybrid_message, toy_message = pickle.load( - fp, - ) + # with open(os.path.join(libpath, "models", "models_message"), "rb") as fp: + # model_message, fluid_message, kinetic_message, hybrid_message, toy_message = pickle.load( + # fp, + # ) except: print("run: struphy --refresh-models") @@ -128,18 +128,18 @@ def struphy(): sys.exit(0) # display subset of models - model_flags = [ - (args.fluid, fluid_message), - (args.kinetic, kinetic_message), - (args.hybrid, hybrid_message), - (args.toy, toy_message), - ] - - for flag, message in model_flags: - if flag: - print(message) - print("For more info on Struphy models, visit https://struphy.pages.mpcdf.de/struphy/sections/models.html") - sys.exit(0) + # model_flags = [ + # (args.fluid, fluid_message), + # (args.kinetic, kinetic_message), + # (args.hybrid, hybrid_message), + # (args.toy, toy_message), + # ] + + # for flag, message in model_flags: + # if flag: + # print(message) + # print("For more info on Struphy models, visit https://struphy.pages.mpcdf.de/struphy/sections/models.html") + # sys.exit(0) # Set default input path if args.set_i: @@ -683,10 +683,10 @@ def add_parser_params(subparsers, list_models, model_message): "params", formatter_class=lambda prog: argparse.RawTextHelpFormatter( prog, - max_help_position=30, + max_help_position=35, ), help="create default parameter file for a model, or show model's options", - description="Creates a default parameter file for a specific model, or shows a model's options.", + description="Create default parameter file (.py) for a specific model.", ) parser_params.add_argument( @@ -698,24 +698,17 @@ def add_parser_params(subparsers, list_models, model_message): ) parser_params.add_argument( - "-f", - "--file", + "-p", + "--params-path", type=str, - metavar="FILE", - help="name of the parameter file (.yml) to be created in the current I/O path (default=params_.yml)", - ) - - parser_params.add_argument( - "-o", - "--options", - help="show model options", - action="store_true", + metavar="PATH", + help="Absolute path to the parameter file (default is getcwd()/params_MODEL.py)", ) parser_params.add_argument( "-y", "--yes", - help="Say yes on prompt to overwrite .yml FILE", + help="Say yes on prompt to overwrite PATH", action="store_true", ) diff --git a/src/struphy/console/params.py b/src/struphy/console/params.py index bf7548b57..d09883405 100644 --- a/src/struphy/console/params.py +++ b/src/struphy/console/params.py @@ -2,7 +2,7 @@ from struphy.models import fluid, hybrid, kinetic, toy -def struphy_params(model_name: str, file, yes=False, options=False): +def struphy_params(model_name: str, params_path: str, yes: bool = False): """Create a model's default parameter file and save in current input path. Parameters @@ -10,14 +10,11 @@ def struphy_params(model_name: str, file, yes=False, options=False): model_name : str The name of the Struphy model. - file : str + params_path : str An alternative file name to the default params_.yml. yes : bool If true, say yes on prompt to overwrite .yml FILE - - show_options : bool - Whether to print to screen all possible options for the model. """ objs = [fluid, kinetic, hybrid, toy] for obj in objs: @@ -28,4 +25,4 @@ def struphy_params(model_name: str, file, yes=False, options=False): pass prompt = not yes - model.generate_default_parameter_file(file_name=file, prompt=prompt) + model.generate_default_parameter_file(path=params_path, prompt=prompt) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 3769d3a9a..8f4561c98 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1702,20 +1702,16 @@ def initialize_coeffs( # special case of const if fb.type == "LogicalConst": vals = fb.values + assert isinstance(vals, (list, tuple)) if self.space_id in {"H1", "L2"}: - assert isinstance(vals, (float, int)) - def f_tmp(e1, e2, e3): - return vals + 0.0 * e1 + return vals[0] + 0.0 * e1 fun = f_tmp else: - assert isinstance(vals, (list, tuple)) assert len(vals) == 3 fun = [] - for i, _v in enumerate(vals): - assert isinstance(_v, (float, int)) or _v is None if vals[0] is not None: fun += [lambda e1, e2, e3: vals[0] + 0.0 * e1] diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 58931822a..23a211c64 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -583,7 +583,7 @@ class TorusModesCos(Perturbation): Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, ms: tuple = (2,), ns: tuple = (1,), amps: tuple = (1e-4,), + def __init__(self, ms: tuple = (2,), ns: tuple = (1,), amps: tuple = (0.1,), pfuns: tuple = ("sin",), pfun_params=None, given_in_basis: GivenInBasis = "0", comp: int = 0,): diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index b07376994..cd477269f 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -6,11 +6,11 @@ from struphy.io.options import FieldsBackground from struphy.initial import perturbations from struphy.kinetic_background import maxwellians -from struphy import struphy +from struphy import main # import model, set verbosity from struphy.models.toy import Maxwell as Model -verbose = False +verbose = True # environment options env = EnvironmentOptions() @@ -19,7 +19,7 @@ units = Units() # time stepping -time = Time() +time_opts = Time() # geometry domain = domains.Cuboid() @@ -41,18 +41,20 @@ # initial conditions (background + perturbation) model.em_fields.b_field.add_background(FieldsBackground()) -model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos()) +model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=0)) +model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=1)) +model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=2)) # optional: exclude variables from saving # model.em_fields.b_field.save_data = False if __name__ == "__main__": # start run - struphy.run(model, - __file__, + main.run(model, + params_path=__file__, env=env, units=units, - time_opts=time, + time_opts=time_opts, domain=domain, equil=equil, grid=grid, diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 6be51e8d9..8b39bb474 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -237,14 +237,15 @@ class FieldsBackground: Type of background. values : tuple[float] - Values for LogicalConst on the unit cube. + Values for LogicalConst on the unit cube. + Can be length 1 for scalar functions; must be length 3 for vector-valued functions. variable : str Name of the function in FluidEquilibrium that should be the background. """ type: BackgroundTypes = "LogicalConst" - values: tuple = (1.5,) + values: tuple = (1.5, 0.7, 2.4) variable: str = None def __post_init__(self): diff --git a/src/struphy/main.py b/src/struphy/main.py index c9fa9ea0a..34c15798c 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -189,10 +189,10 @@ def run( if rank == 0: print("") comm.Barrier() - print(f"Rank {rank}: calling struphy/main.py for model {model_name} ...") + print(f"Rank {rank}: executing main.run() for model {model_name} ...") if size > 32 and rank == 32: - print(f"Ranks > 31: calling struphy/main.py for model {model_name} ...") + print(f"Ranks > 31: executing main.run() for model {model_name} ...") # store geometry vtk if rank == 0: diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 2dd3737cb..41f16a910 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1330,7 +1330,7 @@ def write_parameters_to_file(cls, parameters=None, file=None, save=True, prompt= def generate_default_parameter_file( self, - file_name: str = None, + path: str = None, prompt: bool = True, ): """Generate a parameter file with default options for each species, @@ -1340,31 +1340,26 @@ def generate_default_parameter_file( Parameters ---------- - file_name : str - Alternative filename to params_.py. + path : str + Alternative path to getcwd()/params_MODEL.py. prompt : bool Whether to prompt for overwriting the specified .yml file. """ - - # Read struphy state file - state = read_state() - i_path = state["i_path"] - assert os.path.exists(i_path), f"The path '{i_path}' does not exist. Set path with `struphy --set-i PATH`" - if file_name is None: - file_name = os.path.join(i_path, f"params_{self.__class__.__name__}.py") + if path is None: + path = os.path.join(os.getcwd(), f"params_{self.__class__.__name__}.py") # create new default file try: - file = open(file_name, "x") + file = open(path, "x") except FileExistsError: if not prompt: yn = "Y" else: - yn = input(f"File {file_name} exists, overwrite (Y/n)? ") + yn = input(f"File {path} exists, overwrite (Y/n)? ") if yn in ("", "Y", "y", "yes", "Yes"): - file = open(file_name, "w") + file = open(path, "w") else: print("exiting ...") return @@ -1388,8 +1383,15 @@ def generate_default_parameter_file( for vn, var in species.variables.items(): if isinstance(var, FEECVariable): has_feec = True - init_bckgr_feec = f"model.{sn}.{vn}.add_background(FieldsBackground())\n" - init_pert_feec = f"model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos())\n" + if var.space in ("H1", "L2"): + init_bckgr_feec = f"model.{sn}.{vn}.add_background(FieldsBackground())\n" + init_pert_feec = f"model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos())\n" + else: + init_bckgr_feec = f"model.{sn}.{vn}.add_background(FieldsBackground())\n" + init_pert_feec = f"model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=0))\n\ +model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=1))\n\ +model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=2))\n" + exclude_feec = f"# model.{sn}.{vn}.save_data = False\n" elif isinstance(var, PICVariable): has_pic = True @@ -1406,7 +1408,7 @@ def generate_default_parameter_file( file.write("\n# import model, set verbosity\n") file.write(f"from {self.__module__} import {self.__class__.__name__} as Model\n") - file.write("verbose = False\n") + file.write("verbose = True\n") file.write("\n# environment options\n") file.write("env = EnvironmentOptions()\n") @@ -1468,7 +1470,7 @@ def generate_default_parameter_file( verbose=verbose, \n\ )") - print(f"Default parameter file for '{self.__class__.__name__}' has been created.\n\ + print(f"\nDefault parameter file for '{self.__class__.__name__}' has been created in {path}.\n\ You can now launch with 'struphy run {self.__class__.__name__}' or with 'struphy run -i params_{self.__class__.__name__}.py'") ################### From 7b1d548a98ce9f8a69408bd3047c28eef8aafc7d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 7 Aug 2025 10:01:54 +0200 Subject: [PATCH 069/292] new file models/tests/test_models.py for generic testing --- src/struphy/io/options.py | 1 - src/struphy/io/setup.py | 3 +- src/struphy/models/base.py | 24 ++++-- src/struphy/models/tests/base_test.py | 38 --------- src/struphy/models/tests/test_models.py | 107 ++++++++++++++++++++++++ 5 files changed, 127 insertions(+), 46 deletions(-) delete mode 100644 src/struphy/models/tests/base_test.py create mode 100644 src/struphy/models/tests/test_models.py diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 8b39bb474..f5498a3e9 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -284,7 +284,6 @@ class EnvironmentOptions: out_folders: str = os.getcwd() sim_folder: str = "sim_1" - path_out: str = None restart: bool = False max_runtime: int = 300 save_step: int = 1 diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index 3ad0ccd7c..fdb4a7699 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -4,6 +4,7 @@ import os import shutil import yaml +from types import ModuleType from mpi4py import MPI @@ -14,7 +15,7 @@ from struphy.io.options import Units, Time -def import_parameters_py(params_path: str): +def import_parameters_py(params_path: str) -> ModuleType: """Import a .py parameter file under the module name 'parameters' and return it.""" assert ".py" in params_path spec = importlib.util.spec_from_file_location("parameters", params_path) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 41f16a910..235fb2240 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1357,12 +1357,24 @@ def generate_default_parameter_file( if not prompt: yn = "Y" else: - yn = input(f"File {path} exists, overwrite (Y/n)? ") - if yn in ("", "Y", "y", "yes", "Yes"): - file = open(path, "w") - else: - print("exiting ...") - return + yn = input(f"\nFile {path} exists, overwrite (Y/n)? ") + if yn in ("", "Y", "y", "yes", "Yes"): + file = open(path, "w") + else: + print("exiting ...") + return + except FileNotFoundError: + folder = os.path.join("/", *path.split("/")[:-1]) + if not prompt: + yn = "Y" + else: + yn = input(f"\nFolder {folder} does not exist, create (Y/n)? ") + if yn in ("", "Y", "y", "yes", "Yes"): + os.makedirs(folder) + file = open(path, "x") + else: + print("exiting ...") + return file.write("from struphy.io.options import EnvironmentOptions, Units, Time\n") file.write("from struphy.geometry import domains\n") diff --git a/src/struphy/models/tests/base_test.py b/src/struphy/models/tests/base_test.py deleted file mode 100644 index a87396b76..000000000 --- a/src/struphy/models/tests/base_test.py +++ /dev/null @@ -1,38 +0,0 @@ -from abc import ABCMeta, abstractmethod -from dataclasses import dataclass - -from struphy.io.options import EnvironmentOptions, Units, Time, DerhamOptions -from struphy.topology.grids import TensorProductGrid -from struphy.geometry.base import Domain -from struphy.geometry.domains import Cuboid -from struphy.fields_background.equils import HomogenSlab -from struphy.pic import particles -from struphy.pic.base import Particles -from struphy.propagators.base import Propagator -from struphy.kinetic_background import maxwellians -from struphy.fields_background.base import FluidEquilibrium -from struphy.models.base import StruphyModel - - -@dataclass -class ModelTest(metaclass=ABCMeta): - """Base class for model verification tests.""" - model: StruphyModel = None - units: Units = None - time_opts: Time = None - domain: Domain = None - equil: FluidEquilibrium = None - grid: TensorProductGrid = None - derham_opts: DerhamOptions = None - - @abstractmethod - def set_propagator_opts(self): - """Set propagator options for test.""" - - @abstractmethod - def set_variables_inits(self): - """Set initial conditions for test.""" - - @abstractmethod - def verification_script(self): - """Contains assert statements (and plots) that verify the test.""" \ No newline at end of file diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py new file mode 100644 index 000000000..1487679ee --- /dev/null +++ b/src/struphy/models/tests/test_models.py @@ -0,0 +1,107 @@ +import pytest +import inspect +import os +from types import ModuleType + +from struphy.models import toy, fluid, kinetic, hybrid +from struphy.models.base import StruphyModel +from struphy.io.setup import import_parameters_py +from struphy.io.options import EnvironmentOptions +from struphy import main + + +# available models +toy_models = [] +for name, obj in inspect.getmembers(toy): + if inspect.isclass(obj) and "models.toy" in obj.__module__: + toy_models += [name] +print(f"\n{toy_models = }") + +fluid_models = [] +for name, obj in inspect.getmembers(fluid): + if inspect.isclass(obj) and "models.fluid" in obj.__module__: + fluid_models += [name] +print(f"\n{fluid_models = }") + +kinetic_models = [] +for name, obj in inspect.getmembers(kinetic): + if inspect.isclass(obj) and "models.kinetic" in obj.__module__: + kinetic_models += [name] +print(f"\n{kinetic_models = }") + +hybrid_models = [] +for name, obj in inspect.getmembers(hybrid): + if inspect.isclass(obj) and "models.hybrid" in obj.__module__: + hybrid_models += [name] +print(f"\n{hybrid_models = }") + + +# folder for test simulations +test_folder = os.path.join(os.getcwd(), "struphy_model_tests") + + +# generic function for calling model tests +def call_test(model_name: str, module: ModuleType, verbose=True): + print(f"\n*** Testing '{model_name}':") + model = getattr(module, model_name)() + assert isinstance(model, StruphyModel) + + # generate paramater file for testing + path = os.path.join(test_folder, f"params_{model_name}.py") + model.generate_default_parameter_file(path=path, prompt=False) + del model + print("\nDeleting light-weight instance ...") + + # set environment options + env = EnvironmentOptions(out_folders=test_folder, sim_folder=f"{model_name}_test") + + # read parameters + params_in = import_parameters_py(path) + units = params_in.units + time_opts = params_in.time_opts + domain = params_in.domain + equil = params_in.equil + grid = params_in.grid + derham_opts = params_in.derham_opts + + # test + model = params_in.model + main.run(model, + params_path=path, + env=env, + units=units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) + + +# specific tests +@pytest.mark.parametrize('model_name', toy_models) +def test_toy(model_name: str, verbose=True): + call_test(model_name=model_name, module=toy, verbose=verbose) + +@pytest.mark.parametrize('model_name', fluid_models) +def test_fluid(model_name: str, verbose=True): + call_test(model_name=model_name, module=fluid, verbose=verbose) + +@pytest.mark.parametrize('model_name', kinetic_models) +def test_kinetic(model_name: str, verbose=True): + call_test(model_name=model_name, module=kinetic, verbose=verbose) + +@pytest.mark.parametrize('model_name', hybrid_models) +def test_hybrid(model_name: str, verbose=True): + call_test(model_name=model_name, module=hybrid, verbose=verbose) + + + +if __name__ == '__main__': + test_toy("Maxwell") + test_fluid("LinearMHD") + + + + \ No newline at end of file From 2bd1a8710c16e1ccb4a6f8796500737ba3f52f65 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 7 Aug 2025 10:02:48 +0200 Subject: [PATCH 070/292] minor changes --- src/struphy/console/test.py | 4 ++-- src/struphy/models/tests/test_Maxwell.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index b99f5e482..8fec47ca8 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -4,7 +4,7 @@ def struphy_test( group: str, *, - mpi: int = 2, + mpi: int = 4, fast: bool = False, with_desc: bool = False, Tend: float = None, @@ -22,7 +22,7 @@ def struphy_test( Test identifier: "unit", "models", "fluid", "kinetic", "hybrid", "toy" or a model name. mpi : int - Number of MPI processes used in tests (must be >1, default=2). + Number of MPI processes used in tests (must be >1, default=4). fast : bool Whether to test models just in slab geometry. diff --git a/src/struphy/models/tests/test_Maxwell.py b/src/struphy/models/tests/test_Maxwell.py index 8b22e6c4a..610e5f1c2 100644 --- a/src/struphy/models/tests/test_Maxwell.py +++ b/src/struphy/models/tests/test_Maxwell.py @@ -15,7 +15,7 @@ from struphy import main from struphy.diagnostics.diagn_tools import power_spectrum_2d -test_folder = os.path.join(os.getcwd(), "verification_tests") +test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") @pytest.mark.mpi(min_size=3) From bb41fa9a4783ff6efa3d2dd0b57ad4876240bd8e Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 7 Aug 2025 11:50:06 +0200 Subject: [PATCH 071/292] generate specific default parameter file for LinearMHD --- src/struphy/models/base.py | 11 ++++++++++- src/struphy/models/fluid.py | 16 ++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 235fb2240..934c91a3e 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1345,6 +1345,11 @@ def generate_default_parameter_file( prompt : bool Whether to prompt for overwriting the specified .yml file. + + Returns + ------- + params_path : str + The path of the parameter file. """ if path is None: @@ -1454,7 +1459,7 @@ def generate_default_parameter_file( file.write("model = Model()\n") if has_plasma: - file.write(species_params) + file.write(species_params + "\n") file.write("\n# propagator options\n") for prop in self.propagators.__dict__: @@ -1482,8 +1487,12 @@ def generate_default_parameter_file( verbose=verbose, \n\ )") + file.close() + print(f"\nDefault parameter file for '{self.__class__.__name__}' has been created in {path}.\n\ You can now launch with 'struphy run {self.__class__.__name__}' or with 'struphy run -i params_{self.__class__.__name__}.py'") + + return path ################### # Private methods : diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 15526bf44..f9783d35e 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -143,6 +143,22 @@ def update_scalar_quantities(self): en_Btot = self._tmp_b1.inner(self._tmp_b2) / 2 self.update_scalar("en_B_tot", en_Btot) + + ## default parameters + def generate_default_parameter_file(self, path = None, prompt = True): + params_path = super().generate_default_parameter_file(path, prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "mag_sonic.set_options" in line: + new_file += ["model.propagators.mag_sonic.set_options(b_field=model.em_fields.b_field)\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class LinearExtendedMHDuniform(StruphyModel): From d07d7166974203b9095ba7dcaa3ba1427384f3e8 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 7 Aug 2025 14:25:21 +0200 Subject: [PATCH 072/292] adapt marler propagators PushEta and PushVxB --- src/struphy/models/variables.py | 34 +++- src/struphy/propagators/base.py | 3 +- .../propagators/propagators_markers.py | 158 ++++++++---------- 3 files changed, 103 insertions(+), 92 deletions(-) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index ddf546038..cb1405bc1 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -1,5 +1,6 @@ from abc import ABCMeta, abstractmethod from mpi4py import MPI +import inspect from struphy.initial.base import InitialCondition from struphy.feec.psydac_derham import Derham, SplineFunction @@ -7,6 +8,9 @@ from struphy.initial.perturbations import Perturbation from struphy.geometry.base import Domain from struphy.fields_background.base import FluidEquilibrium +from struphy.pic.base import Particles +from struphy.kinetic_background.base import Maxwellian +from struphy.pic import particles class Variable(metaclass=ABCMeta): @@ -113,7 +117,35 @@ def allocate(self, derham: Derham, domain: Domain = None, equil: FluidEquilibriu class PICVariable(Variable): - pass + def __init__(self, name: str = "a_pic_var", space: str = "Particles6D"): + assert space in ("Particles6D", "Particles5D", "Particles3D", "ParticlesSPH", "DeltaFParticles6D") + self._name = name + self._space = space + + @property + def __name__(self): + return self._name + + @property + def space(self): + return self._space + + @property + def particles(self) -> Particles: + return self._particles + + def add_background(self, background: Maxwellian, verbose=True): + super().add_background(background, verbose=verbose) + + def allocate(self, derham: Derham, domain: Domain = None, equil: FluidEquilibrium = None,): + self._particles = getattr(particles, self.space)( + name=self.__name__, + space_id=self.space, + backgrounds=self.backgrounds, + perturbations=self.perturbations, + domain=domain, + equil=equil, + ) class SPHVariable(Variable): diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 51e8fd25c..6ad02f161 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -11,6 +11,7 @@ from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable from psydac.linalg.stencil import StencilVector from psydac.linalg.block import BlockVector +from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB class Propagator(metaclass=ABCMeta): @@ -162,7 +163,7 @@ def basis_ops(self, basis_ops): self._basis_ops = basis_ops @property - def projected_equil(self): + def projected_equil(self) -> ProjectedFluidEquilibriumWithB: """Fluid equilibrium projected on 3d Derham sequence with commuting projectors.""" assert hasattr( self, diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index bdc1898d8..3572206d3 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -2,6 +2,8 @@ import numpy as np from numpy import array, polynomial, random +from typing import Literal, get_args + from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -17,6 +19,9 @@ from struphy.pic.pushing.pusher import Pusher from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator +from struphy.models.variables import FEECVariable, PICVariable +from struphy.io.options import check_option +from psydac.linalg.basic import LinearOperator class PushEta(Propagator): @@ -35,42 +40,26 @@ class PushEta(Propagator): Available algorithms: * Explicit from :class:`~struphy.ode.utils.ButcherTableau` - - Parameters - ---------- - particles : Particles6D | ParticlesSPH - Particles object. - - algo : str - Algorithm for solving the ODE (see options below). - - density_field: StencilVector - Storage for density evaluation at each __call__. """ - - @staticmethod - def options(default=False): - dct = {} - dct["algo"] = ["rk4", "forward_euler", "heun2", "rk2", "heun3"] - if default: - dct = descend_options_dict(dct, []) - return dct - - def __init__( - self, - particles: Particles6D | ParticlesSPH, - *, - algo: str = options(default=True)["algo"], - density_field: StencilVector | None = None, - ): - # base class constructor call - super().__init__(particles) - + # variables to be updated + var: PICVariable = None + + ## abstract methods + def set_options(self, + butcher: ButcherTableau = None, + ): + + # use setter for options + self.options = self.Options(self, + butcher=butcher, + ) + + def allocate(self): # get kernel kernel = pusher_kernels.push_eta_stage # define algorithm - butcher = ButcherTableau(algo) + butcher: ButcherTableau = self.options.butcher # temp fix due to refactoring of ButcherTableau: import numpy as np @@ -84,7 +73,7 @@ def __init__( ) self._pusher = Pusher( - particles, + self.var.particles, kernel, args_kernel, self.domain.args_domain, @@ -93,28 +82,12 @@ def __init__( mpi_sort="each", ) - self._eval_density = False - if density_field is not None: - self._eval_density = True - self._density_field = density_field - def __call__(self, dt): self._pusher(dt) # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() - - if self._eval_density: - eval_density = lambda eta1, eta2, eta3: self.particles[0].eval_density( - eta1, - eta2, - eta3, - h1=0.1, - h2=0.1, - h3=0.1, - ) - self.derham.P["3"](eval_density, out=self._density_field) + if self.var.particles.control_variate: + self.var.particles.update_weights() class PushVxB(Propagator): @@ -133,46 +106,51 @@ class PushVxB(Propagator): Available algorithms: ``analytic``, ``implicit``. """ - - @staticmethod - def options(default=False): - dct = {} - dct["algo"] = ["analytic", "implicit"] - if default: - dct = descend_options_dict(dct, []) - return dct - - def __init__( - self, - particles: Particles6D, - *, - algo: str = options(default=True)["algo"], - kappa: float = 1.0, - b2: BlockVector | PolarVector, - b2_add: BlockVector | PolarVector = None, - ): + # variables to be updated + ions: PICVariable = None + + # propagator specific options + OptsAlgo = Literal["analytic", "implicit"] + + ## abstract methods + def set_options(self, + algo: OptsAlgo = "analytic", + kappa: float = 1.0, + b2_var: FEECVariable = None, + ): + + # checks + check_option(algo, self.OptsAlgo) + + # use setter for options + self.options = self.Options(self, + algo=algo, + kappa=kappa, + b2_var=b2_var, + ) + + def allocate(self): # TODO: treat PolarVector as well, but polar splines are being reworked at the moment - assert b2.space == self.derham.Vh["2"] - if b2_add is not None: - assert b2_add.space == self.derham.Vh["2"] - - # base class constructor call - super().__init__(particles) - - # parameters that need to be exposed - self._kappa = kappa - self._b2 = b2 - self._b2_add = b2_add + self._b2 = self.projected_equil.b2 + assert self._b2.space == self.derham.Vh["2"] + + if self.options.b2_var is None: + self._b2_var = None + else: + assert self.options.b2_var.spline.vector.space == self.derham.Vh["2"] + self._b2_var = self.options.b2_var.spline.vector + + # allocate dummy vectors to avoid temporary array allocations self._tmp = self.derham.Vh["2"].zeros() self._b_full = self.derham.Vh["2"].zeros() # define pusher kernel - if algo == "analytic": + if self.options.algo == "analytic": kernel = pusher_kernels.push_vxb_analytic - elif algo == "implicit": + elif self.options.algo == "implicit": kernel = pusher_kernels.push_vxb_implicit else: - raise ValueError(f"{algo = } not supported.") + raise ValueError(f"{self.options.algo = } not supported.") # instantiate Pusher args_kernel = ( @@ -183,7 +161,7 @@ def __init__( ) self._pusher = Pusher( - particles, + self.ions.particles, kernel, args_kernel, self.domain.args_domain, @@ -191,24 +169,24 @@ def __init__( ) # transposed extraction operator PolarVector --> BlockVector (identity map in case of no polar splines) - self._E2T = self.derham.extraction_ops["2"].transpose() + self._E2T: LinearOperator = self.derham.extraction_ops["2"].transpose() def __call__(self, dt): # sum up total magnetic field tmp = self._b2.copy(out=self._tmp) - if self._b2_add is not None: - tmp += self._b2_add + if self._b2_var is not None: + tmp += self._b2_var # extract coefficients to tensor product space - b_full = self._E2T.dot(tmp, out=self._b_full) + b_full: BlockVector = self._E2T.dot(tmp, out=self._b_full) b_full.update_ghost_regions() # call pusher kernel - self._pusher(self._kappa * dt) + self._pusher(self.options.kappa * dt) # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() + if self.ions.particles.control_variate: + self.ions.particles.update_weights() class PushVinEfield(Propagator): From 02346a37bae128fe499c2930612994fb447d9cea Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 7 Aug 2025 14:47:15 +0200 Subject: [PATCH 073/292] adapt toy model VLasov --- src/struphy/models/base.py | 5 +- src/struphy/models/toy.py | 112 +++++++----------- src/struphy/propagators/base.py | 5 +- .../propagators/propagators_markers.py | 3 + 4 files changed, 54 insertions(+), 71 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 934c91a3e..1273934c0 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1409,9 +1409,10 @@ def generate_default_parameter_file( model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=1))\n\ model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=2))\n" - exclude_feec = f"# model.{sn}.{vn}.save_data = False\n" + exclude = f"# model.{sn}.{vn}.save_data = False\n" elif isinstance(var, PICVariable): has_pic = True + exclude = f"# model.....save_data = False\n" elif isinstance(var, SPHVariable): has_sph = True @@ -1471,7 +1472,7 @@ def generate_default_parameter_file( file.write(init_pert_feec) file.write("\n# optional: exclude variables from saving\n") - file.write(exclude_feec) + file.write(exclude) file.write('\nif __name__ == "__main__":\n') file.write(" # start run\n") diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 043b36bef..bdc64cf08 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -114,83 +114,59 @@ class Vlasov(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + @dataclass + class KineticIons(KineticSpecies): + var: PICVariable = PICVariable(name="ions", space="Particles6D") + + ## propagators + + class Propagators: + def __init__(self): + self.push_vxb = propagators_markers.PushVxB() + self.push_eta = propagators_markers.PushEta() - dct["kinetic"]["ions"] = "Particles6D" - return dct + ## abstract methods - @staticmethod + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species, variables + self.kinetic_ions = self.KineticIons() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.push_vxb.set_variables( + ions = self.kinetic_ions.var, + ) + + self.propagators.push_eta.set_variables( + var = self.kinetic_ions.var, + ) + + # define scalars for update_scalar_quantities + self.add_scalar("en_f", compute="from_particles", species="ions") + + @property def bulk_species(): return "ions" - @staticmethod + @property def velocity_scale(): return "cyclotron" - - @staticmethod - def propagators_dct(): - return { - propagators_markers.PushVxB: ["ions"], - propagators_markers.PushEta: ["ions"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - from mpi4py.MPI import IN_PLACE, SUM - - # prelim - ions_params = self.kinetic["ions"]["params"] - - # project magnetic background - self._b_eq = self.derham.P["2"]( - [ - self.equil.b2_1, - self.equil.b2_2, - self.equil.b2_3, - ] - ) - - # set keyword arguments for propagators - self._kwargs[propagators_markers.PushVxB] = { - "algo": ions_params["options"]["PushVxB"]["algo"], - "kappa": 1.0, - "b2": self._b_eq, - "b2_add": None, - } - - self._kwargs[propagators_markers.PushEta] = {"algo": ions_params["options"]["PushEta"]["algo"]} - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation - self.add_scalar("en_f", compute="from_particles", species="ions") - - # MPI operations needed for scalar variables - self._mpi_sum = SUM - self._mpi_in_place = IN_PLACE - self._tmp = np.empty(1, dtype=float) + + def allocate_helpers(self): + self._tmp = np.empty(1, dtype=float) def update_scalar_quantities(self): - self._tmp[0] = self.pointer["ions"].markers_wo_holes[:, 6].dot( - self.pointer["ions"].markers_wo_holes[:, 3] ** 2 - + self.pointer["ions"].markers_wo_holes[:, 4] ** 2 - + self.pointer["ions"].markers_wo_holes[:, 5] ** 2, - ) / (2 * self.pointer["ions"].Np) - - # self.derham.comm.Allreduce( - # self._mpi_in_place, self._tmp, op=self._mpi_sum) + particles = self.kinetic_ions.var.particles + self._tmp[0] = particles.markers_wo_holes[:, 6].dot( + particles.markers_wo_holes[:, 3] ** 2 + + particles.markers_wo_holes[:, 4] ** 2 + + particles.markers_wo_holes[:, 5] ** 2, + ) / (2 * particles.Np) self.update_scalar("en_f", self._tmp[0]) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 6ad02f161..1493f2f65 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -44,7 +44,10 @@ def __call__(self, dt): """ def set_variables(self, **vars): - """Update variables in __dict__ and set self.variables with user-defined instance variables (allocated).""" + """Define the variables to be updated by the propagator. + + Update variables in __dict__ and set self.variables with user-defined instance variables (allocated). + """ assert len(vars) == len(self.__dict__), f"Variables must be passed in the following order: {self.__dict__}, but is {vars}." for ((k, v), (kp, vp)) in zip(vars.items(), self.__dict__.items()): assert k == kp, f"Variable name '{k}' not equal to '{kp}'; variables must be passed in the order {self.__dict__}." diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 3572206d3..963b9f4d9 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -3,6 +3,7 @@ import numpy as np from numpy import array, polynomial, random from typing import Literal, get_args +from dataclasses import dataclass from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -24,6 +25,7 @@ from psydac.linalg.basic import LinearOperator +@dataclass class PushEta(Propagator): r"""For each marker :math:`p`, solves @@ -90,6 +92,7 @@ def __call__(self, dt): self.var.particles.update_weights() +@dataclass class PushVxB(Propagator): r"""For each marker :math:`p`, solves From 6f4316f0b81db5222875dc73070ad6bc850f45b5 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 7 Aug 2025 16:48:37 +0200 Subject: [PATCH 074/292] work on the allocate() method for PICVariable --- src/struphy/io/options.py | 5 + src/struphy/models/base.py | 126 +++-------------- src/struphy/models/species.py | 51 ++++++- src/struphy/models/toy.py | 6 +- src/struphy/models/variables.py | 128 ++++++++++++++++-- .../propagators/propagators_markers.py | 3 + 6 files changed, 196 insertions(+), 123 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index f5498a3e9..cb8eb3b38 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -27,6 +27,11 @@ OptsGenSolver = Literal["pbicgstab", "bicgstab"] OptsMassPrecond = Literal["MassMatrixPreconditioner", None] +# markers +OptsMarkerBC = Literal["periodic", "reflect"] +OptsRecontructBC = Literal["periodic", "mirror", "fixed"] +OptsLoading = Literal["pseudo_random", "tesselation"] + ## Option classes diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 1273934c0..329732f73 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -598,109 +598,15 @@ def allocate_variables(self): for k, v in spec.variables.items(): assert isinstance(v, FEECVariable) v.allocate(derham=self.derham, domain=self.domain, equil=self.equil,) - - # marker arrays and plasma parameters of kinetic species + + # allocate memory for marker arrays of kinetic variables if self.kinetic_species: - for species, val in self.kinetic.items(): - assert any([key in val["params"]["markers"] for key in ["Np", "ppc", "ppb"]]) - - bckgr_params = val["params"].get("background", None) - pert_params = val["params"].get("perturbation", None) - boxes_per_dim = val["params"].get("boxes_per_dim", None) - mpi_dims_mask = val["params"].get("dims_mask", None) - weights_params = val["params"].get("weights", None) - - if self.derham is None: - domain_decomp = None - else: - domain_array = self.derham.domain_array - nprocs = self.derham.domain_decomposition.nprocs - domain_decomp = (domain_array, nprocs) - - kinetic_class = getattr(particles, val["space"]) - - val["obj"] = kinetic_class( - comm_world=self.comm_world, - clone_config=self.clone_config, - **val["params"]["markers"], - weights_params=weights_params, - domain_decomp=domain_decomp, - mpi_dims_mask=mpi_dims_mask, - boxes_per_dim=boxes_per_dim, - name=species, - equation_params=self.equation_params[species], - domain=self.domain, - equil=self.equil, - projected_equil=self.projected_equil, - bckgr_params=bckgr_params, - pert_params=pert_params, - ) - - obj = val["obj"] - assert isinstance(obj, Particles) - - self._pointer[species] = obj - - # for storing markers - val["kinetic_data"] = {} - - # for storing the distribution function - if "f" in val["params"]["save_data"]: - slices = val["params"]["save_data"]["f"]["slices"] - n_bins = val["params"]["save_data"]["f"]["n_bins"] - ranges = val["params"]["save_data"]["f"]["ranges"] - - val["kinetic_data"]["f"] = {} - val["kinetic_data"]["df"] = {} - val["bin_edges"] = {} - if len(slices) > 0: - for i, sli in enumerate(slices): - assert ((len(sli) - 2) / 3).is_integer() - assert len(slices[i].split("_")) == len(ranges[i]) == len(n_bins[i]), ( - f"Number of slices names ({len(slices[i].split('_'))}), number of bins ({len(n_bins[i])}), and number of ranges ({len(ranges[i])}) are inconsistent with each other!\n\n" - ) - val["bin_edges"][sli] = [] - dims = (len(sli) - 2) // 3 + 1 - for j in range(dims): - val["bin_edges"][sli] += [ - np.linspace( - ranges[i][j][0], - ranges[i][j][1], - n_bins[i][j] + 1, - ), - ] - val["kinetic_data"]["f"][sli] = np.zeros( - n_bins[i], - dtype=float, - ) - val["kinetic_data"]["df"][sli] = np.zeros( - n_bins[i], - dtype=float, - ) - - # for storing an sph evaluation of the density n - if "n_sph" in val["params"]["save_data"]: - plot_pts = val["params"]["save_data"]["n_sph"]["plot_pts"] - - val["kinetic_data"]["n_sph"] = [] - val["plot_pts"] = [] - for i, pts in enumerate(plot_pts): - assert len(pts) == 3 - eta1 = np.linspace(0.0, 1.0, pts[0]) - eta2 = np.linspace(0.0, 1.0, pts[1]) - eta3 = np.linspace(0.0, 1.0, pts[2]) - ee1, ee2, ee3 = np.meshgrid( - eta1, - eta2, - eta3, - indexing="ij", - ) - val["plot_pts"] += [(ee1, ee2, ee3)] - val["kinetic_data"]["n_sph"] += [np.zeros(ee1.shape, dtype=float)] - - # other data (wave-particle power exchange, etc.) - # TODO - + for species, spec in self.kinetic_species.items(): + assert isinstance(spec, KineticSpecies) + for k, v in spec.variables.items(): + assert isinstance(v, (PICVariable, SPHVariable)) + v.allocate(derham=self.derham, domain=self.domain, equil=self.equil,) + # TODO: allocate memory for FE coeffs of diagnostics # if self.params.diagnostic_fields is not None: # for key, val in self.diagnostics.items(): @@ -1385,7 +1291,8 @@ def generate_default_parameter_file( file.write("from struphy.geometry import domains\n") file.write("from struphy.fields_background import equils\n") - species_params = "\n# species parameters" + species_params = "\n# species parameters\n" + kinetic_params = "" has_plasma = False has_feec = False has_pic = False @@ -1395,8 +1302,12 @@ def generate_default_parameter_file( if isinstance(species, (FluidSpecies, KineticSpecies)): has_plasma = True - species_params += f"\nmodel.{sn}.set_phys_params()" - + species_params += f"model.{sn}.set_phys_params()\n" + if isinstance(species, KineticSpecies): + kinetic_params += f"model.{sn}.set_markers()\n\ +model.{sn}.set_sorting()\n\ +model.{sn}.set_save_data()\n" + for vn, var in species.variables.items(): if isinstance(var, FEECVariable): has_feec = True @@ -1460,7 +1371,10 @@ def generate_default_parameter_file( file.write("model = Model()\n") if has_plasma: - file.write(species_params + "\n") + file.write(species_params) + + if has_pic: + file.write(kinetic_params) file.write("\n# propagator options\n") for prop in self.propagators.__dict__: diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 5dac4caed..37cd82dfd 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -9,6 +9,7 @@ from struphy.initial.base import InitialCondition from struphy.kinetic_background.base import KineticBackground from struphy.io.options import Units +from struphy.io.options import OptsLoading, OptsMarkerBC, OptsRecontructBC from struphy.physics.physics import ConstantsOfNature from struphy.models.variables import Variable @@ -20,7 +21,8 @@ class Species(metaclass=ABCMeta): def __post_init__(self): for k, v in self.__dict__.items(): if isinstance(v, Variable): - v._species = self.__class__.__name__ + # v._species = self.__class__.__name__ + v._species = self @property def variables(self): @@ -92,7 +94,52 @@ class FluidSpecies(Species): class KineticSpecies(Species): """Single kinetic species in 3d + vdim phase space.""" - pass + def set_markers(self, + Np: int = 100, + ppc: int = None, + ppb: int = None, + bc: tuple[OptsMarkerBC] = ("periodic", "periodic", "periodic"), + bc_refill = None, + bc_sph: tuple[OptsRecontructBC] = ("periodic", "periodic", "periodic"), + bufsize: float = 1.0, + loading: OptsLoading = "pseudo_random", + loading_params: dict = None, + control_variate: bool = False, + reject_weights: dict = None, + ): + """Set marker parameters for loading, pushing, weight calculation + and kernel density reconstruction.""" + self.Np = Np + self.ppc = ppc + self.ppb = ppb + self.bc = bc + self.bc_refill = bc_refill + self.bc_sph = bc_sph + self.bufsize = bufsize + self.loading = loading + self.loading_params = loading_params + self.control_variate = control_variate + self.reject_weights = reject_weights + + def set_sorting(self, + boxes_per_dim: tuple = (16, 1, 1), + box_bufsize: float = 2.0, + dims_maks: tuple = (True, True, True), + ): + """For sorting markers in memory.""" + self.boxes_per_dim = boxes_per_dim + self.box_bufsize = box_bufsize + self.dims_mask = dims_maks + + def set_save_data(self, + n_markers: int | float = 3, + f_binned: dict = None, + n_sph: dict = None, + ): + """Saving marker orits, binned data and kernel density reconstructions.""" + self.n_markers = n_markers + self.f_binned = f_binned + self.n_sph = n_sph class DiagnosticSpecies(Species): diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index bdc64cf08..8afb0dcb6 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -150,11 +150,11 @@ def __init__(self): self.add_scalar("en_f", compute="from_particles", species="ions") @property - def bulk_species(): - return "ions" + def bulk_species(self): + return self.kinetic_ions @property - def velocity_scale(): + def velocity_scale(self): return "cyclotron" def allocate_helpers(self): diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index cb1405bc1..50e1ff170 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -8,9 +8,12 @@ from struphy.initial.perturbations import Perturbation from struphy.geometry.base import Domain from struphy.fields_background.base import FluidEquilibrium +from struphy.fields_background.projected_equils import ProjectedFluidEquilibrium from struphy.pic.base import Particles from struphy.kinetic_background.base import Maxwellian from struphy.pic import particles +from struphy.models.species import Species, KineticSpecies +from struphy.utils.clone_config import CloneConfig class Variable(metaclass=ABCMeta): @@ -45,7 +48,7 @@ def save_data(self, new): self._save_data = new @property - def species(self): + def species(self) -> Species: if not hasattr(self, "_species"): self._species = None return self._species @@ -60,7 +63,7 @@ def add_background(self, background, verbose=True): self._backgrounds += [background] if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nVariable '{self.__name__}' of species '{self.species}' - added background '{background.__class__.__name__}' with:") + print(f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added background '{background.__class__.__name__}' with:") for k, v in background.__dict__.items(): print(f' {k}: {v}') @@ -73,7 +76,7 @@ def add_perturbation(self, perturbation: Perturbation, verbose=True): self._perturbations += [perturbation] if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nVariable '{self.__name__}' of species '{self.species}' - added perturbation '{perturbation.__class__.__name__}' with:") + print(f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added perturbation '{perturbation.__class__.__name__}' with:") for k, v in perturbation.__dict__.items(): print(f' {k}: {v}') @@ -118,9 +121,10 @@ def allocate(self, derham: Derham, domain: Domain = None, equil: FluidEquilibriu class PICVariable(Variable): def __init__(self, name: str = "a_pic_var", space: str = "Particles6D"): - assert space in ("Particles6D", "Particles5D", "Particles3D", "ParticlesSPH", "DeltaFParticles6D") + assert space in ("Particles6D", "Particles5D", "Particles3D", "DeltaFParticles6D") self._name = name self._space = space + self._kinetic_data = None @property def __name__(self): @@ -134,18 +138,118 @@ def space(self): def particles(self) -> Particles: return self._particles + @property + def kinetic_data(self): + return self._kinetic_data + def add_background(self, background: Maxwellian, verbose=True): super().add_background(background, verbose=verbose) - def allocate(self, derham: Derham, domain: Domain = None, equil: FluidEquilibrium = None,): - self._particles = getattr(particles, self.space)( - name=self.__name__, - space_id=self.space, - backgrounds=self.backgrounds, - perturbations=self.perturbations, - domain=domain, - equil=equil, + def allocate(self, + clone_config: CloneConfig = None, + derham: Derham = None, + domain: Domain = None, + equil: FluidEquilibrium = None, + projected_equil: ProjectedFluidEquilibrium = None, + ): + + assert isinstance(self.species, KineticSpecies) + + if derham is None: + domain_decomp = None + else: + domain_array = derham.domain_array + nprocs = derham.domain_decomposition.nprocs + domain_decomp = (domain_array, nprocs) + + kinetic_class = getattr(particles, self.space) + + self._particles: Particles = kinetic_class( + comm_world=MPI.COMM_WORLD, + clone_config=clone_config, + Np=self.species.Np, + ppc=self.species.ppc, + domain_decomp=domain_decomp, + mpi_dims_mask=self.species.dims_mask, + ppb=self.species.ppb, + boxes_per_dim=self.species.boxes_per_dim, + box_bufsize=self.species.box_bufsize, + bc=self.species.bc, + bc_refill=self.species.bc_refill, + control_variate=self.species.control_variate, + name=self.species.__class__.__name__, + # bc_sph=self.species.bc_sph, + loading=self.species.loading, + loading_params=self.species.loading_params, + weights_params=self.species.reject_weights, + bufsize=self.species.bufsize, + domain=domain, + equil=equil, + projected_equil=projected_equil, + bckgr_params=self.backgrounds, + pert_params=self.perturbations, + equation_params=self.species.equation_params, + ) + + # for storing markers + self._kinetic_data = {} + + # for storing the distribution function + if "f" in val["params"]["save_data"]: + slices = val["params"]["save_data"]["f"]["slices"] + n_bins = val["params"]["save_data"]["f"]["n_bins"] + ranges = val["params"]["save_data"]["f"]["ranges"] + + val["kinetic_data"]["f"] = {} + val["kinetic_data"]["df"] = {} + val["bin_edges"] = {} + if len(slices) > 0: + for i, sli in enumerate(slices): + assert ((len(sli) - 2) / 3).is_integer() + assert len(slices[i].split("_")) == len(ranges[i]) == len(n_bins[i]), ( + f"Number of slices names ({len(slices[i].split('_'))}), number of bins ({len(n_bins[i])}), and number of ranges ({len(ranges[i])}) are inconsistent with each other!\n\n" ) + val["bin_edges"][sli] = [] + dims = (len(sli) - 2) // 3 + 1 + for j in range(dims): + val["bin_edges"][sli] += [ + np.linspace( + ranges[i][j][0], + ranges[i][j][1], + n_bins[i][j] + 1, + ), + ] + val["kinetic_data"]["f"][sli] = np.zeros( + n_bins[i], + dtype=float, + ) + val["kinetic_data"]["df"][sli] = np.zeros( + n_bins[i], + dtype=float, + ) + + # for storing an sph evaluation of the density n + if "n_sph" in val["params"]["save_data"]: + plot_pts = val["params"]["save_data"]["n_sph"]["plot_pts"] + + val["kinetic_data"]["n_sph"] = [] + val["plot_pts"] = [] + for i, pts in enumerate(plot_pts): + assert len(pts) == 3 + eta1 = np.linspace(0.0, 1.0, pts[0]) + eta2 = np.linspace(0.0, 1.0, pts[1]) + eta3 = np.linspace(0.0, 1.0, pts[2]) + ee1, ee2, ee3 = np.meshgrid( + eta1, + eta2, + eta3, + indexing="ij", + ) + val["plot_pts"] += [(ee1, ee2, ee3)] + val["kinetic_data"]["n_sph"] += [np.zeros(ee1.shape, dtype=float)] + + # other data (wave-particle power exchange, etc.) + # TODO class SPHVariable(Variable): diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 963b9f4d9..78a66beb4 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -51,6 +51,9 @@ def set_options(self, butcher: ButcherTableau = None, ): + if butcher is None: + butcher = ButcherTableau() + # use setter for options self.options = self.Options(self, butcher=butcher, From 86ad11662598057da05379834be67aa250a5bfcc Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 7 Aug 2025 16:58:35 +0200 Subject: [PATCH 075/292] work on storing pic data --- src/struphy/models/variables.py | 43 +++++++++++++++------------------ 1 file changed, 20 insertions(+), 23 deletions(-) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 50e1ff170..5d3903616 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -1,6 +1,6 @@ from abc import ABCMeta, abstractmethod from mpi4py import MPI -import inspect +import numpy as np from struphy.initial.base import InitialCondition from struphy.feec.psydac_derham import Derham, SplineFunction @@ -124,7 +124,7 @@ def __init__(self, name: str = "a_pic_var", space: str = "Particles6D"): assert space in ("Particles6D", "Particles5D", "Particles3D", "DeltaFParticles6D") self._name = name self._space = space - self._kinetic_data = None + self._kinetic_data = {} @property def __name__(self): @@ -191,49 +191,46 @@ def allocate(self, equation_params=self.species.equation_params, ) - # for storing markers - self._kinetic_data = {} - - # for storing the distribution function - if "f" in val["params"]["save_data"]: - slices = val["params"]["save_data"]["f"]["slices"] - n_bins = val["params"]["save_data"]["f"]["n_bins"] - ranges = val["params"]["save_data"]["f"]["ranges"] + # for storing the binned distribution function + if self.species.f_binned is not None: + slices = self.species.f_binned["slices"] + n_bins = self.species.f_binned["n_bins"] + ranges = self.species.f_binned["ranges"] - val["kinetic_data"]["f"] = {} - val["kinetic_data"]["df"] = {} - val["bin_edges"] = {} + self.kinetic_data["f"] = {} + self.kinetic_data["df"] = {} + self.kinetic_data["bin_edges"] = {} if len(slices) > 0: for i, sli in enumerate(slices): assert ((len(sli) - 2) / 3).is_integer() assert len(slices[i].split("_")) == len(ranges[i]) == len(n_bins[i]), ( f"Number of slices names ({len(slices[i].split('_'))}), number of bins ({len(n_bins[i])}), and number of ranges ({len(ranges[i])}) are inconsistent with each other!\n\n" ) - val["bin_edges"][sli] = [] + self.kinetic_data["bin_edges"][sli] = [] dims = (len(sli) - 2) // 3 + 1 for j in range(dims): - val["bin_edges"][sli] += [ + self.kinetic_data["bin_edges"][sli] += [ np.linspace( ranges[i][j][0], ranges[i][j][1], n_bins[i][j] + 1, ), ] - val["kinetic_data"]["f"][sli] = np.zeros( + self.kinetic_data["f"][sli] = np.zeros( n_bins[i], dtype=float, ) - val["kinetic_data"]["df"][sli] = np.zeros( + self.kinetic_data["df"][sli] = np.zeros( n_bins[i], dtype=float, ) # for storing an sph evaluation of the density n - if "n_sph" in val["params"]["save_data"]: - plot_pts = val["params"]["save_data"]["n_sph"]["plot_pts"] + if self.species.n_sph is not None: + plot_pts = self.species.n_sph["plot_pts"] - val["kinetic_data"]["n_sph"] = [] - val["plot_pts"] = [] + self.kinetic_data["n_sph"] = [] + self.kinetic_data["plot_pts"] = [] for i, pts in enumerate(plot_pts): assert len(pts) == 3 eta1 = np.linspace(0.0, 1.0, pts[0]) @@ -245,8 +242,8 @@ def allocate(self, eta3, indexing="ij", ) - val["plot_pts"] += [(ee1, ee2, ee3)] - val["kinetic_data"]["n_sph"] += [np.zeros(ee1.shape, dtype=float)] + self.kinetic_data["plot_pts"] += [(ee1, ee2, ee3)] + self.kinetic_data["n_sph"] += [np.zeros(ee1.shape, dtype=float)] # other data (wave-particle power exchange, etc.) # TODO From 4f9ec66142105c1312067c9ff867489805f8f69f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 11 Aug 2025 10:03:55 +0200 Subject: [PATCH 076/292] Re-factor Maxwllian3D and base class: - remove __init__ from base class - Maxwellian moments must be passed as tuples of the form (float | callable, Perturbation) - adapt .py parameter file --- src/struphy/kinetic_background/base.py | 151 ++++-------------- src/struphy/kinetic_background/maxwellians.py | 54 +++---- src/struphy/main.py | 4 +- src/struphy/models/base.py | 6 + src/struphy/models/variables.py | 20 ++- src/struphy/pic/base.py | 49 +++--- 6 files changed, 105 insertions(+), 179 deletions(-) diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 3538674ec..39013b89f 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -1,13 +1,12 @@ "Base classes for kinetic backgrounds." -import copy from abc import ABCMeta, abstractmethod - import numpy as np +from typing import Callable from struphy.fields_background.base import FluidEquilibrium from struphy.fields_background.equils import set_defaults -from struphy.initial import perturbations +from struphy.initial.base import Perturbation from struphy.initial.utilities import Noise from struphy.kinetic_background import moment_functions @@ -329,40 +328,6 @@ class Maxwellian(KineticBackground): and the thermal velocities :math:`v_{\mathrm{th},i}(\boldsymbol{\eta})`. """ - def __init__( - self, - maxw_params: dict = None, - pert_params: dict = None, - equil: FluidEquilibrium = None, - ): - # Set background parameters - if maxw_params is None: - maxw_params = {} - assert isinstance(maxw_params, dict) - self._maxw_params = set_defaults( - maxw_params, - self.default_maxw_params(), - ) - - # check if fluid background is needed - for key, val in self.maxw_params.items(): - if val == "fluid_background": - assert equil is not None - - # parameters for perturbation - if pert_params is None: - pert_params = {} - assert isinstance(pert_params, dict) - self._pert_params = pert_params - - # Fluid equilibrium - self._equil = equil - - @classmethod - def default_maxw_params(cls): - """Default parameters dictionary defining constant moments of the Maxwellian.""" - pass - @abstractmethod def vth(self, *etas): """Thermal velocities (0-forms). @@ -379,21 +344,17 @@ def vth(self, *etas): pass @property - def maxw_params(self): - """Parameters dictionary defining constant moments of the Maxwellian.""" - return self._maxw_params - - @property - def pert_params(self): - """Parameters dictionary defining the perturbations.""" - return self._pert_params - - @property - def equil(self): - """One of :mod:`~struphy.fields_background.equils` - in case that moments are to be set in that way, None otherwise. - """ - return self._equil + @abstractmethod + def maxw_params(self) -> dict: + """Parameters dictionary defining moments of the Maxwellian.""" + + def check_maxw_params(self): + for k, v in self.maxw_params.items(): + assert isinstance(k, str) + assert isinstance(v, tuple), f"Maxwallian parameter {k} must be tuple, but is {v}" + assert len(v) == 2 + assert isinstance(v[0], (float, int, Callable[..., float])) + assert isinstance(v[1], (Perturbation, None)) @classmethod def gaussian(self, v, u=0.0, vth=1.0, polar=False, volume_form=False): @@ -528,6 +489,10 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name="n"): assert isinstance(eta2, np.ndarray) assert isinstance(eta3, np.ndarray) assert eta1.shape == eta2.shape == eta3.shape + + params = self.maxw_params[name] + assert isinstance(params, tuple) + assert len(params) == 2 # flat evaluation for markers if eta1.ndim == 1: @@ -565,74 +530,26 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name="n"): else: out = 0.0 * etas[0] - # correspondence name -> equilibrium attribute - dct = { - "n": "n0", - "u1": "u_cart_1", - "u2": "u_cart_2", - "u3": "u_cart_3", - "vth1": "vth0", - "vth2": "vth0", - "vth3": "vth0", - "u_para": "u_para0", - "u_perp": None, - "vth_para": "vth0", - "vth_perp": "vth0", - } - - # fluid background - if self.maxw_params[name] == "fluid_background": - if dct[name] is not None: - out += getattr(self.equil, dct[name])(*etas) - if name in ("n") or "vth" in name: - assert np.all(out > 0.0), f"{name} must be positive!" - else: - print(f'Moment evaluation with "fluid_background" not implemented for {name}.') - - # when using moment functions, see test https://gitlab.mpcdf.mpg.de/struphy/struphy/-/blob/devel/src/struphy/kinetic_background/tests/test_maxwellians.py?ref_type=heads#L1760 - elif isinstance(self.maxw_params[name], dict): - mom_funcs = copy.deepcopy(self.maxw_params[name]) - for typ, params in mom_funcs.items(): - assert params["given_in_basis"] == "0", "Moment functions must be passed as 0-forms to Maxwellians." - params.pop("given_in_basis") - nfun = getattr(moment_functions, typ)(**params) - if eta1.ndim == 1: - out += nfun(eta1, eta2, eta3) - else: - out += nfun(*etas) - - # constant background + # evaluate background + background = params[0] + if isinstance(background, (float, int)): + out += background else: + assert callable(background) if eta1.ndim == 1: - out += self.maxw_params[name] + out += background(eta1, eta2, eta3) else: - out += self.maxw_params[name] - - # add possible perturbations - if name in self.pert_params: - pp_copy = copy.deepcopy(self.pert_params) - for pert, params in pp_copy[name].items(): - if pert == "Noise": - noise = Noise(**params) - if eta1.ndim == 1: - out += noise(eta1, eta2, eta3) - else: - out += noise(*etas) - else: - assert params["given_in_basis"] == "0", ( - "Moment perturbations must be passed as 0-forms to Maxwellians." - ) - params.pop("given_in_basis") - - perturbation = getattr(perturbations, pert)( - **params, - ) - - if eta1.ndim == 1: - out += perturbation(eta1, eta2, eta3) - else: - out += perturbation(*etas) - + out += background(*etas) + + # add perturbation + perturbation = params[1] + if perturbation is not None: + assert isinstance(perturbation, Perturbation) + if eta1.ndim == 1: + out += perturbation(eta1, eta2, eta3) + else: + out += perturbation(*etas) + return out diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index 12c27e06a..91dd7fff7 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -1,11 +1,13 @@ "Maxwellian (Gaussian) distributions in velocity space." import numpy as np +from typing import Callable from struphy.fields_background.base import FluidEquilibrium from struphy.fields_background.equils import set_defaults from struphy.kinetic_background import moment_functions from struphy.kinetic_background.base import CanonicalMaxwellian, Maxwellian +from struphy.initial.base import Perturbation class Maxwellian3D(Maxwellian): @@ -13,40 +15,34 @@ class Maxwellian3D(Maxwellian): Parameters ---------- - maxw_params : dict - Parameters for the kinetic background. + n : float | Callable + Fluid density; callables can come from FluidEquilibrium. - pert_params : dict + perturbation : dict Parameters for the kinetic perturbation added to the background. - - equil : FluidEquilibrium - One of :mod:`~struphy.fields_background.equils`. """ - @classmethod - def default_maxw_params(cls): - """Default parameters dictionary defining constant moments of the Maxwellian.""" - return { - "n": 1.0, - "u1": 0.0, - "u2": 0.0, - "u3": 0.0, - "vth1": 1.0, - "vth2": 1.0, - "vth3": 1.0, - } - def __init__( self, - maxw_params: dict = None, - pert_params: dict = None, - equil: FluidEquilibrium = None, + n: tuple[float | Callable[..., float], Perturbation] = (1.0, None), + u1: tuple[float | Callable[..., float], Perturbation] = (0.0, None), + u2: tuple[float | Callable[..., float], Perturbation] = (0.0, None), + u3: tuple[float | Callable[..., float], Perturbation] = (0.0, None), + vth1: tuple[float | Callable[..., float], Perturbation] = (1.0, None), + vth2: tuple[float | Callable[..., float], Perturbation] = (1.0, None), + vth3: tuple[float | Callable[..., float], Perturbation] = (1.0, None), ): - super().__init__( - maxw_params=maxw_params, - pert_params=pert_params, - equil=equil, - ) + + self._maxw_params = {} + self._maxw_params["n"] = n + self._maxw_params["u1"] = u1 + self._maxw_params["u2"] = u2 + self._maxw_params["u3"] = u3 + self._maxw_params["vth1"] = vth1 + self._maxw_params["vth2"] = vth2 + self._maxw_params["vth3"] = vth3 + + self.check_maxw_params() # factors multiplied onto the defined moments n, u and vth (can be set via setter) self._moment_factors = { @@ -55,6 +51,10 @@ def __init__( "vth": [1.0, 1.0, 1.0], } + @property + def maxw_params(self): + return self._maxw_params + @property def coords(self): """Coordinates of the Maxwellian6D, :math:`(v_1, v_2, v_3)`.""" diff --git a/src/struphy/main.py b/src/struphy/main.py index 34c15798c..4aa846158 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -174,6 +174,9 @@ def run( model._projected_equil = None print("GRID:\nNo grid specified - no Derham complex set up.") + # equation paramters + model.setup_equation_params(units=model.units, verbose=verbose) + # allocate variables model.allocate_variables() model.allocate_helpers() @@ -183,7 +186,6 @@ def run( # plasma parameters model.compute_plasma_params(verbose=verbose) - model.setup_equation_params(units=model.units, verbose=verbose) if rank < 32: if rank == 0: diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 329732f73..62f388b52 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1323,6 +1323,9 @@ def generate_default_parameter_file( exclude = f"# model.{sn}.{vn}.save_data = False\n" elif isinstance(var, PICVariable): has_pic = True + init_pert_pic = f"perturbation = perturbations.TorusModesCos()\n" + init_bckgr_pic = f"model.{sn}.{vn}.add_background(maxwellians.Maxwellian3D(perturbation=perturbation))\n" + exclude = f"# model.....save_data = False\n" elif isinstance(var, SPHVariable): has_sph = True @@ -1384,6 +1387,9 @@ def generate_default_parameter_file( if has_feec: file.write(init_bckgr_feec) file.write(init_pert_feec) + if has_pic: + file.write(init_pert_pic) + file.write(init_bckgr_pic) file.write("\n# optional: exclude variables from saving\n") file.write(exclude) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 5d3903616..59f1038ac 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -1,3 +1,8 @@ + +# for type checking (cyclic imports) +from __future__ import annotations +from typing import TYPE_CHECKING + from abc import ABCMeta, abstractmethod from mpi4py import MPI import numpy as np @@ -12,9 +17,10 @@ from struphy.pic.base import Particles from struphy.kinetic_background.base import Maxwellian from struphy.pic import particles -from struphy.models.species import Species, KineticSpecies from struphy.utils.clone_config import CloneConfig +if TYPE_CHECKING: + from struphy.models.species import Species, KineticSpecies class Variable(metaclass=ABCMeta): """Single variable (unknown) of a Species.""" @@ -142,6 +148,12 @@ def particles(self) -> Particles: def kinetic_data(self): return self._kinetic_data + @property + def species(self) -> KineticSpecies: + if not hasattr(self, "_species"): + self._species = None + return self._species + def add_background(self, background: Maxwellian, verbose=True): super().add_background(background, verbose=verbose) @@ -153,7 +165,7 @@ def allocate(self, projected_equil: ProjectedFluidEquilibrium = None, ): - assert isinstance(self.species, KineticSpecies) + #assert isinstance(self.species, KineticSpecies) if derham is None: domain_decomp = None @@ -186,8 +198,8 @@ def allocate(self, domain=domain, equil=equil, projected_equil=projected_equil, - bckgr_params=self.backgrounds, - pert_params=self.perturbations, + backgrounds=self.backgrounds, + perturbations=self.perturbations, equation_params=self.species.equation_params, ) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index b408aea6c..565be60c3 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -17,9 +17,10 @@ from struphy.fields_background.projected_equils import ProjectedFluidEquilibrium from struphy.geometry.base import Domain from struphy.geometry.utilities import TransformedPformComponent -from struphy.initial import perturbations +from struphy.initial.base import Perturbation from struphy.io.output_handling import DataContainer -from struphy.kinetic_background import maxwellians +from struphy.kinetic_background.maxwellians import Maxwellian3D +from struphy.kinetic_background.base import Maxwellian from struphy.pic import sampling_kernels, sobol_seq from struphy.pic.pushing.pusher_args_kernels import MarkerArguments from struphy.pic.pushing.pusher_utilities_kernels import reflect @@ -129,10 +130,10 @@ class Particles(metaclass=ABCMeta): projected_equil : ProjectedFluidEquilibrium Struphy fluid equilibrium projected into a discrete Derham complex. - bckgr_params : dict + backgrounds : Maxwellian | list Kinetic background parameters. - pert_params : dict + perturbations : Perturbation | list Kinetic perturbation parameters. equation_params : dict @@ -165,8 +166,8 @@ def __init__( domain: Domain = None, equil: FluidEquilibrium = None, projected_equil: ProjectedFluidEquilibrium = None, - bckgr_params: dict = None, - pert_params: dict = None, + backgrounds: Maxwellian | list = None, + perturbations: Perturbation | list = None, equation_params: dict = None, verbose: bool = False, ): @@ -310,9 +311,8 @@ def __init__( ) # background - if bckgr_params is None: - bckgr_params = {"Maxwellian3D": {}, "pforms": [None, None]} - self._bckgr_params = bckgr_params + if backgrounds is None: + self._backgrounds = Maxwellian3D() # background p-form description in [eta, v] (None means 0-form, "vol" means volume form -> divide by det) if isinstance(bckgr_params, FluidEquilibrium): @@ -325,7 +325,7 @@ def __init__( self._set_background_coordinates() # perturbation parameters - self._pert_params = pert_params + self._perturbations = perturbations # for loading if self.loading_params["moments"] is None and self.type != "sph" and isinstance(self.bckgr_params, dict): @@ -512,14 +512,14 @@ def clone_id(self): return self._clone_id @property - def bckgr_params(self): - """Kinetic background parameters.""" - return self._bckgr_params + def backgrounds(self): + """Kinetic backgrounds.""" + return self._backgrounds @property - def pert_params(self): - """Kinetic perturbation parameters.""" - return self._pert_params + def perturbations(self): + """Kinetic perturbations.""" + return self._perturbations @property def loading_params(self): @@ -965,12 +965,7 @@ def _set_background_function(self): if isinstance(self.bckgr_params, FluidEquilibrium): self._f0 = self.bckgr_params else: - for fi, maxw_params in self.bckgr_params.items(): - if fi[-2] == "_": - fi_type = fi[:-2] - else: - fi_type = fi - + for bckgr in self.backgrounds: # SPH case: f0 is set to a FluidEquilibrium if self.type == "sph": _eq = getattr(equils, fi_type)(**maxw_params) @@ -984,15 +979,9 @@ def _set_background_function(self): # default case else: if self._f0 is None: - self._f0 = getattr(maxwellians, fi_type)( - maxw_params=maxw_params, - equil=self.equil, - ) + self._f0 = bckgr else: - self._f0 = self._f0 + getattr(maxwellians, fi_type)( - maxw_params=maxw_params, - equil=self.equil, - ) + self._f0 = self._f0 + bckgr def _set_background_coordinates(self): if self.type != "sph" and self.f0.coords == "constants_of_motion": From c5349847103eb3b7f980d0aa7c795e4b0c5bd66e Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 12 Aug 2025 10:27:42 +0200 Subject: [PATCH 077/292] set background in Particles according to new framework - no dicts, but the background is passed directly from the params file. --- src/struphy/initial/base.py | 28 ------ src/struphy/kinetic_background/base.py | 6 +- src/struphy/kinetic_background/maxwellians.py | 8 +- src/struphy/models/fluid.py | 2 +- src/struphy/models/species.py | 1 - src/struphy/models/variables.py | 27 +++--- src/struphy/pic/base.py | 93 ++++++++++--------- src/struphy/pic/particles.py | 12 ++- 8 files changed, 80 insertions(+), 97 deletions(-) diff --git a/src/struphy/initial/base.py b/src/struphy/initial/base.py index 3a5efefa6..899053937 100644 --- a/src/struphy/initial/base.py +++ b/src/struphy/initial/base.py @@ -1,9 +1,6 @@ from typing import Callable from abc import ABCMeta, abstractmethod -from struphy.fields_background.base import FluidEquilibrium -from struphy.io.options import FieldsBackground -from struphy.kinetic_background.base import KineticBackground from struphy.io.options import GivenInBasis, check_option @@ -48,29 +45,4 @@ def comp(self) -> int: def comp(self, new: int): assert new in (0, 1, 2) self._comp = new - - - -class InitialCondition: - """Callable initial condition as sum of background + perturbation.""" - def __init__(self, - background: list = None, - perturbation: list = None, - equil: FluidEquilibrium = None,): - - for b in background: - assert isinstance(b, (FieldsBackground, KineticBackground)) - self._background = background - - for p in perturbation: - assert isinstance(p, tuple) - assert len(p) == 2 - assert isinstance(p[0], Callable) - assert isinstance(p[1], tuple) - self._perturbation = perturbation - - self._equil = equil - - def __call__(eta1, eta2, eta3, *v): - return 1 \ No newline at end of file diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 39013b89f..96623260a 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -46,8 +46,8 @@ def is_polar(self): @property @abstractmethod - def volume_form(self): - """Boolean. True if the background is represented as a volume form (thus including the velocity Jacobian).""" + def volume_form(self) -> bool: + """True if the background is represented as a volume form (thus including the velocity Jacobian).""" pass @abstractmethod @@ -354,7 +354,7 @@ def check_maxw_params(self): assert isinstance(v, tuple), f"Maxwallian parameter {k} must be tuple, but is {v}" assert len(v) == 2 assert isinstance(v[0], (float, int, Callable[..., float])) - assert isinstance(v[1], (Perturbation, None)) + assert isinstance(v[1], Perturbation) or v[1] is None @classmethod def gaussian(self, v, u=0.0, vth=1.0, polar=False, volume_form=False): diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index 91dd7fff7..a9371b8cb 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -15,11 +15,9 @@ class Maxwellian3D(Maxwellian): Parameters ---------- - n : float | Callable - Fluid density; callables can come from FluidEquilibrium. - - perturbation : dict - Parameters for the kinetic perturbation added to the background. + n, ui, vthi : tuple + Moments of the Maxwellian as tuples. The first entry defines the background + (float for constant background or callable), the second entry defines a Perturbation (can be None). """ def __init__( diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index f9783d35e..0a2646c95 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -146,7 +146,7 @@ def update_scalar_quantities(self): ## default parameters def generate_default_parameter_file(self, path = None, prompt = True): - params_path = super().generate_default_parameter_file(path, prompt) + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) new_file = [] with open(params_path, "r") as f: for line in f: diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 37cd82dfd..409f4930d 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -6,7 +6,6 @@ from mpi4py import MPI from struphy.fields_background.base import FluidEquilibrium -from struphy.initial.base import InitialCondition from struphy.kinetic_background.base import KineticBackground from struphy.io.options import Units from struphy.io.options import OptsLoading, OptsMarkerBC, OptsRecontructBC diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 59f1038ac..692e37a12 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -7,7 +7,6 @@ from mpi4py import MPI import numpy as np -from struphy.initial.base import InitialCondition from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.io.options import FieldsBackground from struphy.initial.perturbations import Perturbation @@ -15,7 +14,7 @@ from struphy.fields_background.base import FluidEquilibrium from struphy.fields_background.projected_equils import ProjectedFluidEquilibrium from struphy.pic.base import Particles -from struphy.kinetic_background.base import Maxwellian +from struphy.kinetic_background.base import KineticBackground from struphy.pic import particles from struphy.utils.clone_config import CloneConfig @@ -86,12 +85,6 @@ def add_perturbation(self, perturbation: Perturbation, verbose=True): for k, v in perturbation.__dict__.items(): print(f' {k}: {v}') - def define_initial_condition(self): - self._initial_condition = InitialCondition( - background=self.backgrounds, - perturbation=self.perturbations, - ) - class FEECVariable(Variable): def __init__(self, name: str = "a_feec_var", space: str = "H1"): @@ -154,7 +147,18 @@ def species(self) -> KineticSpecies: self._species = None return self._species - def add_background(self, background: Maxwellian, verbose=True): + @property + def n_as_volume_form(self) -> bool: + """Whether the number density n is given as a volume form or scalar function (=default).""" + if not hasattr(self, "_n_as_volume_form"): + self._n_as_volume_form = False + return self._n_as_volume_form + + def add_background(self, + background: KineticBackground, + n_as_volume_form: bool = False, + verbose=True): + self._n_as_volume_form = n_as_volume_form super().add_background(background, verbose=verbose) def allocate(self, @@ -166,6 +170,7 @@ def allocate(self, ): #assert isinstance(self.species, KineticSpecies) + assert isinstance(self.backgrounds, KineticBackground), f"List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." if derham is None: domain_decomp = None @@ -198,8 +203,8 @@ def allocate(self, domain=domain, equil=equil, projected_equil=projected_equil, - backgrounds=self.backgrounds, - perturbations=self.perturbations, + background=self.backgrounds, + # perturbations=self.perturbations, equation_params=self.species.equation_params, ) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 565be60c3..0c1a3c71b 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -19,8 +19,7 @@ from struphy.geometry.utilities import TransformedPformComponent from struphy.initial.base import Perturbation from struphy.io.output_handling import DataContainer -from struphy.kinetic_background.maxwellians import Maxwellian3D -from struphy.kinetic_background.base import Maxwellian +from struphy.kinetic_background.base import KineticBackground from struphy.pic import sampling_kernels, sobol_seq from struphy.pic.pushing.pusher_args_kernels import MarkerArguments from struphy.pic.pushing.pusher_utilities_kernels import reflect @@ -130,8 +129,11 @@ class Particles(metaclass=ABCMeta): projected_equil : ProjectedFluidEquilibrium Struphy fluid equilibrium projected into a discrete Derham complex. - backgrounds : Maxwellian | list + background : KineticBackground Kinetic background parameters. + + n_as_volume_form: bool + Whether the number density n is given as a volume form or scalar function (=default). perturbations : Perturbation | list Kinetic perturbation parameters. @@ -166,8 +168,9 @@ def __init__( domain: Domain = None, equil: FluidEquilibrium = None, projected_equil: ProjectedFluidEquilibrium = None, - backgrounds: Maxwellian | list = None, - perturbations: Perturbation | list = None, + background: KineticBackground = None, + n_as_volume_form: bool = False, + # perturbations: Perturbation | list = None, equation_params: dict = None, verbose: bool = False, ): @@ -311,25 +314,28 @@ def __init__( ) # background - if backgrounds is None: - self._backgrounds = Maxwellian3D() + if background is None: + raise ValueError("A background function must be passed to Particles.") + else: + self._background = background # background p-form description in [eta, v] (None means 0-form, "vol" means volume form -> divide by det) - if isinstance(bckgr_params, FluidEquilibrium): - self._pforms = [None, None] + if isinstance(background, FluidEquilibrium): + self._pforms = (False, False) else: - self._pforms = bckgr_params.pop("pforms", [None, None]) + self._pforms = (n_as_volume_form, + self.background.volume_form,) # set background function self._set_background_function() self._set_background_coordinates() # perturbation parameters - self._perturbations = perturbations + # self._perturbations = perturbations # for loading - if self.loading_params["moments"] is None and self.type != "sph" and isinstance(self.bckgr_params, dict): - self._auto_sampling_params() + # if self.loading_params["moments"] is None and self.type != "sph" and isinstance(self.bckgr_params, dict): + # self._auto_sampling_params() # create buffers for mpi_sort_markers if self.mpi_comm is not None: @@ -343,8 +349,8 @@ def __init__( @classmethod @abstractmethod - def default_bckgr_params(cls): - """Dictionary holding the minimal information of the default background.""" + def default_background(cls): + """The default background (of type Maxwellian).""" pass @abstractmethod @@ -512,14 +518,14 @@ def clone_id(self): return self._clone_id @property - def backgrounds(self): - """Kinetic backgrounds.""" - return self._backgrounds + def background(self) -> KineticBackground: + """Kinetic background.""" + return self._background - @property - def perturbations(self): - """Kinetic perturbations.""" - return self._perturbations + # @property + # def perturbations(self): + # """Kinetic perturbations.""" + # return self._perturbations @property def loading_params(self): @@ -961,27 +967,28 @@ def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): return dom_arr, tuple(nprocs) def _set_background_function(self): - self._f0 = None - if isinstance(self.bckgr_params, FluidEquilibrium): - self._f0 = self.bckgr_params - else: - for bckgr in self.backgrounds: - # SPH case: f0 is set to a FluidEquilibrium - if self.type == "sph": - _eq = getattr(equils, fi_type)(**maxw_params) - if not isinstance(_eq, NumericalFluidEquilibrium): - _eq.domain = self.domain - if self._f0 is None: - self._f0 = _eq - else: - raise NotImplementedError("Summation of fluid backgrounds not yet implemented.") - # self._f0 = self._f0 + (lambda e1, e2, e3: _eq.n0(e1, e2, e3)) - # default case - else: - if self._f0 is None: - self._f0 = bckgr - else: - self._f0 = self._f0 + bckgr + self._f0 = self.background + # self._f0 = None + # if isinstance(self.bckgr_params, FluidEquilibrium): + # self._f0 = self.bckgr_params + # else: + # for bckgr in self.backgrounds: + # # SPH case: f0 is set to a FluidEquilibrium + # if self.type == "sph": + # _eq = getattr(equils, fi_type)(**maxw_params) + # if not isinstance(_eq, NumericalFluidEquilibrium): + # _eq.domain = self.domain + # if self._f0 is None: + # self._f0 = _eq + # else: + # raise NotImplementedError("Summation of fluid backgrounds not yet implemented.") + # # self._f0 = self._f0 + (lambda e1, e2, e3: _eq.n0(e1, e2, e3)) + # # default case + # else: + # if self._f0 is None: + # self._f0 = bckgr + # else: + # self._f0 = self._f0 + bckgr def _set_background_coordinates(self): if self.type != "sph" and self.f0.coords == "constants_of_motion": diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index b80b1be3d..c49cabf88 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -22,8 +22,8 @@ class Particles6D(Particles): """ @classmethod - def default_bckgr_params(cls): - return {"Maxwellian3D": {}} + def default_background(cls): + return maxwellians.Maxwellian3D() def __init__( self, @@ -31,8 +31,10 @@ def __init__( ): kwargs["type"] = "full_f" - if "bckgr_params" not in kwargs: - kwargs["bckgr_params"] = self.default_bckgr_params() + # if "backgrounds" not in kwargs: + # kwargs["backgrounds"] = self.default_background() + # elif kwargs["backgrounds"] is None: + # kwargs["backgrounds"] = self.default_background() # default number of diagnostics and auxiliary columns self._n_cols_diagnostics = kwargs.pop("n_cols_diagn", 0) @@ -41,7 +43,7 @@ def __init__( super().__init__(**kwargs) # call projected mhd equilibrium in case of CanonicalMaxwellian - if "CanonicalMaxwellian" in kwargs["bckgr_params"]: + if isinstance(kwargs["background"], maxwellians.CanonicalMaxwellian): assert isinstance(self.equil, FluidEquilibriumWithB), ( "CanonicalMaxwellian needs background with magnetic field." ) From b23593ed7748938ca917ef9a3b2dc9c14e7c5a8c Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 12 Aug 2025 14:09:32 +0200 Subject: [PATCH 078/292] remove some unneccessary imports --- src/struphy/main.py | 1 - src/struphy/models/base.py | 1 - src/struphy/post_processing/post_processing_tools.py | 3 --- 3 files changed, 5 deletions(-) diff --git a/src/struphy/main.py b/src/struphy/main.py index 4aa846158..42a04e26d 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -11,7 +11,6 @@ import pickle import h5py -from struphy.feec.psydac_derham import SplineFunction from struphy.fields_background.base import FluidEquilibriumWithB from struphy.io.output_handling import DataContainer from struphy.io.setup import setup_folders diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 62f388b52..5c36644ea 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -10,7 +10,6 @@ import numpy as np import yaml from mpi4py import MPI -from psydac.linalg.stencil import StencilVector from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import WeightedMassOperators diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 58dcc3bbb..4cc7f2dd0 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -1,15 +1,12 @@ import os import shutil import h5py -import matplotlib.pyplot as plt import numpy as np import yaml from tqdm import tqdm import pickle -from struphy.feec.psydac_derham import Derham from struphy.kinetic_background import maxwellians -from struphy.models import fluid, hybrid, kinetic, toy from struphy.io.setup import import_parameters_py from struphy.models.base import setup_derham from struphy.feec.psydac_derham import SplineFunction From e24278e14df95770f930cbb49a597183b7800654 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 13 Aug 2025 10:27:00 +0200 Subject: [PATCH 079/292] re-factor marker paramter passing --- src/struphy/io/options.py | 1 + src/struphy/kinetic_background/base.py | 4 +- src/struphy/main.py | 18 +++- src/struphy/models/base.py | 38 +++---- src/struphy/models/species.py | 30 ++++-- src/struphy/models/toy.py | 2 +- src/struphy/models/variables.py | 11 ++- src/struphy/pic/base.py | 131 ++++++++++++------------- src/struphy/pic/utilities.py | 81 +++++++++++++++ 9 files changed, 213 insertions(+), 103 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index cb8eb3b38..0aa11879a 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -31,6 +31,7 @@ OptsMarkerBC = Literal["periodic", "reflect"] OptsRecontructBC = Literal["periodic", "mirror", "fixed"] OptsLoading = Literal["pseudo_random", "tesselation"] +OptsSpatialLoading = Literal["uniform", "disc"] ## Option classes diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 96623260a..b9a1d5459 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -198,8 +198,10 @@ def u(self, *etas): n1 = self._f1.n(*etas) n2 = self._f2.n(*etas) + u1s = self._f1.u(*etas) + u2s = self._f2.u(*etas) - return [(n1 * u1 + n2 * u2) / (n1 + n2) for u1, u2 in zip(self._f1.u(*etas), self._f2.u(*etas))] + return [(n1 * u1 + n2 * u2) / (n1 + n2) for u1, u2 in zip(u1s, u2s)] def __call__(self, *args): """Evaluates the background distribution function f0(etas, v1, ..., vn). diff --git a/src/struphy/main.py b/src/struphy/main.py index 42a04e26d..9bb50894d 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -29,6 +29,8 @@ from struphy.topology.grids import TensorProductGrid from struphy.io.options import DerhamOptions from struphy.post_processing.post_processing_tools import create_femfields, eval_femfields, create_vtk +from struphy.topology import grids +from struphy.io.options import DerhamOptions def run( @@ -165,13 +167,19 @@ def run( model.setup_domain_and_equil(domain, equil) # allocate derham-related objects - if grid is not None: + if derham_opts is not None: + assert grid is not None, f"Derham complex needs a grid." model.allocate_feec(grid, derham_opts) else: - model._derham = None - model._mass_ops = None - model._projected_equil = None - print("GRID:\nNo grid specified - no Derham complex set up.") + if grid is None: + Nel = (16, 16, 16) + print(f"\nNo grid specified - using TensorProductGrid with {Nel = }.") + grid = grids.TensorProductGrid(Nel=Nel) + p = (3, 3, 3) + spl_kind = (False, False, False) + print(f"\nNo Derham options specified - creating Derham with {p = } and {spl_kind = } for projecting equilibrium.") + derham_opts = DerhamOptions(p=p, spl_kind=spl_kind) + model.allocate_feec(grid, derham_opts) # equation paramters model.setup_equation_params(units=model.units, verbose=verbose) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 5c36644ea..231e29e09 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -439,7 +439,8 @@ def setInDict(dataDict, mapList, value): assert key is not None, "Must provide key if option is not a class." setInDict(dct, species + ["options"] + key, option) - def add_scalar(self, name, species=None, compute=None, summands=None): + def add_scalar(self, + name: str, variable: PICVariable | SPHVariable = None, compute=None, summands=None): """ Add a scalar to be saved during the simulation. @@ -447,8 +448,8 @@ def add_scalar(self, name, species=None, compute=None, summands=None): ---------- name : str Dictionary key for the scalar. - species : str, optional - The species associated with the scalar. Required if compute is 'from_particles'. + variable : PICVariable | SPHVariable, optional + The variable associated with the scalar. Required if compute is 'from_particles'. compute : str, optional Type of scalar, determines the compute operations. Options: 'from_particles' or 'from_field'. Default is None. @@ -459,17 +460,14 @@ def add_scalar(self, name, species=None, compute=None, summands=None): assert isinstance(name, str), "name must be a string" if compute == "from_particles": - assert isinstance( - species, - str, - ), "species must be a string when compute is 'from_particles'" + assert isinstance(variable, (PICVariable, SPHVariable)), f"Variable is needed when {compute = }" if not hasattr(self, "_scalar_quantities"): self._scalar_quantities = {} self._scalar_quantities[name] = { "value": np.empty(1, dtype=float), - "species": species, + "variable": variable, "compute": compute, "summands": summands, } @@ -489,7 +487,7 @@ def update_scalar(self, name, value=None): # Ensure the name is a string assert isinstance(name, str) - species = self._scalar_quantities[name]["species"] + variable: PICVariable | SPHVariable = self._scalar_quantities[name]["variable"] summands = self._scalar_quantities[name]["summands"] compute = self._scalar_quantities[name]["compute"] @@ -552,7 +550,7 @@ def update_scalar(self, name, value=None): if "divide_n_mks" in compute_operations: # Initialize the total number of markers - n_mks_tot = np.array([self.pointer[species].Np]) + n_mks_tot = np.array([variable.particles.Np]) value_array /= n_mks_tot # Update the scalar value @@ -672,15 +670,17 @@ def update_markers_to_be_saved(self): Writes markers with IDs that are supposed to be saved into corresponding array. """ - from struphy.pic.base import Particles - - for _, val in self.kinetic_species.items(): - obj = val["obj"] - assert isinstance(obj, Particles) + for name, species in self.kinetic_species.items(): + assert isinstance(species, KineticSpecies) + assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." + for _, var in species.variables.items(): + assert isinstance(var, PICVariable | SPHVariable) + obj = var.particles + assert isinstance(obj, Particles) # allocate array for saving markers if not present if not hasattr(self, "_n_markers_saved"): - n_markers = val["params"]["save_data"].get("n_markers", 0) + n_markers = species.n_markers if isinstance(n_markers, float): if n_markers > 1.0: @@ -694,7 +694,7 @@ def update_markers_to_be_saved(self): f"The number of markers for which data should be stored (={n_markers}) murst be <= than the total number of markers (={obj.Np})" ) if self._n_markers_saved > 0: - val["kinetic_data"]["markers"] = np.zeros( + var.kinetic_data["markers"] = np.zeros( (self._n_markers_saved, obj.markers.shape[1]), dtype=float, ) @@ -705,8 +705,8 @@ def update_markers_to_be_saved(self): obj.markers[:, -1] < self._n_markers_saved, ) n_markers_on_proc = np.count_nonzero(markers_on_proc) - val["kinetic_data"]["markers"][:] = -1.0 - val["kinetic_data"]["markers"][:n_markers_on_proc] = obj.markers[markers_on_proc] + var.kinetic_data["markers"][:] = -1.0 + var.kinetic_data["markers"][:n_markers_on_proc] = obj.markers[markers_on_proc] def update_distr_functions(self): """ diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 409f4930d..808e84e5c 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -11,6 +11,7 @@ from struphy.io.options import OptsLoading, OptsMarkerBC, OptsRecontructBC from struphy.physics.physics import ConstantsOfNature from struphy.models.variables import Variable +from struphy.pic.utilities import LoadingParameters, WeightsParameters class Species(metaclass=ABCMeta): @@ -77,7 +78,7 @@ def setup_equation_params(self, units: Units = None, verbose=False): self._equation_params["kappa"] = om_p * units.t if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f'Set normalization parameters for species {self.__class__.__name__}:') + print(f'\nSet normalization parameters for species {self.__class__.__name__}:') for key, val in self.equation_params.items(): print((key + ":").ljust(25), "{:4.3e}".format(val)) @@ -93,6 +94,7 @@ class FluidSpecies(Species): class KineticSpecies(Species): """Single kinetic species in 3d + vdim phase space.""" + def set_markers(self, Np: int = 100, ppc: int = None, @@ -101,13 +103,19 @@ def set_markers(self, bc_refill = None, bc_sph: tuple[OptsRecontructBC] = ("periodic", "periodic", "periodic"), bufsize: float = 1.0, - loading: OptsLoading = "pseudo_random", - loading_params: dict = None, - control_variate: bool = False, - reject_weights: dict = None, + loading_params: LoadingParameters = None, + weights_params: WeightsParameters = None, ): """Set marker parameters for loading, pushing, weight calculation and kernel density reconstruction.""" + + # defaults + if loading_params is None: + loading_params = LoadingParameters() + + if weights_params is None: + weights_params = WeightsParameters() + self.Np = Np self.ppc = ppc self.ppb = ppb @@ -115,17 +123,19 @@ def set_markers(self, self.bc_refill = bc_refill self.bc_sph = bc_sph self.bufsize = bufsize - self.loading = loading - self.loading_params = loading_params - self.control_variate = control_variate - self.reject_weights = reject_weights + self.loading = loading_params + self.weights_params = weights_params - def set_sorting(self, + def set_sorting_boxes(self, + do_sort: bool = False, + sorting_frequency: int = 0, boxes_per_dim: tuple = (16, 1, 1), box_bufsize: float = 2.0, dims_maks: tuple = (True, True, True), ): """For sorting markers in memory.""" + self.do_sort = do_sort + self.sorting_fequency = sorting_frequency self.boxes_per_dim = boxes_per_dim self.box_bufsize = box_bufsize self.dims_mask = dims_maks diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 8afb0dcb6..ee932ed81 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -147,7 +147,7 @@ def __init__(self): ) # define scalars for update_scalar_quantities - self.add_scalar("en_f", compute="from_particles", species="ions") + self.add_scalar("en_f", compute="from_particles", variable=self.kinetic_ions.var) @property def bulk_species(self): diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 692e37a12..5b2782cd8 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -198,15 +198,24 @@ def allocate(self, # bc_sph=self.species.bc_sph, loading=self.species.loading, loading_params=self.species.loading_params, - weights_params=self.species.reject_weights, + reject_weights=self.species.reject_weights, + threshold=self.species.threshold, bufsize=self.species.bufsize, domain=domain, equil=equil, projected_equil=projected_equil, background=self.backgrounds, + n_as_volume_form=self.n_as_volume_form, # perturbations=self.perturbations, equation_params=self.species.equation_params, ) + + if self.species.do_sort: + sort = True + else: + sort = False + self.particles.draw_markers(sort=sort) + self.particles.initialize_weights() # for storing the binned distribution function if self.species.f_binned is not None: diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 0c1a3c71b..8a31f3bd6 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -114,8 +114,11 @@ class Particles(metaclass=ABCMeta): loading_params : dict Parameterts for loading, see defaults below. - weights_params : dict - Parameterts for initializing weights, see defaults below. + rejct_weights : bool + Whether to reject weights below threshold. + + threshold : float + Threshold for rejecting weights. bufsize : float Size of buffer (as multiple of total size, default=.25) in markers array. @@ -163,7 +166,8 @@ def __init__( name: str = "some_name", loading: str = "pseudo_random", loading_params: dict = None, - weights_params: dict = None, + reject_weights: bool = False, + threshold: float = 0.0, bufsize: float = 0.25, domain: Domain = None, equil: FluidEquilibrium = None, @@ -302,16 +306,9 @@ def __init__( ) self._spatial = self.loading_params["spatial"] - # weights parameters - weights_params_default = { - "reject_weights": False, - "threshold": 0.0, - } - - self._weights_params = set_defaults( - weights_params, - weights_params_default, - ) + # weights + self._reject_weights = reject_weights + self._threshold = threshold # background if background is None: @@ -533,9 +530,14 @@ def loading_params(self): return self._loading_params @property - def weights_params(self): - """Parameters for initializing weights.""" - return self._weights_params + def reject_weights(self): + """Whether to reect weights below threshold.""" + return self._reject_weights + + @property + def threshold(self): + """Threshold for rejecting weights.""" + return self._threshold @property def boxes_per_dim(self): @@ -1121,56 +1123,59 @@ def _initialize_sorting_boxes(self): else: self._sorting_boxes = None - def _auto_sampling_params(self): - """Automatically determine sampling parameters from the background given""" - ns = [] - us = [] - vths = [] + def _generate_sampling_moments(self): + """Automatically determine moments for sampling distribution (Gaussian) from the given background.""" + + self.loading_params["moments"] = tuple([0.0]*self.vdim + [1.0]*self.vdim) + # TODO: reformulate this function with KineticBackground methods + + # ns = [] + # us = [] + # vths = [] - for fi, params in self.bckgr_params.items(): - if fi[-2] == "_": - fi_type = fi[:-2] - else: - fi_type = fi + # for fi, params in self.bckgr_params.items(): + # if fi[-2] == "_": + # fi_type = fi[:-2] + # else: + # fi_type = fi - us.append([]) - vths.append([]) + # us.append([]) + # vths.append([]) - bckgr = getattr(maxwellians, fi_type) - default_maxw_params = bckgr.default_maxw_params() + # bckgr = getattr(maxwellians, fi_type) - for key in default_maxw_params: - if key[0] == "n": - if key in params: - ns += [params[key]] - else: - ns += [1.0] + # for key in default_maxw_params: + # if key[0] == "n": + # if key in params: + # ns += [params[key]] + # else: + # ns += [1.0] - elif key[0] == "u": - if key in params: - us[-1] += [params[key]] - else: - us[-1] += [0.0] + # elif key[0] == "u": + # if key in params: + # us[-1] += [params[key]] + # else: + # us[-1] += [0.0] - elif key[0] == "v": - if key in params: - vths[-1] += [params[key]] - else: - vths[-1] += [1.0] + # elif key[0] == "v": + # if key in params: + # vths[-1] += [params[key]] + # else: + # vths[-1] += [1.0] - assert len(ns) == len(us) == len(vths) + # assert len(ns) == len(us) == len(vths) - ns = np.array(ns) - us = np.array(us) - vths = np.array(vths) + # ns = np.array(ns) + # us = np.array(us) + # vths = np.array(vths) - new_moments = [] + # new_moments = [] - new_moments += [*np.mean(us, axis=0)] - new_moments += [*(np.max(vths, axis=0) + np.max(np.abs(us), axis=0) - np.mean(us, axis=0))] - new_moments = [float(moment) for moment in new_moments] + # new_moments += [*np.mean(us, axis=0)] + # new_moments += [*(np.max(vths, axis=0) + np.max(np.abs(us), axis=0) - np.mean(us, axis=0))] + # new_moments = [float(moment) for moment in new_moments] - self.loading_params["moments"] = new_moments + # self.loading_params["moments"] = new_moments def _set_initial_condition(self, bp_copy=None, pp_copy=None): """Compute callable initial condition from background + perturbation.""" @@ -1662,8 +1667,8 @@ def initialize_weights( *, bckgr_params: dict = None, pert_params: dict = None, - reject_weights: bool = False, - threshold: float = 1e-8, + # reject_weights: bool = False, + # threshold: float = 1e-8, ): r""" Computes the initial weights @@ -1684,12 +1689,6 @@ def initialize_weights( pert_params : dict Kinetic perturbation parameters for initial condition. - - reject_weights : bool - Whether to use ``threshold`` for rejecting weights. - - threshold : float - Minimal value of a weight; below the marker is set to a hole.les. """ if self.loading == "tesselation": @@ -1733,13 +1732,13 @@ def initialize_weights( # compute w0 and save at vdim + 5 self.weights0 = f_init / self.sampling_density - if reject_weights: - reject = self.markers[:, self.index["w0"]] < threshold + if self.reject_weights: + reject = self.markers[:, self.index["w0"]] < self.threshold self._markers[reject] = -1.0 self.update_holes() self.reset_marker_ids() print( - f"\nWeights < {threshold} have been rejected, number of valid markers on process {self.mpi_rank} is {self.n_mks_loc}." + f"\nWeights < {self.threshold} have been rejected, number of valid markers on process {self.mpi_rank} is {self.n_mks_loc}." ) # compute (time-dependent) weights at vdim + 3 diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index fe702e8c2..bd0615670 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -1,6 +1,87 @@ import numpy as np import struphy.pic.utilities_kernels as utils +from struphy.io.options import OptsLoading, OptsSpatialLoading + + +class LoadingParameters: + """Paramters for particle loading. + + Parameters + ---------- + loading : OptsLoading + How to load markers: "pseudo_random" for Monte-Carlo, or "tesselation" for positioning them on a regular grid. + + seed : int + Seed for random generator. If None, no seed is taken. + + moments : tuple + Mean velocities and temperatures for the Gaussian sampling distribution. + If None, these are auto-calculated form the given background. + + spatial : OptsSpatialLoading + Draw uniformly in eta, or draw uniformly on the "disc" image of (eta1, eta2). + + specific_markers : tuple[tuple] + Each entry is a tuple of phase space coordinates (floats) of a specific marker to be initialized. + + n_quad : int + Number of quadrature points for tesselation. + + dir_external : str + Load markers from external .hdf5 file (absolute path). + + dir_particles_abs : str + Load markers from restart .hdf5 file (absolute path). + + dir_particles : str + Load markers from restart .hdf5 file (relative path to output folder). + """ + def __init__(self, + loading: OptsLoading = "pseudo_random", + seed: int = None, + moments: tuple = None, + spatial: OptsSpatialLoading = "uniform", + specific_markers: tuple[tuple] = None, + n_quad: int = 1, + dir_exrernal: str = None, + dir_particles: str = None, + dir_particles_abs: str = None, + ): + + self.loading = loading + self.seed = seed + self.moments = moments + self.spatial = spatial + self.specific_markers = specific_markers + self.n_quad = n_quad + self.dir_external = dir_exrernal + self.dir_particles = dir_particles + self.dir_particles_abs = dir_particles_abs + + +class WeightsParameters: + """Paramters for particle weights. + + Parameters + ---------- + control_variate : bool + Whether to use a control variate for noise reduction. + + rejct_weights : bool + Whether to reject weights below threshold. + + threshold : float + Threshold for rejecting weights. + """ + def __init__(self, + control_variate: bool = False, + reject_weights: bool = False, + threshold: float = 0.0,): + + self.control_variate = control_variate + self.reject_weights = reject_weights + self.threshold = threshold def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): From 8b3cae3687fe2751885e313704b57b1a5afac5e8 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 14 Aug 2025 11:45:06 +0200 Subject: [PATCH 080/292] Vlasov is running --- src/struphy/io/options.py | 2 +- src/struphy/models/base.py | 59 ++++++------ src/struphy/models/species.py | 2 +- src/struphy/models/variables.py | 5 +- src/struphy/pic/base.py | 164 ++++++++++++++------------------ src/struphy/pic/particles.py | 24 ++--- src/struphy/pic/utilities.py | 9 +- 7 files changed, 121 insertions(+), 144 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 0aa11879a..d871e8b3a 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -30,7 +30,7 @@ # markers OptsMarkerBC = Literal["periodic", "reflect"] OptsRecontructBC = Literal["periodic", "mirror", "fixed"] -OptsLoading = Literal["pseudo_random", "tesselation"] +OptsLoading = Literal["pseudo_random", 'sobol_standard', 'sobol_antithetic', 'external', 'restart', "tesselation"] OptsSpatialLoading = Literal["uniform", "disc"] diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 231e29e09..430cc2845 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -35,6 +35,8 @@ from struphy.pic.base import Particles from struphy.propagators.base import Propagator from struphy.kinetic_background import maxwellians +from psydac.linalg.stencil import StencilVector +from struphy.io.output_handling import DataContainer class StruphyModel(metaclass=ABCMeta): @@ -713,21 +715,23 @@ def update_distr_functions(self): Writes distribution functions slices that are supposed to be saved into corresponding array. """ - from struphy.pic.base import Particles - dim_to_int = {"e1": 0, "e2": 1, "e3": 2, "v1": 3, "v2": 4, "v3": 5} - for _, val in self.kinetic_species.items(): - obj = val["obj"] - assert isinstance(obj, Particles) + for name, species in self.kinetic_species.items(): + assert isinstance(species, KineticSpecies) + assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." + for _, var in species.variables.items(): + assert isinstance(var, PICVariable | SPHVariable) + obj = var.particles + assert isinstance(obj, Particles) if obj.n_cols_diagnostics > 0: for i in range(obj.n_cols_diagnostics): str_dn = f"d{i + 1}" dim_to_int[str_dn] = 3 + obj.vdim + 3 + i - if "f" in val["params"]["save_data"]: - for slice_i, edges in val["bin_edges"].items(): + if species.f_binned is not None: + for slice_i, edges in var.kinetic_data["bin_edges"].items(): comps = slice_i.split("_") components = [False] * (3 + obj.vdim + 3 + obj.n_cols_diagnostics) @@ -736,10 +740,10 @@ def update_distr_functions(self): f_slice, df_slice = obj.binning(components, edges) - val["kinetic_data"]["f"][slice_i][:] = f_slice - val["kinetic_data"]["df"][slice_i][:] = df_slice + var.kinetic_data["f"][slice_i][:] = f_slice + var.kinetic_data["df"][slice_i][:] = df_slice - if "n_sph" in val["params"]["save_data"]: + if species.n_sph is not None: h1 = 1 / obj.boxes_per_dim[0] h2 = 1 / obj.boxes_per_dim[1] h3 = 1 / obj.boxes_per_dim[2] @@ -948,9 +952,6 @@ def initialize_from_restart(self, data): The data object that links to the hdf5 files. """ - from struphy.feec.psydac_derham import Derham - from struphy.pic.base import Particles - # initialize em fields if len(self.em_fields) > 0: for key, val in self.em_fields.items(): @@ -986,7 +987,7 @@ def initialize_from_restart(self, data): # important: sets holes attribute of markers! obj.mpi_sort_markers(do_test=True) - def initialize_data_output(self, data, size): + def initialize_data_output(self, data: DataContainer, size): """ Create datasets in hdf5 files according to model unknowns and diagnostics data. @@ -1007,14 +1008,6 @@ def initialize_data_output(self, data, size): Keys of datasets which are saved at the end of a simulation to enable restarts. """ - from psydac.linalg.stencil import StencilVector - - from struphy.feec.psydac_derham import Derham - from struphy.io.output_handling import DataContainer - from struphy.pic.base import Particles - - assert isinstance(data, DataContainer) - # save scalar quantities in group 'scalar/' for key, scalar in self.scalar_quantities.items(): val = scalar["value"] @@ -1082,16 +1075,20 @@ def initialize_data_output(self, data, size): ) # save kinetic data in group 'kinetic/' - for species, val in self.kinetic_species.items(): - obj = val["obj"] - assert isinstance(obj, Particles) + for name, species in self.kinetic_species.items(): + assert isinstance(species, KineticSpecies) + assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." + for _, var in species.variables.items(): + assert isinstance(var, PICVariable | SPHVariable) + obj = var.particles + assert isinstance(obj, Particles) key_spec = "kinetic/" + key key_spec_restart = "restart/" + key data.add_data({key_spec_restart: obj._markers}) - for key1, val1 in val["kinetic_data"].items(): + for key1, val1 in var.kinetic_data.items(): key_dat = key_spec + "/" + key1 # case of "f" and "df" @@ -1103,8 +1100,8 @@ def initialize_data_output(self, data, size): dims = (len(key2) - 2) // 3 + 1 for dim in range(dims): data.file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = ( - val["bin_edges"][key2][dim][:-1] - + (val["bin_edges"][key2][dim][1] - val["bin_edges"][key2][dim][0]) / 2 + var.kinetic_data["bin_edges"][key2][dim][:-1] + + (var.kinetic_data["bin_edges"][key2][dim][1] - var.kinetic_data["bin_edges"][key2][dim][0]) / 2 ) # case of "n_sph" elif isinstance(val1, list): @@ -1112,9 +1109,9 @@ def initialize_data_output(self, data, size): key_n = key_dat + "/view_" + str(i) data.add_data({key_n: v1}) # save 1d point values, not meshgrids, because attrs size is limited - eta1 = val["plot_pts"][i][0][:, 0, 0] - eta2 = val["plot_pts"][i][1][0, :, 0] - eta3 = val["plot_pts"][i][2][0, 0, :] + eta1 = var.kinetic_data["plot_pts"][i][0][:, 0, 0] + eta2 = var.kinetic_data["plot_pts"][i][1][0, :, 0] + eta3 = var.kinetic_data["plot_pts"][i][2][0, 0, :] data.file[key_n].attrs["eta1"] = eta1 data.file[key_n].attrs["eta2"] = eta2 data.file[key_n].attrs["eta3"] = eta3 diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 808e84e5c..9367529e0 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -123,7 +123,7 @@ def set_markers(self, self.bc_refill = bc_refill self.bc_sph = bc_sph self.bufsize = bufsize - self.loading = loading_params + self.loading_params = loading_params self.weights_params = weights_params def set_sorting_boxes(self, diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 5b2782cd8..dc0f57b37 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -193,13 +193,10 @@ def allocate(self, box_bufsize=self.species.box_bufsize, bc=self.species.bc, bc_refill=self.species.bc_refill, - control_variate=self.species.control_variate, name=self.species.__class__.__name__, # bc_sph=self.species.bc_sph, - loading=self.species.loading, loading_params=self.species.loading_params, - reject_weights=self.species.reject_weights, - threshold=self.species.threshold, + weights_params=self.species.weights_params, bufsize=self.species.bufsize, domain=domain, equil=equil, diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 8a31f3bd6..8ad17f537 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -39,6 +39,8 @@ ) from struphy.utils import utils from struphy.utils.clone_config import CloneConfig +from struphy.pic.utilities import LoadingParameters, WeightsParameters +from struphy.io.options import OptsLoading class Particles(metaclass=ABCMeta): @@ -101,24 +103,14 @@ class Particles(metaclass=ABCMeta): type : str Either 'full_f' (default), 'delta_f' or 'sph'. - control_variate : bool - Whether to use a control variate for noise reduction (only if type is 'full_f' or 'sph'). - name : str Name of particle species. - loading : str - Drawing of markers; either 'pseudo_random', 'sobol_standard', - 'sobol_antithetic', 'external' or 'restart'. - - loading_params : dict - Parameterts for loading, see defaults below. + loading_params : LoadingParameters + Parameterts for particle loading. - rejct_weights : bool - Whether to reject weights below threshold. - - threshold : float - Threshold for rejecting weights. + weights_params : WeightsParameters + Parameters for particle weights. bufsize : float Size of buffer (as multiple of total size, default=.25) in markers array. @@ -162,12 +154,9 @@ def __init__( bc: list = None, bc_refill: str = None, type: str = "full_f", - control_variate: bool = False, name: str = "some_name", - loading: str = "pseudo_random", - loading_params: dict = None, - reject_weights: bool = False, - threshold: float = 0.0, + loading_params: LoadingParameters = None, + weights_params: WeightsParameters = None, bufsize: float = 0.25, domain: Domain = None, equil: FluidEquilibrium = None, @@ -194,6 +183,13 @@ def __init__( self._equil = equil self._projected_equil = projected_equil self._equation_params = equation_params + + # defaults + if loading_params is None: + loading_params = LoadingParameters() + + if weights_params is None: + weights_params = WeightsParameters() # check for mpi communicator (i.e. sub_comm of clone) if self.mpi_comm is None: @@ -274,41 +270,20 @@ def __init__( # particle type assert type in ("full_f", "delta_f", "sph") self._type = type - self._control_variate = control_variate # initialize sorting boxes self._verbose = verbose self._initialize_sorting_boxes() # particle loading parameters - assert loading in ( - "pseudo_random", - "sobol_standard", - "sobol_antithetic", - "external", - "restart", - "tesselation", - ) - self._loading = loading - - loading_params_default = { - "seed": None, - "dir_particles": None, - "moments": None, - "spatial": "uniform", - "initial": None, - "n_quad": 1, - } - - self._loading_params = set_defaults( - loading_params, - loading_params_default, - ) - self._spatial = self.loading_params["spatial"] + self._loading = loading_params.loading + self._loading_params = loading_params + self._spatial = loading_params.spatial # weights - self._reject_weights = reject_weights - self._threshold = threshold + self._reject_weights = weights_params.reject_weights + self._threshold = weights_params.threshold + self._control_variate = weights_params.control_variate # background if background is None: @@ -332,7 +307,7 @@ def __init__( # for loading # if self.loading_params["moments"] is None and self.type != "sph" and isinstance(self.bckgr_params, dict): - # self._auto_sampling_params() + self._generate_sampling_moments() # create buffers for mpi_sort_markers if self.mpi_comm is not None: @@ -444,7 +419,7 @@ def type(self): return self._type @property - def loading(self): + def loading(self) -> OptsLoading: """Type of particle loading.""" return self._loading @@ -525,7 +500,7 @@ def background(self) -> KineticBackground: # return self._perturbations @property - def loading_params(self): + def loading_params(self) -> LoadingParameters: """Parameters for marker loading.""" return self._loading_params @@ -1126,7 +1101,9 @@ def _initialize_sorting_boxes(self): def _generate_sampling_moments(self): """Automatically determine moments for sampling distribution (Gaussian) from the given background.""" - self.loading_params["moments"] = tuple([0.0]*self.vdim + [1.0]*self.vdim) + if self.loading_params.moments is None: + self.loading_params.moments = tuple([0.0]*self.vdim + [1.0]*self.vdim) + # TODO: reformulate this function with KineticBackground methods # ns = [] @@ -1178,37 +1155,38 @@ def _generate_sampling_moments(self): # self.loading_params["moments"] = new_moments def _set_initial_condition(self, bp_copy=None, pp_copy=None): - """Compute callable initial condition from background + perturbation.""" - if bp_copy is None: - bp_copy = copy.deepcopy(self.bckgr_params) - if pp_copy is None: - pp_copy = copy.deepcopy(self.pert_params) - - # Get the initialization function and pass the correct arguments - self._f_init = None - for fi, maxw_params in bp_copy.items(): - if fi[-2] == "_": - fi_type = fi[:-2] - else: - fi_type = fi - - pert_params = pp_copy - if pp_copy is not None: - if fi in pp_copy: - pert_params = pp_copy[fi] - - if self._f_init is None: - self._f_init = getattr(maxwellians, fi_type)( - maxw_params=maxw_params, - pert_params=pert_params, - equil=self.equil, - ) - else: - self._f_init = self._f_init + getattr(maxwellians, fi_type)( - maxw_params=maxw_params, - pert_params=pert_params, - equil=self.equil, - ) + self._f_init = self.background + # """Compute callable initial condition from background + perturbation.""" + # if bp_copy is None: + # bp_copy = copy.deepcopy(self.bckgr_params) + # if pp_copy is None: + # pp_copy = copy.deepcopy(self.pert_params) + + # # Get the initialization function and pass the correct arguments + # self._f_init = None + # for fi, maxw_params in bp_copy.items(): + # if fi[-2] == "_": + # fi_type = fi[:-2] + # else: + # fi_type = fi + + # pert_params = pp_copy + # if pp_copy is not None: + # if fi in pp_copy: + # pert_params = pp_copy[fi] + + # if self._f_init is None: + # self._f_init = getattr(maxwellians, fi_type)( + # maxw_params=maxw_params, + # pert_params=pert_params, + # equil=self.equil, + # ) + # else: + # self._f_init = self._f_init + getattr(maxwellians, fi_type)( + # maxw_params=maxw_params, + # pert_params=pert_params, + # equil=self.equil, + # ) def _load_external( self, @@ -1227,7 +1205,7 @@ def _load_external( """ if self.mpi_rank == 0: file = h5py.File( - self.loading_params["dir_external"], + self.loading_params.dir_external, "r", ) print(f"\nLoading markers from file: {file}") @@ -1260,16 +1238,16 @@ def _load_restart(self): o_path = state["o_path"] - if self.loading_params["dir_particles_abs"] is None: + if self.loading_params.dir_particles_abs is None: data_path = os.path.join( o_path, - self.loading_params["dir_particles"], + self.loading_params.dir_particles, ) else: - data_path = self.loading_params["dir_particles_abs"] + data_path = self.loading_params.dir_particles_abs data = DataContainer(data_path, comm=self.mpi_comm) - self._markers[:, :] = data.file["restart/" + self.loading_params["key"]][-1, :, :] + self._markers[:, :] = data.file["restart/" + self.loading_params.restart_key][-1, :, :] def _load_tesselation(self, n_quad: int = 1): """ @@ -1423,13 +1401,13 @@ def draw_markers( else: if self.mpi_rank == 0 and verbose: print("\nLoading fresh markers:") - for key, val in self.loading_params.items(): + for key, val in self.loading_params.__dict__.items(): print((key + " :").ljust(25), val) # 1. standard random number generator (pseudo-random) if self.loading == "pseudo_random": # set seed - _seed = self.loading_params["seed"] + _seed = self.loading_params.seed if _seed is not None: np.random.seed(_seed) @@ -1510,8 +1488,8 @@ def draw_markers( self.velocities = np.array(self.u_init(self.positions)[0]).T else: # inverse transform sampling in velocity space - u_mean = np.array(self.loading_params["moments"][: self.vdim]) - v_th = np.array(self.loading_params["moments"][self.vdim :]) + u_mean = np.array(self.loading_params.moments[: self.vdim]) + v_th = np.array(self.loading_params.moments[self.vdim :]) # Particles6D: (1d Maxwellian, 1d Maxwellian, 1d Maxwellian) if self.vdim == 3: @@ -1560,8 +1538,8 @@ def draw_markers( self.marker_ids = _first_marker_id + np.arange(n_mks_load_loc, dtype=float) # set specific initial condition for some particles - if self.loading_params["initial"] is not None: - specific_markers = self.loading_params["initial"] + if self.loading_params.specific_markers is not None: + specific_markers = self.loading_params.specific_markers counter = 0 for i in range(len(specific_markers)): @@ -1696,7 +1674,7 @@ def initialize_weights( fvol = TransformedPformComponent([self.f_init], "0", "3", domain=self.domain) else: fvol = self.f_init - cell_avg = self.tesselation.cell_averages(fvol, n_quad=self.loading_params["n_quad"]) + cell_avg = self.tesselation.cell_averages(fvol, n_quad=self.loading_params.n_quad) self.weights0 = cell_avg.flatten() else: assert self.domain is not None, "A domain is needed to initialize weights." diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index c49cabf88..0fd34eeed 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -91,16 +91,16 @@ def svol(self, eta1, eta2, eta3, *v): """ # load sampling density svol (normalized to 1 in logical space) maxw_params = { - "n": 1.0, - "u1": self.loading_params["moments"][0], - "u2": self.loading_params["moments"][1], - "u3": self.loading_params["moments"][2], - "vth1": self.loading_params["moments"][3], - "vth2": self.loading_params["moments"][4], - "vth3": self.loading_params["moments"][5], + "n": (1.0, None), + "u1": (self.loading_params.moments[0], None), + "u2": (self.loading_params.moments[1], None), + "u3": (self.loading_params.moments[2], None), + "vth1": (self.loading_params.moments[3], None), + "vth2": (self.loading_params.moments[4], None), + "vth3": (self.loading_params.moments[5], None), } - fun = maxwellians.Maxwellian3D(maxw_params=maxw_params) + fun = maxwellians.Maxwellian3D(**maxw_params) if self.spatial == "uniform": return fun(eta1, eta2, eta3, *v) @@ -400,10 +400,10 @@ def svol(self, eta1, eta2, eta3, *v): # load sampling density svol (normalized to 1 in logical space) maxw_params = { "n": 1.0, - "u_para": self.loading_params["moments"][0], - "u_perp": self.loading_params["moments"][1], - "vth_para": self.loading_params["moments"][2], - "vth_perp": self.loading_params["moments"][3], + "u_para": self.loading_params.moments[0], + "u_perp": self.loading_params.moments[1], + "vth_para": self.loading_params.moments[2], + "vth_perp": self.loading_params.moments[3], } self._svol = maxwellians.GyroMaxwellian2D( diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index bd0615670..453627768 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -5,12 +5,12 @@ class LoadingParameters: - """Paramters for particle loading. + """Parameters for particle loading. Parameters ---------- loading : OptsLoading - How to load markers: "pseudo_random" for Monte-Carlo, or "tesselation" for positioning them on a regular grid. + How to load markers: multiple options for Monte-Carlo, or "tesselation" for positioning them on a regular grid. seed : int Seed for random generator. If None, no seed is taken. @@ -36,6 +36,9 @@ class LoadingParameters: dir_particles : str Load markers from restart .hdf5 file (relative path to output folder). + + restart_key : str + Key in .hdf5 file's restart/ folder where marker array is stored. """ def __init__(self, loading: OptsLoading = "pseudo_random", @@ -47,6 +50,7 @@ def __init__(self, dir_exrernal: str = None, dir_particles: str = None, dir_particles_abs: str = None, + restart_key: str = None, ): self.loading = loading @@ -58,6 +62,7 @@ def __init__(self, self.dir_external = dir_exrernal self.dir_particles = dir_particles self.dir_particles_abs = dir_particles_abs + self.restart_key = restart_key class WeightsParameters: From 1b27f67448b81e5479ebc2ed49f92f51821e3446 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 14 Aug 2025 12:17:52 +0200 Subject: [PATCH 081/292] default parameters for kinetic species --- src/struphy/models/base.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 430cc2845..287c9d944 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1282,7 +1282,7 @@ def generate_default_parameter_file( else: print("exiting ...") return - + file.write("from struphy.io.options import EnvironmentOptions, Units, Time\n") file.write("from struphy.geometry import domains\n") file.write("from struphy.fields_background import equils\n") @@ -1300,9 +1300,11 @@ def generate_default_parameter_file( has_plasma = True species_params += f"model.{sn}.set_phys_params()\n" if isinstance(species, KineticSpecies): - kinetic_params += f"model.{sn}.set_markers()\n\ -model.{sn}.set_sorting()\n\ -model.{sn}.set_save_data()\n" + kinetic_params += f"\nloading_params = LoadingParameters()\n" + kinetic_params += f"weights_params = WeightsParameters()\n" + kinetic_params += f"model.{sn}.set_markers(loading_params=loading_params, weights_params=weights_params)\n" + kinetic_params += f"model.{sn}.set_sorting_boxes()\n" + kinetic_params += f"model.{sn}.set_save_data()\n" for vn, var in species.variables.items(): if isinstance(var, FEECVariable): @@ -1320,7 +1322,10 @@ def generate_default_parameter_file( elif isinstance(var, PICVariable): has_pic = True init_pert_pic = f"perturbation = perturbations.TorusModesCos()\n" - init_bckgr_pic = f"model.{sn}.{vn}.add_background(maxwellians.Maxwellian3D(perturbation=perturbation))\n" + init_bckgr_pic = f"\nmaxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n" + init_bckgr_pic += f"maxwellian_2 = maxwellians.Maxwellian3D(n=(0.1, None))\n" + init_bckgr_pic += f"background = maxwellian_1 + maxwellian_2\n" + init_bckgr_pic += f"model.{sn}.{vn}.add_background(background)\n" exclude = f"# model.....save_data = False\n" elif isinstance(var, SPHVariable): @@ -1332,6 +1337,7 @@ def generate_default_parameter_file( file.write("from struphy.initial import perturbations\n") file.write("from struphy.kinetic_background import maxwellians\n") + file.write("from struphy.pic.utilities import LoadingParameters, WeightsParameters\n") file.write("from struphy import main\n") file.write("\n# import model, set verbosity\n") From 4e3b26440784032fd542c568939a850990052f66 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 15 Aug 2025 10:24:14 +0200 Subject: [PATCH 082/292] new class BoundaryParameters for particles; start working on Tutorial 01 --- src/struphy/io/options.py | 4 ++ src/struphy/main.py | 18 +++++-- src/struphy/models/base.py | 5 +- src/struphy/models/species.py | 38 +++++++------- src/struphy/models/variables.py | 7 +-- src/struphy/pic/base.py | 52 ++++++++++--------- src/struphy/pic/utilities.py | 46 +++++++++++++++- .../post_processing/post_processing_tools.py | 2 +- .../propagators/propagators_markers.py | 7 ++- 9 files changed, 120 insertions(+), 59 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index d871e8b3a..a0abe7bae 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -298,6 +298,10 @@ class EnvironmentOptions: def __post_init__(self): self.path_out: str = os.path.join(self.out_folders, self.sim_folder) + + def print(self): + for k, v in self.__dict__.items(): + print(f"{k}:".ljust(20), v) def check_option(opt, options): diff --git a/src/struphy/main.py b/src/struphy/main.py index 9bb50894d..c0081767c 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -31,6 +31,10 @@ from struphy.post_processing.post_processing_tools import create_femfields, eval_femfields, create_vtk from struphy.topology import grids from struphy.io.options import DerhamOptions +from struphy.post_processing.post_processing_tools import (post_process_markers, + post_process_f, + post_process_n_sph) +from struphy.post_processing.orbits import orbits_tools def run( @@ -450,7 +454,11 @@ def pproc( if "kinetic" in file.keys(): exist_kinetic = {"markers": False, "f": False, "n_sph": False} + kinetic_species = [] for name in file["kinetic"].keys(): + print(f"{name = }") + kinetic_species += [name] + # check for saved markers if "markers" in file["kinetic"][name]: exist_kinetic["markers"] = True @@ -539,14 +547,14 @@ def pproc( # markers if exist_kinetic["markers"]: - pproc.post_process_markers(path, path_kinetics_species, species, kinetic_kinds[n], step) + post_process_markers(path, path_kinetics_species, species, kinetic_kinds[n], step) if guiding_center: assert kinetic_kinds[n] == "Particles6D" - orbits_pproc.post_process_orbit_guiding_center(path, path_kinetics_species, species) + orbits_tools.post_process_orbit_guiding_center(path, path_kinetics_species, species) if classify: - orbits_pproc.post_process_orbit_classification(path_kinetics_species, species) + orbits_tools.post_process_orbit_classification(path_kinetics_species, species) # distribution function if exist_kinetic["f"]: @@ -555,11 +563,11 @@ def pproc( else: compute_bckgr = False - pproc.post_process_f(path, path_kinetics_species, species, step, compute_bckgr=compute_bckgr) + post_process_f(path, path_kinetics_species, species, step, compute_bckgr=compute_bckgr) # sph density if exist_kinetic["n_sph"]: - pproc.post_process_n_sph(path, path_kinetics_species, species, step, compute_bckgr=compute_bckgr) + post_process_n_sph(path, path_kinetics_species, species, step, compute_bckgr=compute_bckgr) class SimData: diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 287c9d944..6b06c72a4 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1302,7 +1302,8 @@ def generate_default_parameter_file( if isinstance(species, KineticSpecies): kinetic_params += f"\nloading_params = LoadingParameters()\n" kinetic_params += f"weights_params = WeightsParameters()\n" - kinetic_params += f"model.{sn}.set_markers(loading_params=loading_params, weights_params=weights_params)\n" + kinetic_params += f"boundary_params = BoundaryParameters()\n" + kinetic_params += f"model.{sn}.set_markers(loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params)\n" kinetic_params += f"model.{sn}.set_sorting_boxes()\n" kinetic_params += f"model.{sn}.set_save_data()\n" @@ -1337,7 +1338,7 @@ def generate_default_parameter_file( file.write("from struphy.initial import perturbations\n") file.write("from struphy.kinetic_background import maxwellians\n") - file.write("from struphy.pic.utilities import LoadingParameters, WeightsParameters\n") + file.write("from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters\n") file.write("from struphy import main\n") file.write("\n# import model, set verbosity\n") diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 9367529e0..d8e8555ab 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -8,10 +8,11 @@ from struphy.fields_background.base import FluidEquilibrium from struphy.kinetic_background.base import KineticBackground from struphy.io.options import Units -from struphy.io.options import OptsLoading, OptsMarkerBC, OptsRecontructBC from struphy.physics.physics import ConstantsOfNature from struphy.models.variables import Variable -from struphy.pic.utilities import LoadingParameters, WeightsParameters +from struphy.pic.utilities import (LoadingParameters, + WeightsParameters, + BoundaryParameters,) class Species(metaclass=ABCMeta): @@ -96,18 +97,24 @@ class KineticSpecies(Species): """Single kinetic species in 3d + vdim phase space.""" def set_markers(self, - Np: int = 100, - ppc: int = None, - ppb: int = None, - bc: tuple[OptsMarkerBC] = ("periodic", "periodic", "periodic"), - bc_refill = None, - bc_sph: tuple[OptsRecontructBC] = ("periodic", "periodic", "periodic"), - bufsize: float = 1.0, loading_params: LoadingParameters = None, weights_params: WeightsParameters = None, + boundary_params: BoundaryParameters = None, + bufsize: float = 1.0, ): - """Set marker parameters for loading, pushing, weight calculation - and kernel density reconstruction.""" + """Set marker parameters for loading, weight calculation, kernel density reconstruction + and boundary conditions. + + Parameters + ---------- + loading_params : LoadingParameters + + weights_params : WeightsParameters + + boundary_params : BoundaryParameters + + bufsize : float + Size of buffer (as multiple of total size, default=.25) in markers array.""" # defaults if loading_params is None: @@ -116,15 +123,10 @@ def set_markers(self, if weights_params is None: weights_params = WeightsParameters() - self.Np = Np - self.ppc = ppc - self.ppb = ppb - self.bc = bc - self.bc_refill = bc_refill - self.bc_sph = bc_sph - self.bufsize = bufsize + self.boundary_params = boundary_params self.loading_params = loading_params self.weights_params = weights_params + self.bufsize = bufsize def set_sorting_boxes(self, do_sort: bool = False, diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index dc0f57b37..bec807bca 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -184,19 +184,14 @@ def allocate(self, self._particles: Particles = kinetic_class( comm_world=MPI.COMM_WORLD, clone_config=clone_config, - Np=self.species.Np, - ppc=self.species.ppc, domain_decomp=domain_decomp, mpi_dims_mask=self.species.dims_mask, - ppb=self.species.ppb, boxes_per_dim=self.species.boxes_per_dim, box_bufsize=self.species.box_bufsize, - bc=self.species.bc, - bc_refill=self.species.bc_refill, name=self.species.__class__.__name__, - # bc_sph=self.species.bc_sph, loading_params=self.species.loading_params, weights_params=self.species.weights_params, + boundary_params=self.species.boundary_params, bufsize=self.species.bufsize, domain=domain, equil=equil, diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 8ad17f537..9971dc48d 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -39,7 +39,10 @@ ) from struphy.utils import utils from struphy.utils.clone_config import CloneConfig -from struphy.pic.utilities import LoadingParameters, WeightsParameters +from struphy.pic.utilities import (LoadingParameters, + WeightsParameters, + BoundaryParameters, + ) from struphy.io.options import OptsLoading @@ -71,12 +74,6 @@ class Particles(metaclass=ABCMeta): clone_config : CloneConfig Manages the configuration for clone-based (copied grids) parallel processing using MPI. - Np : int - Number of particles. - - ppc : int - Particles per cell. Cells are defined from ``domain_array``. - domain_decomp : tuple The first entry is a domain_array (see :attr:`~struphy.feec.psydac_derham.Derham.domain_array`) and the second entry is the number of MPI processes in each direction. @@ -85,21 +82,12 @@ class Particles(metaclass=ABCMeta): True if the dimension is to be used in the domain decomposition (=default for each dimension). If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. - ppb : int - Particles per sorting box. Boxes are defined from ``boxes_per_dim``. - boxes_per_dim : tuple Number of boxes in each logical direction (n_eta1, n_eta2, n_eta3). box_bufsize : float Between 0 and 1, relative buffer size for box array (default = 0.25). - bc : list - Either 'remove', 'reflect', 'periodic' or 'refill' in each direction. - - bc_refill : list - Either 'inner' or 'outer'. - type : str Either 'full_f' (default), 'delta_f' or 'sph'. @@ -111,6 +99,9 @@ class Particles(metaclass=ABCMeta): weights_params : WeightsParameters Parameters for particle weights. + + boundary_params : BoundaryParameters + Parameters for particle boundary conditions. bufsize : float Size of buffer (as multiple of total size, default=.25) in markers array. @@ -144,26 +135,21 @@ def __init__( self, comm_world: Intracomm = None, clone_config: CloneConfig = None, - Np: int = None, - ppc: int = None, domain_decomp: tuple = None, mpi_dims_mask: tuple | list = None, - ppb: int = 10, boxes_per_dim: tuple | list = None, box_bufsize: float = 2.0, - bc: list = None, - bc_refill: str = None, type: str = "full_f", name: str = "some_name", loading_params: LoadingParameters = None, weights_params: WeightsParameters = None, + boundary_params: BoundaryParameters = None, bufsize: float = 0.25, domain: Domain = None, equil: FluidEquilibrium = None, projected_equil: ProjectedFluidEquilibrium = None, background: KineticBackground = None, n_as_volume_form: bool = False, - # perturbations: Perturbation | list = None, equation_params: dict = None, verbose: bool = False, ): @@ -179,6 +165,9 @@ def __init__( # other parameters self._name = name + self._loading_params = loading_params + self._weights_params = weights_params + self._boundary_params = boundary_params self._domain = domain self._equil = equil self._projected_equil = projected_equil @@ -190,6 +179,9 @@ def __init__( if weights_params is None: weights_params = WeightsParameters() + + if boundary_params is None: + boundary_params = BoundaryParameters() # check for mpi communicator (i.e. sub_comm of clone) if self.mpi_comm is None: @@ -230,6 +222,9 @@ def __init__( print(f"{self.mpi_rank = }, {n_boxes = }") # total number of markers (Np) and particles per cell (ppc) + Np = self.loading_params.Np + ppc = self.loading_params.ppc + ppb = self.loading_params.ppb if Np is not None: self._Np = int(Np) self._ppc = self.Np / n_cells @@ -250,6 +245,8 @@ def __init__( self._allocate_marker_array() # boundary conditions + bc = boundary_params.bc + bc_refill = boundary_params.bc_refill if bc is None: bc = ["periodic", "periodic", "periodic"] @@ -277,7 +274,6 @@ def __init__( # particle loading parameters self._loading = loading_params.loading - self._loading_params = loading_params self._spatial = loading_params.spatial # weights @@ -501,8 +497,16 @@ def background(self) -> KineticBackground: @property def loading_params(self) -> LoadingParameters: - """Parameters for marker loading.""" return self._loading_params + + @property + def weights_params(self) -> WeightsParameters: + return self._weights_params + + @property + def boundary_params(self) -> BoundaryParameters: + """Parameters for marker loading.""" + return self._boundary_params @property def reject_weights(self): diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index 453627768..310ab129a 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -1,7 +1,10 @@ import numpy as np import struphy.pic.utilities_kernels as utils -from struphy.io.options import OptsLoading, OptsSpatialLoading +from struphy.io.options import (OptsLoading, + OptsSpatialLoading, + OptsMarkerBC, + OptsRecontructBC,) class LoadingParameters: @@ -9,6 +12,15 @@ class LoadingParameters: Parameters ---------- + Np : int + Total number of particles to load. + + ppc : int + Particles to load per cell if a grid is defined. Cells are defined from ``domain_array``. + + ppb : int + Particles to load per sorting box. Sorting boxes are defined from ``boxes_per_dim``. + loading : OptsLoading How to load markers: multiple options for Monte-Carlo, or "tesselation" for positioning them on a regular grid. @@ -41,6 +53,9 @@ class LoadingParameters: Key in .hdf5 file's restart/ folder where marker array is stored. """ def __init__(self, + Np: int = 100, + ppc: int = None, + ppb: int = None, loading: OptsLoading = "pseudo_random", seed: int = None, moments: tuple = None, @@ -53,6 +68,9 @@ def __init__(self, restart_key: str = None, ): + self.Np = Np + self.ppc = ppc + self.ppb = ppb self.loading = loading self.seed = seed self.moments = moments @@ -89,6 +107,32 @@ def __init__(self, self.threshold = threshold +class BoundaryParameters: + """Parameters for particle boundary and sph reconstruction boundary conditions. + + Parameters + ---------- + bc : tuple[OptsMarkerBC] + Boundary conditions for particle movement. + Either 'remove', 'reflect', 'periodic' or 'refill' in each direction. + + bc_refill : list + Either 'inner' or 'outer'. + + bc_sph : tuple[OptsRecontructBC] + Boundary conditions for sph kernel reconstruction. + """ + def __init__(self, + bc: tuple[OptsMarkerBC] = ("periodic", "periodic", "periodic"), + bc_refill = None, + bc_sph: tuple[OptsRecontructBC] = ("periodic", "periodic", "periodic"), + ): + self.bc = bc + self.bc_refill = bc_refill + self.bc_sph = bc_sph + + + def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): """ This function is for getting kinetic energy of the case when canonical momentum is used, rather than velocity diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 4cc7f2dd0..8a87d5428 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -430,7 +430,7 @@ def create_vtk( ) -def post_process_markers(path_in, path_out, species, kind, step=1): +def post_process_markers(path_in, path_out, species, kind="Particles6D", step=1): """Computes the Cartesian (x, y, z) coordinates of saved markers during a simulation and writes them to a .npy files and to .txt files. Also saves the weights. diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 78a66beb4..aecdc95df 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -137,8 +137,11 @@ def set_options(self, def allocate(self): # TODO: treat PolarVector as well, but polar splines are being reworked at the moment - self._b2 = self.projected_equil.b2 - assert self._b2.space == self.derham.Vh["2"] + if self.projected_equil is not None: + self._b2 = self.projected_equil.b2 + assert self._b2.space == self.derham.Vh["2"] + else: + self._b2 = self.derham.Vh["2"].zeros() if self.options.b2_var is None: self._b2_var = None From e0b292ccb136171b580fa16a2ce0cd256fe07d37 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 16 Aug 2025 10:46:11 +0200 Subject: [PATCH 083/292] started new particels tutorial --- doc/tutorials/tutorial_01_particles.ipynb | 192 ++++++++++++++++++ src/struphy/main.py | 45 +++- src/struphy/models/base.py | 12 +- .../post_processing/post_processing_tools.py | 21 +- 4 files changed, 248 insertions(+), 22 deletions(-) create mode 100644 doc/tutorials/tutorial_01_particles.ipynb diff --git a/doc/tutorials/tutorial_01_particles.ipynb b/doc/tutorials/tutorial_01_particles.ipynb new file mode 100644 index 000000000..81b36ccb1 --- /dev/null +++ b/doc/tutorials/tutorial_01_particles.ipynb @@ -0,0 +1,192 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d34c79c5", + "metadata": {}, + "source": [ + "# 1 - Parameter files\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3ecab659", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.io.options import EnvironmentOptions, Units, Time\n", + "from struphy.geometry import domains\n", + "from struphy.fields_background import equils\n", + "from struphy.topology import grids\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.io.options import FieldsBackground\n", + "from struphy.initial import perturbations\n", + "from struphy.kinetic_background import maxwellians\n", + "from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters\n", + "from struphy import main\n", + "\n", + "# import model, set verbosity\n", + "from struphy.models.toy import Vlasov as Model\n", + "verbose = True" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc43d2fc", + "metadata": {}, + "outputs": [], + "source": [ + "# environment options\n", + "env = EnvironmentOptions()\n", + "\n", + "# units\n", + "units = Units()\n", + "\n", + "# time stepping\n", + "time_opts = Time(dt=0.2, Tend=10.0)\n", + "\n", + "# geometry\n", + "domain = domains.Cuboid(l1=-5.0,\n", + " r1 = 5.0,\n", + " l2 = -7.0,\n", + " r2 = 7.0,\n", + " l3 = -1.0,\n", + " r3 = 1.0,\n", + " )\n", + "\n", + "# fluid equilibrium (can be used as part of initial conditions)\n", + "equil = None\n", + "\n", + "# grid\n", + "grid = None\n", + "\n", + "# derham options\n", + "derham_opts = None" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6c498fb3", + "metadata": {}, + "outputs": [], + "source": [ + "# light-weight model instance\n", + "model = Model()\n", + "\n", + "# species parameters\n", + "model.kinetic_ions.set_phys_params()\n", + "\n", + "loading_params = LoadingParameters(Np=15)\n", + "weights_params = WeightsParameters()\n", + "boundary_params = BoundaryParameters(bc=('reflect', 'reflect', 'periodic'))\n", + "model.kinetic_ions.set_markers(loading_params=loading_params, \n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params)\n", + "model.kinetic_ions.set_sorting_boxes()\n", + "model.kinetic_ions.set_save_data(n_markers=1.0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "be6875e7", + "metadata": {}, + "outputs": [], + "source": [ + "# propagator options\n", + "model.propagators.push_vxb.set_options()\n", + "model.propagators.push_eta.set_options()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc0ac424", + "metadata": {}, + "outputs": [], + "source": [ + "# initial conditions (background + perturbation)\n", + "perturbation = None\n", + "\n", + "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", + "model.kinetic_ions.var.add_background(background)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "03764138", + "metadata": {}, + "outputs": [], + "source": [ + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " units=units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dd7cfe62", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "main.pproc(path, physical=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "78b8cbab", + "metadata": {}, + "outputs": [], + "source": [ + "main.load_data(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7af3facd", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "env", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/struphy/main.py b/src/struphy/main.py index c0081767c..7a32b1370 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -22,8 +22,9 @@ from struphy.models.species import Species from struphy.models.variables import FEECVariable from struphy.io.options import Units, Time, EnvironmentOptions +from struphy.io.setup import import_parameters_py from struphy.geometry.base import Domain -from struphy.geometry.domains import Cuboid +from struphy.geometry import domains from struphy.fields_background.base import FluidEquilibrium from struphy.fields_background.equils import HomogenSlab from struphy.topology.grids import TensorProductGrid @@ -44,7 +45,7 @@ def run( env: EnvironmentOptions = EnvironmentOptions(), units: Units = Units(), time_opts: Time = Time(), - domain: Domain = Cuboid(), + domain: Domain = domains.Cuboid(), equil: FluidEquilibrium = HomogenSlab(), grid: TensorProductGrid = None, derham_opts: DerhamOptions = None, @@ -132,6 +133,8 @@ def run( pickle.dump(grid, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "derham_opts.bin"), 'wb') as f: pickle.dump(derham_opts, f, pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_out, "model.bin"), 'wb') as f: + pickle.dump(model, f, pickle.HIGHEST_PROTOCOL) # config clones if comm is None: @@ -423,6 +426,18 @@ def pproc( if MPI.COMM_WORLD.Get_rank() == 0: print(f"\n*** Start post-processing of {path}:") + + # load light-weight model instance from simulation + try: + params_in = import_parameters_py(os.path.join(path, "parameters.py")) + model: StruphyModel = params_in.model + domain: Domain = params_in.domain + except FileNotFoundError: + with open(os.path.join(path, "model.bin"), 'rb') as f: + model: StruphyModel = pickle.load(f) + with open(os.path.join(path, "domain.bin"), 'rb') as f: + domain_dct: Domain = pickle.load(f) + domain = getattr(domains, domain_dct["name"])(**domain_dct["params_map"]) # create post-processing folder path_pproc = os.path.join(path, "post_processing") @@ -455,9 +470,10 @@ def pproc( if "kinetic" in file.keys(): exist_kinetic = {"markers": False, "f": False, "n_sph": False} kinetic_species = [] + kinetic_kinds = [] for name in file["kinetic"].keys(): - print(f"{name = }") kinetic_species += [name] + kinetic_kinds += [next(iter(model.species[name].variables.values())).space] # check for saved markers if "markers" in file["kinetic"][name]: @@ -547,7 +563,7 @@ def pproc( # markers if exist_kinetic["markers"]: - post_process_markers(path, path_kinetics_species, species, kinetic_kinds[n], step) + post_process_markers(path, path_kinetics_species, species, domain, kinetic_kinds[n], step,) if guiding_center: assert kinetic_kinds[n] == "Particles6D" @@ -619,9 +635,11 @@ def load_data(path: str) -> SimData: # load time grid simdata.t_grid = np.load(os.path.join(path_pproc, "t_grid.npy")) - # load point data + # data paths path_fields = os.path.join(path_pproc, "fields_data") + path_kinetic = os.path.join(path_pproc, "kinetic_data") + # load point data if os.path.exists(path_fields): # grids @@ -645,6 +663,23 @@ def load_data(path: str) -> SimData: simdata.feec_species[spec] += [var] simdata.arrays[spec][var] = pickle.load(f) + if os.path.exists(path_kinetic): + + # species folders + species = next(os.walk(path_kinetic))[1] + print(f"{species = }") + for spec in species: + simdata.pic_species[spec] = [] + # simdata.arrays[spec] = {} + path_spec = os.path.join(path_kinetic, spec) + wlk = os.walk(path_spec) + print(f"{next(wlk) = }") + files = next(wlk)[2] + for file in files: + print(f"{file = }") + with open(os.path.join(path_spec, file), "r") as f: + simdata.orbits[spec][var] = pickle.load(f) + print("\nThe following data has been loaded:") print(f"{simdata.time_grid_size = }") print(f"{simdata.spline_grid_resolution = }") diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 6b06c72a4..d1a228f68 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1078,23 +1078,23 @@ def initialize_data_output(self, data: DataContainer, size): for name, species in self.kinetic_species.items(): assert isinstance(species, KineticSpecies) assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." - for _, var in species.variables.items(): + for varname, var in species.variables.items(): assert isinstance(var, PICVariable | SPHVariable) obj = var.particles assert isinstance(obj, Particles) - key_spec = "kinetic/" + key - key_spec_restart = "restart/" + key + key_spec = os.path.join("kinetic", name) + key_spec_restart = os.path.join("restart", name) data.add_data({key_spec_restart: obj._markers}) for key1, val1 in var.kinetic_data.items(): - key_dat = key_spec + "/" + key1 + key_dat = os.path.join(key_spec, key1) # case of "f" and "df" if isinstance(val1, dict): for key2, val2 in val1.items(): - key_f = key_dat + "/" + key2 + key_f = os.path.join(key_dat, key2) data.add_data({key_f: val2}) dims = (len(key2) - 2) // 3 + 1 @@ -1106,7 +1106,7 @@ def initialize_data_output(self, data: DataContainer, size): # case of "n_sph" elif isinstance(val1, list): for i, v1 in enumerate(val1): - key_n = key_dat + "/view_" + str(i) + key_n = os.path.join(key_dat, "view_", str(i)) data.add_data({key_n: v1}) # save 1d point values, not meshgrids, because attrs size is limited eta1 = var.kinetic_data["plot_pts"][i][0][:, 0, 0] diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 8a87d5428..fd2f65f1b 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -13,6 +13,7 @@ from struphy.io.options import EnvironmentOptions, Units, Time from struphy.topology.grids import TensorProductGrid from struphy.geometry import domains +from struphy.geometry.base import Domain class ParamsIn: @@ -430,7 +431,7 @@ def create_vtk( ) -def post_process_markers(path_in, path_out, species, kind="Particles6D", step=1): +def post_process_markers(path_in: str, path_out: str, species: str, domain: Domain, kind: str = "Particles6D", step: int = 1,): """Computes the Cartesian (x, y, z) coordinates of saved markers during a simulation and writes them to a .npy files and to .txt files. Also saves the weights. @@ -481,6 +482,9 @@ def post_process_markers(path_in, path_out, species, kind="Particles6D", step=1) species : str Name of the species for which the post processing should be performed. + + domain : Domain + Domain object. kind : str Name of the kinetic kind (Particles6D, Particles5D or Particles3D). @@ -489,17 +493,12 @@ def post_process_markers(path_in, path_out, species, kind="Particles6D", step=1) Whether to do post-processing at every time step (step=1, default), every second time step (step=2), etc. """ - # get # of MPI processes from meta.txt file - with open(os.path.join(path_in, "meta.txt"), "r") as f: - lines = f.readlines() - - nproc = lines[4].split()[-1] - - with open(os.path.join(path_in, "parameters.yml"), "r") as f: - params = yaml.load(f, Loader=yaml.FullLoader) + print(f"{domain = }") - # create domain for calculating markers' physical coordinates - domain = setup_domain_and_equil(params)[0] + # get # of MPI processes from meta.txt file + with open(os.path.join(path_in, "meta.yml"), "r") as f: + meta = yaml.load(f, Loader=yaml.FullLoader) + nproc = meta["MPI processes"] # open hdf5 files and get names and number of saved markers of kinetic species files = [ From 397282c69a1f0a891ebf8b8b4137555794ee82a5 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 25 Aug 2025 18:01:42 +0200 Subject: [PATCH 084/292] new way of saving during pproc --- doc/tutorials/tutorial_01_particles.ipynb | 103 ++++++++++++++++++--- src/struphy/main.py | 49 +++++++--- src/struphy/models/tests/test_LinearMHD.py | 4 +- src/struphy/models/tests/test_Maxwell.py | 6 +- 4 files changed, 131 insertions(+), 31 deletions(-) diff --git a/doc/tutorials/tutorial_01_particles.ipynb b/doc/tutorials/tutorial_01_particles.ipynb index 81b36ccb1..dc4bc05b5 100644 --- a/doc/tutorials/tutorial_01_particles.ipynb +++ b/doc/tutorials/tutorial_01_particles.ipynb @@ -5,8 +5,35 @@ "id": "d34c79c5", "metadata": {}, "source": [ - "# 1 - Parameter files\n", - "\n" + "# Parameter files\n", + "\n", + "Struphy parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", + "For each `MODEL`, the default parameter file can be generated from the console via\n", + "\n", + "```\n", + "struphy params MODEL\n", + "```\n", + "\n", + "This will create a file `params_MODEL.py` in the current working directory. To run the model type\n", + "\n", + "```\n", + "python params_MODEL.py\n", + "```\n", + "\n", + "The user should modify the parameter file to launch a specific simulation.\n", + "As an example, in what follows we discuss the parameter file of the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov) and run some simple examples. The file can be generated from\n", + "\n", + "```\n", + "struphy params Vlasov\n", + "```\n", + "\n", + "To see its contents, open the file in your preferred editor or type\n", + "\n", + "```\n", + "cat params_Vlasov.py\n", + "```\n", + "\n", + "## Part 1: imports" ] }, { @@ -49,13 +76,13 @@ "time_opts = Time(dt=0.2, Tend=10.0)\n", "\n", "# geometry\n", - "domain = domains.Cuboid(l1=-5.0,\n", - " r1 = 5.0,\n", - " l2 = -7.0,\n", - " r2 = 7.0,\n", - " l3 = -1.0,\n", - " r3 = 1.0,\n", - " )\n", + "l1 = -5.0\n", + "r1 = 5.0\n", + "l2 = -7.0\n", + "r2 = 7.0\n", + "l3 = -1.0\n", + "r3 = 1.0\n", + "domain = domains.Cuboid(l1=l1, r1=r1, l2=l2, r2=r2, l3=l3, r3=r3)\n", "\n", "# fluid equilibrium (can be used as part of initial conditions)\n", "equil = None\n", @@ -136,6 +163,14 @@ " )" ] }, + { + "cell_type": "markdown", + "id": "f9ce5099", + "metadata": {}, + "source": [ + "## Post processing" + ] + }, { "cell_type": "code", "execution_count": null, @@ -149,6 +184,14 @@ "main.pproc(path, physical=True)" ] }, + { + "cell_type": "markdown", + "id": "e317f88d", + "metadata": {}, + "source": [ + "## Viewing particle orbits" + ] + }, { "cell_type": "code", "execution_count": null, @@ -156,7 +199,18 @@ "metadata": {}, "outputs": [], "source": [ - "main.load_data(path)" + "simdata = main.load_data(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a329eb38", + "metadata": {}, + "outputs": [], + "source": [ + "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", + " print(f\"{k = }, {type(v) = }\")" ] }, { @@ -165,7 +219,34 @@ "id": "7af3facd", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "from matplotlib import pyplot as plt\n", + "\n", + "fig = plt.figure()\n", + "ax = fig.gca()\n", + "\n", + "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", + "\n", + "time = 0.\n", + "dt = time_opts.dt\n", + "Tend = time_opts.Tend\n", + "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", + " # print(k, v)\n", + " alpha = (Tend - time)/Tend\n", + " for i, particle in enumerate(v):\n", + " ax.scatter(particle[1], particle[2], c=colors[i % 4], alpha=alpha)\n", + " time += dt\n", + " \n", + "ax.plot([l1, l1], [l2, r2], 'k')\n", + "ax.plot([r1, r1], [l2, r2], 'k')\n", + "ax.plot([l1, r1], [l2, l2], 'k')\n", + "ax.plot([l1, r1], [r2, r2], 'k')\n", + "ax.set_xlabel('x')\n", + "ax.set_ylabel('y')\n", + "ax.set_xlim(-6.5, 6.5)\n", + "ax.set_ylim(-9, 9)\n", + "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" + ] } ], "metadata": { diff --git a/src/struphy/main.py b/src/struphy/main.py index 7a32b1370..6952ba621 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -600,6 +600,7 @@ def __init__(self, path: str): self.pic_species = {} self.sph_species = {} self.arrays = {} + self.orbits = {} self.grids_log: list[np.ndarray] = None self.grids_phy: list[np.ndarray] = None self.t_grid: np.ndarray = None @@ -651,17 +652,19 @@ def load_data(path: str) -> SimData: # species folders species = next(os.walk(path_fields))[1] for spec in species: - simdata.feec_species[spec] = [] - simdata.arrays[spec] = {} + simdata.feec_species[spec] = {} + # simdata.arrays[spec] = {} path_spec = os.path.join(path_fields, spec) wlk = os.walk(path_spec) files = next(wlk)[2] + print(f"\nFiles in {path_spec}: {files}") for file in files: if ".bin" in file: var = file.split(".")[0] with open(os.path.join(path_spec, file), "rb") as f: - simdata.feec_species[spec] += [var] - simdata.arrays[spec][var] = pickle.load(f) + # try: + simdata.feec_species[spec][var] = pickle.load(f) + # simdata.arrays[spec][var] = pickle.load(f) if os.path.exists(path_kinetic): @@ -669,23 +672,39 @@ def load_data(path: str) -> SimData: species = next(os.walk(path_kinetic))[1] print(f"{species = }") for spec in species: - simdata.pic_species[spec] = [] - # simdata.arrays[spec] = {} + simdata.pic_species[spec] = {} path_spec = os.path.join(path_kinetic, spec) wlk = os.walk(path_spec) - print(f"{next(wlk) = }") - files = next(wlk)[2] - for file in files: - print(f"{file = }") - with open(os.path.join(path_spec, file), "r") as f: - simdata.orbits[spec][var] = pickle.load(f) + sub_folders = next(wlk)[1] + # print(f"{sub_folders = }") + for folder in sub_folders: + # simdata.pic_species[spec][folder] = {} + tmp = {} + path_dat = os.path.join(path_spec, folder) + sub_wlk = os.walk(path_dat) + files = next(sub_wlk)[2] + for file in files: + # print(f"{file = }") + if ".npy" in file: + var = file.split(".")[0] + tmp[var] = np.load(os.path.join(path_dat, file)) + # sort dict + simdata.pic_species[spec][folder] = dict(sorted(tmp.items())) print("\nThe following data has been loaded:") print(f"{simdata.time_grid_size = }") print(f"{simdata.spline_grid_resolution = }") - print(f"{simdata.feec_species = }") - print(f"{simdata.pic_species = }") - print(f"{simdata.sph_species = }") + print(f"simdata.feec_species:") + for k, v in simdata.feec_species.items(): + print(f" {k}:") + for kk, vv in v.items(): + print(f" {kk}") + print(f"simdata.pic_species:") + for k, v in simdata.pic_species.items(): + print(f" {k}:") + for kk, vv in v.items(): + print(f" {kk}") + print(f"simdata.sph_species:") return simdata diff --git a/src/struphy/models/tests/test_LinearMHD.py b/src/struphy/models/tests/test_LinearMHD.py index 4b28813e6..00918b02c 100644 --- a/src/struphy/models/tests/test_LinearMHD.py +++ b/src/struphy/models/tests/test_LinearMHD.py @@ -87,7 +87,7 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): simdata = main.load_data(env.path_out) # first fft - u_of_t = simdata.arrays["mhd"]["velocity_log"] + u_of_t = simdata.feec_species["mhd"]["velocity_log"] Bsquare = (B0x**2 + B0y**2 + B0z**2) p0 = beta * Bsquare / 2 @@ -121,7 +121,7 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): assert np.abs(coeffs[0][0] - v_alfven) < 0.07 # second fft - p_of_t = simdata.arrays["mhd"]["pressure_log"] + p_of_t = simdata.feec_species["mhd"]["pressure_log"] _1, _2, _3, coeffs = power_spectrum_2d(p_of_t, "pressure_log", diff --git a/src/struphy/models/tests/test_Maxwell.py b/src/struphy/models/tests/test_Maxwell.py index 610e5f1c2..c366cf284 100644 --- a/src/struphy/models/tests/test_Maxwell.py +++ b/src/struphy/models/tests/test_Maxwell.py @@ -79,7 +79,7 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): simdata = main.load_data(env.path_out) # fft - E_of_t = simdata.arrays["em_fields"]["e_field_log"] + E_of_t = simdata.feec_species["em_fields"]["e_field_log"] _1, _2, _3, coeffs = power_spectrum_2d(E_of_t, "e_field_log", grids=simdata.grids_log, @@ -262,5 +262,5 @@ def to_E_theta(X, Y, E_x, E_y): if __name__ == "__main__": - # test_light_wave_1d(algo="explicit", do_plot=True) - test_coaxial(do_plot=True) \ No newline at end of file + test_light_wave_1d(algo="explicit", do_plot=True) + # test_coaxial(do_plot=True) \ No newline at end of file From 7930d7dade3e471d4d91e29d42c3fdd6d9b37ad3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 25 Aug 2025 21:21:55 +0200 Subject: [PATCH 085/292] new tutorial 01: parameter files --- .../tutorial_01_parameter_files.ipynb | 330 ++++++++++++++++++ 1 file changed, 330 insertions(+) create mode 100644 doc/tutorials/tutorial_01_parameter_files.ipynb diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/doc/tutorials/tutorial_01_parameter_files.ipynb new file mode 100644 index 000000000..287108ff0 --- /dev/null +++ b/doc/tutorials/tutorial_01_parameter_files.ipynb @@ -0,0 +1,330 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d34c79c5", + "metadata": {}, + "source": [ + "# Parameter files\n", + "\n", + "Struphy parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", + "For each `MODEL`, the default parameter file can be generated from the console via\n", + "\n", + "```\n", + "struphy params MODEL\n", + "```\n", + "\n", + "This will create a file `params_MODEL.py` in the current working directory. To run the model type\n", + "\n", + "```\n", + "python params_MODEL.py\n", + "```\n", + "\n", + "The user should modify the parameter file to launch a specific simulation.\n", + "As an example, let us discuss the parameter file of the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov) and run some simple examples. \n", + "The file can be generated from\n", + "\n", + "```\n", + "struphy params Vlasov\n", + "```\n", + "\n", + "To see its contents, open the file in your preferred editor or type\n", + "\n", + "```\n", + "cat params_Vlasov.py\n", + "```\n", + "\n", + "## Part 1: Imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3ecab659", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.io.options import EnvironmentOptions, Units, Time\n", + "from struphy.geometry import domains\n", + "from struphy.fields_background import equils\n", + "from struphy.topology import grids\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.io.options import FieldsBackground\n", + "from struphy.initial import perturbations\n", + "from struphy.kinetic_background import maxwellians\n", + "from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters\n", + "from struphy import main\n", + "\n", + "# import model, set verbosity\n", + "from struphy.models.toy import Vlasov as Model\n", + "verbose = True" + ] + }, + { + "cell_type": "markdown", + "id": "5cf6d9c7", + "metadata": {}, + "source": [ + "All Struphy parameter files import the modules listed above, even though some of them might not be needed in a specific model. The last import imports the model itself, always under the alias `Model`.\n", + "\n", + "## Part 2: Generic options\n", + "\n", + "The following lines refer to options that can be set for any model. These are:\n", + "\n", + "* Environment options (paths, saving, domain cloning)\n", + "* Model units\n", + "* Time options\n", + "* Problem geometry (mapped domain)\n", + "* Static background (equilibrium)\n", + "* Grid\n", + "* Derham complex" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc43d2fc", + "metadata": {}, + "outputs": [], + "source": [ + "# environment options\n", + "env = EnvironmentOptions()\n", + "\n", + "# units\n", + "units = Units()\n", + "\n", + "# time stepping\n", + "time_opts = Time(dt=0.2, Tend=10.0)\n", + "\n", + "# geometry\n", + "l1 = -5.0\n", + "r1 = 5.0\n", + "l2 = -7.0\n", + "r2 = 7.0\n", + "l3 = -1.0\n", + "r3 = 1.0\n", + "domain = domains.Cuboid(l1=l1, r1=r1, l2=l2, r2=r2, l3=l3, r3=r3)\n", + "\n", + "# fluid equilibrium (can be used as part of initial conditions)\n", + "equil = None\n", + "\n", + "# grid\n", + "grid = None\n", + "\n", + "# derham options\n", + "derham_opts = None" + ] + }, + { + "cell_type": "markdown", + "id": "74e6f739", + "metadata": {}, + "source": [ + "## Part 3: Model instance\n", + "\n", + "Here, a light-weight instance of the model is created, without allocating memory. The light-weight instance is used to set model-specific parameters for the model's species." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6c498fb3", + "metadata": {}, + "outputs": [], + "source": [ + "# light-weight model instance\n", + "model = Model()\n", + "\n", + "# species parameters\n", + "model.kinetic_ions.set_phys_params()\n", + "\n", + "loading_params = LoadingParameters(Np=15)\n", + "weights_params = WeightsParameters()\n", + "boundary_params = BoundaryParameters(bc=('reflect', 'reflect', 'periodic'))\n", + "model.kinetic_ions.set_markers(loading_params=loading_params, \n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params)\n", + "model.kinetic_ions.set_sorting_boxes()\n", + "model.kinetic_ions.set_save_data(n_markers=1.0)" + ] + }, + { + "cell_type": "markdown", + "id": "b0f65b0a", + "metadata": {}, + "source": [ + "## Part 4: Propagator options" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "be6875e7", + "metadata": {}, + "outputs": [], + "source": [ + "# propagator options\n", + "model.propagators.push_vxb.set_options()\n", + "model.propagators.push_eta.set_options()" + ] + }, + { + "cell_type": "markdown", + "id": "b1ef8b97", + "metadata": {}, + "source": [ + "## Part 5: Initial conditions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc0ac424", + "metadata": {}, + "outputs": [], + "source": [ + "# initial conditions (background + perturbation)\n", + "perturbation = None\n", + "\n", + "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", + "model.kinetic_ions.var.add_background(background)" + ] + }, + { + "cell_type": "markdown", + "id": "879978af", + "metadata": {}, + "source": [ + "## Part 6: Run the model\n", + "\n", + "In the final part of the parameter file, the actual run command is invoked. This command will allocate memory and run the specified simulation. The run command is not executed when the parameter file is imported in another Python script." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "03764138", + "metadata": {}, + "outputs": [], + "source": [ + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " units=units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "f9ce5099", + "metadata": {}, + "source": [ + "## Post processing" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dd7cfe62", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "main.pproc(path, physical=True)" + ] + }, + { + "cell_type": "markdown", + "id": "e317f88d", + "metadata": {}, + "source": [ + "## Viewing particle orbits" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "78b8cbab", + "metadata": {}, + "outputs": [], + "source": [ + "simdata = main.load_data(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a329eb38", + "metadata": {}, + "outputs": [], + "source": [ + "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", + " print(f\"{k = }, {type(v) = }\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7af3facd", + "metadata": {}, + "outputs": [], + "source": [ + "from matplotlib import pyplot as plt\n", + "\n", + "fig = plt.figure()\n", + "ax = fig.gca()\n", + "\n", + "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", + "\n", + "time = 0.\n", + "dt = time_opts.dt\n", + "Tend = time_opts.Tend\n", + "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", + " # print(k, v)\n", + " alpha = (Tend - time)/Tend\n", + " for i, particle in enumerate(v):\n", + " ax.scatter(particle[1], particle[2], c=colors[i % 4], alpha=alpha)\n", + " time += dt\n", + " \n", + "ax.plot([l1, l1], [l2, r2], 'k')\n", + "ax.plot([r1, r1], [l2, r2], 'k')\n", + "ax.plot([l1, r1], [l2, l2], 'k')\n", + "ax.plot([l1, r1], [r2, r2], 'k')\n", + "ax.set_xlabel('x')\n", + "ax.set_ylabel('y')\n", + "ax.set_xlim(-6.5, 6.5)\n", + "ax.set_ylim(-9, 9)\n", + "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "env", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 24bd43c650fe4acffb26185c920891bb2e661684 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 25 Aug 2025 21:27:41 +0200 Subject: [PATCH 086/292] improve tutorial 01 --- doc/tutorials/tutorial_01_parameter_files.ipynb | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/doc/tutorials/tutorial_01_parameter_files.ipynb index 287108ff0..2bb7a8e6a 100644 --- a/doc/tutorials/tutorial_01_parameter_files.ipynb +++ b/doc/tutorials/tutorial_01_parameter_files.ipynb @@ -20,7 +20,7 @@ "python params_MODEL.py\n", "```\n", "\n", - "The user should modify the parameter file to launch a specific simulation.\n", + "The user can modify the parameter file to launch a specific simulation.\n", "As an example, let us discuss the parameter file of the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov) and run some simple examples. \n", "The file can be generated from\n", "\n", @@ -77,7 +77,9 @@ "* Problem geometry (mapped domain)\n", "* Static background (equilibrium)\n", "* Grid\n", - "* Derham complex" + "* Derham complex\n", + "\n", + "Check the respective classes for possible options." ] }, { @@ -122,7 +124,16 @@ "source": [ "## Part 3: Model instance\n", "\n", - "Here, a light-weight instance of the model is created, without allocating memory. The light-weight instance is used to set model-specific parameters for the model's species." + "Here, a light-weight instance of the model is created, without allocating memory. The light-weight instance is used to set model-specific parameters for the model's species.\n", + "\n", + "Check the functions \n", + "\n", + "* `Species.set_phys_params()`\n", + "* `KineticSpecies.set_markers()`\n", + "* `KineticSpecies.set_sorting_boxes()`\n", + "* `KineticSpecies.set_save_data()`\n", + "\n", + " for possible options." ] }, { From 9b151a9a28a0ba1766b3c919228e49c46f42de86 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 25 Aug 2025 21:34:53 +0200 Subject: [PATCH 087/292] improve tutorial 01 more --- doc/tutorials/tutorial_01_parameter_files.ipynb | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/doc/tutorials/tutorial_01_parameter_files.ipynb index 2bb7a8e6a..639552474 100644 --- a/doc/tutorials/tutorial_01_parameter_files.ipynb +++ b/doc/tutorials/tutorial_01_parameter_files.ipynb @@ -164,7 +164,9 @@ "id": "b0f65b0a", "metadata": {}, "source": [ - "## Part 4: Propagator options" + "## Part 4: Propagator options\n", + "\n", + "Check the method `set_options()` of each propagator for possible options." ] }, { @@ -184,7 +186,9 @@ "id": "b1ef8b97", "metadata": {}, "source": [ - "## Part 5: Initial conditions" + "## Part 5: Initial conditions\n", + "\n", + "Use the methods `Variable.add_background()` and `Variable.add_perturbation()` to set initial conditions for each variable of a species. Variables that are not specified are intialized as zero." ] }, { From b1b93c3908fe2b30537db5db2c89b1dc63b163c9 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 26 Aug 2025 11:48:36 +0200 Subject: [PATCH 088/292] finish tutorial 01 --- .../tutorial_01_parameter_files.ipynb | 59 +++++++++++++++---- src/struphy/main.py | 2 - .../post_processing/post_processing_tools.py | 6 +- 3 files changed, 49 insertions(+), 18 deletions(-) diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/doc/tutorials/tutorial_01_parameter_files.ipynb index 639552474..e803e3746 100644 --- a/doc/tutorials/tutorial_01_parameter_files.ipynb +++ b/doc/tutorials/tutorial_01_parameter_files.ipynb @@ -5,7 +5,7 @@ "id": "d34c79c5", "metadata": {}, "source": [ - "# Parameter files\n", + "# Parameter files and `struphy.main`\n", "\n", "Struphy parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", "For each `MODEL`, the default parameter file can be generated from the console via\n", @@ -34,7 +34,7 @@ "cat params_Vlasov.py\n", "```\n", "\n", - "## Part 1: Imports" + "## Parameters part 1: Imports" ] }, { @@ -67,7 +67,7 @@ "source": [ "All Struphy parameter files import the modules listed above, even though some of them might not be needed in a specific model. The last import imports the model itself, always under the alias `Model`.\n", "\n", - "## Part 2: Generic options\n", + "## Parameters part 2: Generic options\n", "\n", "The following lines refer to options that can be set for any model. These are:\n", "\n", @@ -122,7 +122,7 @@ "id": "74e6f739", "metadata": {}, "source": [ - "## Part 3: Model instance\n", + "## Parameters part 3: Model instance\n", "\n", "Here, a light-weight instance of the model is created, without allocating memory. The light-weight instance is used to set model-specific parameters for the model's species.\n", "\n", @@ -164,7 +164,7 @@ "id": "b0f65b0a", "metadata": {}, "source": [ - "## Part 4: Propagator options\n", + "## Parameters part 4: Propagator options\n", "\n", "Check the method `set_options()` of each propagator for possible options." ] @@ -186,7 +186,7 @@ "id": "b1ef8b97", "metadata": {}, "source": [ - "## Part 5: Initial conditions\n", + "## Parameters part 5: Initial conditions\n", "\n", "Use the methods `Variable.add_background()` and `Variable.add_perturbation()` to set initial conditions for each variable of a species. Variables that are not specified are intialized as zero." ] @@ -210,9 +210,9 @@ "id": "879978af", "metadata": {}, "source": [ - "## Part 6: Run the model\n", + "## Parameters part 6: `main.run`\n", "\n", - "In the final part of the parameter file, the actual run command is invoked. This command will allocate memory and run the specified simulation. The run command is not executed when the parameter file is imported in another Python script." + "In the final part of the parameter file, the `main.run` command is invoked. This command will allocate memory and run the specified simulation. The run command is not executed when the parameter file is imported in another Python script." ] }, { @@ -240,7 +240,9 @@ "id": "f9ce5099", "metadata": {}, "source": [ - "## Post processing" + "## Post processing: `main.pproc`\n", + "\n", + "Aside from `run`, the Struphy `main` module has also a `pproc` routine for post-processing raw simulation data:" ] }, { @@ -251,8 +253,8 @@ "outputs": [], "source": [ "import os\n", - "\n", "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "\n", "main.pproc(path, physical=True)" ] }, @@ -261,7 +263,9 @@ "id": "e317f88d", "metadata": {}, "source": [ - "## Viewing particle orbits" + "## Loading data: `main.load_data`\n", + "\n", + "After post-processing, the generated data can be loaded via `main.load_data`:" ] }, { @@ -274,6 +278,25 @@ "simdata = main.load_data(path)" ] }, + { + "cell_type": "markdown", + "id": "a9468479", + "metadata": {}, + "source": [ + "`main.load_data` returns a `SimData` object, which you can inspect to get further info on possible data to load:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0531679e", + "metadata": {}, + "outputs": [], + "source": [ + "for k, v in simdata.__dict__.items():\n", + " print(k, type(v))" + ] + }, { "cell_type": "code", "execution_count": null, @@ -285,6 +308,16 @@ " print(f\"{k = }, {type(v) = }\")" ] }, + { + "cell_type": "markdown", + "id": "08544a91", + "metadata": {}, + "source": [ + "## Plotting particle orbits\n", + "\n", + "In this example, for the species `kinetic_ions` some particle orbits have been saved. Let us plot them:" + ] + }, { "cell_type": "code", "execution_count": null, @@ -303,10 +336,10 @@ "dt = time_opts.dt\n", "Tend = time_opts.Tend\n", "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", - " # print(k, v)\n", + " # print(f\"{v[0] = }\")\n", " alpha = (Tend - time)/Tend\n", " for i, particle in enumerate(v):\n", - " ax.scatter(particle[1], particle[2], c=colors[i % 4], alpha=alpha)\n", + " ax.scatter(particle[0], particle[1], c=colors[i % 4], alpha=alpha)\n", " time += dt\n", " \n", "ax.plot([l1, l1], [l2, r2], 'k')\n", diff --git a/src/struphy/main.py b/src/struphy/main.py index 6952ba621..2fcc74191 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -599,8 +599,6 @@ def __init__(self, path: str): self.feec_species = {} self.pic_species = {} self.sph_species = {} - self.arrays = {} - self.orbits = {} self.grids_log: list[np.ndarray] = None self.grids_phy: list[np.ndarray] = None self.t_grid: np.ndarray = None diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index fd2f65f1b..77c01d432 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -581,10 +581,10 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma change_out_order=True, ) - # move ids to first column and save - temp = np.roll(temp, 1, axis=1) - + # save numpy np.save(file_npy, temp) + # move ids to first column and save txt + temp = np.roll(temp, 1, axis=1) np.savetxt(file_txt, temp[:, (0, 1, 2, 3, -1)], fmt="%12.6f", delimiter=", ") # close hdf5 files From 7ea7aeeac1d20432553ee3a30636939da33244ba Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 27 Aug 2025 11:25:44 +0200 Subject: [PATCH 089/292] - get rid of @dataclass for Species A species has multiple methods to it and its variables have mutable attributes. So a dataclass is a bit misleading. Moreover, the default values are class variables which are shared across all instances (this makes the use of dataclasses dangerous imho). - new method Species.init_variables instead of __post_init__ - use new BaseUnits class to have all imutabke parameters as @dataclass --- doc/sections/tutorials.rst | 13 +- .../tutorial_01_parameter_files.ipynb | 16 +- .../tutorial_02_test_particles.ipynb | 1475 +++++++++++++++++ src/struphy/io/options.py | 33 +- src/struphy/main.py | 8 +- src/struphy/models/base.py | 8 +- src/struphy/models/fluid.py | 18 +- src/struphy/models/species.py | 26 +- src/struphy/models/toy.py | 21 +- src/struphy/models/variables.py | 34 +- src/struphy/propagators/base.py | 4 +- .../tutorial_01_kinetic_particles.ipynb | 0 .../tutorial_01_parameter_files.ipynb | 378 +++++ .../tutorial_01_particles.ipynb | 0 .../tutorial_02_fluid_particles.ipynb | 0 .../tutorial_03_discrete_derham.ipynb | 0 .../tutorial_04_mapped_domains.ipynb | 0 .../tutorial_05_mhd_equilibria.ipynb | 0 .../tutorial_06_poisson.ipynb | 0 .../tutorial_07_heat_equation.ipynb | 0 .../tutorial_08_maxwell.ipynb | 0 .../tutorial_09_vlasov_maxwell.ipynb | 0 .../tutorial_10_linear_mhd.ipynb | 0 .../tutorial_11_data_structures.ipynb | 0 .../tutorial_12_struphy_data_pproc.ipynb | 0 25 files changed, 1948 insertions(+), 86 deletions(-) create mode 100644 doc/tutorials/tutorial_02_test_particles.ipynb rename {doc/tutorials => tutorials_old}/tutorial_01_kinetic_particles.ipynb (100%) create mode 100644 tutorials_old/tutorial_01_parameter_files.ipynb rename {doc/tutorials => tutorials_old}/tutorial_01_particles.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_02_fluid_particles.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_03_discrete_derham.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_04_mapped_domains.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_05_mhd_equilibria.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_06_poisson.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_07_heat_equation.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_08_maxwell.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_09_vlasov_maxwell.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_10_linear_mhd.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_11_data_structures.ipynb (100%) rename {doc/tutorials => tutorials_old}/tutorial_12_struphy_data_pproc.ipynb (100%) diff --git a/doc/sections/tutorials.rst b/doc/sections/tutorials.rst index 9c1ebc8ae..53440ce87 100644 --- a/doc/sections/tutorials.rst +++ b/doc/sections/tutorials.rst @@ -12,15 +12,4 @@ They can thus be used for MPI parallel runs in HPC applications. :maxdepth: 1 :caption: Notebook tutorials: - ../tutorials/tutorial_01_kinetic_particles - ../tutorials/tutorial_02_fluid_particles - ../tutorials/tutorial_03_discrete_derham - ../tutorials/tutorial_04_mapped_domains - ../tutorials/tutorial_05_mhd_equilibria - ../tutorials/tutorial_06_poisson - ../tutorials/tutorial_07_heat_equation - ../tutorials/tutorial_08_maxwell - ../tutorials/tutorial_09_vlasov_maxwell - ../tutorials/tutorial_10_linear_mhd - ../tutorials/tutorial_11_data_structures - ../tutorials/tutorial_12_struphy_data_pproc + ../tutorials/tutorial_01_parameter_files diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/doc/tutorials/tutorial_01_parameter_files.ipynb index e803e3746..53c292b21 100644 --- a/doc/tutorials/tutorial_01_parameter_files.ipynb +++ b/doc/tutorials/tutorial_01_parameter_files.ipynb @@ -5,10 +5,10 @@ "id": "d34c79c5", "metadata": {}, "source": [ - "# Parameter files and `struphy.main`\n", + "# 01 - Parameters and `struphy.main`\n", "\n", "Struphy parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", - "For each `MODEL`, the default parameter file can be generated from the console via\n", + "For each [MODEL](https://struphy.pages.mpcdf.de/struphy/sections/models.html), the default parameter file can be generated from the console via\n", "\n", "```\n", "struphy params MODEL\n", @@ -34,7 +34,7 @@ "cat params_Vlasov.py\n", "```\n", "\n", - "## Parameters part 1: Imports" + "## Part 1: Imports" ] }, { @@ -67,7 +67,7 @@ "source": [ "All Struphy parameter files import the modules listed above, even though some of them might not be needed in a specific model. The last import imports the model itself, always under the alias `Model`.\n", "\n", - "## Parameters part 2: Generic options\n", + "## Part 2: Generic options\n", "\n", "The following lines refer to options that can be set for any model. These are:\n", "\n", @@ -122,7 +122,7 @@ "id": "74e6f739", "metadata": {}, "source": [ - "## Parameters part 3: Model instance\n", + "## Part 3: Model instance\n", "\n", "Here, a light-weight instance of the model is created, without allocating memory. The light-weight instance is used to set model-specific parameters for the model's species.\n", "\n", @@ -164,7 +164,7 @@ "id": "b0f65b0a", "metadata": {}, "source": [ - "## Parameters part 4: Propagator options\n", + "## Part 4: Propagator options\n", "\n", "Check the method `set_options()` of each propagator for possible options." ] @@ -186,7 +186,7 @@ "id": "b1ef8b97", "metadata": {}, "source": [ - "## Parameters part 5: Initial conditions\n", + "## Part 5: Initial conditions\n", "\n", "Use the methods `Variable.add_background()` and `Variable.add_perturbation()` to set initial conditions for each variable of a species. Variables that are not specified are intialized as zero." ] @@ -210,7 +210,7 @@ "id": "879978af", "metadata": {}, "source": [ - "## Parameters part 6: `main.run`\n", + "## Part 6: `main.run`\n", "\n", "In the final part of the parameter file, the `main.run` command is invoked. This command will allocate memory and run the specified simulation. The run command is not executed when the parameter file is imported in another Python script." ] diff --git a/doc/tutorials/tutorial_02_test_particles.ipynb b/doc/tutorials/tutorial_02_test_particles.ipynb new file mode 100644 index 000000000..1ec402de3 --- /dev/null +++ b/doc/tutorials/tutorial_02_test_particles.ipynb @@ -0,0 +1,1475 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "724e47a7", + "metadata": {}, + "source": [ + "# 2 - Test particles\n", + "\n", + "Let us explore some options of the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov), which has already been introdused in Tutorial 1. In particular, we will\n", + "\n", + "1. Change the geometry\n", + "2. Change the loading of the markers\n", + "3. Set a static background magnetic field\n", + "\n", + "## Particles in a cylinder\n", + "\n", + "We use the same setup as in Tutorial 1 but change the domain $\\Omega$ to a cylinder. We explore two options for drawing markers in posittion space:\n", + "\n", + "- uniform in logical space $[0, 1]^3 = F^{-1}(\\Omega)$\n", + "- uniform on the cylinder $\\Omega$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "541d473c", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.io.options import EnvironmentOptions, BaseUnits, Time\n", + "from struphy.geometry import domains\n", + "from struphy.fields_background import equils\n", + "from struphy.topology import grids\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.io.options import FieldsBackground\n", + "from struphy.initial import perturbations\n", + "from struphy.kinetic_background import maxwellians\n", + "from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters\n", + "from struphy import main\n", + "\n", + "# import model, set verbosity\n", + "from struphy.models.toy import Vlasov as Model\n", + "verbose = False" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b6629b5e", + "metadata": {}, + "outputs": [], + "source": [ + "# environment options\n", + "env = EnvironmentOptions()\n", + "env_2 = EnvironmentOptions(sim_folder=\"sim_2\")\n", + "\n", + "# units\n", + "base_units = BaseUnits()\n", + "\n", + "# time stepping\n", + "time_opts = Time(dt=0.2, Tend=0.2)\n", + "\n", + "# geometry\n", + "a1 = 0.\n", + "a2 = 5.\n", + "Lz = 1.\n", + "domain = domains.HollowCylinder(a1=a1, a2=a2, Lz=Lz)\n", + "\n", + "# fluid equilibrium (can be used as part of initial conditions)\n", + "equil = None\n", + "\n", + "# grid\n", + "grid = None\n", + "\n", + "# derham options\n", + "derham_opts = None" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c13b1c6c", + "metadata": {}, + "outputs": [], + "source": [ + "# light-weight model instance\n", + "model = Model()\n", + "model_2 = Model()\n", + "\n", + "# species parameters\n", + "model.kinetic_ions.set_phys_params()\n", + "model_2.kinetic_ions.set_phys_params()\n", + "\n", + "loading_params = LoadingParameters(Np=1000)\n", + "loading_params_2 = LoadingParameters(Np=1000, spatial=\"disc\")\n", + "\n", + "weights_params = WeightsParameters()\n", + "boundary_params = BoundaryParameters(bc=('reflect', 'periodic', 'periodic'))\n", + "\n", + "model.kinetic_ions.set_markers(loading_params=loading_params, \n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params)\n", + "model_2.kinetic_ions.set_markers(loading_params=loading_params_2, \n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params)\n", + "\n", + "model.kinetic_ions.set_sorting_boxes()\n", + "model_2.kinetic_ions.set_sorting_boxes()\n", + "\n", + "model.kinetic_ions.set_save_data(n_markers=1.0)\n", + "model_2.kinetic_ions.set_save_data(n_markers=1.0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "585c3808", + "metadata": {}, + "outputs": [], + "source": [ + "# propagator options\n", + "model.propagators.push_vxb.set_options()\n", + "model.propagators.push_eta.set_options()\n", + "\n", + "model_2.propagators.push_vxb.set_options()\n", + "model_2.propagators.push_eta.set_options()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9394a942", + "metadata": {}, + "outputs": [], + "source": [ + "# initial conditions (background + perturbation)\n", + "perturbation = None\n", + "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", + "\n", + "model.kinetic_ions.var.add_background(background)\n", + "model_2.kinetic_ions.var.add_background(background)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aaa18727", + "metadata": {}, + "outputs": [], + "source": [ + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " base_units=base_units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24b65c51", + "metadata": {}, + "outputs": [], + "source": [ + "main.run(model_2, \n", + " params_path=None, \n", + " env=env_2, \n", + " base_units=base_units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "679e992b", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "path_2 = os.path.join(os.getcwd(), \"sim_2\")\n", + "\n", + "main.pproc(path)\n", + "main.pproc(path_2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a98937e5", + "metadata": {}, + "outputs": [], + "source": [ + "simdata = main.load_data(path)\n", + "simdata_2 = main.load_data(path_2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0aaffb2e", + "metadata": {}, + "outputs": [], + "source": [ + "from matplotlib import pyplot as plt\n", + "\n", + "fig = plt.figure(figsize=(10, 6)) \n", + "\n", + "pushed_pos = simdata.pic_species[\"kinetic_ions\"][\"orbits\"][\"kinetic_ions_0\"]\n", + "pushed_pos_uni = simdata_2.pic_species[\"kinetic_ions\"][\"orbits\"][\"kinetic_ions_0\"]\n", + "\n", + "plt.subplot(1, 2, 1)\n", + "plt.scatter(pushed_pos[:, 0], pushed_pos[:, 1], s=2.)\n", + "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "ax = plt.gca()\n", + "ax.add_patch(circle1)\n", + "ax.set_aspect('equal')\n", + "plt.xlabel('x')\n", + "plt.ylabel('y')\n", + "plt.title('Draw uniform in logical space')\n", + "\n", + "plt.subplot(1, 2, 2)\n", + "plt.scatter(pushed_pos_uni[:, 0], pushed_pos_uni[:, 1], s=2.)\n", + "circle2 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "ax = plt.gca()\n", + "ax.add_patch(circle2)\n", + "ax.set_aspect('equal')\n", + "plt.xlabel('x')\n", + "plt.ylabel('y')\n", + "plt.title('Draw uniform on disc');" + ] + }, + { + "cell_type": "markdown", + "id": "cbb0dfd5", + "metadata": {}, + "source": [ + "## Reflecting boundary conditions \n", + "\n", + "We shall set reflecting boundary conditions in radial direction, which in Struphy is always the logical direction $\\eta_1$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5aa44526", + "metadata": {}, + "outputs": [], + "source": [ + "# instantiate Particle object\n", + "Np = 1000\n", + "bc = ['remove', 'periodic', 'periodic']\n", + "loading_params = {'seed': None}\n", + "\n", + "particles = Particles6D(Np=Np, \n", + " bc=bc, \n", + " loading_params=loading_params)\n", + "\n", + "# instantiate another Particle object\n", + "name = 'test_uni'\n", + "loading_params = {'seed': None, 'spatial': 'disc'}\n", + "particles_uni = Particles6D(Np=Np, \n", + " bc=bc, \n", + " loading_params=loading_params)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2d5cd7a4", + "metadata": {}, + "outputs": [], + "source": [ + "particles.draw_markers()\n", + "particles_uni.draw_markers()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b74f24b3", + "metadata": {}, + "outputs": [], + "source": [ + "# positions on the physical domain Omega\n", + "pushed_pos = domain(particles.positions).T\n", + "pushed_pos_uni = domain(particles_uni.positions).T" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ab1a5b6d", + "metadata": {}, + "outputs": [], + "source": [ + "fig = plt.figure(figsize=(10, 6)) \n", + "\n", + "plt.subplot(1, 2, 1)\n", + "plt.scatter(pushed_pos[:, 0], pushed_pos[:, 1], s=2.)\n", + "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "ax = plt.gca()\n", + "ax.add_patch(circle1)\n", + "ax.set_aspect('equal')\n", + "plt.xlabel('x')\n", + "plt.ylabel('y')\n", + "plt.title('Draw uniform in logical space')\n", + "\n", + "plt.subplot(1, 2, 2)\n", + "plt.scatter(pushed_pos_uni[:, 0], pushed_pos_uni[:, 1], s=2.)\n", + "circle2 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "ax = plt.gca()\n", + "ax.add_patch(circle2)\n", + "ax.set_aspect('equal')\n", + "plt.xlabel('x')\n", + "plt.ylabel('y')\n", + "plt.title('Draw uniform on disc');" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8d5aa753", + "metadata": {}, + "outputs": [], + "source": [ + "# instantiate Particle object\n", + "Np = 15\n", + "bc = ['reflect', 'periodic', 'periodic']\n", + "loading_params = {'seed': None}\n", + "\n", + "particles = Particles6D(Np=Np, \n", + " bc=bc, \n", + " domain=domain,\n", + " loading_params=loading_params)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "be9cdacd", + "metadata": {}, + "outputs": [], + "source": [ + "particles.draw_markers()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ee8e446e", + "metadata": {}, + "outputs": [], + "source": [ + "# positions on the physical domain Omega\n", + "pushed_pos = domain(particles.positions).T\n", + "pushed_pos" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "226f69fa", + "metadata": {}, + "outputs": [], + "source": [ + "fig = plt.figure() \n", + "ax = fig.gca()\n", + "\n", + "for n, pos in enumerate(pushed_pos):\n", + " ax.scatter(pos[0], pos[1], c=colors[n % 4])\n", + " ax.arrow(pos[0], pos[1], particles.velocities[n, 0], particles.velocities[n, 1], color=colors[n % 4], head_width=.2)\n", + "\n", + "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "\n", + "ax.add_patch(circle1)\n", + "ax.set_aspect('equal')\n", + "ax.set_xlabel('x')\n", + "ax.set_ylabel('y')\n", + "ax.set_title('Initial conditions');" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4f56e79b", + "metadata": {}, + "outputs": [], + "source": [ + "# pass simulation parameters to Propagator class\n", + "PushEta.domain = domain" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1f6d91c1", + "metadata": {}, + "outputs": [], + "source": [ + "# instantiate Propagator object\n", + "prop_eta = PushEta(particles)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2fc52d09", + "metadata": {}, + "outputs": [], + "source": [ + "# time stepping\n", + "Tend = 10. \n", + "dt = .2\n", + "Nt = int(Tend / dt)\n", + "\n", + "pos = np.zeros((Nt + 1, Np, 3), dtype=float)\n", + "alpha = np.ones(Nt + 1, dtype=float)\n", + "\n", + "pos[0] = pushed_pos\n", + "\n", + "time = 0.\n", + "n = 0\n", + "while time < (Tend -dt):\n", + " time += dt\n", + " n += 1\n", + " \n", + " # advance in time\n", + " prop_eta(dt)\n", + " \n", + " # positions on the physical domain Omega\n", + " pos[n] = domain(particles.positions).T\n", + " \n", + " # scaling for plotting\n", + " alpha[n] = (Tend - time)/Tend" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "895cc320", + "metadata": {}, + "outputs": [], + "source": [ + "# make scatter plot for each particle in xy-plane\n", + "for i in range(Np):\n", + " ax.scatter(pos[:, i, 0], pos[:, i, 1], c=colors[i % 4], alpha=alpha)\n", + "\n", + "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "\n", + "ax.add_patch(circle1)\n", + "ax.set_aspect('equal')\n", + "ax.set_xlabel('x')\n", + "ax.set_ylabel('y')\n", + "ax.set_title(f'{math.ceil(Tend/dt)} time steps (full color at t=0)');\n", + "\n", + "fig" + ] + }, + { + "cell_type": "markdown", + "id": "bbc72273", + "metadata": {}, + "source": [ + "## Particles in a cylinder with a magnetic field\n", + "\n", + "Let $\\Omega \\subset \\mathbb R^3$ be a cylinder as before. Now, we search for trajectories $(\\mathbf x_p, \\mathbf v_p): [0,T] \\to \\Omega \\times \\mathbb R^3$, $p = 0, \\ldots, N-1$ that satisfy\n", + "\n", + "$$\n", + "\\begin{align}\n", + " \\dot{\\mathbf x}_p &= \\mathbf v_p\\,,\\qquad && \\mathbf x_p(0) = \\mathbf x_{p0}\\,,\n", + " \\\\[2mm]\n", + " \\dot{\\mathbf v}_p &= \\mathbf v_p \\times \\mathbf B_0(\\mathbf x_p) \\qquad && \\mathbf v_p(0) = \\mathbf v_{p0}\\,,\n", + " \\end{align}\n", + "$$\n", + "\n", + "where $\\mathbf B_0$ is a given magnetic field from an [MHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils.html#mhd-equilibria). \n", + "In addition to the Propagator [PushEta](https://struphy.pages.mpcdf.de/struphy/sections/subsections/propagators_markers.html#struphy.propagators.propagators_markers.PushEta) for the position update, we shall use [PushVxB](https://struphy.pages.mpcdf.de/struphy/sections/subsections/propagators_markers.html#struphy.propagators.propagators_markers.PushVxB) for the velocity update." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7ae9ae17", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.geometry.domains import HollowCylinder\n", + "\n", + "a1 = 0.\n", + "a2 = 5.\n", + "Lz = 1.\n", + "domain = HollowCylinder(a1=a1, a2=a2, Lz=Lz)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "336bc5bd", + "metadata": {}, + "outputs": [], + "source": [ + "# instantiate Particle object\n", + "Np = 20\n", + "bc = ['remove', 'periodic', 'periodic']\n", + "loading_params = {'seed': None}\n", + "\n", + "particles = Particles6D(Np=Np, \n", + " bc=bc, \n", + " loading_params=loading_params)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "64860359", + "metadata": {}, + "outputs": [], + "source": [ + "particles.draw_markers()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ed244a1e", + "metadata": {}, + "outputs": [], + "source": [ + "# positions on the physical domain Omega\n", + "pushed_pos = domain(particles.positions).T\n", + "pushed_pos" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3b046601", + "metadata": {}, + "outputs": [], + "source": [ + "fig = plt.figure() \n", + "ax = fig.gca()\n", + "\n", + "for n, pos in enumerate(pushed_pos):\n", + " ax.scatter(pos[0], pos[1], c=colors[n % 4])\n", + " ax.arrow(pos[0], pos[1], particles.velocities[n, 0], particles.velocities[n, 1], color=colors[n % 4], head_width=.2)\n", + "\n", + "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "\n", + "ax.add_patch(circle1)\n", + "ax.set_aspect('equal')\n", + "ax.set_xlabel('x')\n", + "ax.set_ylabel('y')\n", + "ax.set_title('Initial conditions');" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "739f3b41", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.propagators.propagators_markers import PushVxB\n", + "\n", + "# default parameters of Propagator\n", + "opts_vxB = PushVxB.options(default=True)\n", + "print(opts_vxB)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ac87beba", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.fields_background.equils import HomogenSlab\n", + "\n", + "B0x = 0.\n", + "B0y = 0.\n", + "B0z = 1.\n", + "equil = HomogenSlab(B0x=B0x, B0y=B0y, B0z=B0z)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c0a30f29", + "metadata": {}, + "outputs": [], + "source": [ + "# set domain for Cartesian MHD equilibrium\n", + "equil.domain = domain" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "93970eb6", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.fields_background.projected_equils import ProjectedMHDequilibrium\n", + "from struphy.feec.psydac_derham import Derham\n", + "\n", + "# instantiate Derham object\n", + "Nel = [16, 16, 32]\n", + "p = [1, 1, 3]\n", + "spl_kind = [False, True, True]\n", + "derham = Derham(Nel=Nel, p=p, spl_kind=spl_kind)\n", + "\n", + "# instantiate a projected MHD equilibrium object\n", + "proj_equil = ProjectedMHDequilibrium(equil, derham)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4b754e07", + "metadata": {}, + "outputs": [], + "source": [ + "# pass simulation parameters to Propagator class\n", + "PushEta.domain = domain\n", + "PushVxB.domain = domain\n", + "PushVxB.derham = derham" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b7f71e9b", + "metadata": {}, + "outputs": [], + "source": [ + "# instantiate Propagator object\n", + "prop_eta = PushEta(particles)\n", + "prop_vxB = PushVxB(particles, b2=proj_equil.b2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5294ebca", + "metadata": {}, + "outputs": [], + "source": [ + "# time stepping\n", + "Tend = 10. - 1e-6\n", + "dt = .2\n", + "Nt = int(Tend / dt)\n", + "\n", + "pos = []\n", + "alpha = np.ones(Nt + 1, dtype=float)\n", + "\n", + "marker_col = {}\n", + "for marker in particles.markers_wo_holes:\n", + " m_id = int(marker[-1])\n", + " marker_col[m_id] = colors[int(m_id) % 4]\n", + "ids_wo_holes = []\n", + "\n", + "time = 0.\n", + "n = 0\n", + "while time < (Tend - dt):\n", + " time += dt\n", + " n += 1\n", + " \n", + " # advance in time\n", + " prop_vxB(dt/2)\n", + " prop_eta(dt)\n", + " prop_vxB(dt/2)\n", + " \n", + " # positions on the physical domain Omega (can change shape when particles are lost)\n", + " pos += [domain(particles.positions).T]\n", + "\n", + " # id's of non-holes\n", + " ids_wo_holes += [np.int64(particles.markers_wo_holes[:, -1])]\n", + " \n", + " # scaling for plotting\n", + " alpha[n] = (Tend - time)/Tend" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cb2ec764", + "metadata": {}, + "outputs": [], + "source": [ + "# make scatter plot for each particle in xy-plane\n", + "for po, ids, alph in zip(pos, ids_wo_holes, alpha):\n", + " cs = []\n", + " for ii in ids:\n", + " cs += [marker_col[ii]]\n", + " ax.scatter(po[:, 0], po[:, 1], c=cs, alpha=alph)\n", + "\n", + "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "\n", + "ax.add_patch(circle1)\n", + "ax.set_aspect('equal')\n", + "ax.set_xlabel('x')\n", + "ax.set_ylabel('y')\n", + "ax.set_title(f'{math.ceil(Tend/dt)} time steps (full color at t=0)');\n", + "\n", + "fig" + ] + }, + { + "cell_type": "markdown", + "id": "f10f537b", + "metadata": {}, + "source": [ + "## Particles in a Tokamak equilibrium\n", + "\n", + "We use the same Propagators from the previous example but load a more complicated [MHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils.html#mhd-equilibria), namely from an ASDEX-Upgrade equilibrium stored in an EQDSK file.\n", + "\n", + "Let us instatiate an [EQDSKequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils_sub.html#struphy.fields_background.mhd_equil.equils.EQDSKequilibrium) with many of its default parameters, except for the density:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f1c0a4f2", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.fields_background.equils import EQDSKequilibrium\n", + "\n", + "n1 = 0.\n", + "n2 = 0.\n", + "na = 1.\n", + "equil = EQDSKequilibrium(n1=n1, n2=n2, na=na)\n", + "equil.params" + ] + }, + { + "cell_type": "markdown", + "id": "60c488fa", + "metadata": {}, + "source": [ + "Since [EQDSKequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils_sub.html#struphy.fields_background.mhd_equil.equils.EQDSKequilibrium) is an [AxisymmMHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils_sub.html#struphy.fields_background.mhd_equil.base.AxisymmMHDequilibrium), which in turn is a [CartesianMHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils_sub.html#struphy.fields_background.mhd_equil.base.CartesianMHDequilibrium), we are free to choose any mapping for the simulation (e.g. a Cuboid for Cartesian coordinates). In order to be conforming to the boundary of the equilibrium, we shall choose the [Tokamak](https://struphy.pages.mpcdf.de/struphy/sections/subsections/domains_avail.html#struphy.geometry.domains.Tokamak) mapping:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f3e56de4", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.geometry.domains import Tokamak\n", + "\n", + "Nel = (28, 72)\n", + "p = (3, 3)\n", + "psi_power = 0.6\n", + "psi_shifts = (1e-6, 1.)\n", + "domain = Tokamak(equilibrium=equil, \n", + " Nel=Nel,\n", + " p=p,\n", + " psi_power=psi_power,\n", + " psi_shifts=psi_shifts)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1cb1f4cc", + "metadata": {}, + "outputs": [], + "source": [ + "equil.domain = domain" + ] + }, + { + "cell_type": "markdown", + "id": "c3135576", + "metadata": {}, + "source": [ + "The [Tokamak](https://struphy.pages.mpcdf.de/struphy/sections/subsections/domains_avail.html#struphy.geometry.domains.Tokamak) domain is a [PoloidalSplineTorus](https://struphy.pages.mpcdf.de/struphy/sections/subsections/domains_base.html#struphy.geometry.base.PoloidalSplineTorus), hence\n", + "\n", + "$$\n", + " \\begin{align*}\n", + " x &= R \\cos(\\phi)\\,,\n", + " \\\\\n", + " y &= -R \\sin(\\phi)\\,,\n", + " \\\\\n", + " z &= Z\\,,\n", + " \\end{align*}\n", + "$$\n", + "\n", + "between Cartesian $(x, y, z)$- and Tokamak $(R, Z, \\phi)$-coordinates holds, where $(R, Z)$ spans a poloidal plane. Moreover, the Tokamak coordinates are related to general torus coordinates $(r, \\theta, \\phi)$ via a polar mapping in the poloidal plane:\n", + "\n", + "$$\n", + " \\begin{align*}\n", + " R &= R_0 + r \\cos(\\theta)\\,,\n", + " \\\\\n", + " Z &= r \\sin(\\theta)\\,,\n", + " \\\\\n", + " \\phi &= \\phi\\,.\n", + " \\end{align*}\n", + "$$\n", + "\n", + "The torus coordinates are related to Struphy logical coordinates $\\boldsymbol \\eta = (\\eta_1, \\eta_2, \\eta_3) \\in [0, 1]^3$ as \n", + "\n", + "$$\n", + " \\begin{align*}\n", + " r &= a_1 + (a_2 - a_1) \\eta_1\\,,\n", + " \\\\\n", + " \\theta &= 2\\pi \\eta_2\\,,\n", + " \\\\\n", + " \\phi &= 2\\pi \\eta_3\\,,\n", + " \\end{align*}\n", + "$$\n", + "\n", + "where $a_2 > a_1 \\geq 0$ are boundaries in the radial $r$-direction.\n", + "This can be seen for instance in the [HollowTorus](https://struphy.pages.mpcdf.de/struphy/sections/subsections/domains_avail.html#struphy.geometry.domains.HollowTorus) mapping (more complicated angle parametrizations $\\theta(\\eta_1, \\eta_2)$ are also available, but not discussed here).\n", + "\n", + "Let us plot the equilibrium magnetic field strength:\n", + "\n", + "1. in the poloidal plane at $\\phi = 0$\n", + "2. in the top view at $z = 0$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cf741b53", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "# logical grid on the unit cube\n", + "e1 = np.linspace(0., 1., 101)\n", + "e2 = np.linspace(0., 1., 101)\n", + "e3 = np.linspace(0., 1., 101)\n", + "\n", + "# move away from the singular point r = 0\n", + "e1[0] += 1e-5" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "85f37793", + "metadata": {}, + "outputs": [], + "source": [ + "# logical coordinates of the poloidal plane at phi = 0\n", + "eta_poloidal = (e1, e2, 0.)\n", + "# logical coordinates of the top view at theta = 0\n", + "eta_topview_1 = (e1, 0., e3)\n", + "# logical coordinates of the top view at theta = pi\n", + "eta_topview_2 = (e1, .5, e3)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c36dc211", + "metadata": {}, + "outputs": [], + "source": [ + "# Cartesian coordinates (squeezed)\n", + "x_pol, y_pol, z_pol = domain(*eta_poloidal, squeeze_out=True)\n", + "x_top1, y_top1, z_top1 = domain(*eta_topview_1, squeeze_out=True)\n", + "x_top2, y_top2, z_top2 = domain(*eta_topview_2, squeeze_out=True)\n", + "\n", + "print(f'{x_pol.shape = }')\n", + "print(f'{x_top1.shape = }')\n", + "print(f'{x_top2.shape = }')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "078039c9", + "metadata": {}, + "outputs": [], + "source": [ + "# generate two axes\n", + "fig, axs = plt.subplots(2, 1, figsize=(8, 16))\n", + "ax = axs[0]\n", + "ax_top = axs[1]\n", + "\n", + "# min/max of field strength\n", + "Bmax = np.max(equil.absB0(*eta_topview_2, squeeze_out=True))\n", + "Bmin = np.min(equil.absB0(*eta_topview_1, squeeze_out=True))\n", + "levels = np.linspace(Bmin, Bmax, 51)\n", + "\n", + "# absolute magnetic field at phi = 0\n", + "im = ax.contourf(x_pol, z_pol, equil.absB0(*eta_poloidal, squeeze_out=True), levels=levels)\n", + "\n", + "# absolute magnetic field at Z = 0\n", + "im_top = ax_top.contourf(x_top1, y_top1, equil.absB0(*eta_topview_1, squeeze_out=True), levels=levels)\n", + "ax_top.contourf(x_top2, y_top2, equil.absB0(*eta_topview_2, squeeze_out=True), levels=levels)\n", + "\n", + "# last closed flux surface, poloidal\n", + "ax.plot(x_pol[-1], z_pol[-1], color='k')\n", + "\n", + "# last closed flux surface, toroidal\n", + "ax_top.plot(x_top1[-1], y_top1[-1], color='k')\n", + "ax_top.plot(x_top2[-1], y_top2[-1], color='k')\n", + "\n", + "# limiter, poloidal\n", + "ax.plot(equil.limiter_pts_R, equil.limiter_pts_Z, 'tab:orange')\n", + "ax.axis('equal')\n", + "ax.set_xlabel('R')\n", + "ax.set_ylabel('Z')\n", + "ax.set_title('abs(B) at $\\phi=0$')\n", + "fig.colorbar(im);\n", + "\n", + "# limiter, toroidal\n", + "limiter_Rmax = np.max(equil.limiter_pts_R)\n", + "limiter_Rmin = np.min(equil.limiter_pts_R)\n", + "\n", + "thetas = 2*np.pi*e2\n", + "limiter_x_max = limiter_Rmax * np.cos(thetas)\n", + "limiter_y_max = - limiter_Rmax * np.sin(thetas)\n", + "limiter_x_min = limiter_Rmin * np.cos(thetas)\n", + "limiter_y_min = - limiter_Rmin * np.sin(thetas)\n", + "\n", + "ax_top.plot(limiter_x_max, limiter_y_max, 'tab:orange')\n", + "ax_top.plot(limiter_x_min, limiter_y_min, 'tab:orange')\n", + "ax_top.axis('equal')\n", + "ax_top.set_xlabel('x')\n", + "ax_top.set_ylabel('y')\n", + "ax_top.set_title('abs(B) at $Z=0$')\n", + "fig.colorbar(im_top);" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "568c0ba6", + "metadata": {}, + "outputs": [], + "source": [ + "# instantiate Particle object\n", + "Np = 4\n", + "bc = ['remove', 'periodic', 'periodic']\n", + "bufsize = 2.\n", + "\n", + "initial = [[.501, 0.001, 0.001, 0., 0.0450, -0.04], # co-passing particle\n", + " [.511, 0.001, 0.001, 0., -0.0450, -0.04], # counter passing particle\n", + " [.521, 0.001, 0.001, 0., 0.0105, -0.04], # co-trapped particle\n", + " [.531, 0.001, 0.001, 0., -0.0155, -0.04]]\n", + "\n", + "loading_params = {'seed': 1608,\n", + " 'initial' : initial}\n", + "\n", + "particles = Particles6D(Np=Np, \n", + " bc=bc, \n", + " loading_params=loading_params,\n", + " bufsize=bufsize)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "43127584", + "metadata": {}, + "outputs": [], + "source": [ + "particles.draw_markers()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4b2b96f2", + "metadata": {}, + "outputs": [], + "source": [ + "# positions on the physical domain Omega (x, y, z)\n", + "pushed_pos = domain(particles.positions).T\n", + "pushed_pos" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3f7851f8", + "metadata": {}, + "outputs": [], + "source": [ + "# compute R-coordinate\n", + "pushed_r = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "72d79e1c", + "metadata": {}, + "outputs": [], + "source": [ + "labels = ['co-passing',\n", + " 'counter passing',\n", + " 'co_trapped',\n", + " 'counter-trapped']\n", + "\n", + "for n, (r, pos) in enumerate(zip(pushed_r, pushed_pos)):\n", + " # poloidal \n", + " ax.scatter(r, pos[2], c=colors[n % 4], label=labels[n])\n", + " ax.arrow(r, pos[2], particles.velocities[n, 0], particles.velocities[n, 2]*10, color=colors[n % 4], head_width=.05)\n", + " # topview\n", + " ax_top.scatter(pos[0], pos[1], c=colors[n % 4], label=labels[n])\n", + " ax_top.arrow(pos[0], pos[1], particles.velocities[n, 0], particles.velocities[n, 1]*10, color=colors[n % 4], head_width=.05)\n", + "\n", + "ax.set_xlabel('R')\n", + "ax.set_ylabel('Z')\n", + "ax.set_title('Initial conditions')\n", + "ax.legend();\n", + "\n", + "ax_top.set_xlabel('x')\n", + "ax_top.set_ylabel('y')\n", + "ax_top.set_title('Initial conditions')\n", + "ax_top.legend();\n", + "\n", + "fig " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4579b2af", + "metadata": {}, + "outputs": [], + "source": [ + "# instantiate Derham object\n", + "Nel = [32, 72, 1]\n", + "p = [3, 3, 1]\n", + "spl_kind = [False, True, True]\n", + "derham = Derham(Nel=Nel, p=p, spl_kind=spl_kind)\n", + "\n", + "# instantiate a projected MHD equilibrium object\n", + "proj_equil = ProjectedMHDequilibrium(equil, derham)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5b8f74f6", + "metadata": {}, + "outputs": [], + "source": [ + "# pass simulation parameters to Propagator class\n", + "PushEta.domain = domain\n", + "PushVxB.domain = domain\n", + "PushVxB.derham = derham" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "81b357b2", + "metadata": {}, + "outputs": [], + "source": [ + "# instantiate Propagator object\n", + "prop_eta = PushEta(particles)\n", + "prop_vxB = PushVxB(particles, b2=proj_equil.b2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "430a160b", + "metadata": {}, + "outputs": [], + "source": [ + "# time stepping\n", + "Tend = 3000. - 1e-6\n", + "dt = .2\n", + "Nt = int(Tend / dt)\n", + "\n", + "pos = np.zeros((Nt + 2, Np, 3), dtype=float)\n", + "r = np.zeros((Nt + 2, Np), dtype=float)\n", + "\n", + "pos[0] = pushed_pos\n", + "r[0] = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)\n", + "\n", + "time = 0.\n", + "n = 0\n", + "while time < Tend:\n", + " time += dt\n", + " n += 1\n", + " \n", + " # advance in time\n", + " prop_vxB(dt/2)\n", + " prop_eta(dt)\n", + " prop_vxB(dt/2)\n", + " \n", + " # positions on the physical domain Omega\n", + " pushed_pos = domain(particles.positions).T\n", + " \n", + " # compute R-ccordinate\n", + " pos[n] = pushed_pos\n", + " r[n] = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ac6b4341", + "metadata": {}, + "outputs": [], + "source": [ + "# make scatter plot for each particle \n", + "for i in range(pos.shape[1]):\n", + " # poloidal \n", + " ax.scatter(r[:, i], pos[:, i, 2], c=colors[i % 4], s=1)\n", + " # top view\n", + " ax_top.scatter(pos[:, i, 0], pos[:, i, 1], c=colors[i % 4], s=1)\n", + "\n", + "ax.set_title(f'{math.ceil(Tend/dt)} time steps')\n", + "ax_top.set_title(f'{math.ceil(Tend/dt)} time steps');\n", + "\n", + "fig" + ] + }, + { + "cell_type": "markdown", + "id": "602f4408", + "metadata": {}, + "source": [ + "## Guiding-centers in a Tokamak equilibrium\n", + "\n", + "We now use the Propagators []() and []() in ASDEX-Upgrade equilibrium from the previous example.\n", + "\n", + "For this we need to instantiate the [Particles5D]() class." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "64caf9e5", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.pic.particles import Particles5D\n", + "\n", + "# instantiate Particle object\n", + "Np = 4\n", + "bc = ['remove', 'periodic', 'periodic']\n", + "bufsize = 2.\n", + "\n", + "initial = [[.501, 0.001, 0.001, -1.935 , 1.72], # co-passing particle\n", + " [.501, 0.001, 0.001, 1.935 , 1.72], # couner-passing particle\n", + " [.501, 0.001, 0.001, -0.6665, 1.72], # co-trapped particle\n", + " [.501, 0.001, 0.001, 0.4515, 1.72]] # counter-trapped particle\n", + "\n", + "loading_params = {'seed': 1608,\n", + " 'initial' : initial}\n", + "\n", + "particles = Particles5D(proj_equil,\n", + " Np=Np, \n", + " bc=bc, \n", + " loading_params=loading_params,\n", + " bufsize=bufsize)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "847979e0", + "metadata": {}, + "outputs": [], + "source": [ + "particles.draw_markers()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6016bfe0", + "metadata": {}, + "outputs": [], + "source": [ + "# positions on the physical domain Omega (x, y, z)\n", + "pushed_pos = domain(particles.positions).T\n", + "pushed_pos" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a5548daa", + "metadata": {}, + "outputs": [], + "source": [ + "# compute R-coordinate\n", + "pushed_r = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "53fa4d3b", + "metadata": {}, + "outputs": [], + "source": [ + "particles.velocities" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "041005d2", + "metadata": {}, + "outputs": [], + "source": [ + "# generate two axes\n", + "fig, axs = plt.subplots(2, 1, figsize=(8, 16))\n", + "ax = axs[0]\n", + "ax_top = axs[1]\n", + "\n", + "# min/max of field strength\n", + "Bmax = np.max(equil.absB0(*eta_topview_2, squeeze_out=True))\n", + "Bmin = np.min(equil.absB0(*eta_topview_1, squeeze_out=True))\n", + "levels = np.linspace(Bmin, Bmax, 51)\n", + "\n", + "# absolute magnetic field at phi = 0\n", + "im = ax.contourf(x_pol, z_pol, equil.absB0(*eta_poloidal, squeeze_out=True), levels=levels)\n", + "\n", + "# absolute magnetic field at Z = 0\n", + "im_top = ax_top.contourf(x_top1, y_top1, equil.absB0(*eta_topview_1, squeeze_out=True), levels=levels)\n", + "ax_top.contourf(x_top2, y_top2, equil.absB0(*eta_topview_2, squeeze_out=True), levels=levels)\n", + "\n", + "# last closed flux surface, poloidal\n", + "ax.plot(x_pol[-1], z_pol[-1], color='k')\n", + "\n", + "# last closed flux surface, toroidal\n", + "ax_top.plot(x_top1[-1], y_top1[-1], color='k')\n", + "ax_top.plot(x_top2[-1], y_top2[-1], color='k')\n", + "\n", + "# limiter, poloidal\n", + "ax.plot(equil.limiter_pts_R, equil.limiter_pts_Z, 'tab:orange')\n", + "ax.axis('equal')\n", + "ax.set_xlabel('R')\n", + "ax.set_ylabel('Z')\n", + "ax.set_title('abs(B) at $\\phi=0$')\n", + "fig.colorbar(im);\n", + "\n", + "# limiter, toroidal\n", + "limiter_Rmax = np.max(equil.limiter_pts_R)\n", + "limiter_Rmin = np.min(equil.limiter_pts_R)\n", + "\n", + "thetas = 2*np.pi*e2\n", + "limiter_x_max = limiter_Rmax * np.cos(thetas)\n", + "limiter_y_max = - limiter_Rmax * np.sin(thetas)\n", + "limiter_x_min = limiter_Rmin * np.cos(thetas)\n", + "limiter_y_min = - limiter_Rmin * np.sin(thetas)\n", + "\n", + "ax_top.plot(limiter_x_max, limiter_y_max, 'tab:orange')\n", + "ax_top.plot(limiter_x_min, limiter_y_min, 'tab:orange')\n", + "ax_top.axis('equal')\n", + "ax_top.set_xlabel('x')\n", + "ax_top.set_ylabel('y')\n", + "ax_top.set_title('abs(B) at $Z=0$')\n", + "fig.colorbar(im_top);" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2ab72f71", + "metadata": {}, + "outputs": [], + "source": [ + "labels = ['co-passing',\n", + " 'counter passing',\n", + " 'co_trapped',\n", + " 'counter-trapped']\n", + "\n", + "for n, (r, pos) in enumerate(zip(pushed_r, pushed_pos)):\n", + " # poloidal \n", + " ax.scatter(r, pos[2], c=colors[n % 4], label=labels[n])\n", + " # topview\n", + " ax_top.scatter(pos[0], pos[1], c=colors[n % 4], label=labels[n])\n", + " ax_top.arrow(pos[0], pos[1], 0., particles.velocities[n, 0]/5, color=colors[n % 4], head_width=.05)\n", + "\n", + "ax.set_xlabel('R')\n", + "ax.set_ylabel('Z')\n", + "ax.set_title('Initial conditions')\n", + "ax.legend();\n", + "\n", + "ax_top.set_xlabel('x')\n", + "ax_top.set_ylabel('y')\n", + "ax_top.set_title('Initial conditions')\n", + "ax_top.legend();\n", + "\n", + "fig " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "82bb30be", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.propagators.propagators_markers import PushGuidingCenterBxEstar, PushGuidingCenterParallel\n", + "\n", + "# default parameters of Propagator\n", + "opts_BxE = PushGuidingCenterBxEstar.options(default=True)\n", + "print(opts_BxE)\n", + "\n", + "opts_para = PushGuidingCenterParallel.options(default=True)\n", + "print(opts_para)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9bd84960", + "metadata": {}, + "outputs": [], + "source": [ + "# pass simulation parameters to Propagator class\n", + "PushGuidingCenterBxEstar.domain = domain\n", + "PushGuidingCenterParallel.domain = domain\n", + "\n", + "PushGuidingCenterBxEstar.derham = derham\n", + "PushGuidingCenterParallel.derham = derham\n", + "\n", + "PushGuidingCenterBxEstar.projected_equil = proj_equil\n", + "PushGuidingCenterParallel.projected_equil = proj_equil" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "56b19093", + "metadata": {}, + "outputs": [], + "source": [ + "# natural constants\n", + "mH = 1.67262192369e-27 # proton mass (kg)\n", + "e = 1.602176634e-19 # elementary charge (C)\n", + "mu0 = 1.25663706212e-6 # magnetic constant (N/A^2)\n", + "\n", + "# epsilon equation parameter\n", + "A = 1. # mass number in units of proton mass\n", + "Z = 1 # signed charge number in units of elementary charge\n", + "unit_x = 1. # length scale unit in m\n", + "unit_B = 1. # magnetic field unit in T\n", + "unit_n = 1e20 # number density unit in m^(-3)\n", + "unit_v = unit_B / np.sqrt(unit_n * A * mH * mu0) # Alfvén velocity unit\n", + "unit_t = unit_x / unit_v # time unit\n", + "\n", + "# cyclotron frequency and epsilon parameter\n", + "om_c = Z*e * unit_B / (A*mH)\n", + "epsilon = 1./(om_c * unit_t)\n", + "\n", + "print(f'{unit_x = }')\n", + "print(f'{unit_B = }')\n", + "print(f'{unit_n = }')\n", + "print(f'{unit_v = }')\n", + "print(f'{unit_t = }')\n", + "print(f'{epsilon = }')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8cd49aa4", + "metadata": {}, + "outputs": [], + "source": [ + "# instantiate Propagator object\n", + "opts_BxE['algo']['tol'] = 1e-5\n", + "opts_para['algo']['tol'] = 1e-5\n", + "prop_BxE = PushGuidingCenterBxEstar(particles, epsilon=epsilon, algo=opts_BxE['algo'])\n", + "prop_para = PushGuidingCenterParallel(particles, epsilon=epsilon, algo=opts_para['algo'])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18a1defc", + "metadata": {}, + "outputs": [], + "source": [ + "# time stepping\n", + "Tend = 100. - 1e-6\n", + "dt = .1\n", + "Nt = int(Tend / dt)\n", + "\n", + "pos = np.zeros((Nt + 2, Np, 3), dtype=float)\n", + "r = np.zeros((Nt + 2, Np), dtype=float)\n", + "\n", + "pos[0] = pushed_pos\n", + "r[0] = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)\n", + "\n", + "time = 0.\n", + "n = 0\n", + "while time < Tend:\n", + " time += dt\n", + " n += 1\n", + "\n", + " # advance in time\n", + " prop_BxE(dt/2)\n", + " prop_para(dt)\n", + " prop_BxE(dt/2)\n", + " \n", + " # positions on the physical domain Omega\n", + " pushed_pos = domain(particles.positions).T\n", + " \n", + " # compute R-coordinate\n", + " pos[n] = pushed_pos\n", + " r[n] = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bff57f86", + "metadata": {}, + "outputs": [], + "source": [ + "# make scatter plot for each particle in xy-plane\n", + "for i in range(pos.shape[1]):\n", + " # poloidal \n", + " ax.scatter(r[:, i], pos[:, i, 2], c=colors[i % 4], s=1)\n", + " # top view\n", + " ax_top.scatter(pos[:, i, 0], pos[:, i, 1], c=colors[i % 4], s=1)\n", + "\n", + "ax.set_title(f'{math.ceil(Tend/dt)} time steps')\n", + "ax_top.set_title(f'{math.ceil(Tend/dt)} time steps');\n", + "\n", + "fig" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "env", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index a0abe7bae..94af30877 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -13,6 +13,7 @@ # derham PolarRegularity = Literal[-1, 1] +OptsFEECSpace = Literal["H1", "Hcurl", "Hdiv", "L2", "H1vec"] OptsVecSpace = Literal["Hcurl", "Hdiv", "H1vec"] # fields background @@ -28,6 +29,7 @@ OptsMassPrecond = Literal["MassMatrixPreconditioner", None] # markers +OptsPICSpace = Literal["Particles6D", "DeltaFParticles6D", "Particles5D", "Particles3D"] OptsMarkerBC = Literal["periodic", "reflect"] OptsRecontructBC = Literal["periodic", "mirror", "fixed"] OptsLoading = Literal["pseudo_random", 'sobol_standard', 'sobol_antithetic', 'external', 'restart', "tesselation"] @@ -60,7 +62,8 @@ def __post_init__(self): check_option(self.split_algo, SplitAlgos) -class Units: +@dataclass +class BaseUnits: """ Base units are passed to __init__, other units derive from these. @@ -79,17 +82,25 @@ class Units: Unit of internal energy in keV. Only in effect if the velocity scale is set to 'thermal'. """ + x: float = 1.0 + B: float = 1.0 + n: float = 1.0 + kBT: float = None + + +class Units: + """ + Colllects base units and derives other units from these. + """ - def __init__(self, - x: float = 1.0, - B: float = 1.0, - n: float = 1.0, - kBT: float = None,): + def __init__(self, base: BaseUnits = None): + if base is None: + base = BaseUnits() - self._x = x - self._B = B - self._n = n * 1e20 - self._kBT = kBT + self._x = base.x + self._B = base.B + self._n = base.n * 1e20 + self._kBT = base.kBT @property def x(self): @@ -299,7 +310,7 @@ class EnvironmentOptions: def __post_init__(self): self.path_out: str = os.path.join(self.out_folders, self.sim_folder) - def print(self): + def __repr__(self): for k, v in self.__dict__.items(): print(f"{k}:".ljust(20), v) diff --git a/src/struphy/main.py b/src/struphy/main.py index 2fcc74191..f52264175 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -21,7 +21,7 @@ from struphy.pic.base import Particles from struphy.models.species import Species from struphy.models.variables import FEECVariable -from struphy.io.options import Units, Time, EnvironmentOptions +from struphy.io.options import BaseUnits, Units, Time, EnvironmentOptions from struphy.io.setup import import_parameters_py from struphy.geometry.base import Domain from struphy.geometry import domains @@ -43,7 +43,7 @@ def run( *, params_path: str = None, env: EnvironmentOptions = EnvironmentOptions(), - units: Units = Units(), + base_units: BaseUnits = BaseUnits(), time_opts: Time = Time(), domain: Domain = domains.Cuboid(), equil: FluidEquilibrium = HomogenSlab(), @@ -106,6 +106,9 @@ def run( restart=restart, verbose=verbose,) + # add derived units + units = Units(base_units) + # save parameter file if rank == 0: # save python param file @@ -628,6 +631,7 @@ def load_data(path: str) -> SimData: path_pproc = os.path.join(path, "post_processing") assert os.path.exists(path_pproc), f"Path {path_pproc} does not exist, run 'pproc' first?" print("\n*** Loading post-processed simulation data:") + print(f"{path = }") simdata = SimData(path) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index d1a228f68..7876c2380 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -26,7 +26,7 @@ from struphy.utils.utils import dict_to_yaml, read_state from struphy.models.species import Species, FieldSpecies, FluidSpecies, KineticSpecies, DiagnosticSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable -from struphy.io.options import Units, Time, DerhamOptions +from struphy.io.options import BaseUnits, Units, Time, DerhamOptions from struphy.topology.grids import TensorProductGrid from struphy.geometry.base import Domain from struphy.geometry.domains import Cuboid @@ -1283,7 +1283,7 @@ def generate_default_parameter_file( print("exiting ...") return - file.write("from struphy.io.options import EnvironmentOptions, Units, Time\n") + file.write("from struphy.io.options import EnvironmentOptions, BaseUnits, Time\n") file.write("from struphy.geometry import domains\n") file.write("from struphy.fields_background import equils\n") @@ -1349,7 +1349,7 @@ def generate_default_parameter_file( file.write("env = EnvironmentOptions()\n") file.write("\n# units\n") - file.write("units = Units()\n") + file.write("base_units = BaseUnits()\n") file.write("\n# time stepping\n") file.write("time_opts = Time()\n") @@ -1402,7 +1402,7 @@ def generate_default_parameter_file( file.write(" main.run(model, \n\ params_path=__file__, \n\ env=env, \n\ - units=units, \n\ + base_units=base_units, \n\ time_opts=time_opts, \n\ domain=domain, \n\ equil=equil, \n\ diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 0a2646c95..f5829debe 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -45,15 +45,17 @@ class LinearMHD(StruphyModel): """ ## species - @dataclass class EMFields(FieldSpecies): - b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - @dataclass class MHD(FluidSpecies): - density: FEECVariable = FEECVariable(name="density", space="L2") - velocity: FEECVariable = FEECVariable(name="velocity", space="Hdiv") - pressure: FEECVariable = FEECVariable(name="pressure", space="L2") + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="Hdiv") + self.pressure = FEECVariable(space="L2") + self.init_variables() ## propagators @@ -76,12 +78,12 @@ def __init__(self): self.propagators = self.Propagators() # 3. assign variables to propagators - self.propagators.shear_alf.set_variables( + self.propagators.shear_alf.assign_variables( u = self.mhd.velocity, b = self.em_fields.b_field, ) - self.propagators.mag_sonic.set_variables( + self.propagators.mag_sonic.assign_variables( n = self.mhd.density, u = self.mhd.velocity, p = self.mhd.pressure, diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index d8e8555ab..1b4e4711c 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -18,20 +18,21 @@ class Species(metaclass=ABCMeta): """Single species of a StruphyModel.""" + @abstractmethod + def __init__(self): + self.init_variables() + # set species attribute for each variable - def __post_init__(self): + def init_variables(self): + self._variables = {} for k, v in self.__dict__.items(): if isinstance(v, Variable): - # v._species = self.__class__.__name__ + v._name = k v._species = self + self._variables[k] = v @property - def variables(self): - if not hasattr(self, "_variables"): - self._variables = {} - for k, v in self.__dict__.items(): - if isinstance(v, Variable): - self._variables[k] = v + def variables(self) -> dict: return self._variables @property @@ -51,11 +52,11 @@ def set_phys_params(self, charge_number: int = 1, mass_number: int = 1): @property def equation_params(self) -> dict: - if not hasattr(self, "_equation_params"): - self.setup_equation_params() + # if not hasattr(self, "_equation_params"): + # self.setup_equation_params() return self._equation_params - def setup_equation_params(self, units: Units = None, verbose=False): + def setup_equation_params(self, units: Units, verbose=False): """Set the following equation parameters: * alpha = plasma-frequenca / cyclotron frequency @@ -64,9 +65,6 @@ def setup_equation_params(self, units: Units = None, verbose=False): """ Z = self.charge_number A = self.mass_number - - if units is None: - units = Units() con = ConstantsOfNature() diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index ee932ed81..db8f1d9ca 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -37,10 +37,11 @@ class Maxwell(StruphyModel): """ ## species - @dataclass class EMFields(FieldSpecies): - e_field: FEECVariable = FEECVariable(name="e_field", space="Hcurl") - b_field: FEECVariable = FEECVariable(name="b_field", space="Hdiv") + def __init__(self): + self.e_field = FEECVariable(space="Hcurl") + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() ## propagators @@ -61,7 +62,7 @@ def __init__(self): self.propagators = self.Propagators() # 3. assign variables to propagators - self.propagators.maxwell.set_variables( + self.propagators.maxwell.assign_variables( e = self.em_fields.e_field, b = self.em_fields.b_field, ) @@ -113,10 +114,12 @@ class Vlasov(StruphyModel): :ref:`Model info `: """ - - @dataclass + ## species + class KineticIons(KineticSpecies): - var: PICVariable = PICVariable(name="ions", space="Particles6D") + def __init__(self): + self.var = PICVariable(space="Particles6D") + self.init_variables() ## propagators @@ -138,11 +141,11 @@ def __init__(self): self.propagators = self.Propagators() # 3. assign variables to propagators - self.propagators.push_vxb.set_variables( + self.propagators.push_vxb.assign_variables( ions = self.kinetic_ions.var, ) - self.propagators.push_eta.set_variables( + self.propagators.push_eta.assign_variables( var = self.kinetic_ions.var, ) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index bec807bca..e5d39e93e 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -8,7 +8,7 @@ import numpy as np from struphy.feec.psydac_derham import Derham, SplineFunction -from struphy.io.options import FieldsBackground +from struphy.io.options import (FieldsBackground, OptsFEECSpace, OptsPICSpace, check_option,) from struphy.initial.perturbations import Perturbation from struphy.geometry.base import Domain from struphy.fields_background.base import FluidEquilibrium @@ -19,7 +19,7 @@ from struphy.utils.clone_config import CloneConfig if TYPE_CHECKING: - from struphy.models.species import Species, KineticSpecies + from struphy.models.species import Species, KineticSpecies, FieldSpecies, FluidSpecies class Variable(metaclass=ABCMeta): """Single variable (unknown) of a Species.""" @@ -57,6 +57,12 @@ def species(self) -> Species: if not hasattr(self, "_species"): self._species = None return self._species + + @property + def __name__(self): + if not hasattr(self, "_name"): + self._name = None + return self._name def add_background(self, background, verbose=True): """Type inference of added background done in sub class.""" @@ -87,15 +93,10 @@ def add_perturbation(self, perturbation: Perturbation, verbose=True): class FEECVariable(Variable): - def __init__(self, name: str = "a_feec_var", space: str = "H1"): - assert space in ("H1", "Hcurl", "Hdiv", "L2", "H1vec") - self._name = name + def __init__(self, space: OptsFEECSpace = "H1"): + check_option(space, OptsFEECSpace) self._space = space - @property - def __name__(self): - return self._name - @property def space(self): return self._space @@ -104,6 +105,12 @@ def space(self): def spline(self) -> SplineFunction: return self._spline + @property + def species(self) -> FieldSpecies | FluidSpecies: + if not hasattr(self, "_species"): + self._species = None + return self._species + def add_background(self, background: FieldsBackground, verbose=True): super().add_background(background, verbose=verbose) @@ -119,16 +126,11 @@ def allocate(self, derham: Derham, domain: Domain = None, equil: FluidEquilibriu class PICVariable(Variable): - def __init__(self, name: str = "a_pic_var", space: str = "Particles6D"): - assert space in ("Particles6D", "Particles5D", "Particles3D", "DeltaFParticles6D") - self._name = name + def __init__(self, space: OptsPICSpace = "Particles6D"): + check_option(space, OptsPICSpace) self._space = space self._kinetic_data = {} - @property - def __name__(self): - return self._name - @property def space(self): return self._space diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 1493f2f65..6f802f2fd 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -43,8 +43,8 @@ def __call__(self, dt): Time step size. """ - def set_variables(self, **vars): - """Define the variables to be updated by the propagator. + def assign_variables(self, **vars): + """Assign the variables to be updated by the propagator. Update variables in __dict__ and set self.variables with user-defined instance variables (allocated). """ diff --git a/doc/tutorials/tutorial_01_kinetic_particles.ipynb b/tutorials_old/tutorial_01_kinetic_particles.ipynb similarity index 100% rename from doc/tutorials/tutorial_01_kinetic_particles.ipynb rename to tutorials_old/tutorial_01_kinetic_particles.ipynb diff --git a/tutorials_old/tutorial_01_parameter_files.ipynb b/tutorials_old/tutorial_01_parameter_files.ipynb new file mode 100644 index 000000000..e803e3746 --- /dev/null +++ b/tutorials_old/tutorial_01_parameter_files.ipynb @@ -0,0 +1,378 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d34c79c5", + "metadata": {}, + "source": [ + "# Parameter files and `struphy.main`\n", + "\n", + "Struphy parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", + "For each `MODEL`, the default parameter file can be generated from the console via\n", + "\n", + "```\n", + "struphy params MODEL\n", + "```\n", + "\n", + "This will create a file `params_MODEL.py` in the current working directory. To run the model type\n", + "\n", + "```\n", + "python params_MODEL.py\n", + "```\n", + "\n", + "The user can modify the parameter file to launch a specific simulation.\n", + "As an example, let us discuss the parameter file of the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov) and run some simple examples. \n", + "The file can be generated from\n", + "\n", + "```\n", + "struphy params Vlasov\n", + "```\n", + "\n", + "To see its contents, open the file in your preferred editor or type\n", + "\n", + "```\n", + "cat params_Vlasov.py\n", + "```\n", + "\n", + "## Parameters part 1: Imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3ecab659", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.io.options import EnvironmentOptions, Units, Time\n", + "from struphy.geometry import domains\n", + "from struphy.fields_background import equils\n", + "from struphy.topology import grids\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.io.options import FieldsBackground\n", + "from struphy.initial import perturbations\n", + "from struphy.kinetic_background import maxwellians\n", + "from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters\n", + "from struphy import main\n", + "\n", + "# import model, set verbosity\n", + "from struphy.models.toy import Vlasov as Model\n", + "verbose = True" + ] + }, + { + "cell_type": "markdown", + "id": "5cf6d9c7", + "metadata": {}, + "source": [ + "All Struphy parameter files import the modules listed above, even though some of them might not be needed in a specific model. The last import imports the model itself, always under the alias `Model`.\n", + "\n", + "## Parameters part 2: Generic options\n", + "\n", + "The following lines refer to options that can be set for any model. These are:\n", + "\n", + "* Environment options (paths, saving, domain cloning)\n", + "* Model units\n", + "* Time options\n", + "* Problem geometry (mapped domain)\n", + "* Static background (equilibrium)\n", + "* Grid\n", + "* Derham complex\n", + "\n", + "Check the respective classes for possible options." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc43d2fc", + "metadata": {}, + "outputs": [], + "source": [ + "# environment options\n", + "env = EnvironmentOptions()\n", + "\n", + "# units\n", + "units = Units()\n", + "\n", + "# time stepping\n", + "time_opts = Time(dt=0.2, Tend=10.0)\n", + "\n", + "# geometry\n", + "l1 = -5.0\n", + "r1 = 5.0\n", + "l2 = -7.0\n", + "r2 = 7.0\n", + "l3 = -1.0\n", + "r3 = 1.0\n", + "domain = domains.Cuboid(l1=l1, r1=r1, l2=l2, r2=r2, l3=l3, r3=r3)\n", + "\n", + "# fluid equilibrium (can be used as part of initial conditions)\n", + "equil = None\n", + "\n", + "# grid\n", + "grid = None\n", + "\n", + "# derham options\n", + "derham_opts = None" + ] + }, + { + "cell_type": "markdown", + "id": "74e6f739", + "metadata": {}, + "source": [ + "## Parameters part 3: Model instance\n", + "\n", + "Here, a light-weight instance of the model is created, without allocating memory. The light-weight instance is used to set model-specific parameters for the model's species.\n", + "\n", + "Check the functions \n", + "\n", + "* `Species.set_phys_params()`\n", + "* `KineticSpecies.set_markers()`\n", + "* `KineticSpecies.set_sorting_boxes()`\n", + "* `KineticSpecies.set_save_data()`\n", + "\n", + " for possible options." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6c498fb3", + "metadata": {}, + "outputs": [], + "source": [ + "# light-weight model instance\n", + "model = Model()\n", + "\n", + "# species parameters\n", + "model.kinetic_ions.set_phys_params()\n", + "\n", + "loading_params = LoadingParameters(Np=15)\n", + "weights_params = WeightsParameters()\n", + "boundary_params = BoundaryParameters(bc=('reflect', 'reflect', 'periodic'))\n", + "model.kinetic_ions.set_markers(loading_params=loading_params, \n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params)\n", + "model.kinetic_ions.set_sorting_boxes()\n", + "model.kinetic_ions.set_save_data(n_markers=1.0)" + ] + }, + { + "cell_type": "markdown", + "id": "b0f65b0a", + "metadata": {}, + "source": [ + "## Parameters part 4: Propagator options\n", + "\n", + "Check the method `set_options()` of each propagator for possible options." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "be6875e7", + "metadata": {}, + "outputs": [], + "source": [ + "# propagator options\n", + "model.propagators.push_vxb.set_options()\n", + "model.propagators.push_eta.set_options()" + ] + }, + { + "cell_type": "markdown", + "id": "b1ef8b97", + "metadata": {}, + "source": [ + "## Parameters part 5: Initial conditions\n", + "\n", + "Use the methods `Variable.add_background()` and `Variable.add_perturbation()` to set initial conditions for each variable of a species. Variables that are not specified are intialized as zero." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cc0ac424", + "metadata": {}, + "outputs": [], + "source": [ + "# initial conditions (background + perturbation)\n", + "perturbation = None\n", + "\n", + "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", + "model.kinetic_ions.var.add_background(background)" + ] + }, + { + "cell_type": "markdown", + "id": "879978af", + "metadata": {}, + "source": [ + "## Parameters part 6: `main.run`\n", + "\n", + "In the final part of the parameter file, the `main.run` command is invoked. This command will allocate memory and run the specified simulation. The run command is not executed when the parameter file is imported in another Python script." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "03764138", + "metadata": {}, + "outputs": [], + "source": [ + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " units=units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "f9ce5099", + "metadata": {}, + "source": [ + "## Post processing: `main.pproc`\n", + "\n", + "Aside from `run`, the Struphy `main` module has also a `pproc` routine for post-processing raw simulation data:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dd7cfe62", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "\n", + "main.pproc(path, physical=True)" + ] + }, + { + "cell_type": "markdown", + "id": "e317f88d", + "metadata": {}, + "source": [ + "## Loading data: `main.load_data`\n", + "\n", + "After post-processing, the generated data can be loaded via `main.load_data`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "78b8cbab", + "metadata": {}, + "outputs": [], + "source": [ + "simdata = main.load_data(path)" + ] + }, + { + "cell_type": "markdown", + "id": "a9468479", + "metadata": {}, + "source": [ + "`main.load_data` returns a `SimData` object, which you can inspect to get further info on possible data to load:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0531679e", + "metadata": {}, + "outputs": [], + "source": [ + "for k, v in simdata.__dict__.items():\n", + " print(k, type(v))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a329eb38", + "metadata": {}, + "outputs": [], + "source": [ + "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", + " print(f\"{k = }, {type(v) = }\")" + ] + }, + { + "cell_type": "markdown", + "id": "08544a91", + "metadata": {}, + "source": [ + "## Plotting particle orbits\n", + "\n", + "In this example, for the species `kinetic_ions` some particle orbits have been saved. Let us plot them:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7af3facd", + "metadata": {}, + "outputs": [], + "source": [ + "from matplotlib import pyplot as plt\n", + "\n", + "fig = plt.figure()\n", + "ax = fig.gca()\n", + "\n", + "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", + "\n", + "time = 0.\n", + "dt = time_opts.dt\n", + "Tend = time_opts.Tend\n", + "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", + " # print(f\"{v[0] = }\")\n", + " alpha = (Tend - time)/Tend\n", + " for i, particle in enumerate(v):\n", + " ax.scatter(particle[0], particle[1], c=colors[i % 4], alpha=alpha)\n", + " time += dt\n", + " \n", + "ax.plot([l1, l1], [l2, r2], 'k')\n", + "ax.plot([r1, r1], [l2, r2], 'k')\n", + "ax.plot([l1, r1], [l2, l2], 'k')\n", + "ax.plot([l1, r1], [r2, r2], 'k')\n", + "ax.set_xlabel('x')\n", + "ax.set_ylabel('y')\n", + "ax.set_xlim(-6.5, 6.5)\n", + "ax.set_ylim(-9, 9)\n", + "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "env", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/doc/tutorials/tutorial_01_particles.ipynb b/tutorials_old/tutorial_01_particles.ipynb similarity index 100% rename from doc/tutorials/tutorial_01_particles.ipynb rename to tutorials_old/tutorial_01_particles.ipynb diff --git a/doc/tutorials/tutorial_02_fluid_particles.ipynb b/tutorials_old/tutorial_02_fluid_particles.ipynb similarity index 100% rename from doc/tutorials/tutorial_02_fluid_particles.ipynb rename to tutorials_old/tutorial_02_fluid_particles.ipynb diff --git a/doc/tutorials/tutorial_03_discrete_derham.ipynb b/tutorials_old/tutorial_03_discrete_derham.ipynb similarity index 100% rename from doc/tutorials/tutorial_03_discrete_derham.ipynb rename to tutorials_old/tutorial_03_discrete_derham.ipynb diff --git a/doc/tutorials/tutorial_04_mapped_domains.ipynb b/tutorials_old/tutorial_04_mapped_domains.ipynb similarity index 100% rename from doc/tutorials/tutorial_04_mapped_domains.ipynb rename to tutorials_old/tutorial_04_mapped_domains.ipynb diff --git a/doc/tutorials/tutorial_05_mhd_equilibria.ipynb b/tutorials_old/tutorial_05_mhd_equilibria.ipynb similarity index 100% rename from doc/tutorials/tutorial_05_mhd_equilibria.ipynb rename to tutorials_old/tutorial_05_mhd_equilibria.ipynb diff --git a/doc/tutorials/tutorial_06_poisson.ipynb b/tutorials_old/tutorial_06_poisson.ipynb similarity index 100% rename from doc/tutorials/tutorial_06_poisson.ipynb rename to tutorials_old/tutorial_06_poisson.ipynb diff --git a/doc/tutorials/tutorial_07_heat_equation.ipynb b/tutorials_old/tutorial_07_heat_equation.ipynb similarity index 100% rename from doc/tutorials/tutorial_07_heat_equation.ipynb rename to tutorials_old/tutorial_07_heat_equation.ipynb diff --git a/doc/tutorials/tutorial_08_maxwell.ipynb b/tutorials_old/tutorial_08_maxwell.ipynb similarity index 100% rename from doc/tutorials/tutorial_08_maxwell.ipynb rename to tutorials_old/tutorial_08_maxwell.ipynb diff --git a/doc/tutorials/tutorial_09_vlasov_maxwell.ipynb b/tutorials_old/tutorial_09_vlasov_maxwell.ipynb similarity index 100% rename from doc/tutorials/tutorial_09_vlasov_maxwell.ipynb rename to tutorials_old/tutorial_09_vlasov_maxwell.ipynb diff --git a/doc/tutorials/tutorial_10_linear_mhd.ipynb b/tutorials_old/tutorial_10_linear_mhd.ipynb similarity index 100% rename from doc/tutorials/tutorial_10_linear_mhd.ipynb rename to tutorials_old/tutorial_10_linear_mhd.ipynb diff --git a/doc/tutorials/tutorial_11_data_structures.ipynb b/tutorials_old/tutorial_11_data_structures.ipynb similarity index 100% rename from doc/tutorials/tutorial_11_data_structures.ipynb rename to tutorials_old/tutorial_11_data_structures.ipynb diff --git a/doc/tutorials/tutorial_12_struphy_data_pproc.ipynb b/tutorials_old/tutorial_12_struphy_data_pproc.ipynb similarity index 100% rename from doc/tutorials/tutorial_12_struphy_data_pproc.ipynb rename to tutorials_old/tutorial_12_struphy_data_pproc.ipynb From aca8c91539fd374e0499f1aa910e37cc5decbb27 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 2 Sep 2025 10:13:19 +0200 Subject: [PATCH 090/292] improve tutorial 02 --- doc/sections/tutorials.rst | 1 + .../tutorial_01_parameter_files.ipynb | 10 +- .../tutorial_02_test_particles.ipynb | 279 +++++++++--------- src/struphy/main.py | 3 +- src/struphy/models/base.py | 5 +- src/struphy/models/species.py | 2 - src/struphy/models/toy.py | 2 +- src/struphy/models/variables.py | 4 +- 8 files changed, 149 insertions(+), 157 deletions(-) diff --git a/doc/sections/tutorials.rst b/doc/sections/tutorials.rst index 53440ce87..c1b4a414b 100644 --- a/doc/sections/tutorials.rst +++ b/doc/sections/tutorials.rst @@ -13,3 +13,4 @@ They can thus be used for MPI parallel runs in HPC applications. :caption: Notebook tutorials: ../tutorials/tutorial_01_parameter_files + ../tutorials/tutorial_02_test_particles diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/doc/tutorials/tutorial_01_parameter_files.ipynb index 53c292b21..e50877b5a 100644 --- a/doc/tutorials/tutorial_01_parameter_files.ipynb +++ b/doc/tutorials/tutorial_01_parameter_files.ipynb @@ -5,7 +5,7 @@ "id": "d34c79c5", "metadata": {}, "source": [ - "# 01 - Parameters and `struphy.main`\n", + "# 1 - Parameters and `struphy.main`\n", "\n", "Struphy parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", "For each [MODEL](https://struphy.pages.mpcdf.de/struphy/sections/models.html), the default parameter file can be generated from the console via\n", @@ -44,7 +44,7 @@ "metadata": {}, "outputs": [], "source": [ - "from struphy.io.options import EnvironmentOptions, Units, Time\n", + "from struphy.io.options import EnvironmentOptions, BaseUnits, Time\n", "from struphy.geometry import domains\n", "from struphy.fields_background import equils\n", "from struphy.topology import grids\n", @@ -93,7 +93,7 @@ "env = EnvironmentOptions()\n", "\n", "# units\n", - "units = Units()\n", + "base_units = BaseUnits()\n", "\n", "# time stepping\n", "time_opts = Time(dt=0.2, Tend=10.0)\n", @@ -225,7 +225,7 @@ "main.run(model, \n", " params_path=None, \n", " env=env, \n", - " units=units, \n", + " base_units=base_units, \n", " time_opts=time_opts, \n", " domain=domain, \n", " equil=equil, \n", @@ -255,7 +255,7 @@ "import os\n", "path = os.path.join(os.getcwd(), \"sim_1\")\n", "\n", - "main.pproc(path, physical=True)" + "main.pproc(path)" ] }, { diff --git a/doc/tutorials/tutorial_02_test_particles.ipynb b/doc/tutorials/tutorial_02_test_particles.ipynb index 1ec402de3..1ae8f107b 100644 --- a/doc/tutorials/tutorial_02_test_particles.ipynb +++ b/doc/tutorials/tutorial_02_test_particles.ipynb @@ -7,7 +7,7 @@ "source": [ "# 2 - Test particles\n", "\n", - "Let us explore some options of the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov), which has already been introdused in Tutorial 1. In particular, we will\n", + "Let us explore some options of the models [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov), which has already been introdused in Tutorial 1, and [GuidingCenter](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.GuidingCenter). In particular, we will\n", "\n", "1. Change the geometry\n", "2. Change the loading of the markers\n", @@ -15,10 +15,19 @@ "\n", "## Particles in a cylinder\n", "\n", - "We use the same setup as in Tutorial 1 but change the domain $\\Omega$ to a cylinder. We explore two options for drawing markers in posittion space:\n", + "As in Tutorial 1, we re-create the parameter files in the notebook and then run `struphy.main`. The default parameter files can be created from the console via\n", + "\n", + "```\n", + "struphy params Vlasov\n", + "struphy params GuidingCenter \n", + "```\n", + "\n", + "This time, we set the simulation domain $\\Omega$ to be a cylinder. We explore two options for drawing markers in posittion space:\n", "\n", "- uniform in logical space $[0, 1]^3 = F^{-1}(\\Omega)$\n", - "- uniform on the cylinder $\\Omega$" + "- uniform on the cylinder $\\Omega$\n", + "\n", + "We start with the generic imports:" ] }, { @@ -44,6 +53,16 @@ "verbose = False" ] }, + { + "cell_type": "markdown", + "id": "8861c9ef", + "metadata": {}, + "source": [ + "Note that we set `verbose = False` which will be passed to `main.run` to supress screen output during the simulation.\n", + "\n", + "We shall create two simulations, which we store in different output folders, defined throught the environment variables:" + ] + }, { "cell_type": "code", "execution_count": null, @@ -53,8 +72,24 @@ "source": [ "# environment options\n", "env = EnvironmentOptions()\n", - "env_2 = EnvironmentOptions(sim_folder=\"sim_2\")\n", - "\n", + "env_2 = EnvironmentOptions(sim_folder=\"sim_2\")" + ] + }, + { + "cell_type": "markdown", + "id": "5d1ba880", + "metadata": {}, + "source": [ + "The other generic options will be the same for both simulations. Here, we just perform one time step und load a cylindrical geometry:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4dbe46b5", + "metadata": {}, + "outputs": [], + "source": [ "# units\n", "base_units = BaseUnits()\n", "\n", @@ -77,6 +112,14 @@ "derham_opts = None" ] }, + { + "cell_type": "markdown", + "id": "32615bbd", + "metadata": {}, + "source": [ + "For each simulation, we must create the light-weight model instance and set parameters. For the second simulation, in the parameters for particle loading we choose `spatial=\"disc\"` in order to draw uniformly on the cross section of the cylinder: " + ] + }, { "cell_type": "code", "execution_count": null, @@ -112,6 +155,14 @@ "model_2.kinetic_ions.set_save_data(n_markers=1.0)" ] }, + { + "cell_type": "markdown", + "id": "06b41735", + "metadata": {}, + "source": [ + "Propagator options and initial conditions shall be the same in both simulations:" + ] + }, { "cell_type": "code", "execution_count": null, @@ -142,6 +193,14 @@ "model_2.kinetic_ions.var.add_background(background)" ] }, + { + "cell_type": "markdown", + "id": "6ec98182", + "metadata": {}, + "source": [ + "Let us now run the first simulation:" + ] + }, { "cell_type": "code", "execution_count": null, @@ -162,6 +221,14 @@ " )" ] }, + { + "cell_type": "markdown", + "id": "6099f2e8", + "metadata": {}, + "source": [ + "And now the second simulation:" + ] + }, { "cell_type": "code", "execution_count": null, @@ -182,6 +249,14 @@ " )" ] }, + { + "cell_type": "markdown", + "id": "5834fb2f", + "metadata": {}, + "source": [ + "We now post-process both runs, load the generated data and plot the initial particle positions on a cross section of the cylinder:" + ] + }, { "cell_type": "code", "execution_count": null, @@ -250,224 +325,138 @@ "source": [ "## Reflecting boundary conditions \n", "\n", + "Let us now run \n", + "\n", "We shall set reflecting boundary conditions in radial direction, which in Struphy is always the logical direction $\\eta_1$." ] }, { "cell_type": "code", "execution_count": null, - "id": "5aa44526", + "id": "aef8e67a", "metadata": {}, "outputs": [], "source": [ - "# instantiate Particle object\n", - "Np = 1000\n", - "bc = ['remove', 'periodic', 'periodic']\n", - "loading_params = {'seed': None}\n", - "\n", - "particles = Particles6D(Np=Np, \n", - " bc=bc, \n", - " loading_params=loading_params)\n", - "\n", - "# instantiate another Particle object\n", - "name = 'test_uni'\n", - "loading_params = {'seed': None, 'spatial': 'disc'}\n", - "particles_uni = Particles6D(Np=Np, \n", - " bc=bc, \n", - " loading_params=loading_params)" + "# time stepping\n", + "time_opts = Time(dt=0.2, Tend=10.0)" ] }, { "cell_type": "code", "execution_count": null, - "id": "2d5cd7a4", + "id": "b23d8ff8", "metadata": {}, "outputs": [], "source": [ - "particles.draw_markers()\n", - "particles_uni.draw_markers()" + "# light-weight model instance\n", + "model = Model()\n", + "\n", + "# species parameters\n", + "model.kinetic_ions.set_phys_params()\n", + "loading_params = LoadingParameters(Np=15, spatial=\"disc\")\n", + "model.kinetic_ions.set_markers(loading_params=loading_params, \n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params)\n", + "model.kinetic_ions.set_sorting_boxes()\n", + "model.kinetic_ions.set_save_data(n_markers=1.0)" ] }, { "cell_type": "code", "execution_count": null, - "id": "b74f24b3", + "id": "f9f1874f", "metadata": {}, "outputs": [], "source": [ - "# positions on the physical domain Omega\n", - "pushed_pos = domain(particles.positions).T\n", - "pushed_pos_uni = domain(particles_uni.positions).T" + "# propagator options\n", + "model.propagators.push_vxb.set_options()\n", + "model.propagators.push_eta.set_options()" ] }, { "cell_type": "code", "execution_count": null, - "id": "ab1a5b6d", + "id": "5c3c84e6", "metadata": {}, "outputs": [], "source": [ - "fig = plt.figure(figsize=(10, 6)) \n", - "\n", - "plt.subplot(1, 2, 1)\n", - "plt.scatter(pushed_pos[:, 0], pushed_pos[:, 1], s=2.)\n", - "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", - "ax = plt.gca()\n", - "ax.add_patch(circle1)\n", - "ax.set_aspect('equal')\n", - "plt.xlabel('x')\n", - "plt.ylabel('y')\n", - "plt.title('Draw uniform in logical space')\n", + "# initial conditions (background + perturbation)\n", + "perturbation = None\n", + "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", "\n", - "plt.subplot(1, 2, 2)\n", - "plt.scatter(pushed_pos_uni[:, 0], pushed_pos_uni[:, 1], s=2.)\n", - "circle2 = plt.Circle((0, 0), a2, color='k', fill=False)\n", - "ax = plt.gca()\n", - "ax.add_patch(circle2)\n", - "ax.set_aspect('equal')\n", - "plt.xlabel('x')\n", - "plt.ylabel('y')\n", - "plt.title('Draw uniform on disc');" + "model.kinetic_ions.var.add_background(background)" ] }, { "cell_type": "code", "execution_count": null, - "id": "8d5aa753", + "id": "0ec1aaeb", "metadata": {}, "outputs": [], "source": [ - "# instantiate Particle object\n", - "Np = 15\n", - "bc = ['reflect', 'periodic', 'periodic']\n", - "loading_params = {'seed': None}\n", - "\n", - "particles = Particles6D(Np=Np, \n", - " bc=bc, \n", - " domain=domain,\n", - " loading_params=loading_params)" + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " base_units=base_units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" ] }, { "cell_type": "code", "execution_count": null, - "id": "be9cdacd", + "id": "d90b4117", "metadata": {}, "outputs": [], "source": [ - "particles.draw_markers()" + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "main.pproc(path)" ] }, { "cell_type": "code", "execution_count": null, - "id": "ee8e446e", + "id": "3e8b5c2b", "metadata": {}, "outputs": [], "source": [ - "# positions on the physical domain Omega\n", - "pushed_pos = domain(particles.positions).T\n", - "pushed_pos" + "simdata = main.load_data(path)" ] }, { "cell_type": "code", "execution_count": null, - "id": "226f69fa", + "id": "80bb8873", "metadata": {}, "outputs": [], "source": [ - "fig = plt.figure() \n", + "fig = plt.figure()\n", "ax = fig.gca()\n", "\n", - "for n, pos in enumerate(pushed_pos):\n", - " ax.scatter(pos[0], pos[1], c=colors[n % 4])\n", - " ax.arrow(pos[0], pos[1], particles.velocities[n, 0], particles.velocities[n, 1], color=colors[n % 4], head_width=.2)\n", - "\n", - "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", - "\n", - "ax.add_patch(circle1)\n", - "ax.set_aspect('equal')\n", - "ax.set_xlabel('x')\n", - "ax.set_ylabel('y')\n", - "ax.set_title('Initial conditions');" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4f56e79b", - "metadata": {}, - "outputs": [], - "source": [ - "# pass simulation parameters to Propagator class\n", - "PushEta.domain = domain" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1f6d91c1", - "metadata": {}, - "outputs": [], - "source": [ - "# instantiate Propagator object\n", - "prop_eta = PushEta(particles)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2fc52d09", - "metadata": {}, - "outputs": [], - "source": [ - "# time stepping\n", - "Tend = 10. \n", - "dt = .2\n", - "Nt = int(Tend / dt)\n", - "\n", - "pos = np.zeros((Nt + 1, Np, 3), dtype=float)\n", - "alpha = np.ones(Nt + 1, dtype=float)\n", - "\n", - "pos[0] = pushed_pos\n", + "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", "\n", "time = 0.\n", - "n = 0\n", - "while time < (Tend -dt):\n", + "dt = time_opts.dt\n", + "Tend = time_opts.Tend\n", + "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", + " # print(f\"{v[0] = }\")\n", + " alpha = (Tend - time)/Tend\n", + " for i, particle in enumerate(v):\n", + " ax.scatter(particle[0], particle[1], c=colors[i % 4], alpha=alpha)\n", " time += dt\n", - " n += 1\n", - " \n", - " # advance in time\n", - " prop_eta(dt)\n", " \n", - " # positions on the physical domain Omega\n", - " pos[n] = domain(particles.positions).T\n", - " \n", - " # scaling for plotting\n", - " alpha[n] = (Tend - time)/Tend" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "895cc320", - "metadata": {}, - "outputs": [], - "source": [ - "# make scatter plot for each particle in xy-plane\n", - "for i in range(Np):\n", - " ax.scatter(pos[:, i, 0], pos[:, i, 1], c=colors[i % 4], alpha=alpha)\n", - "\n", "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", "\n", "ax.add_patch(circle1)\n", "ax.set_aspect('equal')\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('y')\n", - "ax.set_title(f'{math.ceil(Tend/dt)} time steps (full color at t=0)');\n", - "\n", - "fig" + "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" ] }, { diff --git a/src/struphy/main.py b/src/struphy/main.py index f52264175..7b64d8444 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -195,7 +195,7 @@ def run( model.setup_equation_params(units=model.units, verbose=verbose) # allocate variables - model.allocate_variables() + model.allocate_variables(verbose=verbose) model.allocate_helpers() # pass info to propagators @@ -303,6 +303,7 @@ def run( data.file.close() end_simulation = time.time() if rank == 0: + print(f"\nTime steps done: {time_state["index"][0]}") print( "wall-clock time of simulation [sec]: ", end_simulation - start_simulation, diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 7876c2380..95e076217 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -578,7 +578,7 @@ def add_time_state(self, time_state): if isinstance(prop, Propagator): prop.add_time_state(time_state) - def allocate_variables(self): + def allocate_variables(self, verbose: bool = False): """ Allocate memory for model variables and set initial conditions. """ @@ -604,7 +604,8 @@ def allocate_variables(self): assert isinstance(spec, KineticSpecies) for k, v in spec.variables.items(): assert isinstance(v, (PICVariable, SPHVariable)) - v.allocate(derham=self.derham, domain=self.domain, equil=self.equil,) + v.allocate(derham=self.derham, domain=self.domain, equil=self.equil, + verbose=verbose) # TODO: allocate memory for FE coeffs of diagnostics # if self.params.diagnostic_fields is not None: diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 1b4e4711c..192ce5c84 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -52,8 +52,6 @@ def set_phys_params(self, charge_number: int = 1, mass_number: int = 1): @property def equation_params(self) -> dict: - # if not hasattr(self, "_equation_params"): - # self.setup_equation_params() return self._equation_params def setup_equation_params(self, units: Units, verbose=False): diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index db8f1d9ca..8dcbcf435 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -132,7 +132,7 @@ def __init__(self): def __init__(self): if rank == 0: - print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}' ***") # 1. instantiate all species, variables self.kinetic_ions = self.KineticIons() diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index e5d39e93e..863bca007 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -169,6 +169,7 @@ def allocate(self, domain: Domain = None, equil: FluidEquilibrium = None, projected_equil: ProjectedFluidEquilibrium = None, + verbose: bool = False, ): #assert isinstance(self.species, KineticSpecies) @@ -202,13 +203,14 @@ def allocate(self, n_as_volume_form=self.n_as_volume_form, # perturbations=self.perturbations, equation_params=self.species.equation_params, + verbose=verbose, ) if self.species.do_sort: sort = True else: sort = False - self.particles.draw_markers(sort=sort) + self.particles.draw_markers(sort=sort, verbose=verbose) self.particles.initialize_weights() # for storing the binned distribution function From 9e05edee403814a47d960fefa66e47ddbf5dcc41 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 4 Sep 2025 09:08:09 +0200 Subject: [PATCH 091/292] remove "as Model" in params file --- src/struphy/models/base.py | 4 ++-- src/struphy/models/tests/test_LinearMHD.py | 4 ++-- src/struphy/models/tests/test_Maxwell.py | 8 ++++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 95e076217..912b44a14 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1343,7 +1343,7 @@ def generate_default_parameter_file( file.write("from struphy import main\n") file.write("\n# import model, set verbosity\n") - file.write(f"from {self.__module__} import {self.__class__.__name__} as Model\n") + file.write(f"from {self.__module__} import {self.__class__.__name__}\n") file.write("verbose = True\n") file.write("\n# environment options\n") @@ -1375,7 +1375,7 @@ def generate_default_parameter_file( file.write(derham) file.write("\n# light-weight model instance\n") - file.write("model = Model()\n") + file.write(f"model = {self.__class__.__name__}()\n") if has_plasma: file.write(species_params) diff --git a/src/struphy/models/tests/test_LinearMHD.py b/src/struphy/models/tests/test_LinearMHD.py index 00918b02c..f8e17f826 100644 --- a/src/struphy/models/tests/test_LinearMHD.py +++ b/src/struphy/models/tests/test_LinearMHD.py @@ -20,7 +20,7 @@ @pytest.mark.parametrize('algo', ["implicit", "explicit"]) def test_slab_waves_1d(algo: str, do_plot: bool = False): # import model, set verbosity - from struphy.models.fluid import LinearMHD as Model + from struphy.models.fluid import LinearMHD verbose = True # environment options @@ -51,7 +51,7 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): derham_opts = DerhamOptions(p=(1, 1, 3)) # light-weight model instance - model = Model() + model = LinearMHD() # species parameters model.mhd.set_phys_params() diff --git a/src/struphy/models/tests/test_Maxwell.py b/src/struphy/models/tests/test_Maxwell.py index c366cf284..1f190ff7c 100644 --- a/src/struphy/models/tests/test_Maxwell.py +++ b/src/struphy/models/tests/test_Maxwell.py @@ -22,7 +22,7 @@ @pytest.mark.parametrize('algo', ["implicit", "explicit"]) def test_light_wave_1d(algo: str, do_plot: bool = False): # import model, set verbosity - from struphy.models.toy import Maxwell as Model + from struphy.models.toy import Maxwell verbose = True # environment options @@ -48,7 +48,7 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): derham_opts = DerhamOptions(p=(1, 1, 3)) # light-weight model instance - model = Model() + model = Maxwell() # propagator options model.propagators.maxwell.set_options(algo=algo) @@ -102,7 +102,7 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): @pytest.mark.mpi(min_size=4) def test_coaxial(do_plot: bool = False): # import model, set verbosity - from struphy.models.toy import Maxwell as Model + from struphy.models.toy import Maxwell verbose = True # environment options @@ -134,7 +134,7 @@ def test_coaxial(do_plot: bool = False): ) # light-weight model instance - model = Model() + model = Maxwell() # propagator options model.propagators.maxwell.set_options(algo="implicit") From fc8b841849363e20ca024ac116fe7482c0f84411 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 4 Sep 2025 16:04:03 +0200 Subject: [PATCH 092/292] improve tutorials 1 and 2 --- .../tutorial_01_parameter_files.ipynb | 89 ++- .../tutorial_02_test_particles.ipynb | 551 +++++++++--------- src/struphy/main.py | 18 +- src/struphy/pic/base.py | 7 +- .../post_processing/post_processing_tools.py | 10 +- 5 files changed, 358 insertions(+), 317 deletions(-) diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/doc/tutorials/tutorial_01_parameter_files.ipynb index e50877b5a..69b1f56c6 100644 --- a/doc/tutorials/tutorial_01_parameter_files.ipynb +++ b/doc/tutorials/tutorial_01_parameter_files.ipynb @@ -7,8 +7,12 @@ "source": [ "# 1 - Parameters and `struphy.main`\n", "\n", - "Struphy parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", - "For each [MODEL](https://struphy.pages.mpcdf.de/struphy/sections/models.html), the default parameter file can be generated from the console via\n", + "Struphy is based on \"models\". Each model can be viewed as a simulation code for a certain physical model, described py a set partial differntial equations (PDEs).\n", + "The documentation features a [list of currently available models](https://struphy.pages.mpcdf.de/struphy/sections/models.html).\n", + "A model is launched through a parameter file, where all simulation parameters can be specified by the user.\n", + "\n", + "Model parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", + "For each `MODEL`, the default parameter file can be generated from the console via\n", "\n", "```\n", "struphy params MODEL\n", @@ -21,8 +25,7 @@ "```\n", "\n", "The user can modify the parameter file to launch a specific simulation.\n", - "As an example, let us discuss the parameter file of the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov) and run some simple examples. \n", - "The file can be generated from\n", + "For example, the parameter file of the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov) can be generated from\n", "\n", "```\n", "struphy params Vlasov\n", @@ -34,6 +37,8 @@ "cat params_Vlasov.py\n", "```\n", "\n", + "We shall discuss this parameter file in what follows. Parameter files of all models have a similar structure.\n", + "\n", "## Part 1: Imports" ] }, @@ -56,7 +61,7 @@ "from struphy import main\n", "\n", "# import model, set verbosity\n", - "from struphy.models.toy import Vlasov as Model\n", + "from struphy.models.toy import Vlasov\n", "verbose = True" ] }, @@ -65,7 +70,8 @@ "id": "5cf6d9c7", "metadata": {}, "source": [ - "All Struphy parameter files import the modules listed above, even though some of them might not be needed in a specific model. The last import imports the model itself, always under the alias `Model`.\n", + "All parameter files import the modules listed above, even though some of them might not be needed in a specific model. \n", + "The last import imports the model itself. The `verbose` flag controls the screen output during the simulation run.\n", "\n", "## Part 2: Generic options\n", "\n", @@ -75,8 +81,8 @@ "* Model units\n", "* Time options\n", "* Problem geometry (mapped domain)\n", - "* Static background (equilibrium)\n", - "* Grid\n", + "* Static background (equilibrium) \n", + "* Grid \n", "* Derham complex\n", "\n", "Check the respective classes for possible options." @@ -124,37 +130,49 @@ "source": [ "## Part 3: Model instance\n", "\n", - "Here, a light-weight instance of the model is created, without allocating memory. The light-weight instance is used to set model-specific parameters for the model's species.\n", - "\n", - "Check the functions \n", + "A model has a predefined number of \"species\". Each species is a collection of \"variables\", which are the unknowns of the model.\n", + "In the parameter file, a light-weight instance of the model is created, without allocating memory. \n", + "The light-weight instance is used to set some model-specific parameters for each species.\n", "\n", - "* `Species.set_phys_params()`\n", - "* `KineticSpecies.set_markers()`\n", - "* `KineticSpecies.set_sorting_boxes()`\n", - "* `KineticSpecies.set_save_data()`\n", - "\n", - " for possible options." + "For instance, for each species one can set its charge- and mass number:" ] }, { "cell_type": "code", "execution_count": null, - "id": "6c498fb3", + "id": "83dc7f7f", "metadata": {}, "outputs": [], "source": [ "# light-weight model instance\n", - "model = Model()\n", + "model = Vlasov()\n", "\n", "# species parameters\n", - "model.kinetic_ions.set_phys_params()\n", - "\n", + "model.kinetic_ions.set_phys_params(charge_number=3, mass_number=12)" + ] + }, + { + "cell_type": "markdown", + "id": "aa34840a", + "metadata": {}, + "source": [ + "In case of a kinetic species, one can also set parameters regarding marker drawing, box sorting and data saving: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6c498fb3", + "metadata": {}, + "outputs": [], + "source": [ "loading_params = LoadingParameters(Np=15)\n", "weights_params = WeightsParameters()\n", "boundary_params = BoundaryParameters(bc=('reflect', 'reflect', 'periodic'))\n", "model.kinetic_ions.set_markers(loading_params=loading_params, \n", " weights_params=weights_params,\n", " boundary_params=boundary_params)\n", + "\n", "model.kinetic_ions.set_sorting_boxes()\n", "model.kinetic_ions.set_save_data(n_markers=1.0)" ] @@ -166,7 +184,10 @@ "source": [ "## Part 4: Propagator options\n", "\n", - "Check the method `set_options()` of each propagator for possible options." + "A model is a collection of \"propagators\", which perform the time stepping. Each propagator refers to a splitting step of a single time step. For instance, if only a single propagator is present in a model, then all variables of the model are updated by this propagator. If two or more propagators are present, they are executed in sequence, according to the chosen splitting algorithm (Lie-Trotter, Strang, etc.).\n", + "\n", + "Each propagator has options that can be set in the parameter file.\n", + "Check the method `set_options()` of each propagator for its available options." ] }, { @@ -188,7 +209,7 @@ "source": [ "## Part 5: Initial conditions\n", "\n", - "Use the methods `Variable.add_background()` and `Variable.add_perturbation()` to set initial conditions for each variable of a species. Variables that are not specified are intialized as zero." + "One can use the methods `Variable.add_background()` and `Variable.add_perturbation()` to set initial conditions for each variable of a species. Variables that are not specified are intialized as zero." ] }, { @@ -242,7 +263,7 @@ "source": [ "## Post processing: `main.pproc`\n", "\n", - "Aside from `run`, the Struphy `main` module has also a `pproc` routine for post-processing raw simulation data:" + "Aside from `run`, the Struphy `main` module has also a `pproc` routine for post-processing of raw simulation data:" ] }, { @@ -258,6 +279,26 @@ "main.pproc(path)" ] }, + { + "cell_type": "markdown", + "id": "0076c65c", + "metadata": {}, + "source": [ + "One can also post-process directly from the console:\n", + "\n", + "```\n", + "struphy pproc sim_1\n", + "```\n", + "\n", + "Type \n", + "\n", + "```\n", + "struphy pproc -h\n", + "```\n", + "\n", + "for more info on this command." + ] + }, { "cell_type": "markdown", "id": "e317f88d", diff --git a/doc/tutorials/tutorial_02_test_particles.ipynb b/doc/tutorials/tutorial_02_test_particles.ipynb index 1ae8f107b..f72cfb445 100644 --- a/doc/tutorials/tutorial_02_test_particles.ipynb +++ b/doc/tutorials/tutorial_02_test_particles.ipynb @@ -7,7 +7,7 @@ "source": [ "# 2 - Test particles\n", "\n", - "Let us explore some options of the models [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov), which has already been introdused in Tutorial 1, and [GuidingCenter](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.GuidingCenter). In particular, we will\n", + "Let us explore some options of the models [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov), and [GuidingCenter](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.GuidingCenter). In particular, we will\n", "\n", "1. Change the geometry\n", "2. Change the loading of the markers\n", @@ -15,7 +15,7 @@ "\n", "## Particles in a cylinder\n", "\n", - "As in Tutorial 1, we re-create the parameter files in the notebook and then run `struphy.main`. The default parameter files can be created from the console via\n", + "As in Tutorial 1, we shall re-create the parameter files in the notebook for the purpose of discussion. The default parameter files can be created from the console via\n", "\n", "```\n", "struphy params Vlasov\n", @@ -49,7 +49,7 @@ "from struphy import main\n", "\n", "# import model, set verbosity\n", - "from struphy.models.toy import Vlasov as Model\n", + "from struphy.models.toy import Vlasov\n", "verbose = False" ] }, @@ -99,9 +99,43 @@ "# geometry\n", "a1 = 0.\n", "a2 = 5.\n", - "Lz = 1.\n", - "domain = domains.HollowCylinder(a1=a1, a2=a2, Lz=Lz)\n", - "\n", + "Lz = 20.\n", + "domain = domains.HollowCylinder(a1=a1, a2=a2, Lz=Lz)" + ] + }, + { + "cell_type": "markdown", + "id": "ee94df63", + "metadata": {}, + "source": [ + "We can already look a t the simulation domain:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d4174c86", + "metadata": {}, + "outputs": [], + "source": [ + "domain.show()" + ] + }, + { + "cell_type": "markdown", + "id": "361147f5", + "metadata": {}, + "source": [ + "We can leave the equilibrium, grid and Derham complex empty:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4a03b02e", + "metadata": {}, + "outputs": [], + "source": [ "# fluid equilibrium (can be used as part of initial conditions)\n", "equil = None\n", "\n", @@ -117,29 +151,45 @@ "id": "32615bbd", "metadata": {}, "source": [ - "For each simulation, we must create the light-weight model instance and set parameters. For the second simulation, in the parameters for particle loading we choose `spatial=\"disc\"` in order to draw uniformly on the cross section of the cylinder: " + "For each simulation, we must create the light-weight model instance and set parameters:" ] }, { "cell_type": "code", "execution_count": null, - "id": "c13b1c6c", + "id": "853ac716", "metadata": {}, "outputs": [], "source": [ "# light-weight model instance\n", - "model = Model()\n", - "model_2 = Model()\n", + "model = Vlasov()\n", + "model_2 = Vlasov()\n", "\n", "# species parameters\n", "model.kinetic_ions.set_phys_params()\n", - "model_2.kinetic_ions.set_phys_params()\n", - "\n", + "model_2.kinetic_ions.set_phys_params()" + ] + }, + { + "cell_type": "markdown", + "id": "71ad7471", + "metadata": {}, + "source": [ + "For the second simulation, in the parameters for particle loading we choose `spatial=\"disc\"` in order to draw uniformly on the cross section of the cylinder: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c13b1c6c", + "metadata": {}, + "outputs": [], + "source": [ "loading_params = LoadingParameters(Np=1000)\n", "loading_params_2 = LoadingParameters(Np=1000, spatial=\"disc\")\n", "\n", "weights_params = WeightsParameters()\n", - "boundary_params = BoundaryParameters(bc=('reflect', 'periodic', 'periodic'))\n", + "boundary_params = BoundaryParameters()\n", "\n", "model.kinetic_ions.set_markers(loading_params=loading_params, \n", " weights_params=weights_params,\n", @@ -305,7 +355,7 @@ "ax.set_aspect('equal')\n", "plt.xlabel('x')\n", "plt.ylabel('y')\n", - "plt.title('Draw uniform in logical space')\n", + "plt.title('sim_1: draw uniform in logical space')\n", "\n", "plt.subplot(1, 2, 2)\n", "plt.scatter(pushed_pos_uni[:, 0], pushed_pos_uni[:, 1], s=2.)\n", @@ -315,7 +365,7 @@ "ax.set_aspect('equal')\n", "plt.xlabel('x')\n", "plt.ylabel('y')\n", - "plt.title('Draw uniform on disc');" + "plt.title('sim_2: draw uniform on disc');" ] }, { @@ -325,9 +375,8 @@ "source": [ "## Reflecting boundary conditions \n", "\n", - "Let us now run \n", - "\n", - "We shall set reflecting boundary conditions in radial direction, which in Struphy is always the logical direction $\\eta_1$." + "Let us now run for 50 time steps and with 15 particles in the cylinder.\n", + "Moreover, we set reflecting boundary conditions in radial direction, which in Struphy is always the logical direction $\\eta_1$." ] }, { @@ -337,8 +386,9 @@ "metadata": {}, "outputs": [], "source": [ - "# time stepping\n", - "time_opts = Time(dt=0.2, Tend=10.0)" + "time_opts = Time(dt=0.2, Tend=10.0)\n", + "loading_params = LoadingParameters(Np=15, spatial=\"disc\")\n", + "boundary_params = BoundaryParameters(bc=('reflect', 'periodic', 'periodic'))" ] }, { @@ -349,11 +399,11 @@ "outputs": [], "source": [ "# light-weight model instance\n", - "model = Model()\n", + "model = Vlasov()\n", "\n", "# species parameters\n", "model.kinetic_ions.set_phys_params()\n", - "loading_params = LoadingParameters(Np=15, spatial=\"disc\")\n", + "\n", "model.kinetic_ions.set_markers(loading_params=loading_params, \n", " weights_params=weights_params,\n", " boundary_params=boundary_params)\n", @@ -362,24 +412,24 @@ ] }, { - "cell_type": "code", - "execution_count": null, - "id": "f9f1874f", + "cell_type": "markdown", + "id": "31de5e6d", "metadata": {}, - "outputs": [], "source": [ - "# propagator options\n", - "model.propagators.push_vxb.set_options()\n", - "model.propagators.push_eta.set_options()" + "We still have to set the propagator options and the initial conditions:" ] }, { "cell_type": "code", "execution_count": null, - "id": "5c3c84e6", + "id": "f9f1874f", "metadata": {}, "outputs": [], "source": [ + "# propagator options\n", + "model.propagators.push_vxb.set_options()\n", + "model.propagators.push_eta.set_options()\n", + "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", @@ -387,6 +437,14 @@ "model.kinetic_ions.var.add_background(background)" ] }, + { + "cell_type": "markdown", + "id": "fdffdb7c", + "metadata": {}, + "source": [ + "We can now run the simulation, then post-process the data and plot the resulting orbits:" + ] + }, { "cell_type": "code", "execution_count": null, @@ -466,7 +524,9 @@ "source": [ "## Particles in a cylinder with a magnetic field\n", "\n", - "Let $\\Omega \\subset \\mathbb R^3$ be a cylinder as before. Now, we search for trajectories $(\\mathbf x_p, \\mathbf v_p): [0,T] \\to \\Omega \\times \\mathbb R^3$, $p = 0, \\ldots, N-1$ that satisfy\n", + "Let us add a magnetic field to the simulation. This can be done by setting an [MHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils.html#mhd-equilibria):\n", + "\n", + "" ] }, { "cell_type": "code", "execution_count": null, - "id": "7ae9ae17", + "id": "ac87beba", "metadata": {}, "outputs": [], "source": [ - "from struphy.geometry.domains import HollowCylinder\n", - "\n", - "a1 = 0.\n", - "a2 = 5.\n", - "Lz = 1.\n", - "domain = HollowCylinder(a1=a1, a2=a2, Lz=Lz)" + "B0x = 0.\n", + "B0y = 0.\n", + "B0z = 1.\n", + "equil = equils.HomogenSlab(B0x=B0x, B0y=B0y, B0z=B0z)" ] }, { - "cell_type": "code", - "execution_count": null, - "id": "336bc5bd", + "cell_type": "markdown", + "id": "706486a1", "metadata": {}, - "outputs": [], "source": [ - "# instantiate Particle object\n", - "Np = 20\n", - "bc = ['remove', 'periodic', 'periodic']\n", - "loading_params = {'seed': None}\n", - "\n", - "particles = Particles6D(Np=Np, \n", - " bc=bc, \n", - " loading_params=loading_params)" + "In order to project the equilibrium on the spline basis for fast evaluation in the particle kernels, we need a Derham complex:" ] }, { "cell_type": "code", "execution_count": null, - "id": "64860359", + "id": "48fc6e08", "metadata": {}, "outputs": [], "source": [ - "particles.draw_markers()" + "spl_kind = (False, True, True)\n", + "derham_opts = DerhamOptions(spl_kind=spl_kind)" ] }, { - "cell_type": "code", - "execution_count": null, - "id": "ed244a1e", + "cell_type": "markdown", + "id": "f55b6f60", "metadata": {}, - "outputs": [], "source": [ - "# positions on the physical domain Omega\n", - "pushed_pos = domain(particles.positions).T\n", - "pushed_pos" + "Now we create the light-weight instance of the model and set the species options. We shall `remove` particles that hit the boundary in $\\eta_1$ (radial) direction:" ] }, { "cell_type": "code", "execution_count": null, - "id": "3b046601", + "id": "d8b94d0e", "metadata": {}, "outputs": [], "source": [ - "fig = plt.figure() \n", - "ax = fig.gca()\n", + "# light-weight model instance\n", + "model = Vlasov()\n", "\n", - "for n, pos in enumerate(pushed_pos):\n", - " ax.scatter(pos[0], pos[1], c=colors[n % 4])\n", - " ax.arrow(pos[0], pos[1], particles.velocities[n, 0], particles.velocities[n, 1], color=colors[n % 4], head_width=.2)\n", + "# species parameters\n", + "model.kinetic_ions.set_phys_params()\n", "\n", - "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "loading_params = LoadingParameters(Np=20)\n", + "boundary_params = BoundaryParameters(bc=('remove', 'periodic', 'periodic'))\n", + "model.kinetic_ions.set_markers(loading_params=loading_params, \n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params)\n", + "model.kinetic_ions.set_sorting_boxes()\n", + "model.kinetic_ions.set_save_data(n_markers=1.0)\n", "\n", - "ax.add_patch(circle1)\n", - "ax.set_aspect('equal')\n", - "ax.set_xlabel('x')\n", - "ax.set_ylabel('y')\n", - "ax.set_title('Initial conditions');" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "739f3b41", - "metadata": {}, - "outputs": [], - "source": [ - "from struphy.propagators.propagators_markers import PushVxB\n", + "# propagator options\n", + "model.propagators.push_vxb.set_options()\n", + "model.propagators.push_eta.set_options()\n", "\n", - "# default parameters of Propagator\n", - "opts_vxB = PushVxB.options(default=True)\n", - "print(opts_vxB)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "ac87beba", - "metadata": {}, - "outputs": [], - "source": [ - "from struphy.fields_background.equils import HomogenSlab\n", + "# initial conditions (background + perturbation)\n", + "perturbation = None\n", + "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", "\n", - "B0x = 0.\n", - "B0y = 0.\n", - "B0z = 1.\n", - "equil = HomogenSlab(B0x=B0x, B0y=B0y, B0z=B0z)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "c0a30f29", - "metadata": {}, - "outputs": [], - "source": [ - "# set domain for Cartesian MHD equilibrium\n", - "equil.domain = domain" + "model.kinetic_ions.var.add_background(background)" ] }, { - "cell_type": "code", - "execution_count": null, - "id": "93970eb6", + "cell_type": "markdown", + "id": "5026b318", "metadata": {}, - "outputs": [], "source": [ - "from struphy.fields_background.projected_equils import ProjectedMHDequilibrium\n", - "from struphy.feec.psydac_derham import Derham\n", - "\n", - "# instantiate Derham object\n", - "Nel = [16, 16, 32]\n", - "p = [1, 1, 3]\n", - "spl_kind = [False, True, True]\n", - "derham = Derham(Nel=Nel, p=p, spl_kind=spl_kind)\n", - "\n", - "# instantiate a projected MHD equilibrium object\n", - "proj_equil = ProjectedMHDequilibrium(equil, derham)" + "Now the usual procedure: run, post-process, load data and finally plot the orbits:" ] }, { "cell_type": "code", "execution_count": null, - "id": "4b754e07", + "id": "34893626", "metadata": {}, "outputs": [], "source": [ - "# pass simulation parameters to Propagator class\n", - "PushEta.domain = domain\n", - "PushVxB.domain = domain\n", - "PushVxB.derham = derham" + "# run\n", + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " base_units=base_units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" ] }, { "cell_type": "code", "execution_count": null, - "id": "b7f71e9b", + "id": "966bb91e", "metadata": {}, "outputs": [], "source": [ - "# instantiate Propagator object\n", - "prop_eta = PushEta(particles)\n", - "prop_vxB = PushVxB(particles, b2=proj_equil.b2)" + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "main.pproc(path)\n", + "\n", + "simdata = main.load_data(path)" ] }, { "cell_type": "code", "execution_count": null, - "id": "5294ebca", + "id": "fb7eaf92", "metadata": {}, "outputs": [], "source": [ - "# time stepping\n", - "Tend = 10. - 1e-6\n", - "dt = .2\n", - "Nt = int(Tend / dt)\n", - "\n", - "pos = []\n", - "alpha = np.ones(Nt + 1, dtype=float)\n", + "fig = plt.figure()\n", + "ax = fig.gca()\n", "\n", - "marker_col = {}\n", - "for marker in particles.markers_wo_holes:\n", - " m_id = int(marker[-1])\n", - " marker_col[m_id] = colors[int(m_id) % 4]\n", - "ids_wo_holes = []\n", + "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", "\n", "time = 0.\n", - "n = 0\n", - "while time < (Tend - dt):\n", + "dt = time_opts.dt\n", + "Tend = time_opts.Tend\n", + "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", + " # print(f\"{v[0] = }\")\n", + " alpha = (Tend - time)/Tend\n", + " for i, particle in enumerate(v):\n", + " ax.scatter(particle[0], particle[1], c=colors[i % 4], alpha=alpha)\n", " time += dt\n", - " n += 1\n", " \n", - " # advance in time\n", - " prop_vxB(dt/2)\n", - " prop_eta(dt)\n", - " prop_vxB(dt/2)\n", - " \n", - " # positions on the physical domain Omega (can change shape when particles are lost)\n", - " pos += [domain(particles.positions).T]\n", - "\n", - " # id's of non-holes\n", - " ids_wo_holes += [np.int64(particles.markers_wo_holes[:, -1])]\n", - " \n", - " # scaling for plotting\n", - " alpha[n] = (Tend - time)/Tend" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cb2ec764", - "metadata": {}, - "outputs": [], - "source": [ - "# make scatter plot for each particle in xy-plane\n", - "for po, ids, alph in zip(pos, ids_wo_holes, alpha):\n", - " cs = []\n", - " for ii in ids:\n", - " cs += [marker_col[ii]]\n", - " ax.scatter(po[:, 0], po[:, 1], c=cs, alpha=alph)\n", - "\n", "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", "\n", "ax.add_patch(circle1)\n", "ax.set_aspect('equal')\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('y')\n", - "ax.set_title(f'{math.ceil(Tend/dt)} time steps (full color at t=0)');\n", - "\n", - "fig" + "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" ] }, { @@ -716,9 +692,7 @@ "source": [ "## Particles in a Tokamak equilibrium\n", "\n", - "We use the same Propagators from the previous example but load a more complicated [MHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils.html#mhd-equilibria), namely from an ASDEX-Upgrade equilibrium stored in an EQDSK file.\n", - "\n", - "Let us instatiate an [EQDSKequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils_sub.html#struphy.fields_background.mhd_equil.equils.EQDSKequilibrium) with many of its default parameters, except for the density:" + "Let us try a more complicated [MHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils.html#mhd-equilibria), namely from an ASDEX-Upgrade equilibrium stored in an EQDSK file. We instatiate an [EQDSKequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils_sub.html#struphy.fields_background.mhd_equil.equils.EQDSKequilibrium) with many of its default parameters, except for the density:" ] }, { @@ -728,13 +702,10 @@ "metadata": {}, "outputs": [], "source": [ - "from struphy.fields_background.equils import EQDSKequilibrium\n", - "\n", "n1 = 0.\n", "n2 = 0.\n", "na = 1.\n", - "equil = EQDSKequilibrium(n1=n1, n2=n2, na=na)\n", - "equil.params" + "equil = equils.EQDSKequilibrium(n1=n1, n2=n2, na=na)" ] }, { @@ -752,27 +723,15 @@ "metadata": {}, "outputs": [], "source": [ - "from struphy.geometry.domains import Tokamak\n", - "\n", "Nel = (28, 72)\n", "p = (3, 3)\n", "psi_power = 0.6\n", "psi_shifts = (1e-6, 1.)\n", - "domain = Tokamak(equilibrium=equil, \n", - " Nel=Nel,\n", - " p=p,\n", - " psi_power=psi_power,\n", - " psi_shifts=psi_shifts)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1cb1f4cc", - "metadata": {}, - "outputs": [], - "source": [ - "equil.domain = domain" + "domain = domains.Tokamak(equilibrium=equil, \n", + " Nel=Nel,\n", + " p=p,\n", + " psi_power=psi_power,\n", + " psi_shifts=psi_shifts)" ] }, { @@ -888,6 +847,7 @@ "ax_top = axs[1]\n", "\n", "# min/max of field strength\n", + "equil.domain = domain\n", "Bmax = np.max(equil.absB0(*eta_topview_2, squeeze_out=True))\n", "Bmin = np.min(equil.absB0(*eta_topview_1, squeeze_out=True))\n", "levels = np.linspace(Bmin, Bmax, 51)\n", @@ -933,138 +893,173 @@ "fig.colorbar(im_top);" ] }, + { + "cell_type": "markdown", + "id": "f7dc2c1c", + "metadata": {}, + "source": [ + "We now set up a simualtion of 4 specific particle orbits in this equilibrium:" + ] + }, { "cell_type": "code", "execution_count": null, - "id": "568c0ba6", + "id": "65feff34", "metadata": {}, "outputs": [], "source": [ - "# instantiate Particle object\n", - "Np = 4\n", - "bc = ['remove', 'periodic', 'periodic']\n", - "bufsize = 2.\n", + "# light-weight model instance\n", + "model = Vlasov()\n", "\n", - "initial = [[.501, 0.001, 0.001, 0., 0.0450, -0.04], # co-passing particle\n", - " [.511, 0.001, 0.001, 0., -0.0450, -0.04], # counter passing particle\n", - " [.521, 0.001, 0.001, 0., 0.0105, -0.04], # co-trapped particle\n", - " [.531, 0.001, 0.001, 0., -0.0155, -0.04]]\n", + "# species parameters\n", + "model.kinetic_ions.set_phys_params()\n", "\n", - "loading_params = {'seed': 1608,\n", - " 'initial' : initial}\n", + "initial = ((.501, 0.001, 0.001, 0., 0.0450, -0.04), # co-passing particle\n", + " (.511, 0.001, 0.001, 0., -0.0450, -0.04), # counter passing particle\n", + " (.521, 0.001, 0.001, 0., 0.0105, -0.04), # co-trapped particle\n", + " (.531, 0.001, 0.001, 0., -0.0155, -0.04))\n", "\n", - "particles = Particles6D(Np=Np, \n", - " bc=bc, \n", - " loading_params=loading_params,\n", - " bufsize=bufsize)" + "loading_params = LoadingParameters(Np=4, seed=1608, specific_markers=initial)\n", + "boundary_params = BoundaryParameters(bc=('remove', 'periodic', 'periodic'))\n", + "model.kinetic_ions.set_markers(loading_params=loading_params, \n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params,\n", + " bufsize=2.)\n", + "model.kinetic_ions.set_sorting_boxes()\n", + "model.kinetic_ions.set_save_data(n_markers=1.0)\n", + "\n", + "# propagator options\n", + "model.propagators.push_vxb.set_options()\n", + "model.propagators.push_eta.set_options()\n", + "\n", + "# initial conditions (background + perturbation)\n", + "perturbation = None\n", + "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", + "\n", + "model.kinetic_ions.var.add_background(background)" ] }, { "cell_type": "code", "execution_count": null, - "id": "43127584", + "id": "568c0ba6", "metadata": {}, "outputs": [], "source": [ - "particles.draw_markers()" + "# # instantiate Particle object\n", + "# Np = 4\n", + "# bc = ['remove', 'periodic', 'periodic']\n", + "# bufsize = 2.\n", + "\n", + "# initial = [[.501, 0.001, 0.001, 0., 0.0450, -0.04], # co-passing particle\n", + "# [.511, 0.001, 0.001, 0., -0.0450, -0.04], # counter passing particle\n", + "# [.521, 0.001, 0.001, 0., 0.0105, -0.04], # co-trapped particle\n", + "# [.531, 0.001, 0.001, 0., -0.0155, -0.04]]\n", + "\n", + "# loading_params = {'seed': 1608,\n", + "# 'initial' : initial}\n", + "\n", + "# particles = Particles6D(Np=Np, \n", + "# bc=bc, \n", + "# loading_params=loading_params,\n", + "# bufsize=bufsize)" ] }, { - "cell_type": "code", - "execution_count": null, - "id": "4b2b96f2", + "cell_type": "markdown", + "id": "9de13919", "metadata": {}, - "outputs": [], "source": [ - "# positions on the physical domain Omega (x, y, z)\n", - "pushed_pos = domain(particles.positions).T\n", - "pushed_pos" + "We again need a Derham complex for the projection of the equilibirum onto the spline basis:" ] }, { "cell_type": "code", "execution_count": null, - "id": "3f7851f8", + "id": "4579b2af", "metadata": {}, "outputs": [], "source": [ - "# compute R-coordinate\n", - "pushed_r = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)" + "Nel = (32, 72, 1)\n", + "grid = grids.TensorProductGrid(Nel=Nel)\n", + "\n", + "p = (3, 3, 1)\n", + "spl_kind = (False, True, True)\n", + "derham_opts = DerhamOptions(p=p, spl_kind=spl_kind)" ] }, { - "cell_type": "code", - "execution_count": null, - "id": "72d79e1c", + "cell_type": "markdown", + "id": "9e50f541", "metadata": {}, - "outputs": [], "source": [ - "labels = ['co-passing',\n", - " 'counter passing',\n", - " 'co_trapped',\n", - " 'counter-trapped']\n", - "\n", - "for n, (r, pos) in enumerate(zip(pushed_r, pushed_pos)):\n", - " # poloidal \n", - " ax.scatter(r, pos[2], c=colors[n % 4], label=labels[n])\n", - " ax.arrow(r, pos[2], particles.velocities[n, 0], particles.velocities[n, 2]*10, color=colors[n % 4], head_width=.05)\n", - " # topview\n", - " ax_top.scatter(pos[0], pos[1], c=colors[n % 4], label=labels[n])\n", - " ax_top.arrow(pos[0], pos[1], particles.velocities[n, 0], particles.velocities[n, 1]*10, color=colors[n % 4], head_width=.05)\n", - "\n", - "ax.set_xlabel('R')\n", - "ax.set_ylabel('Z')\n", - "ax.set_title('Initial conditions')\n", - "ax.legend();\n", - "\n", - "ax_top.set_xlabel('x')\n", - "ax_top.set_ylabel('y')\n", - "ax_top.set_title('Initial conditions')\n", - "ax_top.legend();\n", - "\n", - "fig " + "We aim to simulate 15000 time steps with a second-order splitting algorithm:" ] }, { "cell_type": "code", "execution_count": null, - "id": "4579b2af", + "id": "86583b41", "metadata": {}, "outputs": [], "source": [ - "# instantiate Derham object\n", - "Nel = [32, 72, 1]\n", - "p = [3, 3, 1]\n", - "spl_kind = [False, True, True]\n", - "derham = Derham(Nel=Nel, p=p, spl_kind=spl_kind)\n", + "time_opts = Time(dt=0.2, Tend=3000, split_algo=\"Strang\")\n", "\n", - "# instantiate a projected MHD equilibrium object\n", - "proj_equil = ProjectedMHDequilibrium(equil, derham)" + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " base_units=base_units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" ] }, { "cell_type": "code", "execution_count": null, - "id": "5b8f74f6", + "id": "aae7877e", "metadata": {}, "outputs": [], "source": [ - "# pass simulation parameters to Propagator class\n", - "PushEta.domain = domain\n", - "PushVxB.domain = domain\n", - "PushVxB.derham = derham" + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "main.pproc(path)\n", + "\n", + "simdata = main.load_data(path)" ] }, { "cell_type": "code", "execution_count": null, - "id": "81b357b2", + "id": "ae248ef1", "metadata": {}, "outputs": [], "source": [ - "# instantiate Propagator object\n", - "prop_eta = PushEta(particles)\n", - "prop_vxB = PushVxB(particles, b2=proj_equil.b2)" + "fig = plt.figure()\n", + "ax = fig.gca()\n", + "\n", + "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", + "\n", + "time = 0.\n", + "dt = time_opts.dt\n", + "Tend = time_opts.Tend\n", + "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", + " # print(f\"{v[0] = }\")\n", + " alpha = (Tend - time)/Tend\n", + " for i, particle in enumerate(v):\n", + " ax.scatter(particle[0], particle[1], c=colors[i % 4], alpha=alpha)\n", + " time += dt\n", + " \n", + "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", + "\n", + "ax.add_patch(circle1)\n", + "ax.set_aspect('equal')\n", + "ax.set_xlabel('x')\n", + "ax.set_ylabel('y')\n", + "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" ] }, { diff --git a/src/struphy/main.py b/src/struphy/main.py index 7b64d8444..61b0bf794 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -10,6 +10,7 @@ import shutil import pickle import h5py +import copy from struphy.fields_background.base import FluidEquilibriumWithB from struphy.io.output_handling import DataContainer @@ -72,6 +73,7 @@ def run( # check model assert hasattr(model, "propagators"), "Attribute 'self.propagators' must be set in model __init__!" model_name = model.__class__.__name__ + model.verbose = verbose if rank == 0: print(f"\n*** Starting run for model '{model_name}':") @@ -128,6 +130,9 @@ def run( pickle.dump(time_opts, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "domain.bin"), 'wb') as f: # WORKAROUND: cannot pickle pyccelized classes at the moment + # p_copy = copy.deepcopy(domain.params_map) + # p_copy.pop("equilibrium") + # print(f"{domain.__class__.__name__ = }, {p_copy = }") tmp_dct = {"name": domain.__class__.__name__, "params_map": domain.params_map} pickle.dump(tmp_dct, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "equil.bin"), 'wb') as f: @@ -176,20 +181,21 @@ def run( # domain and fluid bckground model.setup_domain_and_equil(domain, equil) + # default grid + if grid is None: + Nel = (16, 16, 16) + print(f"\nNo grid specified - using TensorProductGrid with {Nel = }.") + grid = grids.TensorProductGrid(Nel=Nel) + # allocate derham-related objects if derham_opts is not None: - assert grid is not None, f"Derham complex needs a grid." model.allocate_feec(grid, derham_opts) else: - if grid is None: - Nel = (16, 16, 16) - print(f"\nNo grid specified - using TensorProductGrid with {Nel = }.") - grid = grids.TensorProductGrid(Nel=Nel) p = (3, 3, 3) spl_kind = (False, False, False) print(f"\nNo Derham options specified - creating Derham with {p = } and {spl_kind = } for projecting equilibrium.") derham_opts = DerhamOptions(p=p, spl_kind=spl_kind) - model.allocate_feec(grid, derham_opts) + model.allocate_feec(grid, derham_opts) # equation paramters model.setup_equation_params(units=model.units, verbose=verbose) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 9971dc48d..c0f7ba7dc 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -204,7 +204,7 @@ def __init__( # total number of cells (equal to mpi_size if no grid) n_cells = np.sum(np.prod(self.domain_array[:, 2::3], axis=1, dtype=int)) * self.num_clones if verbose: - print(f"{self.mpi_rank = }, {n_cells = }") + print(f"\n{self.mpi_rank = }, {n_cells = }") # total number of boxes if self.boxes_per_dim is None: @@ -219,7 +219,7 @@ def __init__( n_boxes = np.prod(self.boxes_per_dim, dtype=int) * self.num_clones if verbose: - print(f"{self.mpi_rank = }, {n_boxes = }") + print(f"\n{self.mpi_rank = }, {n_boxes = }") # total number of markers (Np) and particles per cell (ppc) Np = self.loading_params.Np @@ -2239,8 +2239,7 @@ def _set_boxes(self): # A particle on box i only sees particles in boxes that belong to neighbours[i] initialize_neighbours(self._neighbours, self.nx, self.ny, self.nz) - if self._verbose: - print(f"{self._rank = }\n{self._neighbours = }") + # print(f"{self._rank = }\n{self._neighbours = }") self._swap_line_1 = np.zeros(self._markers_shape[1]) self._swap_line_2 = np.zeros(self._markers_shape[1]) diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 77c01d432..fc6dfc261 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -544,7 +544,7 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma # loop over time grid for n in tqdm(range(int((nt - 1) / step) + 1)): # clear buffer - temp[:, :] = 0 + temp[:, :] = 0. # create text file for this time step and this species file_npy = os.path.join( @@ -565,6 +565,8 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma # sorting out lost particles ids = temp[:, -1].astype("int") ids_lost_particles = np.setdiff1d(np.arange(n_markers), ids) + ids_removed_particles = np.nonzero(temp[:, 0] == -1.)[0] + ids_lost_particles = np.array(list(set(ids_lost_particles) | set(ids_removed_particles)), dtype=int) lost_particles_mask[:] = False lost_particles_mask[ids_lost_particles] = True @@ -576,10 +578,8 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma assert np.all(sorted(ids) == np.arange(n_markers)) # compute physical positions (x, y, z) - temp[~lost_particles_mask, :3] = domain( - np.array(temp[~lost_particles_mask, :3]), - change_out_order=True, - ) + pos_phys = domain(np.array(temp[~lost_particles_mask, :3]), change_out_order=True) + temp[~lost_particles_mask, :3] = pos_phys # save numpy np.save(file_npy, temp) From 746370e5c90d8910638f122d8feff1a4c8d81d0b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 5 Sep 2025 11:13:26 +0200 Subject: [PATCH 093/292] minor --- src/struphy/fields_background/base.py | 2 +- src/struphy/main.py | 3 --- src/struphy/models/base.py | 2 +- 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index 0d8fe52bb..768cfb7d1 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -40,7 +40,7 @@ def domain(self): @domain.setter def domain(self, new_domain): - assert isinstance(new_domain, Domain) + assert isinstance(new_domain, Domain) or new_domain is None self._domain = new_domain ########################### diff --git a/src/struphy/main.py b/src/struphy/main.py index 61b0bf794..8097bf9ec 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -130,9 +130,6 @@ def run( pickle.dump(time_opts, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "domain.bin"), 'wb') as f: # WORKAROUND: cannot pickle pyccelized classes at the moment - # p_copy = copy.deepcopy(domain.params_map) - # p_copy.pop("equilibrium") - # print(f"{domain.__class__.__name__ = }, {p_copy = }") tmp_dct = {"name": domain.__class__.__name__, "params_map": domain.params_map} pickle.dump(tmp_dct, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "equil.bin"), 'wb') as f: diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 912b44a14..289956852 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -646,7 +646,7 @@ def integrate(self, dt, split_algo="LieTrotter"): # second order in time elif split_algo == "Strang": - assert len(self.propagators) > 1 + assert len(self.prop_list) > 1 for propagator in self.prop_list[:-1]: prop_name = type(propagator).__name__ From 474e87beb24bcf4ab5dee8ccdb8eab6cb3da101b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 5 Sep 2025 14:38:49 +0200 Subject: [PATCH 094/292] * Refactor Domain base class: Replace abstract attributes by regular attributes with setters. Have some arguments in the constructor, in order to avoid calling params_map. The new attribute params holds the original parameters passed to __init__ of the class in domains.py. * Adapt the mapping Tokamak to the new Domain. --- src/struphy/geometry/base.py | 317 +++++++++++++------------------- src/struphy/geometry/domains.py | 77 +++----- 2 files changed, 153 insertions(+), 241 deletions(-) diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index d044a9d0b..8b3803918 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -37,58 +37,42 @@ class Domain(metaclass=ABCMeta): Only right-handed mappings (:math:`\textnormal{det}(DF) > 0`) are admitted. """ - def __init__(self): - # create IGA attributes for IGA mappings - if self.kind_map < 10: - Nel = self.params_map["Nel"] - p = self.params_map["p"] - spl = self.params_map["spl_kind"] - - self._Nel = Nel - self._p = p - self._spl_kind = spl - - self._NbaseN = [Nel + p - kind * p for Nel, p, kind in zip(Nel, p, spl)] - - el_b = [np.linspace(0.0, 1.0, Nel + 1) for Nel in Nel] - - self._T = [bsp.make_knots(el_b, p, kind) for el_b, p, kind in zip(el_b, p, spl)] - - self._indN = [ - (np.indices((Nel, p + 1))[1] + np.arange(Nel)[:, None]) % NbaseN - for Nel, p, NbaseN in zip(Nel, p, self._NbaseN) - ] + def __init__(self, Nel: tuple[int] = None, p: tuple[int] = None, spl_kind: tuple[bool] = None,): + + if Nel is None or p is None or spl_kind is None: + assert self.kind_map >= 10, "Spline mappings must define Nel, p and spl_kind." + Nel = (1, 1, 1) + p = (1, 1, 1) + spl_kind = (True, True, True) + + # create IGA attributes + self._Nel = Nel + self._p = p + self._spl_kind = spl_kind + + self._NbaseN = [Nel + p - kind * p for Nel, p, kind in zip(Nel, p, spl_kind)] + + el_b = [np.linspace(0.0, 1.0, Nel + 1) for Nel in Nel] + + self._T = [bsp.make_knots(el_b, p, kind) for el_b, p, kind in zip(el_b, p, spl_kind)] + + self._indN = [ + (np.indices((Nel, p + 1))[1] + np.arange(Nel)[:, None]) % NbaseN + for Nel, p, NbaseN in zip(Nel, p, self._NbaseN) + ] - # extend to 3d for 2d IGA mappings - if self._kind_map != 0: - self._Nel = (*self._Nel, 0) - self._p = (*self._p, 0) - self._NbaseN = self._NbaseN + [0] + # extend to 3d for 2d IGA mappings + if 0 < self.kind_map < 10: + self._Nel = (*self._Nel, 0) + self._p = (*self._p, 0) + self._NbaseN = self._NbaseN + [0] - self._T = self._T + [np.zeros((1,), dtype=float)] + self._T = self._T + [np.zeros((1,), dtype=float)] - self._indN = self._indN + [np.zeros((1, 1), dtype=int)] + self._indN = self._indN + [np.zeros((1, 1), dtype=int)] # create dummy attributes for analytical mappings - else: - self._Nel = [0, 0, 0] - self._p = [0, 0, 0] - self._spl_kind = [True, True, True] - - self._NbaseN = [0, 0, 0] - - self._T = [ - np.zeros((1,), dtype=float), - np.zeros((1,), dtype=float), - np.zeros((1,), dtype=float), - ] - - self._indN = [ - np.zeros((1, 1), dtype=int), - np.zeros((1, 1), dtype=int), - np.zeros((1, 1), dtype=int), - ] - + if self.kind_map >= 10: self._cx = np.zeros((1, 1, 1), dtype=float) self._cy = np.zeros((1, 1, 1), dtype=float) self._cz = np.zeros((1, 1, 1), dtype=float) @@ -145,38 +129,69 @@ def __init__(self): ) @property - @abstractmethod - def kind_map(self): + def kind_map(self) -> int: """Integer defining the mapping: * <=9: spline mappings * >=10 and <=19: analytical mappings with cubic domain boundary * >=20 and <=29: analytical cylinder and torus mappings * >=30 and <=39: Shafranov mappings (cylinder)""" - pass + if not hasattr(self, "_kind_map"): + raise AttributeError("Must set 'self.kind_map' for mappings.") + return self._kind_map + + @kind_map.setter + def kind_map(self, new): + assert isinstance(new, int) + self._kind_map = new @property - @abstractmethod - def params_map(self): - """Mapping parameters as dictionary.""" - pass + def params(self) -> dict: + """Mapping parameters passed to __init__() of the class in domains.py, as dictionary.""" + if not hasattr(self, "_params"): + self._params = None + return self._params + + @params.setter + def params(self, new): + assert isinstance(new, dict) + self._params = new @property - @abstractmethod - def params_numpy(self): + def params_numpy(self) -> np.ndarray: """Mapping parameters as numpy array (can be empty).""" - pass + if not hasattr(self, "_params_numpy"): + self._params_numpy = np.array([], dtype=float) + return self._params_numpy + + @params_numpy.setter + def params_numpy(self, new): + assert isinstance(new, np.ndarray) + assert new.ndim == 1 + self._params_numpy = new @property - @abstractmethod - def pole(self): + def pole(self) -> bool: """Bool; True if mapping has one polar point.""" - pass + if not hasattr(self, "_pole"): + self._pole = False + return self._pole + + @pole.setter + def pole(self, new): + assert isinstance(new, bool) + self._pole = new @property - @abstractmethod - def periodic_eta3(self): + def periodic_eta3(self) -> bool: r"""Bool; True if mapping is periodic in :math:`\eta_3` coordinate.""" - pass + if not hasattr(self, "_periodic_eta3"): + raise AttributeError("Must specify whether mapping is periodic in eta3.") + return self._periodic_eta3 + + @periodic_eta3.setter + def periodic_eta3(self, new): + assert isinstance(new, bool) + self._periodic_eta3 = new @property def cx(self): @@ -1379,8 +1394,8 @@ def prepare_params_map(params_user, params_default, return_numpy=True): return params_map @staticmethod - def prepare_params_map_new(return_numpy=True, **params): - """Create parameter dictionary and numpy array. + def get_params_dict(return_numpy: bool = True, **params) -> dict: + """Create parameter dictionary and (optional) numpy array for pyccel kernels. Parameters ---------- @@ -1396,7 +1411,7 @@ def prepare_params_map_new(return_numpy=True, **params): Dictionary with default values for missing keys. params_numpy : np.ndarray - Numpy array with parameters in params_map (if parameter is a float or int). + Numpy array with parameters in params (if parameter is a float or int). """ # parameter numpy array for pyccel kernel (order matters!) if return_numpy: @@ -1488,7 +1503,7 @@ def show( ax.scatter(X[0, 32], Y[0, 32], 20, "red", zorder=10) tstr = "" - for key, val in self.params_map.items(): + for key, val in self.params.items(): if key not in {"cx", "cy", "cz"}: tstr += key + ": " + str(val) + "\n" ax.set_title(self.__class__.__name__ + " at $\\eta_3=0$") @@ -1875,26 +1890,6 @@ def __init__(self, **params): super().__init__() - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 - class PoloidalSpline(Domain): r"""Base class for all mappings that use a 2D spline representation @@ -1911,24 +1906,14 @@ class PoloidalSpline(Domain): The full map :math:`F: [0, 1]^3 \to \Omega` is obtained by defining :math:`(R, Z, \eta_3) \mapsto (x, y, z)` in the child class. """ - def __init__(self, **params): - # set default - params_default = { - "Nel": [8, 24], - "p": [2, 3], - "spl_kind": [False, True], - "cx": None, - "cy": None, - } - - params_map = Domain.prepare_params_map( - params, - params_default, - return_numpy=False, - ) + def __init__(self, Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + spl_kind: tuple[bool] = (False, True), + cx: np.ndarray = None, + cy: np.ndarray = None,): # get default control points - if params_map["cx"] is None or params_map["cy"] is None: + if cx is None or cy is None: def X(eta1, eta2): return eta1 * np.cos(2 * np.pi * eta2) + 3.0 @@ -1936,48 +1921,31 @@ def X(eta1, eta2): def Y(eta1, eta2): return eta1 * np.sin(2 * np.pi * eta2) - cx, cy = interp_mapping( - params_map["Nel"], - params_map["p"], - params_map["spl_kind"], - X, - Y, - ) + cx, cy = interp_mapping(Nel, p, spl_kind, X, Y) # make sure that control points at pole are all the same (eta1=0 there) cx[0] = 3.0 cy[0] = 0.0 - # add control points to parameters dictionary - params_map["cx"] = cx - params_map["cy"] = cy - - self._params_map = params_map - # set control point properties - self._cx = self._params_map["cx"] - self._cy = self._params_map["cy"] + self._cx = cx + self._cy = cy # make sure that control points are 2D assert self.cx.ndim == 2 assert self.cy.ndim == 2 # make sure that control points are compatible with given spline data - expected_shape = tuple( - [ - self._params_map["Nel"][n] + (not self._params_map["spl_kind"][n]) * self._params_map["p"][n] - for n in range(2) - ] - ) + expected_shape = tuple([Nel[n] + (not spl_kind[n]) * p[n] for n in range(2)]) assert self.cx.shape == expected_shape assert self.cy.shape == expected_shape # identify polar singularity at eta1=0 if np.all(self.cx[0, :] == self.cx[0, 0]): - self._pole = True + self.pole = True else: - self._pole = False + self.pole = False # reshape control points to 3D self._cx = self.cx[:, :, None] @@ -1985,27 +1953,7 @@ def Y(eta1, eta2): self._cz = np.zeros((1, 1, 1), dtype=float) # init base class - super().__init__() - - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 + super().__init__(Nel=Nel, p=p, spl_kind=spl_kind) class PoloidalSplineStraight(PoloidalSpline): @@ -2079,7 +2027,8 @@ def Y(eta1, eta2): class PoloidalSplineTorus(PoloidalSpline): - r"""TODO + r"""Torus where the poloidal planes are described by a 2D IGA-spline mapping. + .. math:: F: (R, Z, \eta_3) \mapsto (x, y, z) \textnormal{ as } \left\{\begin{aligned} @@ -2089,29 +2038,41 @@ class PoloidalSplineTorus(PoloidalSpline): z &= Z \,. \end{aligned}\right. + + Parameters + ---------- + Nel : tuple[int] + Number of elements in each poloidal direction. + + p : tuple[int] + Spline degree in each poloidal direction. + + spl_kind : tuple[bool] + Kind of spline in each poloidal direction (True=periodic, False=clamped). + + cx, cy : np.ndarray + Control points (spline coefficients) of the poloidal mapping. + If None, a default square-to-disc mapping of radius 1 centered around (x, y) = (3, 0) is interpolated. + + tor_period : int + The toroidal angle is between [0, 2*pi/tor_period). """ - def __init__(self, **params): - self._kind_map = 2 - - # set default - params_default = { - "Nel": [8, 24], - "p": [2, 3], - "spl_kind": [False, True], - "cx": None, - "cy": None, - "tor_period": 3, - } - - params_map = Domain.prepare_params_map( - params, - params_default, - return_numpy=False, - ) + def __init__(self, Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + spl_kind: tuple[bool] = (False, True), + cx: np.ndarray = None, + cy: np.ndarray = None, + tor_period: int = 3, + ): + + # use setters for mapping attributes + self.kind_map = 2 + self.params_numpy = np.array([float(tor_period)]) + self.periodic_eta3 = True # get default control points - if params_map["cx"] is None or params_map["cy"] is None: + if cx is None or cy is None: def X(eta1, eta2): return eta1 * np.cos(2 * np.pi * eta2) + 3.0 @@ -2119,33 +2080,15 @@ def X(eta1, eta2): def Y(eta1, eta2): return eta1 * np.sin(2 * np.pi * eta2) - cx, cy = interp_mapping( - params_map["Nel"], - params_map["p"], - params_map["spl_kind"], - X, - Y, - ) + cx, cy = interp_mapping(Nel, p, spl_kind, X, Y) # make sure that control points at pole are all 0 (eta1=0 there) cx[0] = 3.0 cy[0] = 0.0 - # add control points to parameters dictionary - params_map["cx"] = cx - params_map["cy"] = cy - - self._params_numpy = np.array([float(params_map["tor_period"])]) - self._periodic_eta3 = True - - # remove "tor_period" temporarily from params_map dictionary (is not a parameter of PoloidalSpline) - tor_period = params_map["tor_period"] - params_map.pop("tor_period") - # init base class - super().__init__(**params_map) + super().__init__(Nel=Nel, p=p, spl_kind=spl_kind, cx=cx, cy=cy,) - self._params_map["tor_period"] = tor_period def interp_mapping(Nel, p, spl_kind, X, Y, Z=None): diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index 2e142fa1f..2fe9a1468 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -2,6 +2,7 @@ import numpy as np +from struphy.fields_background.equils import EQDSKequilibrium from struphy.fields_background.base import AxisymmMHDequilibrium from struphy.geometry.base import Domain, PoloidalSplineStraight, PoloidalSplineTorus, Spline from struphy.geometry.utilities import field_line_tracing @@ -66,15 +67,13 @@ def __init__( p_pre: tuple = (3, 3), tor_period: int = 1, ): - from struphy.fields_background.equils import EQDSKequilibrium - # default MHD equilibrium if equilibrium is None: equilibrium = EQDSKequilibrium() else: assert isinstance(equilibrium, AxisymmMHDequilibrium) - params_map = Domain.prepare_params_map_new( + self.params = Domain.get_params_dict( equilibrium=equilibrium, Nel=Nel, p=p, @@ -88,64 +87,34 @@ def __init__( return_numpy=False, ) - # get control points via field tracing beetwenn fluxes [psi_s, psi_e)] - # flux boundaries of mapping - eq_mhd = params_map["equilibrium"] + # get control points via field tracing between fluxes [psi_s, psi_e] + psi0, psi1 = equilibrium.psi_range[0], equilibrium.psi_range[1] - psi0, psi1 = eq_mhd.psi_range[0], eq_mhd.psi_range[1] - - psi_s = psi0 + params_map["psi_shifts"][0] * 0.01 * (psi1 - psi0) - psi_e = psi1 - params_map["psi_shifts"][1] * 0.01 * (psi1 - psi0) + psi_s = psi0 + self.params["psi_shifts"][0] * 0.01 * (psi1 - psi0) + psi_e = psi1 - self.params["psi_shifts"][1] * 0.01 * (psi1 - psi0) cx, cy = field_line_tracing( - eq_mhd.psi, - eq_mhd.psi_axis_RZ[0], - eq_mhd.psi_axis_RZ[1], + equilibrium.psi, + equilibrium.psi_axis_RZ[0], + equilibrium.psi_axis_RZ[1], psi_s, psi_e, - params_map["Nel"], - params_map["p"], - psi_power=params_map["psi_power"], - xi_param=params_map["xi_param"], - Nel_pre=params_map["Nel_pre"], - p_pre=params_map["p_pre"], - r0=params_map["r0"], + self.params["Nel"], + self.params["p"], + psi_power=self.params["psi_power"], + xi_param=self.params["xi_param"], + Nel_pre=self.params["Nel_pre"], + p_pre=self.params["p_pre"], + r0=self.params["r0"], ) - # add control points to parameters dictionary - params_map["cx"] = cx - params_map["cy"] = cy - - # add spline types to parameters dictionary - params_map["spl_kind"] = [False, True] - - # remove temporarily parameters from params_map dictionary which are not parameters of PoloidalSplineTorus - equilibrium = params_map["equilibrium"] - psi_power = params_map["psi_power"] - psi_shifts = params_map["psi_shifts"] - xi_param = params_map["xi_param"] - r0 = params_map["r0"] - Nel_pre = params_map["Nel_pre"] - p_pre = params_map["p_pre"] - - params_map.pop("equilibrium") - params_map.pop("psi_power") - params_map.pop("psi_shifts") - params_map.pop("xi_param") - params_map.pop("r0") - params_map.pop("Nel_pre") - params_map.pop("p_pre") - # init base class - super().__init__(**params_map) - - self._params_map["equilibrium"] = equilibrium - self._params_map["psi_power"] = psi_power - self._params_map["psi_shifts"] = psi_shifts - self._params_map["xi_param"] = xi_param - self._params_map["r0"] = r0 - self._params_map["Nel_pre"] = Nel_pre - self._params_map["p_pre"] = p_pre + super().__init__(Nel=self.params["Nel"], + p=self.params["p"], + spl_kind=(False, True), + cx=cx, + cy=cy, + tor_period=self.params["tor_period"],) class GVECunit(Spline): @@ -796,7 +765,7 @@ def __init__( ): self._kind_map = 20 - self._params_map, self._params_numpy = Domain.prepare_params_map_new( + self._params_map, self._params_numpy = Domain.get_params_dict( a1=a1, a2=a2, Lz=Lz, From d0a3460df517238ffc5ff19374b2139768b1a7ba Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 5 Sep 2025 15:03:45 +0200 Subject: [PATCH 095/292] adapted GVECunit --- src/struphy/fields_background/base.py | 2 +- src/struphy/geometry/base.py | 68 ++++++++++----------------- src/struphy/geometry/domains.py | 41 ++++++---------- 3 files changed, 39 insertions(+), 72 deletions(-) diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index 0d8fe52bb..0a5bce84f 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -1435,5 +1435,5 @@ def numerical_domain(self): pass @property - def domain(self): + def domain(self) -> Domain: return self.numerical_domain diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index 8b3803918..e0b50a46d 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -160,7 +160,7 @@ def params(self, new): def params_numpy(self) -> np.ndarray: """Mapping parameters as numpy array (can be empty).""" if not hasattr(self, "_params_numpy"): - self._params_numpy = np.array([], dtype=float) + self._params_numpy = np.array([0], dtype=float) return self._params_numpy @params_numpy.setter @@ -1812,7 +1812,8 @@ def show( class Spline(Domain): - r"""TODO + r"""3D IGA spline mapping. + .. math:: F: (\eta_1, \eta_2, \eta_3) \mapsto (x, y, z) \textnormal{ as } \left\{\begin{aligned} @@ -1824,40 +1825,27 @@ class Spline(Domain): \end{aligned}\right. """ - def __init__(self, **params): - self._kind_map = 0 - - # set default - params_default = { - "Nel": None, - "p": None, - "spl_kind": None, - "cx": None, - "cy": None, - "cz": None, - } - - params_map = Domain.prepare_params_map( - params, - params_default, - return_numpy=False, - ) + def __init__(self, Nel: tuple[int] = (8, 24, 6), + p: tuple[int] = (2, 3, 1), + spl_kind: tuple[bool] = (False, True, True), + cx: np.ndarray = None, + cy: np.ndarray = None, + cz: np.ndarray = None,): + + self.kind_map = 0 # get default control points from default GVEC equilibrium - if params_map["cx"] is None or params_map["cy"] is None or params_map["cz"] is None: + if cx is None or cy is None or cz is None: from struphy.fields_background.equils import GVECequilibrium - mhd_equil = GVECequilibrium() - - for key in params_map.keys(): - params_map[key] = getattr(mhd_equil.domain, key) - - self._params_map = params_map + cx = mhd_equil.domain.cx + cx = mhd_equil.domain.cy + cx = mhd_equil.domain.cz # assign control points - self._cx = self._params_map["cx"] - self._cy = self._params_map["cy"] - self._cz = self._params_map["cz"] + self._cx = cx + self._cy = cy + self._cz = cz # check dimensions assert self.cx.ndim == 3 @@ -1865,12 +1853,7 @@ def __init__(self, **params): assert self.cz.ndim == 3 # make sure that control points are compatible with given spline data - expected_shape = tuple( - [ - self._params_map["Nel"][n] + (not self._params_map["spl_kind"][n]) * self._params_map["p"][n] - for n in range(3) - ] - ) + expected_shape = tuple([Nel[n] + (not spl_kind[n]) * p[n] for n in range(3)]) assert self.cx.shape == expected_shape assert self.cy.shape == expected_shape @@ -1878,17 +1861,14 @@ def __init__(self, **params): # identify polar singularity at eta1=0 if np.all(self.cx[0, :, 0] == self.cx[0, 0, 0]): - self._pole = True + self.pole = True else: - self._pole = False - - self._periodic_eta3 = self._params_map["spl_kind"][-1] + self.pole = False - # init base class - # Added one element so we are not passing empty numpy arrays - self._params_numpy = np.array([0.0]) + self.periodic_eta3 = spl_kind[-1] - super().__init__() + # base class + super().__init__(Nel=Nel, p=p, spl_kind=spl_kind) class PoloidalSpline(Domain): diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index 2fe9a1468..e72554991 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -4,9 +4,11 @@ from struphy.fields_background.equils import EQDSKequilibrium from struphy.fields_background.base import AxisymmMHDequilibrium -from struphy.geometry.base import Domain, PoloidalSplineStraight, PoloidalSplineTorus, Spline +from struphy.geometry.base import Domain, PoloidalSplineStraight, PoloidalSplineTorus, Spline, interp_mapping from struphy.geometry.utilities import field_line_tracing +import gvec + class Tokamak(PoloidalSplineTorus): r"""Mappings for Tokamak MHD equilibria constructed via :ref:`field-line tracing ` of a poloidal flux function :math:`\psi`. @@ -73,6 +75,7 @@ def __init__( else: assert isinstance(equilibrium, AxisymmMHDequilibrium) + # use the params setter self.params = Domain.get_params_dict( equilibrium=equilibrium, Nel=Nel, @@ -136,25 +139,23 @@ class GVECunit(Spline): """ def __init__(self, gvec_equil=None): - import gvec - + from struphy.fields_background.equils import GVECequilibrium - from struphy.geometry.base import interp_mapping - + if gvec_equil is None: gvec_equil = GVECequilibrium() else: assert isinstance(gvec_equil, GVECequilibrium) - params_map = { - "Nel": gvec_equil.params["Nel"], - "p": gvec_equil.params["p"], - } + # use the params setter + self.params = Domain.get_params_dict(return_numpy=False, gvec_equil=gvec_equil) + Nel = gvec_equil.params["Nel"] + p =gvec_equil.params["p"] if gvec_equil.params["use_nfp"]: - params_map["spl_kind"] = (False, True, False) + spl_kind = (False, True, False) else: - params_map["spl_kind"] = (False, True, True) + spl_kind = (False, True, True) # project mapping to splines _rmin = gvec_equil.params["rmin"] @@ -182,23 +183,9 @@ def Y(e1, e2, e3): def Z(e1, e2, e3): return XYZ(e1, e2, e3)[2] - cx, cy, cz = interp_mapping( - params_map["Nel"], - params_map["p"], - params_map["spl_kind"], - X, - Y, - Z, - ) + cx, cy, cz = interp_mapping(Nel, p, spl_kind, X, Y, Z) - params_map["cx"] = cx - params_map["cy"] = cy - params_map["cz"] = cz - - super().__init__(**params_map) - - self._params_map["rmin"] = _rmin - self._params_map["equilibrium"] = gvec_equil + super().__init__(Nel=Nel, p=p, spl_kind=spl_kind, cx=cx, cy=cy, cz=cz) class DESCunit(Spline): From 7e9e98016989c5ad9b1e2249077c305ba5b1f2e6 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 5 Sep 2025 16:05:58 +0200 Subject: [PATCH 096/292] use locals() to get parameter dict; adapt some more mappings --- src/struphy/geometry/base.py | 141 ++++------------- src/struphy/geometry/domains.py | 267 ++++++++++---------------------- 2 files changed, 105 insertions(+), 303 deletions(-) diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index e0b50a46d..2c8eed6df 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -154,6 +154,10 @@ def params(self) -> dict: @params.setter def params(self, new): assert isinstance(new, dict) + if "self" in new: + new.pop("self") + if "__class__" in new: + new.pop("__class__") self._params = new @property @@ -1349,80 +1353,14 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): return a_out # ================================ - @staticmethod - def prepare_params_map(params_user, params_default, return_numpy=True): - """Sets missing default key-value pairs in dictionary "params_user" according to "params_default". - - Parameters - ---------- - params_user : dict - Dictionary which is compared to the dictionary "params_default" and to which missing defaults are added. - - params_default : dict - Dictionary with default values. - - return_numpy : bool - Whether to return a numpy parameter array in addition to the always returned parameter dictionary. - - Returns - ------- - params_map : dict - Dictionary with same keys as "params_default" and default values for missing keys. - - params_numpy : np.ndarray - Numpy array with parameters in params_map (if parameter is a float or int). - """ - - # check for correct keys in params_user - for key in params_user: - assert key in params_default, f'Unknown key "{key}". Please choose one of {[*params_default]}.' - - # set default values if key is missing - params_map = params_user - - for key, val in params_default.items(): - params_map.setdefault(key, val) - - # parameter numpy array for pyccel kernel (order matters!) - if return_numpy: - params_numpy = [] - for key in params_default.keys(): - params_numpy.append(params_map[key]) - - return params_map, np.array(params_numpy) - else: - return params_map - - @staticmethod - def get_params_dict(return_numpy: bool = True, **params) -> dict: - """Create parameter dictionary and (optional) numpy array for pyccel kernels. - - Parameters - ---------- - return_numpy : bool - Whether to return a numpy parameter array in addition to the always returned parameter dictionary. - - params : kwargs - Mapping parameters. - - Returns - ------- - params : dict - Dictionary with default values for missing keys. - - params_numpy : np.ndarray - Numpy array with parameters in params (if parameter is a float or int). - """ - # parameter numpy array for pyccel kernel (order matters!) - if return_numpy: - params_numpy = [] - for k, v in params.items(): - params_numpy.append(v) - return params, np.array(params_numpy) - else: - return params - - # ================================ + + def get_params_numpy(self) -> np.ndarray: + """Convert parameter dict into numpy array.""" + params_numpy = [] + for k, v in self.params.items(): + params_numpy.append(v) + return np.array(params_numpy) + def show( self, logical=False, @@ -1937,7 +1875,8 @@ def Y(eta1, eta2): class PoloidalSplineStraight(PoloidalSpline): - r"""TODO + r"""Cylinder where the poloidal planes are described by a 2D IGA-spline mapping. + .. math:: F: (R, Z, \eta_3) \mapsto (x, y, z) \textnormal{ as } \left\{\begin{aligned} @@ -1949,27 +1888,17 @@ class PoloidalSplineStraight(PoloidalSpline): \end{aligned}\right. """ - def __init__(self, **params): - self._kind_map = 1 - - # set default - params_default = { - "Nel": [8, 24], - "p": [2, 3], - "spl_kind": [False, True], - "cx": None, - "cy": None, - "Lz": 4.0, - } - - params_map = Domain.prepare_params_map( - params, - params_default, - return_numpy=False, - ) + def __init__(self, Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + spl_kind: tuple[bool] = (False, True), + cx: np.ndarray = None, + cy: np.ndarray = None, + Lz: float = 4.0,): + + self.kind_map = 1 # get default control points - if params_map["cx"] is None or params_map["cy"] is None: + if cx is None or cy is None: def X(eta1, eta2): return eta1 * np.cos(2 * np.pi * eta2) @@ -1977,33 +1906,17 @@ def X(eta1, eta2): def Y(eta1, eta2): return eta1 * np.sin(2 * np.pi * eta2) - cx, cy = interp_mapping( - params_map["Nel"], - params_map["p"], - params_map["spl_kind"], - X, - Y, - ) + cx, cy = interp_mapping(Nel, p, spl_kind, X, Y) # make sure that control points at pole are all 0 (eta1=0 there) cx[0] = 0.0 cy[0] = 0.0 - # add control points to parameters dictionary - params_map["cx"] = cx - params_map["cy"] = cy - - self._params_numpy = np.array([params_map["Lz"]]) - self._periodic_eta3 = False - - # remove "Lz" temporarily from params_map dictionary (is not a parameter of PoloidalSpline) - Lz = params_map["Lz"] - params_map.pop("Lz") + self.params_numpy = np.array([Lz]) + self.periodic_eta3 = False # init base class - super().__init__(**params_map) - - self._params_map["Lz"] = Lz + super().__init__(Nel=Nel, p=p, spl_kind=spl_kind, cx=cx, cy=cy) class PoloidalSplineTorus(PoloidalSpline): diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index e72554991..11da1e4dd 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -1,10 +1,15 @@ "Mapped domains (single patch)." import numpy as np +import copy from struphy.fields_background.equils import EQDSKequilibrium from struphy.fields_background.base import AxisymmMHDequilibrium -from struphy.geometry.base import Domain, PoloidalSplineStraight, PoloidalSplineTorus, Spline, interp_mapping +from struphy.geometry.base import (Domain, + PoloidalSplineStraight, + PoloidalSplineTorus, + Spline, + interp_mapping,) from struphy.geometry.utilities import field_line_tracing import gvec @@ -76,25 +81,13 @@ def __init__( assert isinstance(equilibrium, AxisymmMHDequilibrium) # use the params setter - self.params = Domain.get_params_dict( - equilibrium=equilibrium, - Nel=Nel, - p=p, - psi_power=psi_power, - psi_shifts=psi_shifts, - xi_param=xi_param, - r0=r0, - Nel_pre=Nel_pre, - p_pre=p_pre, - tor_period=tor_period, - return_numpy=False, - ) + self.params = copy.deepcopy(locals()) # get control points via field tracing between fluxes [psi_s, psi_e] psi0, psi1 = equilibrium.psi_range[0], equilibrium.psi_range[1] - psi_s = psi0 + self.params["psi_shifts"][0] * 0.01 * (psi1 - psi0) - psi_e = psi1 - self.params["psi_shifts"][1] * 0.01 * (psi1 - psi0) + psi_s = psi0 + psi_shifts[0] * 0.01 * (psi1 - psi0) + psi_e = psi1 - psi_shifts[1] * 0.01 * (psi1 - psi0) cx, cy = field_line_tracing( equilibrium.psi, @@ -102,22 +95,22 @@ def __init__( equilibrium.psi_axis_RZ[1], psi_s, psi_e, - self.params["Nel"], - self.params["p"], - psi_power=self.params["psi_power"], - xi_param=self.params["xi_param"], - Nel_pre=self.params["Nel_pre"], - p_pre=self.params["p_pre"], - r0=self.params["r0"], + Nel, + p, + psi_power=psi_power, + xi_param=xi_param, + Nel_pre=Nel_pre, + p_pre=p_pre, + r0=r0, ) # init base class - super().__init__(Nel=self.params["Nel"], - p=self.params["p"], + super().__init__(Nel=Nel, + p=p, spl_kind=(False, True), cx=cx, cy=cy, - tor_period=self.params["tor_period"],) + tor_period=tor_period,) class GVECunit(Spline): @@ -148,7 +141,7 @@ def __init__(self, gvec_equil=None): assert isinstance(gvec_equil, GVECequilibrium) # use the params setter - self.params = Domain.get_params_dict(return_numpy=False, gvec_equil=gvec_equil) + self.params = self.params = copy.deepcopy(locals()) Nel = gvec_equil.params["Nel"] p =gvec_equil.params["p"] @@ -209,25 +202,22 @@ class DESCunit(Spline): def __init__(self, desc_equil=None): from struphy.fields_background.equils import DESCequilibrium - from struphy.geometry.base import interp_mapping if desc_equil is None: desc_equil = DESCequilibrium() else: assert isinstance(desc_equil, DESCequilibrium) - # expose to methods - self._desc_equil = desc_equil + # use params setter + self.params = copy.deepcopy(locals()) - params_map = { - "Nel": desc_equil.params["Nel"], - "p": desc_equil.params["p"], - } + Nel = desc_equil.params["Nel"] + p = desc_equil.params["p"] if desc_equil.eq.NFP > 1 and desc_equil.use_nfp: - params_map["spl_kind"] = (False, True, False) + spl_kind = (False, True, False) else: - params_map["spl_kind"] = (False, True, True) + spl_kind = (False, True, True) _rmin = desc_equil.params["rmin"] @@ -245,27 +235,13 @@ def Y(e1, e2, e3): def Z(e1, e2, e3): return desc_equil.desc_eval("Z", e1, e2, e3, nfp=nfp) - cx, cy, cz = interp_mapping( - params_map["Nel"], - params_map["p"], - params_map["spl_kind"], - X, - Y, - Z, - ) - - params_map["cx"] = cx - params_map["cy"] = cy - params_map["cz"] = cz - - super().__init__(**params_map) + cx, cy, cz = interp_mapping(Nel, p, spl_kind, X, Y, Z) - self._params_map["rmin"] = _rmin - self._params_map["equilibrium"] = desc_equil + super().__init__(Nel=Nel, p=p, spl_kind=spl_kind, cx=cx, cy=cy, cz=cz) class IGAPolarCylinder(PoloidalSplineStraight): - r""" A cylinder with the cross section approximated by a spline mapping. + r"""A cylinder with the cross section approximated by a spline mapping. .. math:: @@ -300,52 +276,33 @@ class IGAPolarCylinder(PoloidalSplineStraight): a : 1. # minor radius """ - def __init__(self, **params): - from struphy.geometry.base import interp_mapping + def __init__(self, + Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + a: float = 1.0, + Lz: float = 4.0, + ): - # set default - params_default = {"Nel": [8, 24], "p": [2, 3], "a": 1.0, "Lz": 4.0} - - params_map = Domain.prepare_params_map( - params, - params_default, - return_numpy=False, - ) + # use params setter + self.params = copy.deepcopy(locals()) # get control points def X(eta1, eta2): - return params_map["a"] * eta1 * np.cos(2 * np.pi * eta2) + return a * eta1 * np.cos(2 * np.pi * eta2) def Y(eta1, eta2): - return params_map["a"] * eta1 * np.sin(2 * np.pi * eta2) - - cx, cy = interp_mapping( - params_map["Nel"], - params_map["p"], - [False, True], - X, - Y, - ) + return a * eta1 * np.sin(2 * np.pi * eta2) + + spl_kind = (False, True) + + cx, cy = interp_mapping(Nel, p, spl_kind, X, Y) # make sure that control points at pole are all the same (eta1=0 there) cx[0] = 0.0 cy[0] = 0.0 - # add control points to parameters dictionary - params_map["cx"] = cx - params_map["cy"] = cy - - # add spline types to parameters dictionary - params_map["spl_kind"] = [False, True] - - # remove "a" temporarily from params_map dictionary (is not a parameter of PoloidalSplineStraight) - a = params_map["a"] - params_map.pop("a") - # init base class - super().__init__(**params_map) - - self._params_map["a"] = a + super().__init__(Nel=Nel, p=p, spl_kind=spl_kind, cx=cx, cy=cy, Lz=Lz) class IGAPolarTorus(PoloidalSplineTorus): @@ -364,9 +321,9 @@ class IGAPolarTorus(PoloidalSplineTorus): Parameters ---------- - Nel : list[int] + Nel : tuple[int] Number of cells in (radial, angular) direction used for spline mapping (default: [8, 24]). - p : list[int] + p : tuple[int] Splines degrees in (radial, angular) direction used for spline mapping (default: [2, 3]). a : float Minor radius of torus (default: 1.). @@ -392,84 +349,40 @@ class IGAPolarTorus(PoloidalSplineTorus): sfl : False # whether to use straight field line coordinates (particular theta parametrization) """ - def __init__(self, **params): - from struphy.geometry.base import interp_mapping + def __init__(self, Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + a: float = 1.0, + R0: float = 3.0, + sfl: bool = False, + tor_period: int = 3,): - # set default - params_default = { - "Nel": [8, 24], - "p": [ - 2, - 3, - ], - "a": 1.0, - "R0": 3.0, - "sfl": False, - "tor_period": 3, - } - - params_map = Domain.prepare_params_map( - params, - params_default, - return_numpy=False, - ) + # use params setter + self.params = copy.deepcopy(locals()) # get control points - if params_map["sfl"]: - + if sfl: def theta(eta1, eta2): - return 2 * np.arctan( - np.sqrt( - (1 + params_map["a"] * eta1 / params_map["R0"]) - / (1 - params_map["a"] * eta1 / params_map["R0"]) - ) - * np.tan(np.pi * eta2) - ) + return 2 * np.arctan(np.sqrt((1 + a * eta1 / R0) / (1 - a * eta1 / R0)) * np.tan(np.pi * eta2)) else: - def theta(eta1, eta2): return 2 * np.pi * eta2 def R(eta1, eta2): - return params_map["a"] * eta1 * np.cos(theta(eta1, eta2)) + params_map["R0"] + return a * eta1 * np.cos(theta(eta1, eta2)) + R0 def Z(eta1, eta2): - return params_map["a"] * eta1 * np.sin(theta(eta1, eta2)) - - cx, cy = interp_mapping( - params_map["Nel"], - params_map["p"], - [False, True], - R, - Z, - ) + return a * eta1 * np.sin(theta(eta1, eta2)) - # make sure that control points at pole are all the same (eta1=0 there) - cx[0] = params_map["R0"] - cy[0] = 0.0 + spl_kind = (False, True) - # add control points to parameters dictionary - params_map["cx"] = cx - params_map["cy"] = cy + cx, cy = interp_mapping(Nel, p, spl_kind, R, Z) - # add spline types to parameters dictionary - params_map["spl_kind"] = [False, True] - - # remove "a", "R0" and "sfl" temporarily from params_map dictionary (is not a parameter of PoloidalSplineTorus) - a = params_map["a"] - R0 = params_map["R0"] - sfl = params_map["sfl"] - - params_map.pop("a") - params_map.pop("R0") - params_map.pop("sfl") + # make sure that control points at pole are all the same (eta1=0 there) + cx[0] = R0 + cy[0] = 0.0 # init base class - super().__init__(**params_map) - - self._params_map["a"] = a - self._params_map["R0"] = R0 - self._params_map["sfl"] = sfl + super().__init__(Nel=Nel, p=p, spl_kind=spl_kind, cx=cx, cy=cy, tor_period=tor_period,) class Cuboid(Domain): @@ -514,50 +427,26 @@ class Cuboid(Domain): r3 : 1. # end of z-interval, r3>l3 """ - def __init__(self, **params): + def __init__(self, l1: float = 0.0, + r1: float = 1.0, + l2: float = 0.0, + r2: float = 1.0, + l3: float = 0.0, + r3: float = 1.0, + ): + self._kind_map = 10 - # set default parameters and remove wrong/not needed keys - params_default = { - "l1": 0.0, - "r1": 2.0, - "l2": 0.0, - "r2": 3.0, - "l3": 0.0, - "r3": 6.0, - } - - self._params_map, self._params_numpy = Domain.prepare_params_map( - params, - params_default, - ) + # use params setter + self.params = copy.deepcopy(locals()) # periodicity in eta3-direction and pole at eta1=0 - self._periodic_eta3 = False - self._pole = False + self.params_numpy = self.get_params_numpy() + self.periodic_eta3 = False + self.pole = False super().__init__() - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 - class Orthogonal(Domain): r""" Slab geometry with orthogonal mesh distortion. @@ -752,7 +641,7 @@ def __init__( ): self._kind_map = 20 - self._params_map, self._params_numpy = Domain.get_params_dict( + self._params_map, self._params_numpy = Domain.get_params_numpy( a1=a1, a2=a2, Lz=Lz, From d306f5021c75fa80102b50fbc27c7badfc5453eb Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 6 Sep 2025 11:56:42 +0200 Subject: [PATCH 097/292] adapt rest of mappings --- src/struphy/geometry/base.py | 2 +- src/struphy/geometry/domains.py | 387 ++++++---------------- src/struphy/geometry/tests/test_domain.py | 20 +- 3 files changed, 116 insertions(+), 293 deletions(-) diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index 2c8eed6df..52234f251 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -148,7 +148,7 @@ def kind_map(self, new): def params(self) -> dict: """Mapping parameters passed to __init__() of the class in domains.py, as dictionary.""" if not hasattr(self, "_params"): - self._params = None + self._params = {} return self._params @params.setter diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index 11da1e4dd..394ed8f1b 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -140,8 +140,7 @@ def __init__(self, gvec_equil=None): else: assert isinstance(gvec_equil, GVECequilibrium) - # use the params setter - self.params = self.params = copy.deepcopy(locals()) + # do not set params here because of a pickling error Nel = gvec_equil.params["Nel"] p =gvec_equil.params["p"] @@ -434,14 +433,13 @@ def __init__(self, l1: float = 0.0, l3: float = 0.0, r3: float = 1.0, ): - - self._kind_map = 10 + self.kind_map = 10 # use params setter self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() # periodicity in eta3-direction and pole at eta1=0 - self.params_numpy = self.get_params_numpy() self.periodic_eta3 = False self.pole = False @@ -484,43 +482,23 @@ class Orthogonal(Domain): Lz : 1. # length in z-direction """ - def __init__(self, **params): - self._kind_map = 11 - - # set default parameters and remove wrong/not needed keys - params_default = {"Lx": 2.0, "Ly": 3.0, "alpha": 0.1, "Lz": 6.0} + def __init__(self, + Lx: float = 2.0, + Ly: float = 3.0, + alpha: float = 0.1, + Lz: float = 6.0,): + self.kind_map = 11 - self._params_map, self._params_numpy = Domain.prepare_params_map( - params, - params_default, - ) + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() # periodicity in eta3-direction and pole at eta1=0 - self._periodic_eta3 = False - self._pole = False + self.periodic_eta3 = False + self.pole = False super().__init__() - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 - class Colella(Domain): r""" Slab geometry with Colella mesh distortion. @@ -558,43 +536,22 @@ class Colella(Domain): Lz : 1. # length in third direction """ - def __init__(self, **params): - self._kind_map = 12 - - # set default parameters and remove wrong/not needed keys - params_default = {"Lx": 2.0, "Ly": 3.0, "alpha": 0.1, "Lz": 6.0} + def __init__(self, Lx: float = 2.0, + Ly: float = 3.0, + alpha: float = 0.1, + Lz: float = 6.0,): + self.kind_map = 12 - self._params_map, self._params_numpy = Domain.prepare_params_map( - params, - params_default, - ) + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() # periodicity in eta3-direction and pole at eta1=0 - self._periodic_eta3 = False - self._pole = False + self.periodic_eta3 = False + self.pole = False super().__init__() - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 - class HollowCylinder(Domain): r""" Cylinder with possible hole around the axis. @@ -639,45 +596,22 @@ def __init__( Lz: float = 4.0, poc: int = 1, ): - self._kind_map = 20 + self.kind_map = 20 - self._params_map, self._params_numpy = Domain.get_params_numpy( - a1=a1, - a2=a2, - Lz=Lz, - poc=float(poc), - ) + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() # periodicity in eta3-direction and pole at eta1=0 - self._periodic_eta3 = False + self.periodic_eta3 = False - if self.params_map["a1"] == 0.0: - self._pole = True + if a1 == 0.0: + self.pole = True else: - self._pole = False + self.pole = False super().__init__() - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 - class PoweredEllipticCylinder(Domain): r""" Cylinder with elliptic cross section and radial power law. @@ -715,43 +649,22 @@ class PoweredEllipticCylinder(Domain): s : .5 # power of radial coordinate """ - def __init__(self, **params): - self._kind_map = 21 - - # set default parameters and remove wrong/not needed keys - params_default = {"rx": 1.0, "ry": 2.0, "Lz": 6.0, "s": 0.5} + def __init__(self, rx: float = 1.0, + ry: float = 2.0, + Lz: float = 6.0, + s: float = 0.5,): + self.kind_map = 21 - self._params_map, self._params_numpy = Domain.prepare_params_map( - params, - params_default, - ) + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() # periodicity in eta3-direction and pole at eta1=0 - self._periodic_eta3 = False - self._pole = True + self.periodic_eta3 = False + self.pole = True super().__init__() - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 - class HollowTorus(Domain): r""" Torus with possible hole around the magnetic axis (center of the smaller circle). @@ -808,81 +721,51 @@ class HollowTorus(Domain): tor_period : 2 # toroidal periodicity built into the mapping: phi = 2*pi * eta3 / tor_period """ - def __init__(self, **params): - self._kind_map = 22 - - # set default parameters and remove wrong/not needed keys - params_default = { - "a1": 0.1, - "a2": 1.0, - "R0": 3.0, - "sfl": False, - "pol_period": 1, - "tor_period": 3, - } - - self._params_map, self._params_numpy = Domain.prepare_params_map( - params, - params_default, - ) + def __init__(self, a1: float = 0.1, + a2: float = 1.0, + R0: float = 3.0, + sfl: bool = False, + pol_period: int = 1, + tor_period: int = 3,): + self.kind_map = 22 - assert self.params_map["a2"] <= self.params_map["R0"], ( - f"The minor radius must be smaller or equal than the major radius! {self.params_map['a2']=}, {self.params_map['R0']=}" - ) + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() - if self.params_map["sfl"]: - assert self.params_map["pol_period"] == 1, ( - f"Piece-of-cake is only implemented for torus coordinates, not for straight field line coordinates!" - ) + assert a2 <= R0, (f"The minor radius must be smaller or equal than the major radius! {a2 = }, {R0 = }") + + if sfl: + assert pol_period == 1, (f"Piece-of-cake is only implemented for torus coordinates, not for straight field line coordinates!") # periodicity in eta3-direction and pole at eta1=0 - self._periodic_eta3 = True + self.periodic_eta3 = True - if self.params_map["a1"] == 0.0: - self._pole = True + if a1 == 0.0: + self.pole = True else: - self._pole = False + self.pole = False super().__init__() - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 - def inverse_map(self, x, y, z, bounded=True, change_out_order=False): """Analytical inverse map of HollowTorus""" - mr = np.sqrt(x**2 + y**2) - self.params_map["R0"] + mr = np.sqrt(x**2 + y**2) - self.params["R0"] eta3 = ( np.arctan2(-y, x) - % (2 * np.pi / self.params_map["tor_period"]) + % (2 * np.pi / self.params["tor_period"]) / (2 * np.pi) - * self.params_map["tor_period"] + * self.params["tor_period"] ) eta2 = ( np.arctan2(z, mr) - % (2 * np.pi / self.params_map["pol_period"]) - / (2 * np.pi / self.params_map["pol_period"]) + % (2 * np.pi / self.params["pol_period"]) + / (2 * np.pi / self.params["pol_period"]) ) - eta1 = (z / np.sin(2 * np.pi * eta2 / self.params_map["pol_period"]) - self.params_map["a1"]) / ( - self.params_map["a2"] - self.params_map["a1"] + eta1 = (z / np.sin(2 * np.pi * eta2 / self.params["pol_period"]) - self.params["a1"]) / ( + self.params["a2"] - self.params["a1"] ) if bounded: @@ -936,43 +819,22 @@ class ShafranovShiftCylinder(Domain): delta : .2 # shift factor, should be in [0, 0.1] """ - def __init__(self, **params): - self._kind_map = 30 - - # set default parameters and remove wrong/not needed keys - params_default = {"rx": 1.0, "ry": 1.0, "Lz": 4.0, "delta": 0.2} + def __init__(self, rx: float = 1.0, + ry: float = 1.0, + Lz: float = 4.0, + delta: float = 0.2,): + self.kind_map = 30 - self._params_map, self._params_numpy = Domain.prepare_params_map( - params, - params_default, - ) + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() # periodicity in eta3-direction and pole at eta1=0 - self._periodic_eta3 = False - self._pole = True + self.periodic_eta3 = False + self.pole = True super().__init__() - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 - class ShafranovSqrtCylinder(Domain): r""" Cylinder with square-root Shafranov shift. @@ -1010,43 +872,22 @@ class ShafranovSqrtCylinder(Domain): delta : .2 # shift factor, should be in [0, 0.1] """ - def __init__(self, **params): - self._kind_map = 31 - - # set default parameters and remove wrong/not needed keys - params_default = {"rx": 1.0, "ry": 1.0, "Lz": 4.0, "delta": 0.2} + def __init__(self, rx: float = 1.0, + ry: float = 1.0, + Lz: float = 4.0, + delta: float = 0.2,): + self.kind_map = 31 - self._params_map, self._params_numpy = Domain.prepare_params_map( - params, - params_default, - ) + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() # periodicity in eta3-direction and pole at eta1=0 - self._periodic_eta3 = False - self._pole = True + self.periodic_eta3 = False + self.pole = True super().__init__() - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 - class ShafranovDshapedCylinder(Domain): r""" Cylinder with D-shaped cross section and quadratic Shafranov shift. @@ -1093,47 +934,21 @@ class ShafranovDshapedCylinder(Domain): kappa_gs : 2. # Kappa: ellipticity (elongation) """ - def __init__(self, **params): - self._kind_map = 32 - - # set default parameters and remove wrong/not needed keys - params_default = { - "R0": 2.0, - "Lz": 3.0, - "delta_x": 0.1, - "delta_y": 0.0, - "delta_gs": 0.33, - "epsilon_gs": 0.32, - "kappa_gs": 1.7, - } - - self._params_map, self._params_numpy = Domain.prepare_params_map( - params, - params_default, - ) + def __init__(self, R0: float = 2.0, + Lz: float = 3.0, + delta_x: float = 0.1, + delta_y: float = 0.0, + delta_gs: float = 0.33, + epsilon_gs: float = 0.32, + kappa_gs: float = 1.7,): + self.kind_map = 32 + + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() # periodicity in eta3-direction and pole at eta1=0 - self._periodic_eta3 = False - self._pole = True + self.periodic_eta3 = False + self.pole = True super().__init__() - - @property - def kind_map(self): - return self._kind_map - - @property - def params_map(self): - return self._params_map - - @property - def params_numpy(self): - return self._params_numpy - - @property - def pole(self): - return self._pole - - @property - def periodic_eta3(self): - return self._periodic_eta3 diff --git a/src/struphy/geometry/tests/test_domain.py b/src/struphy/geometry/tests/test_domain.py index e98a10e65..e511c1900 100644 --- a/src/struphy/geometry/tests/test_domain.py +++ b/src/struphy/geometry/tests/test_domain.py @@ -160,6 +160,7 @@ def test_evaluation_mappings(mapping): import numpy as np + from struphy.geometry.base import Domain from struphy.geometry import domains # arrays: @@ -176,8 +177,9 @@ def test_evaluation_mappings(mapping): print() print("Domain object set.") + assert isinstance(domain, Domain) print("domain's kind_map :", domain.kind_map) - print("domain's params_map :", domain.params_map) + print("domain's params_map :", domain.params) # point-wise evaluation: print("pointwise evaluation, shape:", domain(0.5, 0.5, 0.5, squeeze_out=True).shape) @@ -318,6 +320,7 @@ def test_pullback(): import numpy as np + from struphy.geometry.base import Domain from struphy.geometry import domains # arrays: @@ -339,8 +342,9 @@ def fun(x, y, z): print() print("Domain object set.") + assert isinstance(domain, Domain) print("domain's kind_map :", domain.kind_map) - print("domain's params_map :", domain.params_map) + print("domain's params_map :", domain.params) for p_str in domain.dict_transformations["pull"]: print("component:", p_str) @@ -476,6 +480,7 @@ def test_pushforward(): import numpy as np + from struphy.geometry.base import Domain from struphy.geometry import domains # arrays: @@ -497,8 +502,9 @@ def fun(e1, e2, e3): print() print("Domain object set.") + assert isinstance(domain, Domain) print("domain's kind_map :", domain.kind_map) - print("domain's params_map :", domain.params_map) + print("domain's params_map :", domain.params) for p_str in domain.dict_transformations["push"]: print("component:", p_str) @@ -634,6 +640,7 @@ def test_transform(): import numpy as np + from struphy.geometry.base import Domain from struphy.geometry import domains # arrays: @@ -655,8 +662,9 @@ def fun(e1, e2, e3): print() print("Domain object set.") + assert isinstance(domain, Domain) print("domain's kind_map :", domain.kind_map) - print("domain's params_map :", domain.params_map) + print("domain's params_map :", domain.params) for p_str in domain.dict_transformations["tran"]: print("component:", p_str) @@ -832,7 +840,7 @@ def fun(e1, e2, e3): # print('Domain object set.') # # print('domain\'s kind_map :', domain.kind_map) -# print('domain\'s params_map :', domain.params_map) +# print('domain\'s params_map :', domain.params) # # for p_str in domain.keys_transform: # @@ -913,7 +921,7 @@ def fun(e1, e2, e3): if __name__ == "__main__": test_prepare_arg() - test_evaluation_mappings("Cuboid") + test_evaluation_mappings("GVECunit") test_pullback() test_pushforward() test_transform() From e81a899e161ed16004fda20b416a67921a5beb81 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 6 Sep 2025 17:58:04 +0200 Subject: [PATCH 098/292] replace Domain.params_map -> Domain.params --- doc/sections/subsections/domains_avail.rst | 2 +- .../tutorial_04_mapped_domains.ipynb | 8 +-- .../control_variates/control_variate.py | 8 +-- .../fB_massless_control_variate.py | 12 ++-- .../fnB_massless_control_variate.py | 12 ++-- .../massless_control_variate.py | 8 +-- .../legacy/mass_matrices_3d_pre.py | 54 +++++++------- .../legacy/massless_operators/fB_arrays.py | 8 +-- .../fB_massless_linear_operators.py | 16 ++--- .../pro_local/mhd_operators_3d_local.py | 70 +++++++++---------- .../shape_function_projectors_L2.py | 18 ++--- .../shape_function_projectors_local.py | 18 ++--- src/struphy/fields_background/base.py | 2 +- src/struphy/geometry/tests/test_domain.py | 10 +-- .../initial/tests/test_init_perturbations.py | 4 +- .../tests/test_maxwellians.py | 2 +- src/struphy/models/base.py | 4 +- src/struphy/pic/base.py | 2 +- src/struphy/pic/particles.py | 4 +- 19 files changed, 132 insertions(+), 130 deletions(-) diff --git a/doc/sections/subsections/domains_avail.rst b/doc/sections/subsections/domains_avail.rst index b38b4951a..6fcade7e2 100644 --- a/doc/sections/subsections/domains_avail.rst +++ b/doc/sections/subsections/domains_avail.rst @@ -8,5 +8,5 @@ Available domains .. automodule:: struphy.geometry.domains :members: - :exclude-members: kind_map, params_map, params_numpy, pole, periodic_eta3 + :exclude-members: kind_map, params, params_numpy, pole, periodic_eta3 :show-inheritance: diff --git a/doc/tutorials/tutorial_04_mapped_domains.ipynb b/doc/tutorials/tutorial_04_mapped_domains.ipynb index ebc633b2b..3127a3546 100644 --- a/doc/tutorials/tutorial_04_mapped_domains.ipynb +++ b/doc/tutorials/tutorial_04_mapped_domains.ipynb @@ -45,7 +45,7 @@ "metadata": {}, "outputs": [], "source": [ - "for key, val in domain.params_map.items():\n", + "for key, val in domain.params.items():\n", " print(key, '=', val)" ] }, @@ -213,7 +213,7 @@ "metadata": {}, "outputs": [], "source": [ - "for key, val in domain.params_map.items():\n", + "for key, val in domain.params.items():\n", " print(key, '=', val)" ] }, @@ -268,7 +268,7 @@ "metadata": {}, "outputs": [], "source": [ - "for key, val in domain.params_map.items():\n", + "for key, val in domain.params.items():\n", " if 'cx' not in key and 'cy' not in key:\n", " print(key, '=', val)" ] @@ -352,7 +352,7 @@ "metadata": {}, "outputs": [], "source": [ - "for key, val in domain.params_map.items():\n", + "for key, val in domain.params.items():\n", " if 'cx' not in key and 'cy' not in key and 'cz' not in key:\n", " print(key, '=', val)" ] diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py index e775bc16a..ca5134edc 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py @@ -83,7 +83,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): self.mat_jh1, kind_fun_eq[0], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -102,7 +102,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): self.mat_jh2, kind_fun_eq[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -121,7 +121,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): self.mat_jh3, kind_fun_eq[2], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -154,7 +154,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): self.mat_nh, kind_fun_eq[3], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py index 4c7441c28..562df0691 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py @@ -269,7 +269,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -299,7 +299,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -329,7 +329,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -359,7 +359,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -398,7 +398,7 @@ def quadrature_density(gather, domain): gather.n_quad, gather.gather_quadrature, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -432,7 +432,7 @@ def quadrature_grid(gather, domain): gather.Nel, gather.gather_grid, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py index a936871ea..f0279fbcf 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py @@ -496,7 +496,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -526,7 +526,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -556,7 +556,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -586,7 +586,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -625,7 +625,7 @@ def quadrature_density(gather, Nel, pts1, pts2, pts3, n_quad, domain): n_quad, gather, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -659,7 +659,7 @@ def quadrature_grid(gather, Nel, domain): Nel, gather, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py index be31711ed..09deb07f1 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py @@ -495,7 +495,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -524,7 +524,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -553,7 +553,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -582,7 +582,7 @@ def vv_right( tensor_space_FEM.T[2], particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], diff --git a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py index 80515ff51..1feccebd9 100644 --- a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py +++ b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py @@ -33,13 +33,13 @@ def get_M0_PRE(tensor_space_FEM, domain): # spaces_pre[2].set_extraction_operators() spaces_pre[0].assemble_M0( - lambda eta: (domain.params_map[1] - domain.params_map[0]) * np.ones(eta.shape, dtype=float) + lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) ) spaces_pre[1].assemble_M0( - lambda eta: (domain.params_map[3] - domain.params_map[2]) * np.ones(eta.shape, dtype=float) + lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) ) spaces_pre[2].assemble_M0( - lambda eta: (domain.params_map[5] - domain.params_map[4]) * np.ones(eta.shape, dtype=float) + lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) ) c_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] @@ -70,23 +70,23 @@ def get_M1_PRE(tensor_space_FEM, domain): # spaces_pre[2].set_extraction_operators() spaces_pre[0].assemble_M0( - lambda eta: (domain.params_map[1] - domain.params_map[0]) * np.ones(eta.shape, dtype=float) + lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) ) spaces_pre[1].assemble_M0( - lambda eta: (domain.params_map[3] - domain.params_map[2]) * np.ones(eta.shape, dtype=float) + lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) ) spaces_pre[2].assemble_M0( - lambda eta: (domain.params_map[5] - domain.params_map[4]) * np.ones(eta.shape, dtype=float) + lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) ) spaces_pre[0].assemble_M1( - lambda eta: 1 / (domain.params_map[1] - domain.params_map[0]) * np.ones(eta.shape, dtype=float) + lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) ) spaces_pre[1].assemble_M1( - lambda eta: 1 / (domain.params_map[3] - domain.params_map[2]) * np.ones(eta.shape, dtype=float) + lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) ) spaces_pre[2].assemble_M1( - lambda eta: 1 / (domain.params_map[5] - domain.params_map[4]) * np.ones(eta.shape, dtype=float) + lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) ) c11_pre = [spaces_pre[0].M1.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] @@ -129,23 +129,23 @@ def get_M2_PRE(tensor_space_FEM, domain): # spaces_pre[2].set_extraction_operators() spaces_pre[0].assemble_M0( - lambda eta: (domain.params_map[1] - domain.params_map[0]) * np.ones(eta.shape, dtype=float) + lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) ) spaces_pre[1].assemble_M0( - lambda eta: (domain.params_map[3] - domain.params_map[2]) * np.ones(eta.shape, dtype=float) + lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) ) spaces_pre[2].assemble_M0( - lambda eta: (domain.params_map[5] - domain.params_map[4]) * np.ones(eta.shape, dtype=float) + lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) ) spaces_pre[0].assemble_M1( - lambda eta: 1 / (domain.params_map[1] - domain.params_map[0]) * np.ones(eta.shape, dtype=float) + lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) ) spaces_pre[1].assemble_M1( - lambda eta: 1 / (domain.params_map[3] - domain.params_map[2]) * np.ones(eta.shape, dtype=float) + lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) ) spaces_pre[2].assemble_M1( - lambda eta: 1 / (domain.params_map[5] - domain.params_map[4]) * np.ones(eta.shape, dtype=float) + lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) ) c11_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M1.toarray()[:, 0], spaces_pre[2].M1.toarray()[:, 0]] @@ -188,13 +188,13 @@ def get_M3_PRE(tensor_space_FEM, domain): # spaces_pre[2].set_extraction_operators() spaces_pre[0].assemble_M1( - lambda eta: 1 / (domain.params_map[1] - domain.params_map[0]) * np.ones(eta.shape, dtype=float) + lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) ) spaces_pre[1].assemble_M1( - lambda eta: 1 / (domain.params_map[3] - domain.params_map[2]) * np.ones(eta.shape, dtype=float) + lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) ) spaces_pre[2].assemble_M1( - lambda eta: 1 / (domain.params_map[5] - domain.params_map[4]) * np.ones(eta.shape, dtype=float) + lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) ) c_pre = [spaces_pre[0].M1.toarray()[:, 0], spaces_pre[1].M1.toarray()[:, 0], spaces_pre[2].M1.toarray()[:, 0]] @@ -224,21 +224,21 @@ def get_Mv_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M0(lambda eta: domain.params_map[0] ** 3 * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M0(lambda eta: domain.params_map[1] * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M0(lambda eta: domain.params_map[2] * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M0(lambda eta: domain.params[0] ** 3 * np.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: domain.params[1] * np.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: domain.params[2] * np.ones(eta.shape, dtype=float)) c11_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] - spaces_pre[0].assemble_M0(lambda eta: domain.params_map[0] * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M0(lambda eta: domain.params_map[1] ** 3 * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M0(lambda eta: domain.params_map[2] * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M0(lambda eta: domain.params[0] * np.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: domain.params[1] ** 3 * np.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: domain.params[2] * np.ones(eta.shape, dtype=float)) c22_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] - spaces_pre[0].assemble_M0(lambda eta: domain.params_map[0] * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M0(lambda eta: domain.params_map[1] * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M0(lambda eta: domain.params_map[2] ** 3 * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M0(lambda eta: domain.params[0] * np.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: domain.params[1] * np.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: domain.params[2] ** 3 * np.ones(eta.shape, dtype=float)) c33_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py index 6cd356c0c..7705da576 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py @@ -679,7 +679,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): # evaluate Jacobian matrix mapping_fast.df_all( DOMAIN.kind_map, - DOMAIN.params_map, + DOMAIN.params, DOMAIN.T[0], DOMAIN.T[1], DOMAIN.T[2], @@ -768,7 +768,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): TENSOR_SPACE_FEM.pts[2][ie3, q3], 1, DOMAIN.kind_map, - DOMAIN.params_map, + DOMAIN.params, DOMAIN.T[0], DOMAIN.T[1], DOMAIN.T[2], @@ -784,7 +784,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): TENSOR_SPACE_FEM.pts[2][ie3, q3], 2, DOMAIN.kind_map, - DOMAIN.params_map, + DOMAIN.params, DOMAIN.T[0], DOMAIN.T[1], DOMAIN.T[2], @@ -800,7 +800,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): TENSOR_SPACE_FEM.pts[2][ie3, q3], 3, DOMAIN.kind_map, - DOMAIN.params_map, + DOMAIN.params, DOMAIN.T[0], DOMAIN.T[1], DOMAIN.T[2], diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py index 351130dcd..a26a22679 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py @@ -396,7 +396,7 @@ def gather(self, index, dt, acc, tensor_space_FEM, domain, particles_loc, Np_loc acc.gather3_loc, acc.mid_particles, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -431,7 +431,7 @@ def gather(self, index, dt, acc, tensor_space_FEM, domain, particles_loc, Np_loc acc.gather3_loc, acc.stage1_out_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -466,7 +466,7 @@ def gather(self, index, dt, acc, tensor_space_FEM, domain, particles_loc, Np_loc acc.gather3_loc, acc.stage2_out_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -501,7 +501,7 @@ def gather(self, index, dt, acc, tensor_space_FEM, domain, particles_loc, Np_loc acc.gather3_loc, acc.stage3_out_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -559,7 +559,7 @@ def scatter(self, index, acc, tensor_space_FEM, domain, particles_loc, Np_loc, N tensor_space_FEM.n_quad, particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -593,7 +593,7 @@ def scatter(self, index, acc, tensor_space_FEM, domain, particles_loc, Np_loc, N tensor_space_FEM.n_quad, particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -627,7 +627,7 @@ def scatter(self, index, acc, tensor_space_FEM, domain, particles_loc, Np_loc, N tensor_space_FEM.n_quad, particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -661,7 +661,7 @@ def scatter(self, index, acc, tensor_space_FEM, domain, particles_loc, Np_loc, N tensor_space_FEM.n_quad, particles_loc, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py index b8f7b1624..aef91b87d 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py @@ -635,7 +635,7 @@ def projection_Q_0form(self, domain): mat_eq, 11, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -691,7 +691,7 @@ def projection_Q_0form(self, domain): mat_eq, 11, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -747,7 +747,7 @@ def projection_Q_0form(self, domain): mat_eq, 11, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -944,7 +944,7 @@ def projection_Q_2form(self, domain): mat_eq, 11, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1000,7 +1000,7 @@ def projection_Q_2form(self, domain): mat_eq, 11, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1056,7 +1056,7 @@ def projection_Q_2form(self, domain): mat_eq, 11, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1231,7 +1231,7 @@ def projection_W_0form(self, domain): mat_eq, 12, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1429,7 +1429,7 @@ def projection_T_0form(self, domain): mat_eq, 23, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1475,7 +1475,7 @@ def projection_T_0form(self, domain): mat_eq, 22, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1524,7 +1524,7 @@ def projection_T_0form(self, domain): mat_eq, 23, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1581,7 +1581,7 @@ def projection_T_0form(self, domain): mat_eq, 21, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1630,7 +1630,7 @@ def projection_T_0form(self, domain): mat_eq, 22, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1687,7 +1687,7 @@ def projection_T_0form(self, domain): mat_eq, 21, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1984,7 +1984,7 @@ def projection_T_1form(self, domain): mat_eq, 23, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2030,7 +2030,7 @@ def projection_T_1form(self, domain): mat_eq, 22, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2079,7 +2079,7 @@ def projection_T_1form(self, domain): mat_eq, 23, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2125,7 +2125,7 @@ def projection_T_1form(self, domain): mat_eq, 21, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2174,7 +2174,7 @@ def projection_T_1form(self, domain): mat_eq, 22, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2220,7 +2220,7 @@ def projection_T_1form(self, domain): mat_eq, 21, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2517,7 +2517,7 @@ def projection_T_2form(self, domain): mat_eq, 23, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2563,7 +2563,7 @@ def projection_T_2form(self, domain): mat_eq, 22, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2612,7 +2612,7 @@ def projection_T_2form(self, domain): mat_eq, 23, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2658,7 +2658,7 @@ def projection_T_2form(self, domain): mat_eq, 21, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2707,7 +2707,7 @@ def projection_T_2form(self, domain): mat_eq, 22, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -2753,7 +2753,7 @@ def projection_T_2form(self, domain): mat_eq, 21, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3015,7 +3015,7 @@ def projection_S_0form(self, domain): mat_eq, 31, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3071,7 +3071,7 @@ def projection_S_0form(self, domain): mat_eq, 31, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3127,7 +3127,7 @@ def projection_S_0form(self, domain): mat_eq, 31, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3324,7 +3324,7 @@ def projection_S_2form(self, domain): mat_eq, 31, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3380,7 +3380,7 @@ def projection_S_2form(self, domain): mat_eq, 31, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3436,7 +3436,7 @@ def projection_S_2form(self, domain): mat_eq, 31, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3606,7 +3606,7 @@ def projection_K_3form(self, domain): mat_eq, 41, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3760,7 +3760,7 @@ def projection_N_0form(self, domain): mat_eq, 51, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3816,7 +3816,7 @@ def projection_N_0form(self, domain): mat_eq, 51, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -3872,7 +3872,7 @@ def projection_N_0form(self, domain): mat_eq, 51, domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py index d74e64df8..c1ae9e9f6 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py @@ -558,7 +558,7 @@ def heavy_test(self, test1, test2, test3, acc, particles_loc, Np, domain): self.NbaseD, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -570,7 +570,7 @@ def heavy_test(self, test1, test2, test3, acc, particles_loc, Np, domain): domain.cz, ) - # ker_loc.kernel_1_heavy(self.pts[0][0], self.pts[1][0], self.pts[2][0], self.wts[0][0], self.wts[1][0], self.wts[2][0], test1, test2, test3, acc.oneform_temp1, acc.oneform_temp2, acc.oneform_temp3, Np, self.n_quad, self.p, self.Nel, self.p_shape, self.p_size, particles_loc, self.lambdas_1_11, self.lambdas_1_12, self.lambdas_1_13, self.lambdas_1_21, self.lambdas_1_22, self.lambdas_1_23, self.lambdas_1_31, self.lambdas_1_32, self.lambdas_1_33, self.num_cell, self.coeff_i[0], self.coeff_i[1], self.coeff_i[2], self.coeff_h[0], self.coeff_h[1], self.coeff_h[2], self.NbaseN, self.NbaseD, particles_loc.shape[1], domain.kind_map, domain.params_map, domain.T[0], domain.T[1], domain.T[2], domain.p, domain.Nel, domain.NbaseN, domain.cx, domain.cy, domain.cz) + # ker_loc.kernel_1_heavy(self.pts[0][0], self.pts[1][0], self.pts[2][0], self.wts[0][0], self.wts[1][0], self.wts[2][0], test1, test2, test3, acc.oneform_temp1, acc.oneform_temp2, acc.oneform_temp3, Np, self.n_quad, self.p, self.Nel, self.p_shape, self.p_size, particles_loc, self.lambdas_1_11, self.lambdas_1_12, self.lambdas_1_13, self.lambdas_1_21, self.lambdas_1_22, self.lambdas_1_23, self.lambdas_1_31, self.lambdas_1_32, self.lambdas_1_33, self.num_cell, self.coeff_i[0], self.coeff_i[1], self.coeff_i[2], self.coeff_h[0], self.coeff_h[1], self.coeff_h[2], self.NbaseN, self.NbaseD, particles_loc.shape[1], domain.kind_map, domain.params, domain.T[0], domain.T[1], domain.T[2], domain.p, domain.Nel, domain.NbaseN, domain.cx, domain.cy, domain.cz) def potential_pi_0(self, particles_loc, Np, domain, mpi_comm): """ @@ -601,7 +601,7 @@ def potential_pi_0(self, particles_loc, Np, domain, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -648,7 +648,7 @@ def S_pi_0(self, particles_loc, Np, domain): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -742,7 +742,7 @@ def S_pi_1(self, particles_loc, Np, domain): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -799,7 +799,7 @@ def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -851,7 +851,7 @@ def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -903,7 +903,7 @@ def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -955,7 +955,7 @@ def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py index bd0347d81..d361075e3 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py @@ -654,7 +654,7 @@ def heavy_test(self, test1, test2, test3, acc, particles_loc, Np, domain): self.NbaseD, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -695,7 +695,7 @@ def potential_pi_0(self, particles_loc, Np, domain, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -742,7 +742,7 @@ def S_pi_0(self, particles_loc, Np, domain): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -831,7 +831,7 @@ def S_pi_1(self, particles_loc, Np, domain): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -909,7 +909,7 @@ def S_pi_01(self, particles_loc, Np, domain): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -966,7 +966,7 @@ def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1018,7 +1018,7 @@ def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1070,7 +1070,7 @@ def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], @@ -1122,7 +1122,7 @@ def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): self.related, particles_loc.shape[1], domain.kind_map, - domain.params_map, + domain.params, domain.T[0], domain.T[1], domain.T[2], diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index 0a5bce84f..c3cebe6d6 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -946,7 +946,7 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): print(key, ": ", val) print("\nMapping parameters:") - for key, val in self.domain.params_map.items(): + for key, val in self.domain.params.items(): if key not in {"cx", "cy", "cz"}: print(key, ": ", val) diff --git a/src/struphy/geometry/tests/test_domain.py b/src/struphy/geometry/tests/test_domain.py index e511c1900..5f56aa426 100644 --- a/src/struphy/geometry/tests/test_domain.py +++ b/src/struphy/geometry/tests/test_domain.py @@ -179,7 +179,7 @@ def test_evaluation_mappings(mapping): assert isinstance(domain, Domain) print("domain's kind_map :", domain.kind_map) - print("domain's params_map :", domain.params) + print("domain's params :", domain.params) # point-wise evaluation: print("pointwise evaluation, shape:", domain(0.5, 0.5, 0.5, squeeze_out=True).shape) @@ -344,7 +344,7 @@ def fun(x, y, z): assert isinstance(domain, Domain) print("domain's kind_map :", domain.kind_map) - print("domain's params_map :", domain.params) + print("domain's params :", domain.params) for p_str in domain.dict_transformations["pull"]: print("component:", p_str) @@ -504,7 +504,7 @@ def fun(e1, e2, e3): assert isinstance(domain, Domain) print("domain's kind_map :", domain.kind_map) - print("domain's params_map :", domain.params) + print("domain's params :", domain.params) for p_str in domain.dict_transformations["push"]: print("component:", p_str) @@ -664,7 +664,7 @@ def fun(e1, e2, e3): assert isinstance(domain, Domain) print("domain's kind_map :", domain.kind_map) - print("domain's params_map :", domain.params) + print("domain's params :", domain.params) for p_str in domain.dict_transformations["tran"]: print("component:", p_str) @@ -840,7 +840,7 @@ def fun(e1, e2, e3): # print('Domain object set.') # # print('domain\'s kind_map :', domain.kind_map) -# print('domain\'s params_map :', domain.params) +# print('domain\'s params :', domain.params) # # for p_str in domain.keys_transform: # diff --git a/src/struphy/initial/tests/test_init_perturbations.py b/src/struphy/initial/tests/test_init_perturbations.py index 7ff965254..c2fcfc66d 100644 --- a/src/struphy/initial/tests/test_init_perturbations.py +++ b/src/struphy/initial/tests/test_init_perturbations.py @@ -26,6 +26,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False from struphy.feec.psydac_derham import Derham from struphy.geometry import domains + from struphy.geometry.base import Domain from struphy.initial import perturbations comm = MPI.COMM_WORLD @@ -34,6 +35,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False # Domain domain_class = getattr(domains, mapping[0]) domain = domain_class(**mapping[1]) + assert isinstance(domain, Domain) # Derham derham = Derham(Nel, p, spl_kind, comm=comm) @@ -60,7 +62,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False ls = [0, 0] pfuns = ["sin", "sin"] - pmap = domain.params_map + pmap = domain.params if isinstance(domain, domains.Cuboid): Lx = pmap["r1"] - pmap["l1"] Ly = pmap["r2"] - pmap["l2"] diff --git a/src/struphy/kinetic_background/tests/test_maxwellians.py b/src/struphy/kinetic_background/tests/test_maxwellians.py index 67b8e4b2c..0c061ed58 100644 --- a/src/struphy/kinetic_background/tests/test_maxwellians.py +++ b/src/struphy/kinetic_background/tests/test_maxwellians.py @@ -1480,7 +1480,7 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): energy = 1 / 2 * v_para**2 + mu * absB # shifted canonical toroidal momentum - a1 = mhd_equil.domain.params_map["a1"] + a1 = mhd_equil.domain.params["a1"] R0 = mhd_equil.params["R0"] B0 = mhd_equil.params["B0"] diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 221e5247f..01eaf8323 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -112,7 +112,7 @@ def __init__( print("\nDOMAIN:") print(f"type:".ljust(25), self.domain.__class__.__name__) - for key, val in self.domain.params_map.items(): + for key, val in self.domain.params.items(): if key not in {"cx", "cy", "cz"}: print((key + ":").ljust(25), val) @@ -2164,7 +2164,7 @@ def _compute_plasma_params(self, verbose=True): if val["space"] != "ParticlesSPH" and tmp.coords == "constants_of_motion": # call parameters - a1 = self.domain.params_map["a1"] + a1 = self.domain.params["a1"] r = eta1mg * (1 - a1) + a1 psi = self.equil.psi_r(r) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 55d11bc06..851c79b5c 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2025,7 +2025,7 @@ def particle_refilling(self): if kind == "inner": outside_inds = np.nonzero(self._is_outside_left)[0] self.markers[outside_inds, 0] = 1e-4 - r_loss = self._domain.params_map["a1"] + r_loss = self.domain.params["a1"] else: outside_inds = np.nonzero(self._is_outside_right)[0] diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index b80b1be3d..a651a6b71 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -202,7 +202,7 @@ def save_constants_of_motion(self): ) / (2) # eval psi at etas - a1 = self.equil.domain.params_map["a1"] + a1 = self.equil.domain.params["a1"] R0 = self.equil.params["R0"] B0 = self.equil.params["B0"] @@ -516,7 +516,7 @@ def save_constants_of_motion(self): ) # eval psi at etas - a1 = self.equil.domain.params_map["a1"] + a1 = self.equil.domain.params["a1"] R0 = self.equil.params["R0"] B0 = self.equil.params["B0"] From 5eca46c21561998f4be357281d61117bb78ba215 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 8 Sep 2025 07:35:38 +0200 Subject: [PATCH 099/292] fix unit tests --- doc/tutorials/tutorial_04_mapped_domains.ipynb | 18 ------------------ src/struphy/feec/tests/test_l2_projectors.py | 2 +- 2 files changed, 1 insertion(+), 19 deletions(-) diff --git a/doc/tutorials/tutorial_04_mapped_domains.ipynb b/doc/tutorials/tutorial_04_mapped_domains.ipynb index 3127a3546..4ef0a0ef0 100644 --- a/doc/tutorials/tutorial_04_mapped_domains.ipynb +++ b/doc/tutorials/tutorial_04_mapped_domains.ipynb @@ -357,24 +357,6 @@ " print(key, '=', val)" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The mapping parameters contain the spline mapping data and the related [GVECequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/mhd_equils.html#struphy.fields_background.mhd_equil.equils.GVECequilibrium), which is specified through the parameter file.\n", - "Let us check the parameters of the default GVEC equilibrium:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for key, val in domain.params_map['equilibrium'].params.items():\n", - " print(key, val)" - ] - }, { "cell_type": "markdown", "metadata": {}, diff --git a/src/struphy/feec/tests/test_l2_projectors.py b/src/struphy/feec/tests/test_l2_projectors.py index 2bac876d8..e376d2d84 100644 --- a/src/struphy/feec/tests/test_l2_projectors.py +++ b/src/struphy/feec/tests/test_l2_projectors.py @@ -35,7 +35,7 @@ def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_desc, do_plo dom_types = [] dom_classes = [] for key, val in inspect.getmembers(domains): - if inspect.isclass(val) and key != "Domain" and "AxisymmMHDequilibrium" not in key: + if inspect.isclass(val) and val.__module__ == domains.__name__ and "AxisymmMHDequilibrium" not in key: dom_types += [key] dom_classes += [val] From 390f62c3b4121acdfba78570c70f6cc3e108ef3f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 8 Sep 2025 09:03:04 +0200 Subject: [PATCH 100/292] use same params.setter structure in equils as in domains (remove method Equilibrium.set_params) --- .../tutorial_05_mhd_equilibria.ipynb | 9 +- src/struphy/fields_background/base.py | 17 +- src/struphy/fields_background/equils.py | 396 +++++++----------- src/struphy/geometry/domains.py | 3 +- 4 files changed, 156 insertions(+), 269 deletions(-) diff --git a/doc/tutorials/tutorial_05_mhd_equilibria.ipynb b/doc/tutorials/tutorial_05_mhd_equilibria.ipynb index cdee68326..ad7d89358 100644 --- a/doc/tutorials/tutorial_05_mhd_equilibria.ipynb +++ b/doc/tutorials/tutorial_05_mhd_equilibria.ipynb @@ -30,11 +30,6 @@ "import numpy as np" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [] - }, { "cell_type": "markdown", "metadata": {}, @@ -135,7 +130,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "env", "language": "python", "name": "python3" }, @@ -149,7 +144,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.12.3" } }, "nbformat": 4, diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index c3cebe6d6..afe3b4286 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -22,13 +22,20 @@ class FluidEquilibrium(metaclass=ABCMeta): """ @property - def params(self): - """Parameters dictionary.""" + def params(self) -> dict: + """Parameters passed to __init__() of the class in equils.py, as dictionary.""" + if not hasattr(self, "_params"): + self._params = {} return self._params - def set_params(self, **params): - """Generates self.params dictionary from keyword arguments.""" - self._params = params + @params.setter + def params(self, new): + assert isinstance(new, dict) + if "self" in new: + new.pop("self") + if "__class__" in new: + new.pop("__class__") + self._params = new @property def domain(self): diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 07dd72c40..cedabc82a 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -5,9 +5,15 @@ import sys import warnings from time import time +import copy import numpy as np +from scipy.integrate import quad, odeint +from scipy.interpolate import UnivariateSpline, RectBivariateSpline +from scipy.optimize import fsolve, minimize +import struphy +from struphy.fields_background.mhd_equil.eqdsk import readeqdsk from struphy.fields_background.base import ( AxisymmMHDequilibrium, CartesianFluidEquilibrium, @@ -73,13 +79,8 @@ def __init__( beta: float = 0.1, n0: float = 1.0, ): - self.set_params( - B0x=B0x, - B0y=B0y, - B0z=B0z, - beta=beta, - n0=n0, - ) + # use params setter + self.params = copy.deepcopy(locals()) # =============================================================== # profiles on physical domain @@ -206,17 +207,8 @@ def __init__( na: float = 1.0, beta: float = 0.1, ): - self.set_params( - a=a, - R0=R0, - B0=B0, - q0=q0, - q1=q1, - n1=n1, - n2=n2, - na=na, - beta=beta, - ) + # use params setter + self.params = copy.deepcopy(locals()) # =============================================================== # profiles for a sheared slab geometry @@ -427,29 +419,22 @@ class ShearFluid(CartesianMHDequilibrium): B0z : 0. # magnetic field in z """ - def __init__(self, **params): - params_default = { - "a": 1.0, - "b": 1.0, - "c": 1.0, - "z1": 0.25, - "z2": 0.75, - "delta": 0.06666666, - "na": 1.0, - "nb": 0.25, - "pa": 1.0, - "pb": 0.0, - "B0x": 1.0, - "B0y": 0.0, - "B0z": 0.0, - } - - self._params = set_defaults(params, params_default) - - @property - def params(self): - """Parameters dictionary.""" - return self._params + def __init__(self, + a: float = 1.0, + b: float = 1.0, + c: float = 1.0, + z1: float = 0.25, + z2: float = 0.75, + delta: float = 0.06666666, + na: float = 1.0, + nb: float = 0.25, + pa: float = 1.0, + pb: float = 0.0, + B0x: float = 1.0, + B0y: float = 0.0, + B0z: float = 0.0,): + # use params setter + self.params = copy.deepcopy(locals()) # =============================================================== # profiles for a sheared slab geometry @@ -644,18 +629,8 @@ def __init__( p0: float = 1.0e-8, beta: float = 0.1, ): - self.set_params( - a=a, - R0=R0, - B0=B0, - q0=q0, - q1=q1, - n1=n1, - n2=n2, - na=na, - p0=p0, - beta=beta, - ) + # use params setter + self.params = copy.deepcopy(locals()) # inverse cylindrical coordinate transformation (x, y, z) --> (r, theta, phi) self.r = lambda x, y, z: np.sqrt(x**2 + y**2) @@ -935,31 +910,24 @@ class AdhocTorus(AxisymmMHDequilibrium): psi_nel : 50 # number of cells to be used for interpolation of poloidal flux function (only needed if q_kind=1) """ - def __init__(self, **params): - from scipy.integrate import quad - from scipy.interpolate import UnivariateSpline - - # parameters - params_default = { - "a": 1.0, - "R0": 3.0, - "B0": 2.0, - "q_kind": 0, - "q0": 1.71, - "q1": 1.87, - "n1": 2.0, - "n2": 1.0, - "na": 0.2, - "p_kind": 1, - "p0": 1.0, - "p1": 0.1, - "p2": 0.1, - "beta": 0.179, - "psi_k": 3, - "psi_nel": 50, - } - - self._params = set_defaults(params, params_default) + def __init__(self, a: float = 1.0, + R0: float = 3.0, + B0: float = 2.0, + q_kind: int = 0, + q0: float = 1.71, + q1: float = 1.87, + n1: float = 2.0, + n2: float = 1.0, + na: float = 0.2, + p_kind: int = 1, + p0: float = 1.0, + p1: float = 0.1, + p2: float = 0.1, + beta: float = 0.179, + psi_k: int = 3, + psi_nel: int = 50,): + # use params setter + self.params = copy.deepcopy(locals()) # plasma boundary contour ths = np.linspace(0.0, 2 * np.pi, 201) @@ -1017,11 +985,6 @@ def dp_dr(r): ext=3, ) - @property - def params(self): - """Parameters dictionary.""" - return self._params - @property def boundary_pts_R(self): """R-coordinates of plasma boundary contour.""" @@ -1440,30 +1403,22 @@ class AdhocTorusQPsi(AxisymmMHDequilibrium): psi_nel : 50 # number of cells to be used for interpolation of poloidal flux functionq_kind=1) """ - def __init__(self, **params): - from scipy.integrate import odeint - from scipy.interpolate import UnivariateSpline - from scipy.optimize import fsolve - - # parameters - params_default = { - "a": 0.361925, - "R0": 1.0, - "B0": 1.0, - "q0": 0.6, - "q1": 2.5, - "q0p": 0.78, - "q1p": 5.00, - "n1": 2.0, - "n2": 1.0, - "na": 0.2, - "beta": 4.0, - "p1": 0.25, - "psi_k": 3, - "psi_nel": 50, - } - - self._params = set_defaults(params, params_default) + def __init__(self, a: float = 0.361925, + R0: float = 1.0, + B0: float = 1.0, + q0: float = 0.6, + q1: float = 2.5, + q0p: float =0.78, + q1p: float =5.00, + n1: float = 2.0, + n2: float = 1.0, + na: float = 0.2, + beta: float = 4.0, + p1: float = 0.25, + psi_k: int = 3, + psi_nel: int = 50,): + # use params setter + self.params = copy.deepcopy(locals()) # plasma boundary contour ths = np.linspace(0.0, 2 * np.pi, 201) @@ -1513,11 +1468,6 @@ def fun(psi1): ext=3, ) - @property - def params(self): - """Parameters dictionary describing the equilibrium.""" - return self._params - @property def boundary_pts_R(self): """R-coordinates of plasma boundary contour.""" @@ -1770,13 +1720,10 @@ def __init__( n1: float = 2.0, n2: float = 1.0, na: float = 0.2, - units=None, + units: dict = None, ): - from scipy.interpolate import RectBivariateSpline, UnivariateSpline - from scipy.optimize import minimize - - import struphy - from struphy.fields_background.mhd_equil.eqdsk import readeqdsk + # use params setter + self.params = copy.deepcopy(locals()) # default input file if file is None: @@ -1797,23 +1744,10 @@ def __init__( self._units = units - self.set_params( - rel_path=rel_path, - file=file, - data_type=data_type, - p_for_psi=p_for_psi, - psi_resolution=psi_resolution, - p_for_flux=p_for_flux, - flux_resolution=flux_resolution, - n1=n1, - n2=n2, - na=na, - ) - if self.params["rel_path"]: - _path = struphy.__path__[0] + "/fields_background/mhd_equil/eqdsk/data/" + self.params["file"] + _path = struphy.__path__[0] + "/fields_background/mhd_equil/eqdsk/data/" + file else: - _path = self.params["file"] + _path = file eqdsk = readeqdsk.Geqdsk() eqdsk.openFile(_path, data_type=self.params["data_type"]) @@ -1923,11 +1857,6 @@ def units(self): """All Struphy units.""" return self._units - @property - def params(self): - """Parameters describing the equilibrium.""" - return self._params - @property def boundary_pts_R(self): """R-coordinates of plasma boundary contour.""" @@ -2153,7 +2082,28 @@ class GVECequilibrium(NumericalMHDequilibrium): p0 : 1. """ - def __init__(self, units=None, **params): + def __init__(self, + rel_path: bool = True, + # dat_file: str = "run_01/CIRCTOK_State_0000_00000000.dat", + # dat_file: str = "run_02/W7X_State_0000_00000000.dat", + dat_file: str = "run_03/NEO-SPITZER_State_0000_00003307.dat", + # param_file: str = "run_01/parameter.ini", + # param_file: str = "run_02/parameter-w7x.ini", + param_file: str = "run_03/parameter-fig8.ini", + use_boozer: bool = False, + use_nfp: bool = True, + rmin: float = 0.01, + Nel: tuple[int] = (16, 16, 16), + p: tuple[int] = (3, 3, 3), + density_profile: str = "pressure", + p0: float = 0.1, + n0: float = 0.2, + n1: float = 0.0, + units: dict = None,): + + # use params setter + self.params = copy.deepcopy(locals()) + # install if necessary gvec_spec = importlib.util.find_spec("gvec") if gvec_spec is None: @@ -2165,9 +2115,6 @@ def __init__(self, units=None, **params): print(f"{exc.value.code = }") import gvec - from mpi4py import MPI - - import struphy from struphy.geometry.domains import GVECunit # no rescaling if units are not provided @@ -2184,27 +2131,6 @@ def __init__(self, units=None, **params): self._units = units - params_default = { - "rel_path": True, - # "dat_file": "run_01/CIRCTOK_State_0000_00000000.dat", - # "dat_file": "run_02/W7X_State_0000_00000000.dat", - "dat_file": "run_03/NEO-SPITZER_State_0000_00003307.dat", - # "param_file": "run_01/parameter.ini", - # "param_file": "run_02/parameter-w7x.ini", - "param_file": "run_03/parameter-fig8.ini", - "use_boozer": False, - "use_nfp": True, - "rmin": 0.01, - "Nel": (16, 16, 16), - "p": (3, 3, 3), - "density_profile": "pressure", - "p0": 0.1, - "n0": 0.2, - "n1": 0.0, - } - - self._params = set_defaults(params, params_default) - assert self.params["dat_file"][-4:] == ".dat" assert self.params["param_file"][-4:] == ".ini" @@ -2224,8 +2150,8 @@ def __init__(self, units=None, **params): self.params["param_file"], ) else: - dat_file = params["dat_file"] - param_file = params["param_file"] + dat_file = self.params["dat_file"] + param_file = self.params["param_file"] # gvec object self._state = gvec.State(param_file, dat_file) @@ -2253,11 +2179,6 @@ def units(self): """All Struphy units.""" return self._units - @property - def params(self): - """Parameters describing the equilibrium.""" - return self._params - def bv(self, *etas, squeeze_out=False): """Contra-variant (vector field) magnetic field on logical cube [0, 1]^3 in Tesla / meter.""" # evaluate @@ -2413,8 +2334,6 @@ class DESCequilibrium(NumericalMHDequilibrium): Parameters ---------- - units : dict - All Struphy units. If None, no rescaling of EQDSK output is performed. eq_name : str Name of existing DESC equilibrium object (.h5 or binary). rel_path : bool @@ -2429,6 +2348,8 @@ class DESCequilibrium(NumericalMHDequilibrium): Number of cells in each direction used for interpolation of the mapping (default: (16, 16, 16)). p : tuple[int] Spline degree in each direction used for interpolation of the mapping (default: (3, 3, 3)). + units : dict + All Struphy units. If None, no rescaling of EQDSK output is performed. T_kelvin : maximum of temperature in Kelvin (default: 100000). @@ -2447,8 +2368,18 @@ class DESCequilibrium(NumericalMHDequilibrium): T_kelvin : 100000 # maximum temperature in Kelvin used to set density """ - def __init__(self, units=None, **params): - import os + def __init__(self, + eq_name: str = None, + rel_path: bool = False, + use_pest: bool = False, + use_nfp: bool = True, + rmin: float = 0.01, + Nel: tuple[int] = (16, 16, 50), + p: tuple[int] = (3, 3, 3), + T_kelvin: float = 100000., + units: dict = None,): + # use params setter + self.params = copy.deepcopy(locals()) t = time() # install if necessary @@ -2460,16 +2391,9 @@ def __init__(self, units=None, **params): sys.exit(1) import desc - print(f"DESC import: {time() - t} seconds") - from mpi4py import MPI - - import struphy from struphy.geometry.domains import DESCunit - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - # no rescaling if units are not provided if units is None: units = {} @@ -2484,27 +2408,14 @@ def __init__(self, units=None, **params): self._units = units - params_default = { - "eq_name": None, - "rel_path": False, - "use_pest": False, - "use_nfp": True, - "rmin": 0.01, - "Nel": (16, 16, 50), - "p": (3, 3, 3), - "T_kelvin": 100000, - } - - self._params = set_defaults(params, params_default) - - if self._params["rel_path"]: + if self.params["rel_path"]: eq_name = os.path.join( struphy.__path__[0], "fields_background/mhd_equil/desc", - self._params["eq_name"], + self.params["eq_name"], ) else: - eq_name = self._params["eq_name"] + eq_name = self.params["eq_name"] t = time() # desc object @@ -2514,11 +2425,11 @@ def __init__(self, units=None, **params): self._eq = desc.io.load(eq_name) print(f"Eq. load: {time() - t} seconds") - self._rmin = params["rmin"] - self._use_nfp = params["use_nfp"] + self._rmin = self.params["rmin"] + self._use_nfp = self.params["use_nfp"] # straight field line coords - if self._params["use_pest"]: + if self.params["use_pest"]: raise ValueError( "PEST coordinates not yet implemented in desc interface.", ) @@ -2563,11 +2474,6 @@ def units(self): """All Struphy units.""" return self._units - @property - def params(self): - """Parameters describing the equilibrium.""" - return self._params - def bv(self, *etas, squeeze_out=False): """Contra-variant (vector field) magnetic field on logical cube [0, 1]^3 in Tesla / meter.""" # check if already cached @@ -2979,23 +2885,15 @@ class ConstantVelocity(CartesianFluidEquilibrium): """ - def __init__(self, **params): - params_default = { - "ux": 1.0, - "uy": 1.0, - "uz": 1.0, - "n": 1.0, - "n1": 0.0, - "density_profile": "affine", - "p0": 1.0, - } - - self._params = set_defaults(params, params_default) - - @property - def params(self): - """Parameters dictionary.""" - return self._params + def __init__(self, ux: float = 1.0, + uy: float = 1.0, + uz: float = 1.0, + n: float = 1.0, + n1: float = 0.0, + density_profile: str = "affine", + p0: float = 1.0,): + # use params setter + self.params = copy.deepcopy(locals()) # equilibrium ion velocity def u_xyz(self, x, y, z): @@ -3070,15 +2968,15 @@ class HomogenSlabITG(CartesianFluidEquilibriumWithB): eps : .1 """ - def __init__(self, **params): - params_default = {"B0z": 1.0, "Lx": 6.0, "p0": 1.0, "pmin": 0.1, "n0": 1.0, "eps": 0.1} - - self._params = set_defaults(params, params_default) - - @property - def params(self): - """Parameters dictionary.""" - return self._params + def __init__(self, + B0z: float = 1.0, + Lx: float = 6.0, + p0: float = 1.0, + pmin: float = 0.1, + n0: float = 1.0, + eps: float = 0.1,): + # use params setter + self.params = copy.deepcopy(locals()) # =============================================================== # profiles on physical domain @@ -3180,25 +3078,17 @@ class CircularTokamak(AxisymmMHDequilibrium): Bp : 12.5 # poloidal magnetic field """ - def __init__(self, **params): - # parameters - params_default = { - "a": 1.0, - "R0": 2.0, - "B0": 10.0, - "Bp": 12.5, - } - - self._params = set_defaults(params, params_default) + def __init__(self, + a: float = 1.0, + R0: float = 2.0, + B0: float = 10.0, + Bp: float = 12.5,): + # use params setter + self.params = copy.deepcopy(locals()) self._psi0 = 0.0 self._psi1 = self.params["a"] * self.params["R0"] * self.params["Bp"] * 0.5 - @property - def params(self): - """Parameters dictionary.""" - return self._params - # =============================================================== # abstract properties # =============================================================== @@ -3339,15 +3229,11 @@ class CurrentSheet(CartesianMHDequilibrium): """ - def __init__(self, **params): - params_default = {"delta": 0.1, "amp": 1.0} - - self._params = set_defaults(params, params_default) - - @property - def params(self): - """Parameters dictionary.""" - return self._params + def __init__(self, + delta: float = 0.1, + amp: float = 1.0): + # use params setter + self.params = copy.deepcopy(locals()) # =============================================================== # profiles for a straight tokamak equilibrium diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index 394ed8f1b..e73ddf318 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -12,8 +12,6 @@ interp_mapping,) from struphy.geometry.utilities import field_line_tracing -import gvec - class Tokamak(PoloidalSplineTorus): r"""Mappings for Tokamak MHD equilibria constructed via :ref:`field-line tracing ` of a poloidal flux function :math:`\psi`. @@ -133,6 +131,7 @@ class GVECunit(Spline): def __init__(self, gvec_equil=None): + import gvec from struphy.fields_background.equils import GVECequilibrium if gvec_equil is None: From 994413e43ee63a0e6fa9cd9573a0287fc1b0b05f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 8 Sep 2025 09:04:28 +0200 Subject: [PATCH 101/292] formatting --- .../legacy/mass_matrices_3d_pre.py | 72 ++---- src/struphy/fields_background/base.py | 2 +- src/struphy/fields_background/equils.py | 218 ++++++++++-------- src/struphy/geometry/base.py | 115 +++++---- src/struphy/geometry/domains.py | 202 +++++++++------- src/struphy/geometry/tests/test_domain.py | 8 +- 6 files changed, 323 insertions(+), 294 deletions(-) diff --git a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py index 1feccebd9..3b6d2a31b 100644 --- a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py +++ b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py @@ -32,15 +32,9 @@ def get_M0_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M0( - lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[1].assemble_M0( - lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[2].assemble_M0( - lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) - ) + spaces_pre[0].assemble_M0(lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) c_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] @@ -69,25 +63,13 @@ def get_M1_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M0( - lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[1].assemble_M0( - lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[2].assemble_M0( - lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) - ) + spaces_pre[0].assemble_M0(lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) - spaces_pre[0].assemble_M1( - lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[1].assemble_M1( - lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[2].assemble_M1( - lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) - ) + spaces_pre[0].assemble_M1(lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M1(lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M1(lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) c11_pre = [spaces_pre[0].M1.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] c22_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M1.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] @@ -128,25 +110,13 @@ def get_M2_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M0( - lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[1].assemble_M0( - lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[2].assemble_M0( - lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) - ) + spaces_pre[0].assemble_M0(lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) - spaces_pre[0].assemble_M1( - lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[1].assemble_M1( - lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[2].assemble_M1( - lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) - ) + spaces_pre[0].assemble_M1(lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M1(lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M1(lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) c11_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M1.toarray()[:, 0], spaces_pre[2].M1.toarray()[:, 0]] c22_pre = [spaces_pre[0].M1.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M1.toarray()[:, 0]] @@ -187,15 +157,9 @@ def get_M3_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M1( - lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[1].assemble_M1( - lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float) - ) - spaces_pre[2].assemble_M1( - lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float) - ) + spaces_pre[0].assemble_M1(lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M1(lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M1(lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) c_pre = [spaces_pre[0].M1.toarray()[:, 0], spaces_pre[1].M1.toarray()[:, 0], spaces_pre[2].M1.toarray()[:, 0]] diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index afe3b4286..910e56297 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -25,7 +25,7 @@ class FluidEquilibrium(metaclass=ABCMeta): def params(self) -> dict: """Parameters passed to __init__() of the class in equils.py, as dictionary.""" if not hasattr(self, "_params"): - self._params = {} + self._params = {} return self._params @params.setter diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index cedabc82a..af11467fe 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -1,19 +1,18 @@ "Available fluid backgrounds:" +import copy import importlib.util import os import sys import warnings from time import time -import copy import numpy as np -from scipy.integrate import quad, odeint -from scipy.interpolate import UnivariateSpline, RectBivariateSpline +from scipy.integrate import odeint, quad +from scipy.interpolate import RectBivariateSpline, UnivariateSpline from scipy.optimize import fsolve, minimize import struphy -from struphy.fields_background.mhd_equil.eqdsk import readeqdsk from struphy.fields_background.base import ( AxisymmMHDequilibrium, CartesianFluidEquilibrium, @@ -29,6 +28,7 @@ NumericalFluidEquilibriumWithB, NumericalMHDequilibrium, ) +from struphy.fields_background.mhd_equil.eqdsk import readeqdsk from struphy.utils.utils import read_state, subp_run @@ -419,20 +419,22 @@ class ShearFluid(CartesianMHDequilibrium): B0z : 0. # magnetic field in z """ - def __init__(self, - a: float = 1.0, - b: float = 1.0, - c: float = 1.0, - z1: float = 0.25, - z2: float = 0.75, - delta: float = 0.06666666, - na: float = 1.0, - nb: float = 0.25, - pa: float = 1.0, - pb: float = 0.0, - B0x: float = 1.0, - B0y: float = 0.0, - B0z: float = 0.0,): + def __init__( + self, + a: float = 1.0, + b: float = 1.0, + c: float = 1.0, + z1: float = 0.25, + z2: float = 0.75, + delta: float = 0.06666666, + na: float = 1.0, + nb: float = 0.25, + pa: float = 1.0, + pb: float = 0.0, + B0x: float = 1.0, + B0y: float = 0.0, + B0z: float = 0.0, + ): # use params setter self.params = copy.deepcopy(locals()) @@ -910,22 +912,25 @@ class AdhocTorus(AxisymmMHDequilibrium): psi_nel : 50 # number of cells to be used for interpolation of poloidal flux function (only needed if q_kind=1) """ - def __init__(self, a: float = 1.0, - R0: float = 3.0, - B0: float = 2.0, - q_kind: int = 0, - q0: float = 1.71, - q1: float = 1.87, - n1: float = 2.0, - n2: float = 1.0, - na: float = 0.2, - p_kind: int = 1, - p0: float = 1.0, - p1: float = 0.1, - p2: float = 0.1, - beta: float = 0.179, - psi_k: int = 3, - psi_nel: int = 50,): + def __init__( + self, + a: float = 1.0, + R0: float = 3.0, + B0: float = 2.0, + q_kind: int = 0, + q0: float = 1.71, + q1: float = 1.87, + n1: float = 2.0, + n2: float = 1.0, + na: float = 0.2, + p_kind: int = 1, + p0: float = 1.0, + p1: float = 0.1, + p2: float = 0.1, + beta: float = 0.179, + psi_k: int = 3, + psi_nel: int = 50, + ): # use params setter self.params = copy.deepcopy(locals()) @@ -1403,20 +1408,23 @@ class AdhocTorusQPsi(AxisymmMHDequilibrium): psi_nel : 50 # number of cells to be used for interpolation of poloidal flux functionq_kind=1) """ - def __init__(self, a: float = 0.361925, - R0: float = 1.0, - B0: float = 1.0, - q0: float = 0.6, - q1: float = 2.5, - q0p: float =0.78, - q1p: float =5.00, - n1: float = 2.0, - n2: float = 1.0, - na: float = 0.2, - beta: float = 4.0, - p1: float = 0.25, - psi_k: int = 3, - psi_nel: int = 50,): + def __init__( + self, + a: float = 0.361925, + R0: float = 1.0, + B0: float = 1.0, + q0: float = 0.6, + q1: float = 2.5, + q0p: float = 0.78, + q1p: float = 5.00, + n1: float = 2.0, + n2: float = 1.0, + na: float = 0.2, + beta: float = 4.0, + p1: float = 0.25, + psi_k: int = 3, + psi_nel: int = 50, + ): # use params setter self.params = copy.deepcopy(locals()) @@ -2082,28 +2090,29 @@ class GVECequilibrium(NumericalMHDequilibrium): p0 : 1. """ - def __init__(self, - rel_path: bool = True, - # dat_file: str = "run_01/CIRCTOK_State_0000_00000000.dat", - # dat_file: str = "run_02/W7X_State_0000_00000000.dat", - dat_file: str = "run_03/NEO-SPITZER_State_0000_00003307.dat", - # param_file: str = "run_01/parameter.ini", - # param_file: str = "run_02/parameter-w7x.ini", - param_file: str = "run_03/parameter-fig8.ini", - use_boozer: bool = False, - use_nfp: bool = True, - rmin: float = 0.01, - Nel: tuple[int] = (16, 16, 16), - p: tuple[int] = (3, 3, 3), - density_profile: str = "pressure", - p0: float = 0.1, - n0: float = 0.2, - n1: float = 0.0, - units: dict = None,): - + def __init__( + self, + rel_path: bool = True, + # dat_file: str = "run_01/CIRCTOK_State_0000_00000000.dat", + # dat_file: str = "run_02/W7X_State_0000_00000000.dat", + dat_file: str = "run_03/NEO-SPITZER_State_0000_00003307.dat", + # param_file: str = "run_01/parameter.ini", + # param_file: str = "run_02/parameter-w7x.ini", + param_file: str = "run_03/parameter-fig8.ini", + use_boozer: bool = False, + use_nfp: bool = True, + rmin: float = 0.01, + Nel: tuple[int] = (16, 16, 16), + p: tuple[int] = (3, 3, 3), + density_profile: str = "pressure", + p0: float = 0.1, + n0: float = 0.2, + n1: float = 0.0, + units: dict = None, + ): # use params setter self.params = copy.deepcopy(locals()) - + # install if necessary gvec_spec = importlib.util.find_spec("gvec") if gvec_spec is None: @@ -2115,6 +2124,7 @@ def __init__(self, print(f"{exc.value.code = }") import gvec + from struphy.geometry.domains import GVECunit # no rescaling if units are not provided @@ -2368,16 +2378,18 @@ class DESCequilibrium(NumericalMHDequilibrium): T_kelvin : 100000 # maximum temperature in Kelvin used to set density """ - def __init__(self, - eq_name: str = None, - rel_path: bool = False, - use_pest: bool = False, - use_nfp: bool = True, - rmin: float = 0.01, - Nel: tuple[int] = (16, 16, 50), - p: tuple[int] = (3, 3, 3), - T_kelvin: float = 100000., - units: dict = None,): + def __init__( + self, + eq_name: str = None, + rel_path: bool = False, + use_pest: bool = False, + use_nfp: bool = True, + rmin: float = 0.01, + Nel: tuple[int] = (16, 16, 50), + p: tuple[int] = (3, 3, 3), + T_kelvin: float = 100000.0, + units: dict = None, + ): # use params setter self.params = copy.deepcopy(locals()) @@ -2391,6 +2403,7 @@ def __init__(self, sys.exit(1) import desc + print(f"DESC import: {time() - t} seconds") from struphy.geometry.domains import DESCunit @@ -2885,13 +2898,16 @@ class ConstantVelocity(CartesianFluidEquilibrium): """ - def __init__(self, ux: float = 1.0, - uy: float = 1.0, - uz: float = 1.0, - n: float = 1.0, - n1: float = 0.0, - density_profile: str = "affine", - p0: float = 1.0,): + def __init__( + self, + ux: float = 1.0, + uy: float = 1.0, + uz: float = 1.0, + n: float = 1.0, + n1: float = 0.0, + density_profile: str = "affine", + p0: float = 1.0, + ): # use params setter self.params = copy.deepcopy(locals()) @@ -2968,13 +2984,15 @@ class HomogenSlabITG(CartesianFluidEquilibriumWithB): eps : .1 """ - def __init__(self, - B0z: float = 1.0, - Lx: float = 6.0, - p0: float = 1.0, - pmin: float = 0.1, - n0: float = 1.0, - eps: float = 0.1,): + def __init__( + self, + B0z: float = 1.0, + Lx: float = 6.0, + p0: float = 1.0, + pmin: float = 0.1, + n0: float = 1.0, + eps: float = 0.1, + ): # use params setter self.params = copy.deepcopy(locals()) @@ -3078,11 +3096,13 @@ class CircularTokamak(AxisymmMHDequilibrium): Bp : 12.5 # poloidal magnetic field """ - def __init__(self, - a: float = 1.0, - R0: float = 2.0, - B0: float = 10.0, - Bp: float = 12.5,): + def __init__( + self, + a: float = 1.0, + R0: float = 2.0, + B0: float = 10.0, + Bp: float = 12.5, + ): # use params setter self.params = copy.deepcopy(locals()) @@ -3229,9 +3249,7 @@ class CurrentSheet(CartesianMHDequilibrium): """ - def __init__(self, - delta: float = 0.1, - amp: float = 1.0): + def __init__(self, delta: float = 0.1, amp: float = 1.0): # use params setter self.params = copy.deepcopy(locals()) diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index 52234f251..5b9ff359f 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -37,15 +37,19 @@ class Domain(metaclass=ABCMeta): Only right-handed mappings (:math:`\textnormal{det}(DF) > 0`) are admitted. """ - def __init__(self, Nel: tuple[int] = None, p: tuple[int] = None, spl_kind: tuple[bool] = None,): - + def __init__( + self, + Nel: tuple[int] = None, + p: tuple[int] = None, + spl_kind: tuple[bool] = None, + ): if Nel is None or p is None or spl_kind is None: assert self.kind_map >= 10, "Spline mappings must define Nel, p and spl_kind." Nel = (1, 1, 1) p = (1, 1, 1) spl_kind = (True, True, True) - - # create IGA attributes + + # create IGA attributes self._Nel = Nel self._p = p self._spl_kind = spl_kind @@ -138,7 +142,7 @@ def kind_map(self) -> int: if not hasattr(self, "_kind_map"): raise AttributeError("Must set 'self.kind_map' for mappings.") return self._kind_map - + @kind_map.setter def kind_map(self, new): assert isinstance(new, int) @@ -148,9 +152,9 @@ def kind_map(self, new): def params(self) -> dict: """Mapping parameters passed to __init__() of the class in domains.py, as dictionary.""" if not hasattr(self, "_params"): - self._params = {} + self._params = {} return self._params - + @params.setter def params(self, new): assert isinstance(new, dict) @@ -166,7 +170,7 @@ def params_numpy(self) -> np.ndarray: if not hasattr(self, "_params_numpy"): self._params_numpy = np.array([0], dtype=float) return self._params_numpy - + @params_numpy.setter def params_numpy(self, new): assert isinstance(new, np.ndarray) @@ -179,7 +183,7 @@ def pole(self) -> bool: if not hasattr(self, "_pole"): self._pole = False return self._pole - + @pole.setter def pole(self, new): assert isinstance(new, bool) @@ -191,7 +195,7 @@ def periodic_eta3(self) -> bool: if not hasattr(self, "_periodic_eta3"): raise AttributeError("Must specify whether mapping is periodic in eta3.") return self._periodic_eta3 - + @periodic_eta3.setter def periodic_eta3(self, new): assert isinstance(new, bool) @@ -1353,14 +1357,14 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): return a_out # ================================ - + def get_params_numpy(self) -> np.ndarray: """Convert parameter dict into numpy array.""" params_numpy = [] for k, v in self.params.items(): params_numpy.append(v) return np.array(params_numpy) - + def show( self, logical=False, @@ -1751,7 +1755,7 @@ def show( class Spline(Domain): r"""3D IGA spline mapping. - + .. math:: F: (\eta_1, \eta_2, \eta_3) \mapsto (x, y, z) \textnormal{ as } \left\{\begin{aligned} @@ -1763,18 +1767,21 @@ class Spline(Domain): \end{aligned}\right. """ - def __init__(self, Nel: tuple[int] = (8, 24, 6), - p: tuple[int] = (2, 3, 1), - spl_kind: tuple[bool] = (False, True, True), - cx: np.ndarray = None, - cy: np.ndarray = None, - cz: np.ndarray = None,): - + def __init__( + self, + Nel: tuple[int] = (8, 24, 6), + p: tuple[int] = (2, 3, 1), + spl_kind: tuple[bool] = (False, True, True), + cx: np.ndarray = None, + cy: np.ndarray = None, + cz: np.ndarray = None, + ): self.kind_map = 0 # get default control points from default GVEC equilibrium if cx is None or cy is None or cz is None: from struphy.fields_background.equils import GVECequilibrium + mhd_equil = GVECequilibrium() cx = mhd_equil.domain.cx cx = mhd_equil.domain.cy @@ -1824,12 +1831,14 @@ class PoloidalSpline(Domain): The full map :math:`F: [0, 1]^3 \to \Omega` is obtained by defining :math:`(R, Z, \eta_3) \mapsto (x, y, z)` in the child class. """ - def __init__(self, Nel: tuple[int] = (8, 24), - p: tuple[int] = (2, 3), - spl_kind: tuple[bool] = (False, True), - cx: np.ndarray = None, - cy: np.ndarray = None,): - + def __init__( + self, + Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + spl_kind: tuple[bool] = (False, True), + cx: np.ndarray = None, + cy: np.ndarray = None, + ): # get default control points if cx is None or cy is None: @@ -1876,7 +1885,7 @@ def Y(eta1, eta2): class PoloidalSplineStraight(PoloidalSpline): r"""Cylinder where the poloidal planes are described by a 2D IGA-spline mapping. - + .. math:: F: (R, Z, \eta_3) \mapsto (x, y, z) \textnormal{ as } \left\{\begin{aligned} @@ -1888,13 +1897,15 @@ class PoloidalSplineStraight(PoloidalSpline): \end{aligned}\right. """ - def __init__(self, Nel: tuple[int] = (8, 24), - p: tuple[int] = (2, 3), - spl_kind: tuple[bool] = (False, True), - cx: np.ndarray = None, - cy: np.ndarray = None, - Lz: float = 4.0,): - + def __init__( + self, + Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + spl_kind: tuple[bool] = (False, True), + cx: np.ndarray = None, + cy: np.ndarray = None, + Lz: float = 4.0, + ): self.kind_map = 1 # get default control points @@ -1921,7 +1932,7 @@ def Y(eta1, eta2): class PoloidalSplineTorus(PoloidalSpline): r"""Torus where the poloidal planes are described by a 2D IGA-spline mapping. - + .. math:: F: (R, Z, \eta_3) \mapsto (x, y, z) \textnormal{ as } \left\{\begin{aligned} @@ -1931,7 +1942,7 @@ class PoloidalSplineTorus(PoloidalSpline): z &= Z \,. \end{aligned}\right. - + Parameters ---------- Nel : tuple[int] @@ -1942,23 +1953,24 @@ class PoloidalSplineTorus(PoloidalSpline): spl_kind : tuple[bool] Kind of spline in each poloidal direction (True=periodic, False=clamped). - + cx, cy : np.ndarray - Control points (spline coefficients) of the poloidal mapping. + Control points (spline coefficients) of the poloidal mapping. If None, a default square-to-disc mapping of radius 1 centered around (x, y) = (3, 0) is interpolated. - + tor_period : int The toroidal angle is between [0, 2*pi/tor_period). """ - def __init__(self, Nel: tuple[int] = (8, 24), - p: tuple[int] = (2, 3), - spl_kind: tuple[bool] = (False, True), - cx: np.ndarray = None, - cy: np.ndarray = None, - tor_period: int = 3, - ): - + def __init__( + self, + Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + spl_kind: tuple[bool] = (False, True), + cx: np.ndarray = None, + cy: np.ndarray = None, + tor_period: int = 3, + ): # use setters for mapping attributes self.kind_map = 2 self.params_numpy = np.array([float(tor_period)]) @@ -1980,8 +1992,13 @@ def Y(eta1, eta2): cy[0] = 0.0 # init base class - super().__init__(Nel=Nel, p=p, spl_kind=spl_kind, cx=cx, cy=cy,) - + super().__init__( + Nel=Nel, + p=p, + spl_kind=spl_kind, + cx=cx, + cy=cy, + ) def interp_mapping(Nel, p, spl_kind, X, Y, Z=None): diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index e73ddf318..269ebfc81 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -1,15 +1,18 @@ "Mapped domains (single patch)." -import numpy as np import copy -from struphy.fields_background.equils import EQDSKequilibrium +import numpy as np + from struphy.fields_background.base import AxisymmMHDequilibrium -from struphy.geometry.base import (Domain, - PoloidalSplineStraight, - PoloidalSplineTorus, - Spline, - interp_mapping,) +from struphy.fields_background.equils import EQDSKequilibrium +from struphy.geometry.base import ( + Domain, + PoloidalSplineStraight, + PoloidalSplineTorus, + Spline, + interp_mapping, +) from struphy.geometry.utilities import field_line_tracing @@ -72,7 +75,6 @@ def __init__( p_pre: tuple = (3, 3), tor_period: int = 1, ): - if equilibrium is None: equilibrium = EQDSKequilibrium() else: @@ -103,12 +105,14 @@ def __init__( ) # init base class - super().__init__(Nel=Nel, - p=p, - spl_kind=(False, True), - cx=cx, - cy=cy, - tor_period=tor_period,) + super().__init__( + Nel=Nel, + p=p, + spl_kind=(False, True), + cx=cx, + cy=cy, + tor_period=tor_period, + ) class GVECunit(Spline): @@ -130,10 +134,10 @@ class GVECunit(Spline): """ def __init__(self, gvec_equil=None): - import gvec + from struphy.fields_background.equils import GVECequilibrium - + if gvec_equil is None: gvec_equil = GVECequilibrium() else: @@ -142,7 +146,7 @@ def __init__(self, gvec_equil=None): # do not set params here because of a pickling error Nel = gvec_equil.params["Nel"] - p =gvec_equil.params["p"] + p = gvec_equil.params["p"] if gvec_equil.params["use_nfp"]: spl_kind = (False, True, False) else: @@ -274,13 +278,13 @@ class IGAPolarCylinder(PoloidalSplineStraight): a : 1. # minor radius """ - def __init__(self, - Nel: tuple[int] = (8, 24), - p: tuple[int] = (2, 3), - a: float = 1.0, - Lz: float = 4.0, - ): - + def __init__( + self, + Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + a: float = 1.0, + Lz: float = 4.0, + ): # use params setter self.params = copy.deepcopy(locals()) @@ -347,21 +351,25 @@ class IGAPolarTorus(PoloidalSplineTorus): sfl : False # whether to use straight field line coordinates (particular theta parametrization) """ - def __init__(self, Nel: tuple[int] = (8, 24), - p: tuple[int] = (2, 3), - a: float = 1.0, - R0: float = 3.0, - sfl: bool = False, - tor_period: int = 3,): - + def __init__( + self, + Nel: tuple[int] = (8, 24), + p: tuple[int] = (2, 3), + a: float = 1.0, + R0: float = 3.0, + sfl: bool = False, + tor_period: int = 3, + ): # use params setter self.params = copy.deepcopy(locals()) # get control points if sfl: + def theta(eta1, eta2): return 2 * np.arctan(np.sqrt((1 + a * eta1 / R0) / (1 - a * eta1 / R0)) * np.tan(np.pi * eta2)) else: + def theta(eta1, eta2): return 2 * np.pi * eta2 @@ -380,7 +388,14 @@ def Z(eta1, eta2): cy[0] = 0.0 # init base class - super().__init__(Nel=Nel, p=p, spl_kind=spl_kind, cx=cx, cy=cy, tor_period=tor_period,) + super().__init__( + Nel=Nel, + p=p, + spl_kind=spl_kind, + cx=cx, + cy=cy, + tor_period=tor_period, + ) class Cuboid(Domain): @@ -425,13 +440,15 @@ class Cuboid(Domain): r3 : 1. # end of z-interval, r3>l3 """ - def __init__(self, l1: float = 0.0, - r1: float = 1.0, - l2: float = 0.0, - r2: float = 1.0, - l3: float = 0.0, - r3: float = 1.0, - ): + def __init__( + self, + l1: float = 0.0, + r1: float = 1.0, + l2: float = 0.0, + r2: float = 1.0, + l3: float = 0.0, + r3: float = 1.0, + ): self.kind_map = 10 # use params setter @@ -481,11 +498,13 @@ class Orthogonal(Domain): Lz : 1. # length in z-direction """ - def __init__(self, - Lx: float = 2.0, - Ly: float = 3.0, - alpha: float = 0.1, - Lz: float = 6.0,): + def __init__( + self, + Lx: float = 2.0, + Ly: float = 3.0, + alpha: float = 0.1, + Lz: float = 6.0, + ): self.kind_map = 11 # use params setter @@ -535,10 +554,13 @@ class Colella(Domain): Lz : 1. # length in third direction """ - def __init__(self, Lx: float = 2.0, - Ly: float = 3.0, - alpha: float = 0.1, - Lz: float = 6.0,): + def __init__( + self, + Lx: float = 2.0, + Ly: float = 3.0, + alpha: float = 0.1, + Lz: float = 6.0, + ): self.kind_map = 12 # use params setter @@ -648,10 +670,13 @@ class PoweredEllipticCylinder(Domain): s : .5 # power of radial coordinate """ - def __init__(self, rx: float = 1.0, - ry: float = 2.0, - Lz: float = 6.0, - s: float = 0.5,): + def __init__( + self, + rx: float = 1.0, + ry: float = 2.0, + Lz: float = 6.0, + s: float = 0.5, + ): self.kind_map = 21 # use params setter @@ -720,22 +745,27 @@ class HollowTorus(Domain): tor_period : 2 # toroidal periodicity built into the mapping: phi = 2*pi * eta3 / tor_period """ - def __init__(self, a1: float = 0.1, - a2: float = 1.0, - R0: float = 3.0, - sfl: bool = False, - pol_period: int = 1, - tor_period: int = 3,): + def __init__( + self, + a1: float = 0.1, + a2: float = 1.0, + R0: float = 3.0, + sfl: bool = False, + pol_period: int = 1, + tor_period: int = 3, + ): self.kind_map = 22 # use params setter self.params = copy.deepcopy(locals()) self.params_numpy = self.get_params_numpy() - assert a2 <= R0, (f"The minor radius must be smaller or equal than the major radius! {a2 = }, {R0 = }") + assert a2 <= R0, f"The minor radius must be smaller or equal than the major radius! {a2 = }, {R0 = }" if sfl: - assert pol_period == 1, (f"Piece-of-cake is only implemented for torus coordinates, not for straight field line coordinates!") + assert pol_period == 1, ( + f"Piece-of-cake is only implemented for torus coordinates, not for straight field line coordinates!" + ) # periodicity in eta3-direction and pole at eta1=0 self.periodic_eta3 = True @@ -752,17 +782,8 @@ def inverse_map(self, x, y, z, bounded=True, change_out_order=False): mr = np.sqrt(x**2 + y**2) - self.params["R0"] - eta3 = ( - np.arctan2(-y, x) - % (2 * np.pi / self.params["tor_period"]) - / (2 * np.pi) - * self.params["tor_period"] - ) - eta2 = ( - np.arctan2(z, mr) - % (2 * np.pi / self.params["pol_period"]) - / (2 * np.pi / self.params["pol_period"]) - ) + eta3 = np.arctan2(-y, x) % (2 * np.pi / self.params["tor_period"]) / (2 * np.pi) * self.params["tor_period"] + eta2 = np.arctan2(z, mr) % (2 * np.pi / self.params["pol_period"]) / (2 * np.pi / self.params["pol_period"]) eta1 = (z / np.sin(2 * np.pi * eta2 / self.params["pol_period"]) - self.params["a1"]) / ( self.params["a2"] - self.params["a1"] ) @@ -818,10 +839,13 @@ class ShafranovShiftCylinder(Domain): delta : .2 # shift factor, should be in [0, 0.1] """ - def __init__(self, rx: float = 1.0, - ry: float = 1.0, - Lz: float = 4.0, - delta: float = 0.2,): + def __init__( + self, + rx: float = 1.0, + ry: float = 1.0, + Lz: float = 4.0, + delta: float = 0.2, + ): self.kind_map = 30 # use params setter @@ -871,10 +895,13 @@ class ShafranovSqrtCylinder(Domain): delta : .2 # shift factor, should be in [0, 0.1] """ - def __init__(self, rx: float = 1.0, - ry: float = 1.0, - Lz: float = 4.0, - delta: float = 0.2,): + def __init__( + self, + rx: float = 1.0, + ry: float = 1.0, + Lz: float = 4.0, + delta: float = 0.2, + ): self.kind_map = 31 # use params setter @@ -933,13 +960,16 @@ class ShafranovDshapedCylinder(Domain): kappa_gs : 2. # Kappa: ellipticity (elongation) """ - def __init__(self, R0: float = 2.0, - Lz: float = 3.0, - delta_x: float = 0.1, - delta_y: float = 0.0, - delta_gs: float = 0.33, - epsilon_gs: float = 0.32, - kappa_gs: float = 1.7,): + def __init__( + self, + R0: float = 2.0, + Lz: float = 3.0, + delta_x: float = 0.1, + delta_y: float = 0.0, + delta_gs: float = 0.33, + epsilon_gs: float = 0.32, + kappa_gs: float = 1.7, + ): self.kind_map = 32 # use params setter diff --git a/src/struphy/geometry/tests/test_domain.py b/src/struphy/geometry/tests/test_domain.py index 5f56aa426..b63c5b4c6 100644 --- a/src/struphy/geometry/tests/test_domain.py +++ b/src/struphy/geometry/tests/test_domain.py @@ -160,8 +160,8 @@ def test_evaluation_mappings(mapping): import numpy as np - from struphy.geometry.base import Domain from struphy.geometry import domains + from struphy.geometry.base import Domain # arrays: arr1 = np.linspace(0.0, 1.0, 4) @@ -320,8 +320,8 @@ def test_pullback(): import numpy as np - from struphy.geometry.base import Domain from struphy.geometry import domains + from struphy.geometry.base import Domain # arrays: arr1 = np.linspace(0.0, 1.0, 4) @@ -480,8 +480,8 @@ def test_pushforward(): import numpy as np - from struphy.geometry.base import Domain from struphy.geometry import domains + from struphy.geometry.base import Domain # arrays: arr1 = np.linspace(0.0, 1.0, 4) @@ -640,8 +640,8 @@ def test_transform(): import numpy as np - from struphy.geometry.base import Domain from struphy.geometry import domains + from struphy.geometry.base import Domain # arrays: arr1 = np.linspace(0.0, 1.0, 4) From 250ed36313855d6c567c30591d9ab5591695a02f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 8 Sep 2025 15:47:11 +0200 Subject: [PATCH 102/292] tutorial 2 works up to GuidingCenter --- .../tutorial_02_test_particles.ipynb | 230 ++++++++++-------- src/struphy/main.py | 47 +++- 2 files changed, 163 insertions(+), 114 deletions(-) diff --git a/doc/tutorials/tutorial_02_test_particles.ipynb b/doc/tutorials/tutorial_02_test_particles.ipynb index f72cfb445..5ff7d31b6 100644 --- a/doc/tutorials/tutorial_02_test_particles.ipynb +++ b/doc/tutorials/tutorial_02_test_particles.ipynb @@ -344,11 +344,11 @@ "\n", "fig = plt.figure(figsize=(10, 6)) \n", "\n", - "pushed_pos = simdata.pic_species[\"kinetic_ions\"][\"orbits\"][\"kinetic_ions_0\"]\n", - "pushed_pos_uni = simdata_2.pic_species[\"kinetic_ions\"][\"orbits\"][\"kinetic_ions_0\"]\n", + "pushed_pos = simdata.pic_species[\"kinetic_ions\"][\"orbits\"]\n", + "pushed_pos_uni = simdata_2.pic_species[\"kinetic_ions\"][\"orbits\"]\n", "\n", "plt.subplot(1, 2, 1)\n", - "plt.scatter(pushed_pos[:, 0], pushed_pos[:, 1], s=2.)\n", + "plt.scatter(pushed_pos[0, :, 0], pushed_pos[0, :, 1], s=2.)\n", "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle1)\n", @@ -358,7 +358,7 @@ "plt.title('sim_1: draw uniform in logical space')\n", "\n", "plt.subplot(1, 2, 2)\n", - "plt.scatter(pushed_pos_uni[:, 0], pushed_pos_uni[:, 1], s=2.)\n", + "plt.scatter(pushed_pos_uni[0, :, 0], pushed_pos_uni[0, :, 1], s=2.)\n", "circle2 = plt.Circle((0, 0), a2, color='k', fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle2)\n", @@ -486,6 +486,14 @@ "simdata = main.load_data(path)" ] }, + { + "cell_type": "markdown", + "id": "2b1736a1", + "metadata": {}, + "source": [ + "Under `pic_species[\"kinetic_ions\"][\"orbits\"]` one finds a three-dimensional numpy array; the first index refers to the time step, the second index to the particle and the third index to the particel attribute. The first three attributes are the partciel positions, followed by the velocities and the (initial and time-dependent) weights." + ] + }, { "cell_type": "code", "execution_count": null, @@ -493,20 +501,35 @@ "metadata": {}, "outputs": [], "source": [ + "orbits = simdata.pic_species[\"kinetic_ions\"][\"orbits\"]\n", + "print(f\"{orbits.shape = }\")\n", + "\n", + "Nt = orbits.shape[0]\n", + "Np = orbits.shape[1]\n", + "Nattr = orbits.shape[2]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "81afac14", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", "fig = plt.figure()\n", "ax = fig.gca()\n", "\n", "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", "\n", - "time = 0.\n", - "dt = time_opts.dt\n", + "# create alpha for color scaling\n", "Tend = time_opts.Tend\n", - "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", - " # print(f\"{v[0] = }\")\n", - " alpha = (Tend - time)/Tend\n", - " for i, particle in enumerate(v):\n", - " ax.scatter(particle[0], particle[1], c=colors[i % 4], alpha=alpha)\n", - " time += dt\n", + "alpha = np.linspace(1., 0., Nt + 1)\n", + "\n", + "# loop through particles, plot all time steps\n", + "for i in range(Np):\n", + " ax.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], alpha=alpha)\n", " \n", "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", "\n", @@ -514,7 +537,7 @@ "ax.set_aspect('equal')\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('y')\n", - "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" + "ax.set_title(f'{Nt - 1} time steps (full color at t=0)');" ] }, { @@ -654,6 +677,20 @@ "simdata = main.load_data(path)" ] }, + { + "cell_type": "code", + "execution_count": null, + "id": "4cfe0ee8", + "metadata": {}, + "outputs": [], + "source": [ + "orbits = simdata.pic_species[\"kinetic_ions\"][\"orbits\"]\n", + "print(f\"{orbits.shape = }\")\n", + "\n", + "Nt = orbits.shape[0]\n", + "Np = orbits.shape[1]" + ] + }, { "cell_type": "code", "execution_count": null, @@ -666,15 +703,13 @@ "\n", "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", "\n", - "time = 0.\n", - "dt = time_opts.dt\n", + "# create alpha for color scaling\n", "Tend = time_opts.Tend\n", - "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", - " # print(f\"{v[0] = }\")\n", - " alpha = (Tend - time)/Tend\n", - " for i, particle in enumerate(v):\n", - " ax.scatter(particle[0], particle[1], c=colors[i % 4], alpha=alpha)\n", - " time += dt\n", + "alpha = np.linspace(1., 0., Nt + 1)\n", + "\n", + "# loop through particles, plot all time steps\n", + "for i in range(Np):\n", + " ax.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], alpha=alpha)\n", " \n", "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", "\n", @@ -682,7 +717,7 @@ "ax.set_aspect('equal')\n", "ax.set_xlabel('x')\n", "ax.set_ylabel('y')\n", - "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" + "ax.set_title(f'{int(Nt - 1)} time steps (full color at t=0)');" ] }, { @@ -1025,6 +1060,9 @@ "metadata": {}, "outputs": [], "source": [ + "import os\n", + "from struphy import main\n", + "\n", "path = os.path.join(os.getcwd(), \"sim_1\")\n", "main.pproc(path)\n", "\n", @@ -1034,86 +1072,38 @@ { "cell_type": "code", "execution_count": null, - "id": "ae248ef1", + "id": "3be0d0a3", "metadata": {}, "outputs": [], "source": [ - "fig = plt.figure()\n", - "ax = fig.gca()\n", - "\n", - "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", + "orbits = simdata.pic_species[\"kinetic_ions\"][\"orbits\"]\n", + "print(f\"{orbits.shape = }\")\n", "\n", - "time = 0.\n", - "dt = time_opts.dt\n", - "Tend = time_opts.Tend\n", - "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", - " # print(f\"{v[0] = }\")\n", - " alpha = (Tend - time)/Tend\n", - " for i, particle in enumerate(v):\n", - " ax.scatter(particle[0], particle[1], c=colors[i % 4], alpha=alpha)\n", - " time += dt\n", - " \n", - "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", - "\n", - "ax.add_patch(circle1)\n", - "ax.set_aspect('equal')\n", - "ax.set_xlabel('x')\n", - "ax.set_ylabel('y')\n", - "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" + "Nt = orbits.shape[0]\n", + "Np = orbits.shape[1]" ] }, { "cell_type": "code", "execution_count": null, - "id": "430a160b", + "id": "ae248ef1", "metadata": {}, "outputs": [], "source": [ - "# time stepping\n", - "Tend = 3000. - 1e-6\n", - "dt = .2\n", - "Nt = int(Tend / dt)\n", + "import math\n", "\n", - "pos = np.zeros((Nt + 2, Np, 3), dtype=float)\n", - "r = np.zeros((Nt + 2, Np), dtype=float)\n", + "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", "\n", - "pos[0] = pushed_pos\n", - "r[0] = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)\n", + "dt = time_opts.dt\n", + "Tend = time_opts.Tend\n", "\n", - "time = 0.\n", - "n = 0\n", - "while time < Tend:\n", - " time += dt\n", - " n += 1\n", - " \n", - " # advance in time\n", - " prop_vxB(dt/2)\n", - " prop_eta(dt)\n", - " prop_vxB(dt/2)\n", - " \n", - " # positions on the physical domain Omega\n", - " pushed_pos = domain(particles.positions).T\n", - " \n", - " # compute R-ccordinate\n", - " pos[n] = pushed_pos\n", - " r[n] = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)\n", - " " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "ac6b4341", - "metadata": {}, - "outputs": [], - "source": [ - "# make scatter plot for each particle \n", - "for i in range(pos.shape[1]):\n", + "for i in range(Np):\n", + " r = np.sqrt(orbits[:, i, 0]**2 + orbits[:, i, 1]**2)\n", " # poloidal \n", - " ax.scatter(r[:, i], pos[:, i, 2], c=colors[i % 4], s=1)\n", + " ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n", " # top view\n", - " ax_top.scatter(pos[:, i, 0], pos[:, i, 1], c=colors[i % 4], s=1)\n", - "\n", + " ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n", + " \n", "ax.set_title(f'{math.ceil(Tend/dt)} time steps')\n", "ax_top.set_title(f'{math.ceil(Tend/dt)} time steps');\n", "\n", @@ -1127,37 +1117,75 @@ "source": [ "## Guiding-centers in a Tokamak equilibrium\n", "\n", - "We now use the Propagators []() and []() in ASDEX-Upgrade equilibrium from the previous example.\n", - "\n", - "For this we need to instantiate the [Particles5D]() class." + "Let us run a similar test for guiding-centers:" ] }, { "cell_type": "code", "execution_count": null, - "id": "64caf9e5", + "id": "a5182604", "metadata": {}, "outputs": [], "source": [ - "from struphy.pic.particles import Particles5D\n", + "from struphy.models.toy import GuidingCenter\n", "\n", - "# instantiate Particle object\n", - "Np = 4\n", - "bc = ['remove', 'periodic', 'periodic']\n", - "bufsize = 2.\n", + "# light-weight model instance\n", + "model = GuidingCenter()\n", + "\n", + "# species parameters\n", + "model.kinetic_ions.set_phys_params()\n", "\n", "initial = [[.501, 0.001, 0.001, -1.935 , 1.72], # co-passing particle\n", " [.501, 0.001, 0.001, 1.935 , 1.72], # couner-passing particle\n", " [.501, 0.001, 0.001, -0.6665, 1.72], # co-trapped particle\n", - " [.501, 0.001, 0.001, 0.4515, 1.72]] # counter-trapped particle\n", + " [.501, 0.001, 0.001, 0.4515, 1.72]] # counter-trapped particl\n", "\n", - "loading_params = {'seed': 1608,\n", - " 'initial' : initial}\n", + "loading_params = LoadingParameters(Np=4, seed=1608, specific_markers=initial)\n", + "boundary_params = BoundaryParameters(bc=('remove', 'periodic', 'periodic'))\n", + "model.kinetic_ions.set_markers(loading_params=loading_params, \n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params,\n", + " bufsize=2.)\n", + "model.kinetic_ions.set_sorting_boxes()\n", + "model.kinetic_ions.set_save_data(n_markers=1.0)\n", "\n", - "particles = Particles5D(proj_equil,\n", - " Np=Np, \n", - " bc=bc, \n", - " loading_params=loading_params,\n", + "# propagator options\n", + "model.propagators.push_vxb.set_options()\n", + "model.propagators.push_eta.set_options()\n", + "\n", + "# initial conditions (background + perturbation)\n", + "perturbation = None\n", + "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", + "\n", + "model.kinetic_ions.var.add_background(background)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "64caf9e5", + "metadata": {}, + "outputs": [], + "source": [ + "# from struphy.pic.particles import Particles5D\n", + "\n", + "# # instantiate Particle object\n", + "# Np = 4\n", + "# bc = ['remove', 'periodic', 'periodic']\n", + "# bufsize = 2.\n", + "\n", + "# initial = [[.501, 0.001, 0.001, -1.935 , 1.72], # co-passing particle\n", + "# [.501, 0.001, 0.001, 1.935 , 1.72], # couner-passing particle\n", + "# [.501, 0.001, 0.001, -0.6665, 1.72], # co-trapped particle\n", + "# [.501, 0.001, 0.001, 0.4515, 1.72]] # counter-trapped particle\n", + "\n", + "# loading_params = {'seed': 1608,\n", + "# 'initial' : initial}\n", + "\n", + "# particles = Particles5D(proj_equil,\n", + "# Np=Np, \n", + "# bc=bc, \n", + "# loading_params=loading_params,\n", " bufsize=bufsize)" ] }, diff --git a/src/struphy/main.py b/src/struphy/main.py index 8097bf9ec..c7534c1c5 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -130,10 +130,15 @@ def run( pickle.dump(time_opts, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "domain.bin"), 'wb') as f: # WORKAROUND: cannot pickle pyccelized classes at the moment - tmp_dct = {"name": domain.__class__.__name__, "params_map": domain.params_map} + tmp_dct = {"name": domain.__class__.__name__, "params": domain.params} pickle.dump(tmp_dct, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "equil.bin"), 'wb') as f: - pickle.dump(equil, f, pickle.HIGHEST_PROTOCOL) + # WORKAROUND: cannot pickle pyccelized classes at the moment + if equil is not None: + tmp_dct = {"name": equil.__class__.__name__, "params": equil.params} + else: + tmp_dct = {} + pickle.dump(tmp_dct, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "grid.bin"), 'wb') as f: pickle.dump(grid, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "derham_opts.bin"), 'wb') as f: @@ -443,8 +448,11 @@ def pproc( with open(os.path.join(path, "model.bin"), 'rb') as f: model: StruphyModel = pickle.load(f) with open(os.path.join(path, "domain.bin"), 'rb') as f: + # domain: Domain = pickle.load(f) + # print(f"{domain = }") + # print(f"{domain.params = }") domain_dct: Domain = pickle.load(f) - domain = getattr(domains, domain_dct["name"])(**domain_dct["params_map"]) + domain = getattr(domains, domain_dct["name"])(**domain_dct["params"]) # create post-processing folder path_pproc = os.path.join(path, "post_processing") @@ -682,20 +690,33 @@ def load_data(path: str) -> SimData: path_spec = os.path.join(path_kinetic, spec) wlk = os.walk(path_spec) sub_folders = next(wlk)[1] - # print(f"{sub_folders = }") for folder in sub_folders: - # simdata.pic_species[spec][folder] = {} - tmp = {} path_dat = os.path.join(path_spec, folder) sub_wlk = os.walk(path_dat) files = next(sub_wlk)[2] - for file in files: - # print(f"{file = }") - if ".npy" in file: - var = file.split(".")[0] - tmp[var] = np.load(os.path.join(path_dat, file)) - # sort dict - simdata.pic_species[spec][folder] = dict(sorted(tmp.items())) + if "orbits" in folder: + Nt = len(files) // 2 + n = 0 + for file in files: + # print(f"{file = }") + if ".npy" in file: + step = int(file.split(".")[0].split("_")[-1]) + tmp = np.load(os.path.join(path_dat, file)) + if n == 0: + orbits = np.zeros((Nt, *tmp.shape), dtype=float) + orbits[step] = tmp + n += 1 + simdata.pic_species[spec][folder] = orbits + else: + # simdata.pic_species[spec][folder] = {} + tmp = {} + for file in files: + # print(f"{file = }") + if ".npy" in file: + var = file.split(".")[0] + tmp[var] = np.load(os.path.join(path_dat, file)) + # sort dict + simdata.pic_species[spec][folder] = dict(sorted(tmp.items())) print("\nThe following data has been loaded:") print(f"{simdata.time_grid_size = }") From 702fc95d2c2a1e5f3e9a3a24481bd6802340daa5 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 9 Sep 2025 15:08:09 +0200 Subject: [PATCH 103/292] Introduce SimData.orbits and SimData.spline_values --- .../tutorial_02_test_particles.ipynb | 62 ++++--------- src/struphy/main.py | 89 +++++++++++++------ 2 files changed, 79 insertions(+), 72 deletions(-) diff --git a/doc/tutorials/tutorial_02_test_particles.ipynb b/doc/tutorials/tutorial_02_test_particles.ipynb index 5ff7d31b6..b9277aca8 100644 --- a/doc/tutorials/tutorial_02_test_particles.ipynb +++ b/doc/tutorials/tutorial_02_test_particles.ipynb @@ -344,11 +344,14 @@ "\n", "fig = plt.figure(figsize=(10, 6)) \n", "\n", - "pushed_pos = simdata.pic_species[\"kinetic_ions\"][\"orbits\"]\n", - "pushed_pos_uni = simdata_2.pic_species[\"kinetic_ions\"][\"orbits\"]\n", + "orbits = simdata.orbits[\"kinetic_ions\"]\n", + "orbits_uni = simdata_2.orbits[\"kinetic_ions\"]\n", + "\n", + "# orbits = simdata.pic_species[\"kinetic_ions\"][\"orbits\"]\n", + "# orbits_uni = simdata_2.pic_species[\"kinetic_ions\"][\"orbits\"]\n", "\n", "plt.subplot(1, 2, 1)\n", - "plt.scatter(pushed_pos[0, :, 0], pushed_pos[0, :, 1], s=2.)\n", + "plt.scatter(orbits[0, :, 0], orbits[0, :, 1], s=2.)\n", "circle1 = plt.Circle((0, 0), a2, color='k', fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle1)\n", @@ -358,7 +361,7 @@ "plt.title('sim_1: draw uniform in logical space')\n", "\n", "plt.subplot(1, 2, 2)\n", - "plt.scatter(pushed_pos_uni[0, :, 0], pushed_pos_uni[0, :, 1], s=2.)\n", + "plt.scatter(orbits_uni[0, :, 0], orbits_uni[0, :, 1], s=2.)\n", "circle2 = plt.Circle((0, 0), a2, color='k', fill=False)\n", "ax = plt.gca()\n", "ax.add_patch(circle2)\n", @@ -491,7 +494,7 @@ "id": "2b1736a1", "metadata": {}, "source": [ - "Under `pic_species[\"kinetic_ions\"][\"orbits\"]` one finds a three-dimensional numpy array; the first index refers to the time step, the second index to the particle and the third index to the particel attribute. The first three attributes are the partciel positions, followed by the velocities and the (initial and time-dependent) weights." + "Under `simdata.orbits[]` one finds a three-dimensional numpy array; the first index refers to the time step, the second index to the particle and the third index to the particel attribute. The first three attributes are the partciel positions, followed by the velocities and the (initial and time-dependent) weights." ] }, { @@ -501,12 +504,11 @@ "metadata": {}, "outputs": [], "source": [ - "orbits = simdata.pic_species[\"kinetic_ions\"][\"orbits\"]\n", - "print(f\"{orbits.shape = }\")\n", + "orbits = simdata.orbits[\"kinetic_ions\"]\n", "\n", - "Nt = orbits.shape[0]\n", - "Np = orbits.shape[1]\n", - "Nattr = orbits.shape[2]" + "Nt = simdata.Nt[\"kinetic_ions\"]\n", + "Np = simdata.Np[\"kinetic_ions\"]\n", + "Nattr = simdata.Nattr[\"kinetic_ions\"]" ] }, { @@ -684,11 +686,10 @@ "metadata": {}, "outputs": [], "source": [ - "orbits = simdata.pic_species[\"kinetic_ions\"][\"orbits\"]\n", - "print(f\"{orbits.shape = }\")\n", + "orbits = simdata.orbits[\"kinetic_ions\"]\n", "\n", - "Nt = orbits.shape[0]\n", - "Np = orbits.shape[1]" + "Nt = simdata.Nt[\"kinetic_ions\"]\n", + "Np = simdata.Np[\"kinetic_ions\"]" ] }, { @@ -974,32 +975,6 @@ "model.kinetic_ions.var.add_background(background)" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "568c0ba6", - "metadata": {}, - "outputs": [], - "source": [ - "# # instantiate Particle object\n", - "# Np = 4\n", - "# bc = ['remove', 'periodic', 'periodic']\n", - "# bufsize = 2.\n", - "\n", - "# initial = [[.501, 0.001, 0.001, 0., 0.0450, -0.04], # co-passing particle\n", - "# [.511, 0.001, 0.001, 0., -0.0450, -0.04], # counter passing particle\n", - "# [.521, 0.001, 0.001, 0., 0.0105, -0.04], # co-trapped particle\n", - "# [.531, 0.001, 0.001, 0., -0.0155, -0.04]]\n", - "\n", - "# loading_params = {'seed': 1608,\n", - "# 'initial' : initial}\n", - "\n", - "# particles = Particles6D(Np=Np, \n", - "# bc=bc, \n", - "# loading_params=loading_params,\n", - "# bufsize=bufsize)" - ] - }, { "cell_type": "markdown", "id": "9de13919", @@ -1076,11 +1051,10 @@ "metadata": {}, "outputs": [], "source": [ - "orbits = simdata.pic_species[\"kinetic_ions\"][\"orbits\"]\n", - "print(f\"{orbits.shape = }\")\n", + "orbits = simdata.orbits[\"kinetic_ions\"]\n", "\n", - "Nt = orbits.shape[0]\n", - "Np = orbits.shape[1]" + "Nt = simdata.Nt[\"kinetic_ions\"]\n", + "Np = simdata.Np[\"kinetic_ions\"]" ] }, { diff --git a/src/struphy/main.py b/src/struphy/main.py index c7534c1c5..c5e947606 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Optional, TypedDict import os import sysconfig import time @@ -611,13 +611,50 @@ class SimData: """ def __init__(self, path: str): self.path = path - self.feec_species = {} - self.pic_species = {} + self._orbits = {} + self._spline_values = {} self.sph_species = {} self.grids_log: list[np.ndarray] = None self.grids_phy: list[np.ndarray] = None self.t_grid: np.ndarray = None + @property + def orbits(self) -> dict[str, np.ndarray]: + """Keys: species name. Values: 3d arrays indexed by (n, p, a), where 'n' is the time index, 'p' the particle index and 'a' the attribute index.""" + return self._orbits + + @property + def spline_values(self) -> dict[str, dict[str, np.ndarray]]: + """Keys: species name. Values: dicts of variable names with values being 3d arrays on the grid.""" + return self._spline_values + + @property + def Nt(self) -> dict[str, int]: + """Number of available time points (snap shots) for each species.""" + if not hasattr(self, "_Nt"): + self._Nt = {} + for spec, orbs in self.orbits.items(): + self._Nt[spec] = orbs.shape[0] + return self._Nt + + @property + def Np(self) -> dict[str, int]: + """Number of particle orbits for each species.""" + if not hasattr(self, "_Np"): + self._Np = {} + for spec, orbs in self.orbits.items(): + self._Np[spec] = orbs.shape[1] + return self._Np + + @property + def Nattr(self) -> dict[str, int]: + """Number of particle attributes for each species.""" + if not hasattr(self, "_Nattr"): + self._Nattr = {} + for spec, orbs in self.orbits.items(): + self._Nattr[spec] = orbs.shape[2] + return self._Nattr + @property def spline_grid_resolution(self): if self.grids_log is not None: @@ -666,7 +703,7 @@ def load_data(path: str) -> SimData: # species folders species = next(os.walk(path_fields))[1] for spec in species: - simdata.feec_species[spec] = {} + simdata.spline_values[spec] = {} # simdata.arrays[spec] = {} path_spec = os.path.join(path_fields, spec) wlk = os.walk(path_spec) @@ -677,7 +714,7 @@ def load_data(path: str) -> SimData: var = file.split(".")[0] with open(os.path.join(path_spec, file), "rb") as f: # try: - simdata.feec_species[spec][var] = pickle.load(f) + simdata.spline_values[spec][var] = pickle.load(f) # simdata.arrays[spec][var] = pickle.load(f) if os.path.exists(path_kinetic): @@ -686,7 +723,6 @@ def load_data(path: str) -> SimData: species = next(os.walk(path_kinetic))[1] print(f"{species = }") for spec in species: - simdata.pic_species[spec] = {} path_spec = os.path.join(path_kinetic, spec) wlk = os.walk(path_spec) sub_folders = next(wlk)[1] @@ -703,42 +739,39 @@ def load_data(path: str) -> SimData: step = int(file.split(".")[0].split("_")[-1]) tmp = np.load(os.path.join(path_dat, file)) if n == 0: - orbits = np.zeros((Nt, *tmp.shape), dtype=float) - orbits[step] = tmp + simdata.orbits[spec] = np.zeros((Nt, *tmp.shape), dtype=float) + simdata.orbits[spec][step] = tmp n += 1 - simdata.pic_species[spec][folder] = orbits else: - # simdata.pic_species[spec][folder] = {} - tmp = {} - for file in files: - # print(f"{file = }") - if ".npy" in file: - var = file.split(".")[0] - tmp[var] = np.load(os.path.join(path_dat, file)) - # sort dict - simdata.pic_species[spec][folder] = dict(sorted(tmp.items())) + raise NotImplementedError + # # simdata.pic_species[spec][folder] = {} + # tmp = {} + # for file in files: + # # print(f"{file = }") + # if ".npy" in file: + # var = file.split(".")[0] + # tmp[var] = np.load(os.path.join(path_dat, file)) + # # sort dict + # simdata.pic_species[spec][folder] = dict(sorted(tmp.items())) print("\nThe following data has been loaded:") print(f"{simdata.time_grid_size = }") print(f"{simdata.spline_grid_resolution = }") - print(f"simdata.feec_species:") - for k, v in simdata.feec_species.items(): - print(f" {k}:") - for kk, vv in v.items(): - print(f" {kk}") - print(f"simdata.pic_species:") - for k, v in simdata.pic_species.items(): + print(f"simdata.spline_values:") + for k, v in simdata.spline_values.items(): print(f" {k}:") for kk, vv in v.items(): print(f" {kk}") + print(f"simdata.orbits:") + for k, v in simdata.orbits.items(): + print(f" {k}") + # for kk, vv in v.items(): + # print(f" {kk}") print(f"simdata.sph_species:") return simdata - - - if __name__ == "__main__": import argparse import os From 9e99367a9ef1e900806d826f3f94f466f2f48679 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 9 Sep 2025 15:41:02 +0200 Subject: [PATCH 104/292] repair Maxwell and LinearMHD tests --- .../tutorial_01_parameter_files.ipynb | 51 ++++++++----------- src/struphy/main.py | 4 +- src/struphy/models/tests/test_LinearMHD.py | 10 ++-- src/struphy/models/tests/test_Maxwell.py | 42 +++++++-------- .../post_processing/post_processing_tools.py | 13 +++-- 5 files changed, 58 insertions(+), 62 deletions(-) diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/doc/tutorials/tutorial_01_parameter_files.ipynb index 69b1f56c6..fb316d85d 100644 --- a/doc/tutorials/tutorial_01_parameter_files.ipynb +++ b/doc/tutorials/tutorial_01_parameter_files.ipynb @@ -306,7 +306,7 @@ "source": [ "## Loading data: `main.load_data`\n", "\n", - "After post-processing, the generated data can be loaded via `main.load_data`:" + "After post-processing, the generated data can be loaded via `main.load_data`. This function returns a `SimData` object, which you can inspect to get further info on possible data to load." ] }, { @@ -321,42 +321,34 @@ }, { "cell_type": "markdown", - "id": "a9468479", - "metadata": {}, - "source": [ - "`main.load_data` returns a `SimData` object, which you can inspect to get further info on possible data to load:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0531679e", + "id": "08544a91", "metadata": {}, - "outputs": [], "source": [ - "for k, v in simdata.__dict__.items():\n", - " print(k, type(v))" + "## Plotting particle orbits\n", + "\n", + "In this example, for the species `kinetic_ions` some particle orbits have been saved. Under `simdata.orbits[]` one finds a three-dimensional numpy array; the first index refers to the time step, the second index to the particle and the third index to the particel attribute. The first three attributes are the particle positions, followed by the velocities and the (initial and time-dependent) weights." ] }, { "cell_type": "code", "execution_count": null, - "id": "a329eb38", + "id": "0ea08d24", "metadata": {}, "outputs": [], "source": [ - "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", - " print(f\"{k = }, {type(v) = }\")" + "orbits = simdata.orbits[\"kinetic_ions\"]\n", + "\n", + "Nt = simdata.Nt[\"kinetic_ions\"]\n", + "Np = simdata.Np[\"kinetic_ions\"]\n", + "Nattr = simdata.Nattr[\"kinetic_ions\"]" ] }, { "cell_type": "markdown", - "id": "08544a91", + "id": "001437bc", "metadata": {}, "source": [ - "## Plotting particle orbits\n", - "\n", - "In this example, for the species `kinetic_ions` some particle orbits have been saved. Let us plot them:" + "Let us plot the orbits:" ] }, { @@ -367,21 +359,20 @@ "outputs": [], "source": [ "from matplotlib import pyplot as plt\n", + "import numpy as np\n", "\n", "fig = plt.figure()\n", "ax = fig.gca()\n", "\n", "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", "\n", - "time = 0.\n", - "dt = time_opts.dt\n", + "# create alpha for color scaling\n", "Tend = time_opts.Tend\n", - "for k, v in simdata.pic_species[\"kinetic_ions\"][\"orbits\"].items():\n", - " # print(f\"{v[0] = }\")\n", - " alpha = (Tend - time)/Tend\n", - " for i, particle in enumerate(v):\n", - " ax.scatter(particle[0], particle[1], c=colors[i % 4], alpha=alpha)\n", - " time += dt\n", + "alpha = np.linspace(1., 0., Nt + 1)\n", + "\n", + "# loop through particles, plot all time steps\n", + "for i in range(Np):\n", + " ax.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], alpha=alpha)\n", " \n", "ax.plot([l1, l1], [l2, r2], 'k')\n", "ax.plot([r1, r1], [l2, r2], 'k')\n", @@ -391,7 +382,7 @@ "ax.set_ylabel('y')\n", "ax.set_xlim(-6.5, 6.5)\n", "ax.set_ylim(-9, 9)\n", - "ax.set_title(f'{int(Tend/dt)} time steps (full color at t=0)');" + "ax.set_title(f'{int(Nt - 1)} time steps (full color at t=0)');" ] } ], diff --git a/src/struphy/main.py b/src/struphy/main.py index c5e947606..46453c278 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -451,8 +451,8 @@ def pproc( # domain: Domain = pickle.load(f) # print(f"{domain = }") # print(f"{domain.params = }") - domain_dct: Domain = pickle.load(f) - domain = getattr(domains, domain_dct["name"])(**domain_dct["params"]) + domain_dct = pickle.load(f) + domain: Domain = getattr(domains, domain_dct["name"])(**domain_dct["params"]) # create post-processing folder path_pproc = os.path.join(path, "post_processing") diff --git a/src/struphy/models/tests/test_LinearMHD.py b/src/struphy/models/tests/test_LinearMHD.py index f8e17f826..cec893000 100644 --- a/src/struphy/models/tests/test_LinearMHD.py +++ b/src/struphy/models/tests/test_LinearMHD.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from struphy.io.options import EnvironmentOptions, Units, Time +from struphy.io.options import EnvironmentOptions, BaseUnits, Time from struphy.geometry import domains from struphy.fields_background import equils from struphy.topology import grids @@ -28,7 +28,7 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): env = EnvironmentOptions(out_folders=out_folders, sim_folder="slab_waves_1d") # units - units = Units() + base_units = BaseUnits() # time stepping time_opts = Time(dt=0.15, Tend=180.0) @@ -69,7 +69,7 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): main.run(model, params_path=None, env=env, - units=units, + base_units=base_units, time_opts=time_opts, domain=domain, equil=equil, @@ -87,7 +87,7 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): simdata = main.load_data(env.path_out) # first fft - u_of_t = simdata.feec_species["mhd"]["velocity_log"] + u_of_t = simdata.spline_values["mhd"]["velocity_log"] Bsquare = (B0x**2 + B0y**2 + B0z**2) p0 = beta * Bsquare / 2 @@ -121,7 +121,7 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): assert np.abs(coeffs[0][0] - v_alfven) < 0.07 # second fft - p_of_t = simdata.feec_species["mhd"]["pressure_log"] + p_of_t = simdata.spline_values["mhd"]["pressure_log"] _1, _2, _3, coeffs = power_spectrum_2d(p_of_t, "pressure_log", diff --git a/src/struphy/models/tests/test_Maxwell.py b/src/struphy/models/tests/test_Maxwell.py index 1f190ff7c..46401367b 100644 --- a/src/struphy/models/tests/test_Maxwell.py +++ b/src/struphy/models/tests/test_Maxwell.py @@ -5,7 +5,7 @@ from scipy.special import jv, yn from matplotlib import pyplot as plt -from struphy.io.options import EnvironmentOptions, Units, Time +from struphy.io.options import EnvironmentOptions, BaseUnits, Time from struphy.geometry import domains from struphy.fields_background import equils from struphy.topology import grids @@ -30,7 +30,7 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): env = EnvironmentOptions(out_folders=out_folders, sim_folder="light_wave_1d") # units - units = Units() + base_units = BaseUnits() # time stepping time_opts = Time(dt=0.05, Tend=50.0) @@ -57,18 +57,18 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0, seed=123)) model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1, seed=123)) - # start run - main.run(model, - params_path=None, - env=env, - units=units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=verbose, - ) + # # start run + # main.run(model, + # params_path=None, + # env=env, + # base_units=base_units, + # time_opts=time_opts, + # domain=domain, + # equil=equil, + # grid=grid, + # derham_opts=derham_opts, + # verbose=verbose, + # ) # post processing if MPI.COMM_WORLD.Get_rank() == 0: @@ -79,7 +79,7 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): simdata = main.load_data(env.path_out) # fft - E_of_t = simdata.feec_species["em_fields"]["e_field_log"] + E_of_t = simdata.spline_values["em_fields"]["e_field_log"] _1, _2, _3, coeffs = power_spectrum_2d(E_of_t, "e_field_log", grids=simdata.grids_log, @@ -110,7 +110,7 @@ def test_coaxial(do_plot: bool = False): env = EnvironmentOptions(out_folders=out_folders, sim_folder="coaxial") # units - units = Units() + base_units = BaseUnits() # time time_opts = Time(dt=0.05, Tend=10.0) @@ -149,7 +149,7 @@ def test_coaxial(do_plot: bool = False): main.run(model, params_path=None, env=env, - units=units, + base_units=base_units, time_opts=time_opts, domain=domain, equil=equil, @@ -175,8 +175,8 @@ def test_coaxial(do_plot: bool = False): t_grid = simdata.t_grid grids_phy = simdata.grids_phy - e_field_phy = simdata.arrays["em_fields"]["e_field_phy"] - b_field_phy = simdata.arrays["em_fields"]["b_field_phy"] + e_field_phy = simdata.spline_values["em_fields"]["e_field_phy"] + b_field_phy = simdata.spline_values["em_fields"]["b_field_phy"] X = grids_phy[0][:, :, 0] Y = grids_phy[1][:, :, 0] @@ -262,5 +262,5 @@ def to_E_theta(X, Y, E_x, E_y): if __name__ == "__main__": - test_light_wave_1d(algo="explicit", do_plot=True) - # test_coaxial(do_plot=True) \ No newline at end of file + # test_light_wave_1d(algo="explicit", do_plot=True) + test_coaxial(do_plot=True) \ No newline at end of file diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index fc6dfc261..f1a60d888 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -7,6 +7,8 @@ import pickle from struphy.kinetic_background import maxwellians +from struphy.fields_background.base import FluidEquilibrium +from struphy.fields_background import equils from struphy.io.setup import import_parameters_py from struphy.models.base import setup_derham from struphy.feec.psydac_derham import SplineFunction @@ -69,11 +71,14 @@ def get_params_of_run(path: str) -> ParamsIn: with open(os.path.join(path, "domain.bin"), "rb") as f: # WORKAROUND: cannot pickle pyccelized classes at the moment domain_dct = pickle.load(f) - name = domain_dct["name"] - params_map = domain_dct["params_map"] - domain = getattr(domains, name)(**params_map) + domain: Domain = getattr(domains, domain_dct["name"])(**domain_dct["params"]) with open(os.path.join(path, "equil.bin"), "rb") as f: - equil = pickle.load(f) + # WORKAROUND: cannot pickle pyccelized classes at the moment + equil_dct = pickle.load(f) + if equil_dct: + equil: FluidEquilibrium = getattr(equils, equil_dct["name"])(**equil_dct["params"]) + else: + equil = None with open(os.path.join(path, "grid.bin"), "rb") as f: grid = pickle.load(f) with open(os.path.join(path, "derham_opts.bin"), "rb") as f: From bdcaa7f645a9923b62343f0137753a306ad1b2ed Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 9 Sep 2025 17:04:09 +0200 Subject: [PATCH 105/292] Two new abstract methods for Propagator: class Variables: and __init__(self): The method `assign_variables` is thus obsolete. This leads to a new way of how the variables to be updated by a propagator are defined. --- src/struphy/models/fluid.py | 16 +- src/struphy/models/toy.py | 91 ++++++---- src/struphy/propagators/base.py | 36 +++- src/struphy/propagators/propagators_fields.py | 158 +++++++++++++----- .../propagators/propagators_markers.py | 54 ++++-- 5 files changed, 254 insertions(+), 101 deletions(-) diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index f5829debe..254274b5d 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -78,16 +78,12 @@ def __init__(self): self.propagators = self.Propagators() # 3. assign variables to propagators - self.propagators.shear_alf.assign_variables( - u = self.mhd.velocity, - b = self.em_fields.b_field, - ) - - self.propagators.mag_sonic.assign_variables( - n = self.mhd.density, - u = self.mhd.velocity, - p = self.mhd.pressure, - ) + self.propagators.shear_alf.variables.u = self.mhd.velocity + self.propagators.shear_alf.variables.b = self.em_fields.b_field + + self.propagators.mag_sonic.variables.n = self.mhd.density + self.propagators.mag_sonic.variables.u = self.mhd.velocity + self.propagators.mag_sonic.variables.p = self.mhd.pressure # define scalars for update_scalar_quantities self.add_scalar("en_U") diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 8dcbcf435..9ff7d8461 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -35,6 +35,7 @@ class Maxwell(StruphyModel): :ref:`Model info `: """ + ## species class EMFields(FieldSpecies): @@ -62,10 +63,8 @@ def __init__(self): self.propagators = self.Propagators() # 3. assign variables to propagators - self.propagators.maxwell.assign_variables( - e = self.em_fields.e_field, - b = self.em_fields.b_field, - ) + self.propagators.maxwell.variables.e = self.em_fields.e_field + self.propagators.maxwell.variables.b = self.em_fields.b_field # define scalars for update_scalar_quantities self.add_scalar("electric energy") @@ -114,6 +113,7 @@ class Vlasov(StruphyModel): :ref:`Model info `: """ + ## species class KineticIons(KineticSpecies): @@ -141,13 +141,8 @@ def __init__(self): self.propagators = self.Propagators() # 3. assign variables to propagators - self.propagators.push_vxb.assign_variables( - ions = self.kinetic_ions.var, - ) - - self.propagators.push_eta.assign_variables( - var = self.kinetic_ions.var, - ) + self.propagators.push_vxb.variables.ions = self.kinetic_ions.var + self.propagators.push_eta.variables.var = self.kinetic_ions.var # define scalars for update_scalar_quantities self.add_scalar("en_f", compute="from_particles", variable=self.kinetic_ions.var) @@ -208,35 +203,65 @@ class GuidingCenter(StruphyModel): :ref:`Model info `: """ + + ## species + + class KineticIons(KineticSpecies): + def __init__(self): + self.var = PICVariable(space="Particles5D") + self.init_variables() + + ## propagators + + class Propagators: + def __init__(self): + self.push_bxe = propagators_markers.PushGuidingCenterBxEstar() + self.push_parallel = propagators_markers.PushGuidingCenterParallel() - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## abstract methods - dct["kinetic"]["ions"] = "Particles5D" - return dct + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}' ***") + + # 1. instantiate all species, variables + self.kinetic_ions = self.KineticIons() - @staticmethod - def bulk_species(): - return "ions" + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.push_bxe.assign_variables( + ions = self.kinetic_ions.var, + ) + + self.propagators.push_parallel.assign_variables( + var = self.kinetic_ions.var, + ) + + # define scalars for update_scalar_quantities + self.add_scalar("en_f", compute="from_particles", variable=self.kinetic_ions.var) - @staticmethod - def velocity_scale(): + @property + def bulk_species(self): + return self.kinetic_ions + + @property + def velocity_scale(self): return "alfvén" + + def allocate_helpers(self): + self._tmp = np.empty(1, dtype=float) - @staticmethod - def propagators_dct(): - return { - propagators_markers.PushGuidingCenterBxEstar: ["ions"], - propagators_markers.PushGuidingCenterParallel: ["ions"], - } + def update_scalar_quantities(self): + particles = self.kinetic_ions.var.particles + self._tmp[0] = particles.markers_wo_holes[:, 6].dot( + particles.markers_wo_holes[:, 3] ** 2 + + particles.markers_wo_holes[:, 4] ** 2 + + particles.markers_wo_holes[:, 5] ** 2, + ) / (2 * particles.Np) - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + self.update_scalar("en_f", self._tmp[0]) def __init__(self, params, comm, clone_config=None): # initialize base class diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 6f802f2fd..a40a49b48 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -23,6 +23,26 @@ class Propagator(metaclass=ABCMeta): in one of the modules ``propagators_fields.py``, ``propagators_markers.py`` or ``propagators_coupling.py``. Only propagators that update both a FEEC and a PIC species go into ``propagators_coupling.py``. """ + @abstractmethod + def __init__(self): + # variables to be updated + self.variables = self.Variables() + + @abstractmethod + class Variables: + """Define variable names and types to be updated by the propagator.""" + def __init__(self): + self._var1 = None + + @property + def var1(self): + return self._var1 + + @var1.setter + def var1(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles6D" + self._var1 = new @abstractmethod def set_options(self, **kwargs): @@ -71,9 +91,9 @@ def update_feec_variables(self, **new_coeffs): max_diff for all feec variables. """ diffs = {} - assert len(new_coeffs) == len(self.variables), f"Coefficients must be passed in the following order: {self.variables}, but is {new_coeffs}." - for ((k, new), (kp, vp)) in zip(new_coeffs.items(), self.variables.items()): - assert k == kp, f"Variable name '{k}' not equal to '{kp}'; variables must be passed in the order {self.variables}." + assert len(new_coeffs) == len(self.variables.__dict__), f"Coefficients must be passed in the following order: {self.variables}, but is {new_coeffs}." + for ((k, new), (kp, vp)) in zip(new_coeffs.items(), self.variables.__dict__.items()): + assert k in kp, f"Variable name '{k}' not in '{kp}'; variables must be passed in the order {self.variables.__dict__}." assert isinstance(new, (StencilVector, BlockVector)) assert isinstance(vp, FEECVariable) old = vp.spline.vector @@ -90,11 +110,11 @@ def update_feec_variables(self, **new_coeffs): return diffs - @property - def variables(self): - """Dict of Variables to be updated by the propagator. - """ - return self._variables + # @property + # def variables(self): + # """Dict of Variables to be updated by the propagator. + # """ + # return self._variables @property def init_kernels(self): diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index aca2b66a0..82857d7af 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -50,7 +50,6 @@ from struphy.models.variables import FEECVariable, PICVariable, SPHVariable -@dataclass class Maxwell(Propagator): r""":ref:`FEEC ` discretization of the following equations: find :math:`\mathbf E \in H(\textnormal{curl})` and :math:`\mathbf B \in H(\textnormal{div})` such that @@ -63,9 +62,34 @@ class Maxwell(Propagator): :ref:`time_discret`: Crank-Nicolson (implicit mid-point). System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`. """ - # variables to be updated - e: FEECVariable = None - b: FEECVariable = None + class Variables: + def __init__(self): + self._e: FEECVariable = None + self._b: FEECVariable = None + + @property + def e(self) -> FEECVariable: + return self._e + + @e.setter + def e(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._e = new + + @property + def b(self) -> FEECVariable: + return self._b + + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._b = new + + def __init__(self): + # variables to be updated + self.variables = self.Variables() # propagator specific options OptsAlgo = Literal["implicit", "explicit"] @@ -149,8 +173,8 @@ def allocate(self): weak_curl = M1_inv @ curl.T @ M2 # allocate output of vector field - out1 = self.e.spline.vector.space.zeros() - out2 = self.b.spline.vector.space.zeros() + out1 = self.variables.e.spline.vector.space.zeros() + out2 = self.variables.b.spline.vector.space.zeros() def f1(t, y1, y2, out: BlockVector = out1): weak_curl.dot(y2, out=out) @@ -163,18 +187,18 @@ def f2(t, y1, y2, out: BlockVector = out2): out.update_ghost_regions() return out - vector_field = {self.e.spline.vector: f1, self.b.spline.vector: f2} + vector_field = {self.variables.e.spline.vector: f1, self.variables.b.spline.vector: f2} self._ode_solver = ODEsolverFEEC(vector_field, butcher=self.options.butcher) # allocate place-holder vectors to avoid temporary array allocations in __call__ - self._e_tmp1 = self.e.spline.vector.space.zeros() - self._e_tmp2 = self.e.spline.vector.space.zeros() - self._b_tmp1 = self.b.spline.vector.space.zeros() + self._e_tmp1 = self.variables.e.spline.vector.space.zeros() + self._e_tmp2 = self.variables.e.spline.vector.space.zeros() + self._b_tmp1 = self.variables.b.spline.vector.space.zeros() def __call__(self, dt): # current FE coeffs - en = self.e.spline.vector - bn = self.b.spline.vector + en = self.variables.e.spline.vector + bn = self.variables.b.spline.vector if self.options.algo == "implicit": # solve for new e coeffs @@ -419,7 +443,6 @@ def __call__(self, dt): print() -@dataclass class ShearAlfven(Propagator): r""":ref:`FEEC ` discretization of the following equations: find :math:`\mathbf U \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\}` and :math:`\mathbf B \in H(\textnormal{div})` such that @@ -443,9 +466,34 @@ class ShearAlfven(Propagator): where :math:`\alpha \in \{1, 2, v\}` and :math:`\mathbb M^\rho_\alpha` is a weighted mass matrix in :math:`\alpha`-space, the weight being :math:`\rho_0`, the MHD equilibirum density. The solution of the above system is based on the :ref:`Schur complement `. """ - # variables to be updated - u: FEECVariable = None - b: FEECVariable = None + class Variables: + def __init__(self): + self._u: FEECVariable = None + self._b: FEECVariable = None + + @property + def u(self) -> FEECVariable: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new + + @property + def b(self) -> FEECVariable: + return self._b + + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._b = new + + def __init__(self): + # variables to be updated + self.variables = self.Variables() # propagator specific options OptsAlgo = Literal["implicit", "explicit"] @@ -538,8 +586,8 @@ def allocate(self): _f2 = curl @ _T # allocate output of vector field - out1 = self.u.spline.vector.space.zeros() - out2 = self.b.spline.vector.space.zeros() + out1 = self.variables.u.spline.vector.space.zeros() + out2 = self.variables.b.spline.vector.space.zeros() def f1(t, y1, y2, out: BlockVector = out1): _f1.dot(y2, out=out) @@ -552,18 +600,18 @@ def f2(t, y1, y2, out: BlockVector = out2): out.update_ghost_regions() return out - vector_field = {self.u.spline.vector: f1, self.b.spline.vector: f2} + vector_field = {self.variables.u.spline.vector: f1, self.variables.b.spline.vector: f2} self._ode_solver = ODEsolverFEEC(vector_field, butcher=self.options.butcher) # allocate dummy vectors to avoid temporary array allocations - self._u_tmp1 = self.u.spline.vector.space.zeros() - self._u_tmp2 = self.u.spline.vector.space.zeros() - self._b_tmp1 = self.b.spline.vector.space.zeros() + self._u_tmp1 = self.variables.u.spline.vector.space.zeros() + self._u_tmp2 = self.variables.u.spline.vector.space.zeros() + self._b_tmp1 = self.variables.b.spline.vector.space.zeros() def __call__(self, dt): # current FE coeffs - un = self.u.spline.vector - bn = self.b.spline.vector + un = self.variables.u.spline.vector + bn = self.variables.b.spline.vector if self.options.algo == "implicit": # solve for new u coeffs @@ -840,7 +888,6 @@ def __call__(self, dt): print() -@dataclass class Magnetosonic(Propagator): r""" :ref:`FEEC ` discretization of the following equations: @@ -874,10 +921,45 @@ class Magnetosonic(Propagator): \boldsymbol{\rho}^{n+1} = \boldsymbol{\rho}^n - \frac{\Delta t}{2} \mathbb D \mathcal Q^\alpha (\mathbf u^{n+1} + \mathbf u^n) \,. """ - # variables to be updated - n: FEECVariable = None - u: FEECVariable = None - p: FEECVariable = None + class Variables: + def __init__(self): + self._n: FEECVariable = None + self._u: FEECVariable = None + self._p: FEECVariable = None + + @property + def n(self) -> FEECVariable: + return self._n + + @n.setter + def n(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._n = new + + @property + def u(self) -> FEECVariable: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new + + @property + def p(self) -> FEECVariable: + return self._p + + @p.setter + def p(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._p = new + + def __init__(self): + # variables to be updated + self.variables = self.Variables() ## abstract methods def set_options(self, @@ -931,8 +1013,8 @@ def allocate(self): _K = getattr(self.basis_ops, id_K) if id_U is None: - _U = IdentityOperator(self.u.spline.vector.space) - _UT = IdentityOperator(self.u.spline.vector.space) + _U = IdentityOperator(self.variables.u.spline.vector.space) + _UT = IdentityOperator(self.variables.u.spline.vector.space) else: _U = getattr(self.basis_ops, id_U) _UT = _U.T @@ -964,10 +1046,10 @@ def allocate(self): ) # allocate dummy vectors to avoid temporary array allocations - self._u_tmp1 = self.u.spline.vector.space.zeros() - self._u_tmp2 = self.u.spline.vector.space.zeros() - self._p_tmp1 = self.p.spline.vector.space.zeros() - self._n_tmp1 = self.n.spline.vector.space.zeros() + self._u_tmp1 = self.variables.u.spline.vector.space.zeros() + self._u_tmp2 = self.variables.u.spline.vector.space.zeros() + self._p_tmp1 = self.variables.p.spline.vector.space.zeros() + self._n_tmp1 = self.variables.n.spline.vector.space.zeros() self._b_tmp1 = self._b.space.zeros() self._byn1 = self._B.codomain.zeros() @@ -975,9 +1057,9 @@ def allocate(self): def __call__(self, dt): # current FE coeffs - nn = self.n.spline.vector - un = self.u.spline.vector - pn = self.p.spline.vector + nn = self.variables.n.spline.vector + un = self.variables.u.spline.vector + pn = self.variables.p.spline.vector # solve for new u coeffs (no tmps created here) byn1 = self._B.dot(pn, out=self._byn1) diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index aecdc95df..0259e0a6a 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -25,7 +25,6 @@ from psydac.linalg.basic import LinearOperator -@dataclass class PushEta(Propagator): r"""For each marker :math:`p`, solves @@ -43,8 +42,22 @@ class PushEta(Propagator): * Explicit from :class:`~struphy.ode.utils.ButcherTableau` """ - # variables to be updated - var: PICVariable = None + class Variables: + def __init__(self): + self._var: PICVariable = None + + @property + def var(self) -> PICVariable: + return self._var + + @var.setter + def var(self, new): + assert isinstance(new, PICVariable) + self._var = new + + def __init__(self): + # variables to be updated + self.variables = self.Variables() ## abstract methods def set_options(self, @@ -78,7 +91,7 @@ def allocate(self): ) self._pusher = Pusher( - self.var.particles, + self.variables.var.particles, kernel, args_kernel, self.domain.args_domain, @@ -91,11 +104,10 @@ def __call__(self, dt): self._pusher(dt) # update_weights - if self.var.particles.control_variate: - self.var.particles.update_weights() + if self.variables.var.particles.control_variate: + self.variables.var.particles.update_weights() -@dataclass class PushVxB(Propagator): r"""For each marker :math:`p`, solves @@ -112,8 +124,23 @@ class PushVxB(Propagator): Available algorithms: ``analytic``, ``implicit``. """ - # variables to be updated - ions: PICVariable = None + class Variables: + def __init__(self): + self._ions: PICVariable = None + + @property + def ions(self) -> PICVariable: + return self._ions + + @ions.setter + def ions(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles6D" + self._ions = new + + def __init__(self): + # variables to be updated + self.variables = self.Variables() # propagator specific options OptsAlgo = Literal["analytic", "implicit"] @@ -170,7 +197,7 @@ def allocate(self): ) self._pusher = Pusher( - self.ions.particles, + self.variables.ions.particles, kernel, args_kernel, self.domain.args_domain, @@ -194,8 +221,8 @@ def __call__(self, dt): self._pusher(self.options.kappa * dt) # update_weights - if self.ions.particles.control_variate: - self.ions.particles.update_weights() + if self.variables.ions.particles.control_variate: + self.variables.ions.particles.update_weights() class PushVinEfield(Propagator): @@ -365,6 +392,7 @@ def __call__(self, dt): self.particles[0].update_weights() +@dataclass class PushGuidingCenterBxEstar(Propagator): r"""For each marker :math:`p`, solves @@ -394,6 +422,8 @@ class PushGuidingCenterBxEstar(Propagator): * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_1st_order_newton` * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_2nd_order` """ + # variables to be updated + ions: PICVariable = None @staticmethod def options(default=False): From 56e3c9c1e80d6413a610f85970c81ddc117ef83f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 10 Sep 2025 08:01:38 +0200 Subject: [PATCH 106/292] New way of setting Propagator options through three abstract methods: - class Options - def options(self: - @options.setter The benefit is that now self.options is an Options object, which can be well documented. Model Maxwell has been adapted. --- src/struphy/io/options.py | 1 + src/struphy/models/base.py | 2 +- src/struphy/models/fluid.py | 1 + src/struphy/models/tests/test_Maxwell.py | 30 +++--- src/struphy/propagators/base.py | 92 +++++++---------- src/struphy/propagators/propagators_fields.py | 67 +++++++------ .../propagators/propagators_markers.py | 99 +++++++++++-------- 7 files changed, 151 insertions(+), 141 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 94af30877..fb91d4d4d 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -34,6 +34,7 @@ OptsRecontructBC = Literal["periodic", "mirror", "fixed"] OptsLoading = Literal["pseudo_random", 'sobol_standard', 'sobol_antithetic', 'external', 'restart', "tesselation"] OptsSpatialLoading = Literal["uniform", "disc"] +OptsMPIsort = Literal["each", "last", None] ## Option classes diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index b7375a1c6..d28cfef36 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1385,7 +1385,7 @@ def generate_default_parameter_file( file.write("\n# propagator options\n") for prop in self.propagators.__dict__: - file.write(f"model.propagators.{prop}.set_options()\n") + file.write(f"model.propagators.{prop}.options = model.propagators.{prop}.Options()\n") file.write("\n# initial conditions (background + perturbation)\n") if has_feec: diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 254274b5d..680f9cae5 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -43,6 +43,7 @@ class LinearMHD(StruphyModel): :ref:`Model info `: """ + ## species class EMFields(FieldSpecies): diff --git a/src/struphy/models/tests/test_Maxwell.py b/src/struphy/models/tests/test_Maxwell.py index 46401367b..297092f4f 100644 --- a/src/struphy/models/tests/test_Maxwell.py +++ b/src/struphy/models/tests/test_Maxwell.py @@ -51,24 +51,24 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): model = Maxwell() # propagator options - model.propagators.maxwell.set_options(algo=algo) + model.propagators.maxwell.options = model.propagators.maxwell.Options(algo=algo) # initial conditions (background + perturbation) model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0, seed=123)) model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1, seed=123)) # # start run - # main.run(model, - # params_path=None, - # env=env, - # base_units=base_units, - # time_opts=time_opts, - # domain=domain, - # equil=equil, - # grid=grid, - # derham_opts=derham_opts, - # verbose=verbose, - # ) + main.run(model, + params_path=None, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) # post processing if MPI.COMM_WORLD.Get_rank() == 0: @@ -116,8 +116,8 @@ def test_coaxial(do_plot: bool = False): time_opts = Time(dt=0.05, Tend=10.0) # geometry - a1=2.326744 - a2=3.686839 + a1 = 2.326744 + a2 = 3.686839 Lz = 2.0 domain = domains.HollowCylinder(a1=a1, a2=a2, Lz=Lz) @@ -137,7 +137,7 @@ def test_coaxial(do_plot: bool = False): model = Maxwell() # propagator options - model.propagators.maxwell.set_options(algo="implicit") + model.propagators.maxwell.options = model.propagators.maxwell.Options(algo="implicit") # initial conditions (background + perturbation) m = 3 diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index a40a49b48..412244cba 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -3,6 +3,8 @@ from abc import ABCMeta, abstractmethod import numpy as np from mpi4py import MPI +from dataclasses import dataclass +from typing import Literal from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import WeightedMassOperators @@ -12,6 +14,7 @@ from psydac.linalg.stencil import StencilVector from psydac.linalg.block import BlockVector from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB +from struphy.io.options import check_option class Propagator(metaclass=ABCMeta): @@ -23,11 +26,6 @@ class Propagator(metaclass=ABCMeta): in one of the modules ``propagators_fields.py``, ``propagators_markers.py`` or ``propagators_coupling.py``. Only propagators that update both a FEEC and a PIC species go into ``propagators_coupling.py``. """ - @abstractmethod - def __init__(self): - # variables to be updated - self.variables = self.Variables() - @abstractmethod class Variables: """Define variable names and types to be updated by the propagator.""" @@ -43,10 +41,39 @@ def var1(self, new): assert isinstance(new, PICVariable) assert new.space == "Particles6D" self._var1 = new - + + @abstractmethod + def __init__(self): + self.variables = self.Variables() + + @abstractmethod + @dataclass + class Options: + # define specific literals + OptsTemplate = Literal["implicit", "explicit"] + # propagator options + opt1: str = "implicit", + + def __post_init__(self): + # checks + check_option(self.opt1, self.OptsTemplate) + + @property + @abstractmethod + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter @abstractmethod - def set_options(self, **kwargs): - """Set the dynamical options of the propagator (kwargs).""" + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nInstance of propagator '{self.__class__.__name__}' with:") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new @abstractmethod def allocate(): @@ -62,24 +89,6 @@ def __call__(self, dt): dt : float Time step size. """ - - def assign_variables(self, **vars): - """Assign the variables to be updated by the propagator. - - Update variables in __dict__ and set self.variables with user-defined instance variables (allocated). - """ - assert len(vars) == len(self.__dict__), f"Variables must be passed in the following order: {self.__dict__}, but is {vars}." - for ((k, v), (kp, vp)) in zip(vars.items(), self.__dict__.items()): - assert k == kp, f"Variable name '{k}' not equal to '{kp}'; variables must be passed in the order {self.__dict__}." - assert isinstance(v, Variable) - if isinstance(v, (PICVariable, SPHVariable)): - # comm = var.obj.mpi_comm - pass - elif isinstance(v, FEECVariable): - # comm = var.obj.comm - pass - self.__dict__.update(vars) - self._variables = vars def update_feec_variables(self, **new_coeffs): r"""Return max_diff = max(abs(new - old)) for each new_coeffs, @@ -110,12 +119,6 @@ def update_feec_variables(self, **new_coeffs): return diffs - # @property - # def variables(self): - # """Dict of Variables to be updated by the propagator. - # """ - # return self._variables - @property def init_kernels(self): r"""List of initialization kernels for evaluation at @@ -303,27 +306,4 @@ def add_eval_kernel( comps, args_eval, ) - ] - - @property - def options(self): - if not hasattr(self, "_options"): - self._options = self.Options() - return self._options - - @options.setter - def options(self, new): - assert isinstance(new, self.Options) - self._options = new - - class Options: - def __init__(self, prop, **kwargs): - self.__dict__.update(kwargs) - if MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nInstance of propagator '{prop.__class__.__name__}' with:") - for k, v in self.__dict__.items(): - print(f' {k}: {v}') - - # @property - # def kwargs(self): - # return self._kwargs \ No newline at end of file + ] \ No newline at end of file diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 82857d7af..eed1659a7 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -4,6 +4,7 @@ from copy import deepcopy from dataclasses import dataclass from typing import Literal, get_args +import copy import numpy as np import scipy as sc @@ -88,40 +89,46 @@ def b(self, new): self._b = new def __init__(self): - # variables to be updated self.variables = self.Variables() - # propagator specific options - OptsAlgo = Literal["implicit", "explicit"] - - ## abstract methods - def set_options(self, - algo: OptsAlgo = "implicit", - solver: OptsSymmSolver = "pcg", - precond: OptsMassPrecond = "MassMatrixPreconditioner", - solver_params: SolverParameters = None, - butcher: ButcherTableau = None, - ): - - # checks - check_option(algo, self.OptsAlgo) - check_option(solver, OptsSymmSolver) - check_option(precond, OptsMassPrecond) + @dataclass + class Options: + # define specific literals + OptsAlgo = Literal["implicit", "explicit"] + # propagator options + algo: OptsAlgo = "implicit" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + butcher: ButcherTableau = None - # defaults - if solver_params is None: - solver_params = SolverParameters() + def __post_init__(self): + # checks + check_option(self.algo, self.OptsAlgo) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) - if algo == "explicit" and butcher is None: - butcher = ButcherTableau() - - # use setter for options - self.options = self.Options(self, - algo=algo, - solver=solver, - precond=precond, - solver_params=solver_params, - butcher=butcher,) + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.algo == "explicit" and self.butcher is None: + self.butcher = ButcherTableau() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new def allocate(self): # obtain needed matrices diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 0259e0a6a..2843c1974 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -3,7 +3,7 @@ import numpy as np from numpy import array, polynomial, random from typing import Literal, get_args -from dataclasses import dataclass +import copy from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -21,7 +21,7 @@ from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator from struphy.models.variables import FEECVariable, PICVariable -from struphy.io.options import check_option +from struphy.io.options import check_option, OptsMPIsort from psydac.linalg.basic import LinearOperator @@ -67,17 +67,16 @@ def set_options(self, if butcher is None: butcher = ButcherTableau() - # use setter for options - self.options = self.Options(self, - butcher=butcher, - ) + # setter + self.options = copy.deepcopy(locals()) + self.options.pop("self") def allocate(self): # get kernel kernel = pusher_kernels.push_eta_stage # define algorithm - butcher: ButcherTableau = self.options.butcher + butcher: ButcherTableau = self.options["butcher"] # temp fix due to refactoring of ButcherTableau: import numpy as np @@ -155,12 +154,9 @@ def set_options(self, # checks check_option(algo, self.OptsAlgo) - # use setter for options - self.options = self.Options(self, - algo=algo, - kappa=kappa, - b2_var=b2_var, - ) + # setter + self.options = copy.deepcopy(locals()) + self.options.pop("self") def allocate(self): # TODO: treat PolarVector as well, but polar splines are being reworked at the moment @@ -392,7 +388,6 @@ def __call__(self, dt): self.particles[0].update_weights() -@dataclass class PushGuidingCenterBxEstar(Propagator): r"""For each marker :math:`p`, solves @@ -422,32 +417,58 @@ class PushGuidingCenterBxEstar(Propagator): * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_1st_order_newton` * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_2nd_order` """ - # variables to be updated - ions: PICVariable = None - - @staticmethod - def options(default=False): - dct = {} - dct["algo"] = { - "method": [ - "discrete_gradient_2nd_order", + class Variables: + def __init__(self): + self._ions: PICVariable = None + + @property + def ions(self) -> PICVariable: + return self._ions + + @ions.setter + def ions(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles5D" + self._ions = new + + def __init__(self): + # variables to be updated + self.variables = self.Variables() + + # propagator specific options + OptsAlgo = Literal["discrete_gradient_2nd_order", "discrete_gradient_1st_order", - "discrete_gradient_1st_order_newton", - "rk4", - "forward_euler", - "heun2", - "rk2", - "heun3", - ], - "maxiter": 20, - "tol": 1e-7, - "mpi_sort": "each", - "verbose": False, - } - if default: - dct = descend_options_dict(dct, []) + "discrete_gradient_1st_order_newton", "explicit",] + + ## abstract methods + def set_options(self, + algo: OptsAlgo = "discrete_gradient_1st_order", + butcher: ButcherTableau = None, + maxiter: int = 20, + tol: float = 1e-7, + mpi_sort: OptsMPIsort = "each", + verbose: bool = False, + ): + + # checks + check_option(algo, self.OptsAlgo) + check_option(mpi_sort, OptsMPIsort) + + if algo == "explicit" and butcher is None: + butcher = ButcherTableau() + + # use setter for options + self.options = self.Options(self, + algo=algo, + butcher=butcher, + maxiter=maxiter, + tol=tol, + mpi_sort=mpi_sort, + verbose=verbose, + ) - return dct + def allocate(self): + pass def __init__( self, @@ -457,7 +478,7 @@ def __init__( evaluate_e_field: bool = False, b_tilde: BlockVector = None, epsilon: float = 1.0, - algo: dict = options(default=True)["algo"], + algo: dict = None, ): super().__init__(particles) From dc05b3a1fe79a9b6ae7b26500fbf6ab28864979f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 10 Sep 2025 08:26:56 +0200 Subject: [PATCH 107/292] adapt LinearMHD to new options setting --- src/struphy/models/fluid.py | 4 +- src/struphy/models/tests/test_LinearMHD.py | 4 +- src/struphy/propagators/base.py | 4 +- src/struphy/propagators/propagators_fields.py | 136 ++++++++++-------- 4 files changed, 79 insertions(+), 69 deletions(-) diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 680f9cae5..c307e49f9 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -149,8 +149,8 @@ def generate_default_parameter_file(self, path = None, prompt = True): new_file = [] with open(params_path, "r") as f: for line in f: - if "mag_sonic.set_options" in line: - new_file += ["model.propagators.mag_sonic.set_options(b_field=model.em_fields.b_field)\n"] + if "mag_sonic.Options" in line: + new_file += ["model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n"] else: new_file += [line] diff --git a/src/struphy/models/tests/test_LinearMHD.py b/src/struphy/models/tests/test_LinearMHD.py index cec893000..5b994fd64 100644 --- a/src/struphy/models/tests/test_LinearMHD.py +++ b/src/struphy/models/tests/test_LinearMHD.py @@ -57,8 +57,8 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): model.mhd.set_phys_params() # propagator options - model.propagators.shear_alf.set_options(algo=algo) - model.propagators.mag_sonic.set_options(b_field=model.em_fields.b_field) + model.propagators.shear_alf.options = model.propagators.shear_alf.Options(algo=algo) + model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field) # initial conditions (background + perturbation) model.mhd.velocity.add_perturbation(perturbations.Noise(amp=0.1, comp=0, seed=123)) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 412244cba..da092568a 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -49,7 +49,7 @@ def __init__(self): @abstractmethod @dataclass class Options: - # define specific literals + # specific literals OptsTemplate = Literal["implicit", "explicit"] # propagator options opt1: str = "implicit", @@ -70,7 +70,7 @@ def options(self) -> Options: def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nInstance of propagator '{self.__class__.__name__}' with:") + print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): print(f' {k}: {v}') self._options = new diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index eed1659a7..8ebaf991e 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -93,7 +93,7 @@ def __init__(self): @dataclass class Options: - # define specific literals + # specific literals OptsAlgo = Literal["implicit", "explicit"] # propagator options algo: OptsAlgo = "implicit" @@ -499,43 +499,48 @@ def b(self, new): self._b = new def __init__(self): - # variables to be updated self.variables = self.Variables() - # propagator specific options - OptsAlgo = Literal["implicit", "explicit"] - - ## abstract methods - def set_options(self, - u_space: OptsVecSpace = "Hdiv", - algo: OptsAlgo = "implicit", - solver: OptsSymmSolver = "pcg", - precond: OptsMassPrecond = "MassMatrixPreconditioner", - solver_params: SolverParameters = None, - butcher: ButcherTableau = None, - ): + @dataclass + class Options: + # specific literals + OptsAlgo = Literal["implicit", "explicit"] + # propagator options + u_space: OptsVecSpace = "Hdiv" + algo: OptsAlgo = "implicit" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + butcher: ButcherTableau = None - # checks - check_option(u_space, OptsVecSpace) - check_option(algo, self.OptsAlgo) - check_option(solver, OptsSymmSolver) - check_option(precond, OptsMassPrecond) - - # defaults - if solver_params is None: - solver_params = SolverParameters() + def __post_init__(self): + # checks + check_option(self.u_space, OptsVecSpace) + check_option(self.algo, self.OptsAlgo) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) - if algo == "explicit" and butcher is None: - butcher = ButcherTableau() - - # use setter for options - self.options = self.Options(self, - u_space=u_space, - algo=algo, - solver=solver, - precond=precond, - solver_params=solver_params, - butcher=butcher,) + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.algo == "explicit" and self.butcher is None: + self.butcher = ButcherTableau() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new def allocate(self): u_space = self.options.u_space @@ -965,38 +970,42 @@ def p(self, new): self._p = new def __init__(self): - # variables to be updated self.variables = self.Variables() - - ## abstract methods - def set_options(self, - b_field: FEECVariable = None, - u_space: OptsVecSpace = "Hdiv", - solver: OptsGenSolver = "pbicgstab", - precond: OptsMassPrecond = "MassMatrixPreconditioner", - solver_params: SolverParameters = None, - ): - - # checks - check_option(u_space, OptsVecSpace) - check_option(solver, OptsGenSolver) - check_option(precond, OptsMassPrecond) - # defaults - if b_field is None: - b_field = FEECVariable(space="Hdiv") - b_field.allocate(self.derham, self.domain) - if solver_params is None: - solver_params = SolverParameters() + @dataclass + class Options: + b_field: FEECVariable = None + u_space: OptsVecSpace = "Hdiv" + solver: OptsGenSolver = "pbicgstab" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None - # use setter for options - self.options = self.Options(self, - u_space=u_space, - b_field=b_field, - solver=solver, - precond=precond, - solver_params=solver_params, - ) + def __post_init__(self): + # checks + check_option(self.u_space, OptsVecSpace) + check_option(self.solver, OptsGenSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.b_field is None: + self.b_field = FEECVariable(space="Hdiv") + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new def allocate(self): u_space = self.options.u_space @@ -1032,6 +1041,7 @@ def allocate(self): self._MJ = getattr(self.mass_ops, id_MJ) self._DQ = self.derham.div @ getattr(self.basis_ops, id_Q) + self.options.b_field.allocate(self.derham, self.domain) self._b = self.options.b_field.spline.vector # preconditioner From 00a9c5cc36cb9173bc570822d428448d1b31b7e3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 10 Sep 2025 08:54:19 +0200 Subject: [PATCH 108/292] adapt Vlasov to new options --- .../tutorial_01_parameter_files.ipynb | 17 ++-- .../tutorial_02_test_particles.ipynb | 20 ++--- src/struphy/pic/utilities.py | 2 +- .../propagators/propagators_markers.py | 78 ++++++++++++------- 4 files changed, 70 insertions(+), 47 deletions(-) diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/doc/tutorials/tutorial_01_parameter_files.ipynb index fb316d85d..e4c5006a1 100644 --- a/doc/tutorials/tutorial_01_parameter_files.ipynb +++ b/doc/tutorials/tutorial_01_parameter_files.ipynb @@ -7,9 +7,9 @@ "source": [ "# 1 - Parameters and `struphy.main`\n", "\n", - "Struphy is based on \"models\". Each model can be viewed as a simulation code for a certain physical model, described py a set partial differntial equations (PDEs).\n", - "The documentation features a [list of currently available models](https://struphy.pages.mpcdf.de/struphy/sections/models.html).\n", - "A model is launched through a parameter file, where all simulation parameters can be specified by the user.\n", + "Struphy is based on \"models\". Each model is a parallelized simulation code for a certain physical model, described py a set partial differntial equations (PDEs).\n", + "Check the documentation for a [list of currently available models](https://struphy.pages.mpcdf.de/struphy/sections/models.html).\n", + "A model is launched through its parameter file, where all simulation parameters can be specified by the user.\n", "\n", "Model parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", "For each `MODEL`, the default parameter file can be generated from the console via\n", @@ -21,7 +21,7 @@ "This will create a file `params_MODEL.py` in the current working directory. To run the model type\n", "\n", "```\n", - "python params_MODEL.py\n", + "python3 params_MODEL.py\n", "```\n", "\n", "The user can modify the parameter file to launch a specific simulation.\n", @@ -184,10 +184,9 @@ "source": [ "## Part 4: Propagator options\n", "\n", - "A model is a collection of \"propagators\", which perform the time stepping. Each propagator refers to a splitting step of a single time step. For instance, if only a single propagator is present in a model, then all variables of the model are updated by this propagator. If two or more propagators are present, they are executed in sequence, according to the chosen splitting algorithm (Lie-Trotter, Strang, etc.).\n", + "A model is a collection of \"propagators\" which perform the time stepping. Each propagator refers to part of a single time step and can be viewed as a splitting step. For instance, if only a single propagator is present in a model, then all variables of the model are updated by this propagator. If two or more propagators are present, they are executed in sequence, according to the chosen splitting algorithm (Lie-Trotter, Strang, etc.).\n", "\n", - "Each propagator has options that can be set in the parameter file.\n", - "Check the method `set_options()` of each propagator for its available options." + "Each propagator has options that can be set in the parameter file as follows:" ] }, { @@ -198,8 +197,8 @@ "outputs": [], "source": [ "# propagator options\n", - "model.propagators.push_vxb.set_options()\n", - "model.propagators.push_eta.set_options()" + "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n", + "model.propagators.push_eta.options = model.propagators.push_eta.Options()" ] }, { diff --git a/doc/tutorials/tutorial_02_test_particles.ipynb b/doc/tutorials/tutorial_02_test_particles.ipynb index b9277aca8..96b88961d 100644 --- a/doc/tutorials/tutorial_02_test_particles.ipynb +++ b/doc/tutorials/tutorial_02_test_particles.ipynb @@ -221,11 +221,11 @@ "outputs": [], "source": [ "# propagator options\n", - "model.propagators.push_vxb.set_options()\n", - "model.propagators.push_eta.set_options()\n", + "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n", + "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n", "\n", - "model_2.propagators.push_vxb.set_options()\n", - "model_2.propagators.push_eta.set_options()" + "model_2.propagators.push_vxb.options = model_2.propagators.push_vxb.Options()\n", + "model_2.propagators.push_eta.options = model_2.propagators.push_eta.Options()" ] }, { @@ -430,8 +430,8 @@ "outputs": [], "source": [ "# propagator options\n", - "model.propagators.push_vxb.set_options()\n", - "model.propagators.push_eta.set_options()\n", + "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n", + "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n", "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", @@ -627,8 +627,8 @@ "model.kinetic_ions.set_save_data(n_markers=1.0)\n", "\n", "# propagator options\n", - "model.propagators.push_vxb.set_options()\n", - "model.propagators.push_eta.set_options()\n", + "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n", + "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n", "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", @@ -965,8 +965,8 @@ "model.kinetic_ions.set_save_data(n_markers=1.0)\n", "\n", "# propagator options\n", - "model.propagators.push_vxb.set_options()\n", - "model.propagators.push_eta.set_options()\n", + "model.propagators.push_vxb.options = model.propagators.push_vxb.Options()\n", + "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n", "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index 310ab129a..c0047bf92 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -91,7 +91,7 @@ class WeightsParameters: control_variate : bool Whether to use a control variate for noise reduction. - rejct_weights : bool + reject_weights : bool Whether to reject weights below threshold. threshold : float diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 2843c1974..733a7d749 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -4,6 +4,8 @@ from numpy import array, polynomial, random from typing import Literal, get_args import copy +from dataclasses import dataclass +from mpi4py import MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -56,27 +58,38 @@ def var(self, new): self._var = new def __init__(self): - # variables to be updated self.variables = self.Variables() - - ## abstract methods - def set_options(self, - butcher: ButcherTableau = None, - ): - if butcher is None: - butcher = ButcherTableau() + @dataclass + class Options: + butcher: ButcherTableau = None - # setter - self.options = copy.deepcopy(locals()) - self.options.pop("self") + def __post_init__(self): + # defaults + if self.butcher is None: + self.butcher = ButcherTableau() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new def allocate(self): # get kernel kernel = pusher_kernels.push_eta_stage # define algorithm - butcher: ButcherTableau = self.options["butcher"] + butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: import numpy as np @@ -141,22 +154,33 @@ def __init__(self): # variables to be updated self.variables = self.Variables() - # propagator specific options - OptsAlgo = Literal["analytic", "implicit"] - - ## abstract methods - def set_options(self, - algo: OptsAlgo = "analytic", - kappa: float = 1.0, - b2_var: FEECVariable = None, - ): - - # checks - check_option(algo, self.OptsAlgo) + @dataclass + class Options: + # specific literals + OptsAlgo = Literal["analytic", "implicit"] + # propagator options + algo: OptsAlgo = "analytic" + kappa: float = 1.0 + b2_var: FEECVariable = None - # setter - self.options = copy.deepcopy(locals()) - self.options.pop("self") + def __post_init__(self): + # checks + check_option(self.algo, self.OptsAlgo) + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new def allocate(self): # TODO: treat PolarVector as well, but polar splines are being reworked at the moment From ab21da66f455c6e5ac869f560d1ccb8663d145be Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 10 Sep 2025 14:38:38 +0200 Subject: [PATCH 109/292] make toy model GuidingCenter run in new framework --- .../tutorial_02_test_particles.ipynb | 215 ++++++------- src/struphy/feec/psydac_derham.py | 3 + src/struphy/kinetic_background/maxwellians.py | 65 ++-- src/struphy/models/base.py | 16 +- src/struphy/models/species.py | 51 +-- src/struphy/models/toy.py | 98 +----- src/struphy/pic/particles.py | 17 +- src/struphy/propagators/base.py | 17 +- .../propagators/propagators_markers.py | 296 +++++++++--------- 9 files changed, 389 insertions(+), 389 deletions(-) diff --git a/doc/tutorials/tutorial_02_test_particles.ipynb b/doc/tutorials/tutorial_02_test_particles.ipynb index 96b88961d..5fb31d436 100644 --- a/doc/tutorials/tutorial_02_test_particles.ipynb +++ b/doc/tutorials/tutorial_02_test_particles.ipynb @@ -1013,19 +1013,19 @@ "metadata": {}, "outputs": [], "source": [ - "time_opts = Time(dt=0.2, Tend=3000, split_algo=\"Strang\")\n", + "# time_opts = Time(dt=0.2, Tend=3000, split_algo=\"Strang\")\n", "\n", - "main.run(model, \n", - " params_path=None, \n", - " env=env, \n", - " base_units=base_units, \n", - " time_opts=time_opts, \n", - " domain=domain, \n", - " equil=equil, \n", - " grid=grid, \n", - " derham_opts=derham_opts, \n", - " verbose=verbose, \n", - " )" + "# main.run(model, \n", + "# params_path=None, \n", + "# env=env, \n", + "# base_units=base_units, \n", + "# time_opts=time_opts, \n", + "# domain=domain, \n", + "# equil=equil, \n", + "# grid=grid, \n", + "# derham_opts=derham_opts, \n", + "# verbose=verbose, \n", + "# )" ] }, { @@ -1064,24 +1064,24 @@ "metadata": {}, "outputs": [], "source": [ - "import math\n", + "# import math\n", "\n", - "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", + "# colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", "\n", - "dt = time_opts.dt\n", - "Tend = time_opts.Tend\n", + "# dt = time_opts.dt\n", + "# Tend = time_opts.Tend\n", "\n", - "for i in range(Np):\n", - " r = np.sqrt(orbits[:, i, 0]**2 + orbits[:, i, 1]**2)\n", - " # poloidal \n", - " ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n", - " # top view\n", - " ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n", + "# for i in range(Np):\n", + "# r = np.sqrt(orbits[:, i, 0]**2 + orbits[:, i, 1]**2)\n", + "# # poloidal \n", + "# ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n", + "# # top view\n", + "# ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n", " \n", - "ax.set_title(f'{math.ceil(Tend/dt)} time steps')\n", - "ax_top.set_title(f'{math.ceil(Tend/dt)} time steps');\n", + "# ax.set_title(f'{math.ceil(Tend/dt)} time steps')\n", + "# ax_top.set_title(f'{math.ceil(Tend/dt)} time steps');\n", "\n", - "fig" + "# fig" ] }, { @@ -1124,12 +1124,12 @@ "model.kinetic_ions.set_save_data(n_markers=1.0)\n", "\n", "# propagator options\n", - "model.propagators.push_vxb.set_options()\n", - "model.propagators.push_eta.set_options()\n", + "model.propagators.push_bxe.options = model.propagators.push_bxe.Options(tol=1e-5)\n", + "model.propagators.push_parallel.options = model.propagators.push_parallel.Options(tol=1e-5)\n", "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", - "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", + "background = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation))\n", "\n", "model.kinetic_ions.var.add_background(background)" ] @@ -1160,50 +1160,7 @@ "# Np=Np, \n", "# bc=bc, \n", "# loading_params=loading_params,\n", - " bufsize=bufsize)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "847979e0", - "metadata": {}, - "outputs": [], - "source": [ - "particles.draw_markers()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6016bfe0", - "metadata": {}, - "outputs": [], - "source": [ - "# positions on the physical domain Omega (x, y, z)\n", - "pushed_pos = domain(particles.positions).T\n", - "pushed_pos" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "a5548daa", - "metadata": {}, - "outputs": [], - "source": [ - "# compute R-coordinate\n", - "pushed_r = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "53fa4d3b", - "metadata": {}, - "outputs": [], - "source": [ - "particles.velocities" + " # bufsize=bufsize)" ] }, { @@ -1219,6 +1176,7 @@ "ax_top = axs[1]\n", "\n", "# min/max of field strength\n", + "equil.domain = domain\n", "Bmax = np.max(equil.absB0(*eta_topview_2, squeeze_out=True))\n", "Bmin = np.min(equil.absB0(*eta_topview_1, squeeze_out=True))\n", "levels = np.linspace(Bmin, Bmax, 51)\n", @@ -1267,68 +1225,111 @@ { "cell_type": "code", "execution_count": null, - "id": "2ab72f71", + "id": "e792ad7e", "metadata": {}, "outputs": [], "source": [ - "labels = ['co-passing',\n", - " 'counter passing',\n", - " 'co_trapped',\n", - " 'counter-trapped']\n", + "time_opts = Time(dt=0.1, Tend=100, split_algo=\"Strang\")\n", "\n", - "for n, (r, pos) in enumerate(zip(pushed_r, pushed_pos)):\n", - " # poloidal \n", - " ax.scatter(r, pos[2], c=colors[n % 4], label=labels[n])\n", - " # topview\n", - " ax_top.scatter(pos[0], pos[1], c=colors[n % 4], label=labels[n])\n", - " ax_top.arrow(pos[0], pos[1], 0., particles.velocities[n, 0]/5, color=colors[n % 4], head_width=.05)\n", + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " base_units=base_units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ef3b687f", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from struphy import main\n", "\n", - "ax.set_xlabel('R')\n", - "ax.set_ylabel('Z')\n", - "ax.set_title('Initial conditions')\n", - "ax.legend();\n", + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "main.pproc(path)\n", "\n", - "ax_top.set_xlabel('x')\n", - "ax_top.set_ylabel('y')\n", - "ax_top.set_title('Initial conditions')\n", - "ax_top.legend();\n", + "simdata = main.load_data(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9819e15f", + "metadata": {}, + "outputs": [], + "source": [ + "orbits = simdata.orbits[\"kinetic_ions\"]\n", "\n", - "fig " + "Nt = simdata.Nt[\"kinetic_ions\"]\n", + "Np = simdata.Np[\"kinetic_ions\"]" ] }, { "cell_type": "code", "execution_count": null, - "id": "82bb30be", + "id": "05cc5d7d", "metadata": {}, "outputs": [], "source": [ - "from struphy.propagators.propagators_markers import PushGuidingCenterBxEstar, PushGuidingCenterParallel\n", + "import math\n", "\n", - "# default parameters of Propagator\n", - "opts_BxE = PushGuidingCenterBxEstar.options(default=True)\n", - "print(opts_BxE)\n", + "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", "\n", - "opts_para = PushGuidingCenterParallel.options(default=True)\n", - "print(opts_para)" + "dt = time_opts.dt\n", + "Tend = time_opts.Tend\n", + "\n", + "for i in range(Np):\n", + " r = np.sqrt(orbits[:, i, 0]**2 + orbits[:, i, 1]**2)\n", + " # poloidal \n", + " ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n", + " # top view\n", + " ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n", + " \n", + "ax.set_title(f'{math.ceil(Tend/dt)} time steps')\n", + "ax_top.set_title(f'{math.ceil(Tend/dt)} time steps');\n", + "\n", + "fig" ] }, { "cell_type": "code", "execution_count": null, - "id": "9bd84960", + "id": "2ab72f71", "metadata": {}, "outputs": [], "source": [ - "# pass simulation parameters to Propagator class\n", - "PushGuidingCenterBxEstar.domain = domain\n", - "PushGuidingCenterParallel.domain = domain\n", + "labels = ['co-passing',\n", + " 'counter passing',\n", + " 'co_trapped',\n", + " 'counter-trapped']\n", + "\n", + "for n, (r, pos) in enumerate(zip(pushed_r, pushed_pos)):\n", + " # poloidal \n", + " ax.scatter(r, pos[2], c=colors[n % 4], label=labels[n])\n", + " # topview\n", + " ax_top.scatter(pos[0], pos[1], c=colors[n % 4], label=labels[n])\n", + " ax_top.arrow(pos[0], pos[1], 0., particles.velocities[n, 0]/5, color=colors[n % 4], head_width=.05)\n", + "\n", + "ax.set_xlabel('R')\n", + "ax.set_ylabel('Z')\n", + "ax.set_title('Initial conditions')\n", + "ax.legend();\n", "\n", - "PushGuidingCenterBxEstar.derham = derham\n", - "PushGuidingCenterParallel.derham = derham\n", + "ax_top.set_xlabel('x')\n", + "ax_top.set_ylabel('y')\n", + "ax_top.set_title('Initial conditions')\n", + "ax_top.legend();\n", "\n", - "PushGuidingCenterBxEstar.projected_equil = proj_equil\n", - "PushGuidingCenterParallel.projected_equil = proj_equil" + "fig " ] }, { diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 5b49fdb7a..0158dcffc 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1416,6 +1416,9 @@ class SplineFunction: domain : Domain Mapping for pullback/transform of initial condition. + + equil : FluidEquilibrium + Fluid background used for inital condition. """ def __init__( diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index a9371b8cb..dcd1ff3ee 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -3,7 +3,7 @@ import numpy as np from typing import Callable -from struphy.fields_background.base import FluidEquilibrium +from struphy.fields_background.base import FluidEquilibriumWithB from struphy.fields_background.equils import set_defaults from struphy.kinetic_background import moment_functions from struphy.kinetic_background.base import CanonicalMaxwellian, Maxwellian @@ -143,47 +143,47 @@ class GyroMaxwellian2D(Maxwellian): Parameters ---------- + n, u_para, u_perp, vth_para, vth_perp : tuple + Moments of the Maxwellian as tuples. The first entry defines the background + (float for constant background or callable), the second entry defines a Perturbation (can be None). + maxw_params : dict Parameters for the kinetic background. pert_params : dict Parameters for the kinetic perturbation added to the background. - equil : FluidEquilibrium - One of :mod:`~struphy.fields_background.equils`. + equil : FluidEquilibriumWithB + Fluid background. volume_form : bool Whether to represent the Maxwellian as a volume form; if True it is multiplied by the Jacobian determinant |v_perp| of the polar coordinate transofrmation (default = False). """ - - @classmethod - def default_maxw_params(cls): - """Default parameters dictionary defining constant moments of the Maxwellian.""" - return { - "n": 1.0, - "u_para": 0.0, - "u_perp": 0.0, - "vth_para": 1.0, - "vth_perp": 1.0, - } - def __init__( self, - maxw_params: dict = None, - pert_params: dict = None, - equil: FluidEquilibrium = None, + n: tuple[float | Callable[..., float], Perturbation] = (1.0, None), + u_para: tuple[float | Callable[..., float], Perturbation] = (0.0, None), + u_perp: tuple[float | Callable[..., float], Perturbation] = (0.0, None), + vth_para: tuple[float | Callable[..., float], Perturbation] = (1.0, None), + vth_perp: tuple[float | Callable[..., float], Perturbation] = (1.0, None), + equil: FluidEquilibriumWithB = None, volume_form: bool = True, ): - super().__init__( - maxw_params=maxw_params, - pert_params=pert_params, - equil=equil, - ) + + self._maxw_params = {} + self._maxw_params["n"] = n + self._maxw_params["u_para"] = u_para + self._maxw_params["u_perp"] = u_perp + self._maxw_params["vth_para"] = vth_para + self._maxw_params["vth_perp"] = vth_perp + + self.check_maxw_params() # volume form represenation self._volume_form = volume_form + self._equil = equil # factors multiplied onto the defined moments n, u and vth (can be set via setter) self._moment_factors = { @@ -192,6 +192,10 @@ def __init__( "vth": [1.0, 1.0], } + @property + def maxw_params(self): + return self._maxw_params + @property def coords(self): r"""Coordinates of the Maxwellian5D, :math:`(v_\parallel, v_\perp)`.""" @@ -257,9 +261,14 @@ def velocity_jacobian_det(self, eta1, eta2, eta3, *v): return jacobian_det @property - def volume_form(self): + def volume_form(self) -> bool: """Boolean. True if the background is represented as a volume form (thus including the velocity Jacobian |v_perp|).""" return self._volume_form + + @property + def equil(self) -> FluidEquilibriumWithB: + """Fluid background with B-field.""" + return self._equil @property def moment_factors(self): @@ -304,8 +313,8 @@ class CanonicalMaxwellian(CanonicalMaxwellian): pert_params : dict Parameters for the kinetic perturbation added to the background. - equil : FluidEquilibrium - One of :mod:`~struphy.fields_background.equils`. + equil : FluidEquilibriumWithB + Fluid background. volume_form : bool Whether to represent the Maxwellian as a volume form; @@ -326,7 +335,7 @@ def __init__( self, maxw_params: dict = None, pert_params: dict = None, - equil: FluidEquilibrium = None, + equil: FluidEquilibriumWithB = None, volume_form: bool = True, ): # Set background parameters @@ -536,7 +545,7 @@ def __init__( self, maxw_params: dict = None, pert_params: dict = None, - equil: FluidEquilibrium = None, + equil: FluidEquilibriumWithB = None, ): super().__init__( maxw_params=maxw_params, diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index d28cfef36..914fd2acb 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -604,8 +604,12 @@ def allocate_variables(self, verbose: bool = False): assert isinstance(spec, KineticSpecies) for k, v in spec.variables.items(): assert isinstance(v, (PICVariable, SPHVariable)) - v.allocate(derham=self.derham, domain=self.domain, equil=self.equil, - verbose=verbose) + v.allocate(clone_config=self.clone_config, + derham=self.derham, + domain=self.domain, + equil=self.equil, + projected_equil=self.projected_equil, + verbose=verbose,) # TODO: allocate memory for FE coeffs of diagnostics # if self.params.diagnostic_fields is not None: @@ -1324,8 +1328,12 @@ def generate_default_parameter_file( elif isinstance(var, PICVariable): has_pic = True init_pert_pic = f"perturbation = perturbations.TorusModesCos()\n" - init_bckgr_pic = f"\nmaxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n" - init_bckgr_pic += f"maxwellian_2 = maxwellians.Maxwellian3D(n=(0.1, None))\n" + if "6D" in var.space: + init_bckgr_pic = f"\nmaxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n" + init_bckgr_pic += f"maxwellian_2 = maxwellians.Maxwellian3D(n=(0.1, None))\n" + elif "5D" in var.space: + init_bckgr_pic = f"\nmaxwellian_1 = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation))\n" + init_bckgr_pic += f"maxwellian_2 = maxwellians.GyroMaxwellian2D(n=(0.1, None))\n" init_bckgr_pic += f"background = maxwellian_1 + maxwellian_2\n" init_bckgr_pic += f"model.{sn}.{vn}.add_background(background)\n" diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 192ce5c84..4cf94e097 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -50,8 +50,34 @@ def set_phys_params(self, charge_number: int = 1, mass_number: int = 1): self._charge_number = charge_number self._mass_number = mass_number + class EquationParameters: + """Normalization parameters of one species, appearing in scaled equations.""" + + def __init__(self, species, units: Units = None, verbose: bool = False): + if units is None: + units = Units() + + Z = species.charge_number + A = species.mass_number + + con = ConstantsOfNature() + + # relevant frequencies + om_p = np.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_c = Z * con.e * units.B / (A * con.mH) + + # compute equation parameters + self.alpha = om_p / om_c + self.epsilon = 1.0 / (om_c * units.t) + self.kappa = om_p * units.t + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f'\nSet normalization parameters for species {species.__class__.__name__}:') + for key, val in self.__dict__.items(): + print((key + ":").ljust(25), "{:4.3e}".format(val)) + @property - def equation_params(self) -> dict: + def equation_params(self) -> EquationParameters: return self._equation_params def setup_equation_params(self, units: Units, verbose=False): @@ -61,23 +87,9 @@ def setup_equation_params(self, units: Units, verbose=False): * epsilon = 1 / (cyclotron frequency * time unit) * kappa = plasma frequency * time unit """ - Z = self.charge_number - A = self.mass_number - - con = ConstantsOfNature() + self._equation_params = self.EquationParameters(species=self, units=units, verbose=verbose) - # compute equation parameters - self._equation_params = {} - om_p = np.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) - om_c = Z * con.e * units.B / (A * con.mH) - self._equation_params["alpha"] = om_p / om_c - self._equation_params["epsilon"] = 1.0 / (om_c * units.t) - self._equation_params["kappa"] = om_p * units.t - - if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f'\nSet normalization parameters for species {self.__class__.__name__}:') - for key, val in self.equation_params.items(): - print((key + ":").ljust(25), "{:4.3e}".format(val)) + class FieldSpecies(Species): @@ -118,10 +130,13 @@ def set_markers(self, if weights_params is None: weights_params = WeightsParameters() + + if boundary_params is None: + boundary_params = BoundaryParameters() - self.boundary_params = boundary_params self.loading_params = loading_params self.weights_params = weights_params + self.boundary_params = boundary_params self.bufsize = bufsize def set_sorting_boxes(self, diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 9ff7d8461..b3bc527c1 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -231,16 +231,13 @@ def __init__(self): self.propagators = self.Propagators() # 3. assign variables to propagators - self.propagators.push_bxe.assign_variables( - ions = self.kinetic_ions.var, - ) - - self.propagators.push_parallel.assign_variables( - var = self.kinetic_ions.var, - ) + self.propagators.push_bxe.variables.ions = self.kinetic_ions.var + self.propagators.push_parallel.variables.ions = self.kinetic_ions.var # define scalars for update_scalar_quantities - self.add_scalar("en_f", compute="from_particles", variable=self.kinetic_ions.var) + self.add_scalar("en_fv", compute="from_particles", variable=self.kinetic_ions.var) + self.add_scalar("en_fB", compute="from_particles", variable=self.kinetic_ions.var) + self.add_scalar("en_tot", compute="from_particles", variable=self.kinetic_ions.var) @property def bulk_species(self): @@ -251,99 +248,38 @@ def velocity_scale(self): return "alfvén" def allocate_helpers(self): - self._tmp = np.empty(1, dtype=float) - - def update_scalar_quantities(self): - particles = self.kinetic_ions.var.particles - self._tmp[0] = particles.markers_wo_holes[:, 6].dot( - particles.markers_wo_holes[:, 3] ** 2 - + particles.markers_wo_holes[:, 4] ** 2 - + particles.markers_wo_holes[:, 5] ** 2, - ) / (2 * particles.Np) - - self.update_scalar("en_f", self._tmp[0]) - - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - from mpi4py.MPI import IN_PLACE, SUM - - # prelim - ions_params = self.kinetic["ions"]["params"] - epsilon = self.equation_params["ions"]["epsilon"] - - # set keyword arguments for propagators - self._kwargs[propagators_markers.PushGuidingCenterBxEstar] = { - "epsilon": epsilon, - "algo": ions_params["options"]["PushGuidingCenterBxEstar"]["algo"], - } - - self._kwargs[propagators_markers.PushGuidingCenterParallel] = { - "epsilon": epsilon, - "algo": ions_params["options"]["PushGuidingCenterParallel"]["algo"], - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation - self.add_scalar("en_fv") - self.add_scalar("en_fB") - self.add_scalar("en_tot") - - # MPI operations needed for scalar variables - self._mpi_sum = SUM - self._mpi_in_place = IN_PLACE self._en_fv = np.empty(1, dtype=float) self._en_fB = np.empty(1, dtype=float) self._en_tot = np.empty(1, dtype=float) self._n_lost_particles = np.empty(1, dtype=float) def update_scalar_quantities(self): - # particles' kinetic energy + particles = self.kinetic_ions.var.particles - self._en_fv[0] = self.pointer["ions"].markers[~self.pointer["ions"].holes, 5].dot( - self.pointer["ions"].markers[~self.pointer["ions"].holes, 3] ** 2, - ) / (2.0 * self.pointer["ions"].Np) + # particles' kinetic energy + self._en_fv[0] = particles.markers[~particles.holes, 5].dot( + particles.markers[~particles.holes, 3] ** 2, + ) / (2.0 * particles.Np) - self.pointer["ions"].save_magnetic_background_energy() + particles.save_magnetic_background_energy() self._en_tot[0] = ( - self.pointer["ions"] - .markers[~self.pointer["ions"].holes, 5] - .dot( - self.pointer["ions"].markers[~self.pointer["ions"].holes, 8], + particles.markers[~particles.holes, 5].dot( + particles.markers[~particles.holes, 8], ) - / self.pointer["ions"].Np + / particles.Np ) self._en_fB[0] = self._en_tot[0] - self._en_fv[0] - self.derham.comm.Allreduce( - self._mpi_in_place, - self._en_fv, - op=self._mpi_sum, - ) - self.derham.comm.Allreduce( - self._mpi_in_place, - self._en_tot, - op=self._mpi_sum, - ) - self.derham.comm.Allreduce( - self._mpi_in_place, - self._en_fB, - op=self._mpi_sum, - ) - self.update_scalar("en_fv", self._en_fv[0]) self.update_scalar("en_fB", self._en_fB[0]) self.update_scalar("en_tot", self._en_tot[0]) - self._n_lost_particles[0] = self.pointer["ions"].n_lost_markers + self._n_lost_particles[0] = particles.n_lost_markers self.derham.comm.Allreduce( - self._mpi_in_place, + MPI.IN_PLACE, self._n_lost_particles, - op=self._mpi_sum, + op=MPI.SUM, ) diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 825505236..7a7da2054 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -299,20 +299,21 @@ class Particles5D(Particles): **kwargs : dict Parameters for markers, see :class:`~struphy.pic.base.Particles`. """ - @classmethod - def default_bckgr_params(cls): - return {"GyroMaxwellian2D": {}} + def default_background(cls): + return maxwellians.GyroMaxwellian2D() def __init__( self, projected_equil: ProjectedFluidEquilibriumWithB, **kwargs, ): + assert projected_equil is not None, "Particles5D needs a projected MHD equilibrium." + kwargs["type"] = "full_f" - if "bckgr_params" not in kwargs: - kwargs["bckgr_params"] = self.default_bckgr_params() + # if "bckgr_params" not in kwargs: + # kwargs["bckgr_params"] = self.default_bckgr_params() # default number of diagnostics and auxiliary columns self._n_cols_diagnostics = kwargs.pop("n_cols_diagn", 3) @@ -407,7 +408,11 @@ def svol(self, eta1, eta2, eta3, *v): } self._svol = maxwellians.GyroMaxwellian2D( - maxw_params=maxw_params, + n = (1.0, None), + u_para = (self.loading_params.moments[0], None), + u_perp = (self.loading_params.moments[1], None), + vth_para = (self.loading_params.moments[2], None), + vth_perp = (self.loading_params.moments[3], None), volume_form=True, equil=self._magn_bckgr, ) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index da092568a..b1cfcc0a8 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -76,11 +76,11 @@ def options(self, new): self._options = new @abstractmethod - def allocate(): + def allocate(self): """Allocate all data/objects of the instance.""" - + @abstractmethod - def __call__(self, dt): + def __call__(self, dt: float): """Update variables from t -> t + dt. Use ``Propagators.feec_vars_update`` to write to FEEC variables to ``Propagator.feec_vars``. @@ -198,8 +198,9 @@ def projected_equil(self) -> ProjectedFluidEquilibriumWithB: return self._projected_equil @projected_equil.setter - def projected_equil(self, projected_equil): - self._projected_equil = projected_equil + def projected_equil(self, new): + assert isinstance(new, ProjectedFluidEquilibriumWithB) + self._projected_equil = new @property def time_state(self): @@ -247,6 +248,9 @@ def add_init_kernel( else: comps = np.array(comps, dtype=int) + if not hasattr(self, "_init_kernels"): + self._init_kernels = [] + self._init_kernels += [ ( kernel, @@ -298,6 +302,9 @@ def add_eval_kernel( else: comps = np.array(comps, dtype=int) + if not hasattr(self, "_eval_kernels"): + self._eval_kernels = [] + self._eval_kernels += [ ( kernel, diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 733a7d749..b87493005 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -151,7 +151,6 @@ def ions(self, new): self._ions = new def __init__(self): - # variables to be updated self.variables = self.Variables() @dataclass @@ -456,56 +455,54 @@ def ions(self, new): self._ions = new def __init__(self): - # variables to be updated self.variables = self.Variables() - - # propagator specific options - OptsAlgo = Literal["discrete_gradient_2nd_order", - "discrete_gradient_1st_order", - "discrete_gradient_1st_order_newton", "explicit",] - - ## abstract methods - def set_options(self, - algo: OptsAlgo = "discrete_gradient_1st_order", - butcher: ButcherTableau = None, - maxiter: int = 20, - tol: float = 1e-7, - mpi_sort: OptsMPIsort = "each", - verbose: bool = False, - ): - # checks - check_option(algo, self.OptsAlgo) - check_option(mpi_sort, OptsMPIsort) + @dataclass + class Options: + # specific literals + OptsAlgo = Literal["discrete_gradient_2nd_order", + "discrete_gradient_1st_order", + "discrete_gradient_1st_order_newton", + "explicit",] + # propagator options + phi: FEECVariable = None + evaluate_e_field: bool = False + b_tilde: FEECVariable = None + algo: OptsAlgo = "discrete_gradient_1st_order" + butcher: ButcherTableau = None + maxiter: int = 20 + tol: float = 1e-7 + mpi_sort: OptsMPIsort = "each" + verbose: bool = False - if algo == "explicit" and butcher is None: - butcher = ButcherTableau() + def __post_init__(self): + # checks + check_option(self.algo, self.OptsAlgo) + check_option(self.mpi_sort, OptsMPIsort) + + # defaults + if self.phi is None: + self.phi = FEECVariable(space="H1") + + if self.algo == "explicit" and self.butcher is None: + self.butcher = ButcherTableau() - # use setter for options - self.options = self.Options(self, - algo=algo, - butcher=butcher, - maxiter=maxiter, - tol=tol, - mpi_sort=mpi_sort, - verbose=verbose, - ) + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new def allocate(self): - pass - - def __init__( - self, - particles: Particles5D, - *, - phi: StencilVector = None, - evaluate_e_field: bool = False, - b_tilde: BlockVector = None, - epsilon: float = 1.0, - algo: dict = None, - ): - super().__init__(particles) - # magnetic equilibrium field unit_b1 = self.projected_equil.unit_b1 self._gradB1 = self.projected_equil.gradB1 @@ -513,14 +510,13 @@ def __init__( curl_unit_b_dot_b0 = self.projected_equil.curl_unit_b_dot_b0 # magnetic perturbation - self._b_tilde = b_tilde - if self._b_tilde is not None: + if self.options.b_tilde is not None: self._B_dot_b = self.derham.Vh["0"].zeros() self._grad_b_full = self.derham.Vh["1"].zeros() self._PB = getattr(self.basis_ops, "PB") - B_dot_b = self._PB.dot(self._b_tilde, out=self._B_dot_b) + B_dot_b = self._PB.dot(self.options.b_tilde.spline.vector, out=self._B_dot_b) B_dot_b.update_ghost_regions() grad_b_full = self.derham.grad.dot(B_dot_b, out=self._grad_b_full) @@ -533,20 +529,20 @@ def __init__( self._B_dot_b = self._absB0 # allocate electric field - if phi is None: - phi = self.derham.Vh["0"].zeros() - self._evaluate_e_field = False - self._phi = phi - self._evaluate_e_field = evaluate_e_field + self.options.phi.allocate(self.derham, self.domain) + self._phi = self.options.phi.spline.vector + self._evaluate_e_field = self.options.evaluate_e_field self._e_field = self.derham.Vh["1"].zeros() - self._epsilon = epsilon + self._epsilon = self.variables.ions.species.equation_params.epsilon # choose method - if "discrete_gradient" in algo["method"]: + particles = self.variables.ions.particles + + if "discrete_gradient" in self.options.algo: # place for storing data during iteration first_free_idx = particles.args_markers.first_free_idx - if "1st_order" in algo["method"]: + if "1st_order" in self.options.algo: # init kernels self.add_init_kernel( eval_kernels_gc.driftkinetic_hamiltonian, @@ -585,7 +581,7 @@ def __init__( ), ) - if "newton" in algo["method"]: + if "newton" in self.options.algo: # eval kernels self.add_eval_kernel( eval_kernels_gc.driftkinetic_hamiltonian, @@ -701,7 +697,7 @@ def __init__( self._evaluate_e_field, ) - elif "2nd_order" in algo["method"]: + elif "2nd_order" in self.options.algo: # init kernels (evaluate at eta^n and save) self.add_init_kernel( eval_kernels_gc.driftkinetic_hamiltonian, @@ -752,11 +748,6 @@ def __init__( self._evaluate_e_field, ) - else: - raise NotImplementedError( - f"Chosen method {algo['method']} is not implemented.", - ) - # Pusher instance self._pusher = Pusher( particles, @@ -766,14 +757,17 @@ def __init__( alpha_in_kernel=alpha_in_kernel, init_kernels=self.init_kernels, eval_kernels=self.eval_kernels, - maxiter=algo["maxiter"], - tol=algo["tol"], - mpi_sort=algo["mpi_sort"], - verbose=algo["verbose"], + maxiter=self.options.maxiter, + tol=self.options.tol, + mpi_sort=self.options.mpi_sort, + verbose=self.options.verbose, ) else: - butcher = ButcherTableau(algo["method"]) + if self.options.butcher is None: + butcher = ButcherTableau() + else: + butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: import numpy as np @@ -809,8 +803,8 @@ def __init__( self.domain.args_domain, alpha_in_kernel=1.0, n_stages=butcher.n_stages, - mpi_sort=algo["mpi_sort"], - verbose=algo["verbose"], + mpi_sort=self.options.mpi_sort, + verbose=self.options.verbose, ) def __call__(self, dt): @@ -821,8 +815,8 @@ def __call__(self, dt): e_field.update_ghost_regions() # magnetic perturbation - if self._b_tilde is not None: - B_dot_b = self._PB.dot(self._b_tilde, out=self._B_dot_b) + if self.options.b_tilde is not None: + B_dot_b = self._PB.dot(self.options.b_tilde.spline.vector, out=self._B_dot_b) B_dot_b.update_ghost_regions() grad_b_full = self.derham.grad.dot(B_dot_b, out=self._grad_b_full) @@ -835,8 +829,8 @@ def __call__(self, dt): self._pusher(dt) # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() + if self.variables.ions.species.weights_params.control_variate: + self.variables.ions.particles.update_weights() class PushGuidingCenterParallel(Propagator): @@ -880,45 +874,69 @@ class PushGuidingCenterParallel(Propagator): * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_Bstar_discrete_gradient_1st_order_newton` * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_Bstar_discrete_gradient_2nd_order` """ + class Variables: + def __init__(self): + self._ions: PICVariable = None + + @property + def ions(self) -> PICVariable: + return self._ions + + @ions.setter + def ions(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles5D" + self._ions = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsAlgo = Literal["discrete_gradient_2nd_order", + "discrete_gradient_1st_order", + "discrete_gradient_1st_order_newton", + "explicit",] + # propagator options + phi: FEECVariable = None + evaluate_e_field: bool = False + b_tilde: FEECVariable = None + algo: OptsAlgo = "discrete_gradient_1st_order" + butcher: ButcherTableau = None + maxiter: int = 20 + tol: float = 1e-7 + mpi_sort: OptsMPIsort = "each" + verbose: bool = False + + def __post_init__(self): + # checks + check_option(self.algo, self.OptsAlgo) + check_option(self.mpi_sort, OptsMPIsort) + + # defaults + if self.phi is None: + self.phi = FEECVariable(space="H1") + + if self.algo == "explicit" and self.butcher is None: + self.butcher = ButcherTableau() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new - @staticmethod - def options(default=False): - dct = {} - dct["algo"] = { - "method": [ - "discrete_gradient_2nd_order", - "discrete_gradient_1st_order", - "discrete_gradient_1st_order_newton", - "rk4", - "forward_euler", - "heun2", - "rk2", - "heun3", - ], - "maxiter": 20, - "tol": 1e-7, - "mpi_sort": "each", - "verbose": False, - } - if default: - dct = descend_options_dict(dct, []) - - return dct - - def __init__( - self, - particles: Particles5D, - *, - phi: StencilVector = None, - evaluate_e_field: bool = False, - b_tilde: BlockVector = None, - epsilon: float = 1.0, - algo: dict = options(default=True)["algo"], - ): - super().__init__(particles) - - self._epsilon = epsilon - + def allocate(self): # magnetic equilibrium field self._gradB1 = self.projected_equil.gradB1 b2 = self.projected_equil.b2 @@ -927,14 +945,13 @@ def __init__( curl_unit_b_dot_b0 = self.projected_equil.curl_unit_b_dot_b0 # magnetic perturbation - self._b_tilde = b_tilde - if self._b_tilde is not None: + if self.options.b_tilde is not None: self._B_dot_b = self.derham.Vh["0"].zeros() self._grad_b_full = self.derham.Vh["1"].zeros() self._PB = getattr(self.basis_ops, "PB") - B_dot_b = self._PB.dot(self._b_tilde, out=self._B_dot_b) + B_dot_b = self._PB.dot(self.options.b_tilde.spline.vector, out=self._B_dot_b) B_dot_b.update_ghost_regions() grad_b_full = self.derham.grad.dot(B_dot_b, out=self._grad_b_full) @@ -947,19 +964,20 @@ def __init__( self._B_dot_b = self._absB0 # allocate electric field - if phi is None: - phi = self.derham.Vh["0"].zeros() - self._phi = phi - self._evaluate_e_field = evaluate_e_field + self.options.phi.allocate(self.derham, domain=self.domain) + self._phi = self.options.phi.spline.vector + self._evaluate_e_field = self.options.evaluate_e_field self._e_field = self.derham.Vh["1"].zeros() - self._epsilon = epsilon + self._epsilon = self.variables.ions.species.equation_params.epsilon # choose method - if "discrete_gradient" in algo["method"]: + particles = self.variables.ions.particles + + if "discrete_gradient" in self.options.algo: # place for storing data during iteration first_free_idx = particles.args_markers.first_free_idx - if "1st_order" in algo["method"]: + if "1st_order" in self.options.algo: # init kernels self.add_init_kernel( eval_kernels_gc.driftkinetic_hamiltonian, @@ -1002,7 +1020,7 @@ def __init__( ), ) - if "newton" in algo["method"]: + if "newton" in self.options.algo: # eval kernels self.add_eval_kernel( eval_kernels_gc.driftkinetic_hamiltonian, @@ -1117,7 +1135,7 @@ def __init__( self._evaluate_e_field, ) - elif "2nd_order" in algo["method"]: + elif "2nd_order" in self.options.algo: # init kernels (evaluate at eta^n and save) self.add_init_kernel( eval_kernels_gc.driftkinetic_hamiltonian, @@ -1171,11 +1189,6 @@ def __init__( self._evaluate_e_field, ) - else: - raise NotImplementedError( - f"Chosen method {algo['method']} is not implemented.", - ) - # Pusher instance self._pusher = Pusher( particles, @@ -1185,14 +1198,17 @@ def __init__( alpha_in_kernel=alpha_in_kernel, init_kernels=self.init_kernels, eval_kernels=self.eval_kernels, - maxiter=algo["maxiter"], - tol=algo["tol"], - mpi_sort=algo["mpi_sort"], - verbose=algo["verbose"], + maxiter=self.options.maxiter, + tol=self.options.tol, + mpi_sort=self.options.mpi_sort, + verbose=self.options.verbose, ) else: - butcher = ButcherTableau(algo["method"]) + if self.options.butcher is None: + butcher = ButcherTableau() + else: + butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: import numpy as np @@ -1231,8 +1247,8 @@ def __init__( self.domain.args_domain, alpha_in_kernel=1.0, n_stages=butcher.n_stages, - mpi_sort=algo["mpi_sort"], - verbose=algo["verbose"], + mpi_sort=self.options.mpi_sort, + verbose=self.options.verbose, ) def __call__(self, dt): @@ -1243,8 +1259,8 @@ def __call__(self, dt): e_field.update_ghost_regions() # magnetic perturbation - if self._b_tilde is not None: - B_dot_b = self._PB.dot(self._b_tilde, out=self._B_dot_b) + if self.options.b_tilde is not None: + B_dot_b = self._PB.dot(self.options.b_tilde.spline.vector, out=self._B_dot_b) B_dot_b.update_ghost_regions() grad_b_full = self.derham.grad.dot(B_dot_b, out=self._grad_b_full) @@ -1257,8 +1273,8 @@ def __call__(self, dt): self._pusher(dt) # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() + if self.variables.ions.species.weights_params.control_variate: + self.variables.ions.particles.update_weights() class StepStaticEfield(Propagator): From 6694072c4fd80423a134e324ef7c71c82f5e1efe Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 10 Sep 2025 14:49:00 +0200 Subject: [PATCH 110/292] tutorial 2 on test particles is now complete and running. Adaption of GuidingCenter model completed. --- .../tutorial_02_test_particles.ipynb | 224 +++--------------- .../post_processing/post_processing_tools.py | 5 +- 2 files changed, 31 insertions(+), 198 deletions(-) diff --git a/doc/tutorials/tutorial_02_test_particles.ipynb b/doc/tutorials/tutorial_02_test_particles.ipynb index 5fb31d436..3d89d1ec0 100644 --- a/doc/tutorials/tutorial_02_test_particles.ipynb +++ b/doc/tutorials/tutorial_02_test_particles.ipynb @@ -1013,19 +1013,19 @@ "metadata": {}, "outputs": [], "source": [ - "# time_opts = Time(dt=0.2, Tend=3000, split_algo=\"Strang\")\n", + "time_opts = Time(dt=0.2, Tend=3000, split_algo=\"Strang\")\n", "\n", - "# main.run(model, \n", - "# params_path=None, \n", - "# env=env, \n", - "# base_units=base_units, \n", - "# time_opts=time_opts, \n", - "# domain=domain, \n", - "# equil=equil, \n", - "# grid=grid, \n", - "# derham_opts=derham_opts, \n", - "# verbose=verbose, \n", - "# )" + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " base_units=base_units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" ] }, { @@ -1064,24 +1064,24 @@ "metadata": {}, "outputs": [], "source": [ - "# import math\n", + "import math\n", "\n", - "# colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", + "colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']\n", "\n", - "# dt = time_opts.dt\n", - "# Tend = time_opts.Tend\n", + "dt = time_opts.dt\n", + "Tend = time_opts.Tend\n", "\n", - "# for i in range(Np):\n", - "# r = np.sqrt(orbits[:, i, 0]**2 + orbits[:, i, 1]**2)\n", - "# # poloidal \n", - "# ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n", - "# # top view\n", - "# ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n", + "for i in range(Np):\n", + " r = np.sqrt(orbits[:, i, 0]**2 + orbits[:, i, 1]**2)\n", + " # poloidal \n", + " ax.scatter(r, orbits[:, i, 2], c=colors[i % 4], s=1)\n", + " # top view\n", + " ax_top.scatter(orbits[:, i, 0], orbits[:, i, 1], c=colors[i % 4], s=1)\n", " \n", - "# ax.set_title(f'{math.ceil(Tend/dt)} time steps')\n", - "# ax_top.set_title(f'{math.ceil(Tend/dt)} time steps');\n", + "ax.set_title(f'{math.ceil(Tend/dt)} time steps')\n", + "ax_top.set_title(f'{math.ceil(Tend/dt)} time steps');\n", "\n", - "# fig" + "fig" ] }, { @@ -1109,10 +1109,10 @@ "# species parameters\n", "model.kinetic_ions.set_phys_params()\n", "\n", - "initial = [[.501, 0.001, 0.001, -1.935 , 1.72], # co-passing particle\n", - " [.501, 0.001, 0.001, 1.935 , 1.72], # couner-passing particle\n", - " [.501, 0.001, 0.001, -0.6665, 1.72], # co-trapped particle\n", - " [.501, 0.001, 0.001, 0.4515, 1.72]] # counter-trapped particl\n", + "initial = ((.501, 0.001, 0.001, -1.935 , 1.72), # co-passing particle\n", + " (.501, 0.001, 0.001, 1.935 , 1.72), # couner-passing particle\n", + " (.501, 0.001, 0.001, -0.6665, 1.72), # co-trapped particle\n", + " (.501, 0.001, 0.001, 0.4515, 1.72)) # counter-trapped particl\n", "\n", "loading_params = LoadingParameters(Np=4, seed=1608, specific_markers=initial)\n", "boundary_params = BoundaryParameters(bc=('remove', 'periodic', 'periodic'))\n", @@ -1134,35 +1134,6 @@ "model.kinetic_ions.var.add_background(background)" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "64caf9e5", - "metadata": {}, - "outputs": [], - "source": [ - "# from struphy.pic.particles import Particles5D\n", - "\n", - "# # instantiate Particle object\n", - "# Np = 4\n", - "# bc = ['remove', 'periodic', 'periodic']\n", - "# bufsize = 2.\n", - "\n", - "# initial = [[.501, 0.001, 0.001, -1.935 , 1.72], # co-passing particle\n", - "# [.501, 0.001, 0.001, 1.935 , 1.72], # couner-passing particle\n", - "# [.501, 0.001, 0.001, -0.6665, 1.72], # co-trapped particle\n", - "# [.501, 0.001, 0.001, 0.4515, 1.72]] # counter-trapped particle\n", - "\n", - "# loading_params = {'seed': 1608,\n", - "# 'initial' : initial}\n", - "\n", - "# particles = Particles5D(proj_equil,\n", - "# Np=Np, \n", - "# bc=bc, \n", - "# loading_params=loading_params,\n", - " # bufsize=bufsize)" - ] - }, { "cell_type": "code", "execution_count": null, @@ -1299,143 +1270,6 @@ "\n", "fig" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2ab72f71", - "metadata": {}, - "outputs": [], - "source": [ - "labels = ['co-passing',\n", - " 'counter passing',\n", - " 'co_trapped',\n", - " 'counter-trapped']\n", - "\n", - "for n, (r, pos) in enumerate(zip(pushed_r, pushed_pos)):\n", - " # poloidal \n", - " ax.scatter(r, pos[2], c=colors[n % 4], label=labels[n])\n", - " # topview\n", - " ax_top.scatter(pos[0], pos[1], c=colors[n % 4], label=labels[n])\n", - " ax_top.arrow(pos[0], pos[1], 0., particles.velocities[n, 0]/5, color=colors[n % 4], head_width=.05)\n", - "\n", - "ax.set_xlabel('R')\n", - "ax.set_ylabel('Z')\n", - "ax.set_title('Initial conditions')\n", - "ax.legend();\n", - "\n", - "ax_top.set_xlabel('x')\n", - "ax_top.set_ylabel('y')\n", - "ax_top.set_title('Initial conditions')\n", - "ax_top.legend();\n", - "\n", - "fig " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "56b19093", - "metadata": {}, - "outputs": [], - "source": [ - "# natural constants\n", - "mH = 1.67262192369e-27 # proton mass (kg)\n", - "e = 1.602176634e-19 # elementary charge (C)\n", - "mu0 = 1.25663706212e-6 # magnetic constant (N/A^2)\n", - "\n", - "# epsilon equation parameter\n", - "A = 1. # mass number in units of proton mass\n", - "Z = 1 # signed charge number in units of elementary charge\n", - "unit_x = 1. # length scale unit in m\n", - "unit_B = 1. # magnetic field unit in T\n", - "unit_n = 1e20 # number density unit in m^(-3)\n", - "unit_v = unit_B / np.sqrt(unit_n * A * mH * mu0) # Alfvén velocity unit\n", - "unit_t = unit_x / unit_v # time unit\n", - "\n", - "# cyclotron frequency and epsilon parameter\n", - "om_c = Z*e * unit_B / (A*mH)\n", - "epsilon = 1./(om_c * unit_t)\n", - "\n", - "print(f'{unit_x = }')\n", - "print(f'{unit_B = }')\n", - "print(f'{unit_n = }')\n", - "print(f'{unit_v = }')\n", - "print(f'{unit_t = }')\n", - "print(f'{epsilon = }')" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8cd49aa4", - "metadata": {}, - "outputs": [], - "source": [ - "# instantiate Propagator object\n", - "opts_BxE['algo']['tol'] = 1e-5\n", - "opts_para['algo']['tol'] = 1e-5\n", - "prop_BxE = PushGuidingCenterBxEstar(particles, epsilon=epsilon, algo=opts_BxE['algo'])\n", - "prop_para = PushGuidingCenterParallel(particles, epsilon=epsilon, algo=opts_para['algo'])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "18a1defc", - "metadata": {}, - "outputs": [], - "source": [ - "# time stepping\n", - "Tend = 100. - 1e-6\n", - "dt = .1\n", - "Nt = int(Tend / dt)\n", - "\n", - "pos = np.zeros((Nt + 2, Np, 3), dtype=float)\n", - "r = np.zeros((Nt + 2, Np), dtype=float)\n", - "\n", - "pos[0] = pushed_pos\n", - "r[0] = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)\n", - "\n", - "time = 0.\n", - "n = 0\n", - "while time < Tend:\n", - " time += dt\n", - " n += 1\n", - "\n", - " # advance in time\n", - " prop_BxE(dt/2)\n", - " prop_para(dt)\n", - " prop_BxE(dt/2)\n", - " \n", - " # positions on the physical domain Omega\n", - " pushed_pos = domain(particles.positions).T\n", - " \n", - " # compute R-coordinate\n", - " pos[n] = pushed_pos\n", - " r[n] = np.sqrt(pushed_pos[:, 0]**2 + pushed_pos[:, 1]**2)\n", - " " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "bff57f86", - "metadata": {}, - "outputs": [], - "source": [ - "# make scatter plot for each particle in xy-plane\n", - "for i in range(pos.shape[1]):\n", - " # poloidal \n", - " ax.scatter(r[:, i], pos[:, i, 2], c=colors[i % 4], s=1)\n", - " # top view\n", - " ax_top.scatter(pos[:, i, 0], pos[:, i, 1], c=colors[i % 4], s=1)\n", - "\n", - "ax.set_title(f'{math.ceil(Tend/dt)} time steps')\n", - "ax_top.set_title(f'{math.ceil(Tend/dt)} time steps');\n", - "\n", - "fig" - ] } ], "metadata": { diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index f1a60d888..f12be939f 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -524,14 +524,13 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma log_nt = int(np.log10(int(((nt - 1) / step)))) + 1 # directory for .txt files and marker index which will be saved + path_orbits = os.path.join(path_out, "orbits") + if "5D" in kind: - path_orbits = os.path.join(path_out, "guiding_center") save_index = list(range(0, 6)) + [10] + [-1] elif "6D" in kind or "SPH" in kind: - path_orbits = os.path.join(path_out, "orbits") save_index = list(range(0, 7)) + [-1] else: - path_orbits = os.path.join(path_out, "orbits") save_index = list(range(0, 4)) + [-1] try: From ff45240c8bdb8d6f2253f79cd9b402dcdc88b1f1 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 10 Sep 2025 16:37:45 +0200 Subject: [PATCH 111/292] started to transfer model VlasovAmpereOneSpecies --- src/struphy/main.py | 2 +- src/struphy/models/toy.py | 6 - src/struphy/pic/base.py | 31 ----- .../propagators/propagators_coupling.py | 110 ++++++++++++------ 4 files changed, 78 insertions(+), 71 deletions(-) diff --git a/src/struphy/main.py b/src/struphy/main.py index 46453c278..bdfdfccd2 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -180,7 +180,7 @@ def run( Z_bulk=Z_bulk, verbose=verbose,) - # domain and fluid bckground + # domain and fluid background model.setup_domain_and_equil(domain, equil) # default grid diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index b3bc527c1..902882a31 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -32,8 +32,6 @@ class Maxwell(StruphyModel): :ref:`propagators` (called in sequence): 1. :class:`~struphy.propagators.propagators_fields.Maxwell` - - :ref:`Model info `: """ ## species @@ -110,8 +108,6 @@ class Vlasov(StruphyModel): 1. :class:`~struphy.propagators.propagators_markers.PushVxB` 2. :class:`~struphy.propagators.propagators_markers.PushEta` - - :ref:`Model info `: """ ## species @@ -200,8 +196,6 @@ class GuidingCenter(StruphyModel): 1. :class:`~struphy.propagators.propagators_markers.PushGuidingCenterBxEstar` 2. :class:`~struphy.propagators.propagators_markers.PushGuidingCenterParallel` - - :ref:`Model info `: """ ## species diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index dee1fd197..c826e20ea 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1167,37 +1167,6 @@ def _generate_sampling_moments(self): def _set_initial_condition(self, bp_copy=None, pp_copy=None): self._f_init = self.background - # """Compute callable initial condition from background + perturbation.""" - # if bp_copy is None: - # bp_copy = copy.deepcopy(self.bckgr_params) - # if pp_copy is None: - # pp_copy = copy.deepcopy(self.pert_params) - - # # Get the initialization function and pass the correct arguments - # self._f_init = None - # for fi, maxw_params in bp_copy.items(): - # if fi[-2] == "_": - # fi_type = fi[:-2] - # else: - # fi_type = fi - - # pert_params = pp_copy - # if pp_copy is not None: - # if fi in pp_copy: - # pert_params = pp_copy[fi] - - # if self._f_init is None: - # self._f_init = getattr(maxwellians, fi_type)( - # maxw_params=maxw_params, - # pert_params=pert_params, - # equil=self.equil, - # ) - # else: - # self._f_init = self._f_init + getattr(maxwellians, fi_type)( - # maxw_params=maxw_params, - # pert_params=pert_params, - # equil=self.equil, - # ) def _load_external( self, diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index f7a38b5ac..1aa70a9d2 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -1,6 +1,10 @@ "Particle and FEEC variables are updated." import numpy as np +from dataclasses import dataclass +from mpi4py import MPI +from typing import Literal + from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -17,6 +21,9 @@ from struphy.pic.pushing.pusher import Pusher from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator +from struphy.models.variables import FEECVariable, PICVariable +from struphy.linear_algebra.solver import SolverParameters +from struphy.io.options import (check_option, OptsSymmSolver, OptsMassPrecond, OptsGenSolver, OptsVecSpace) class VlasovAmpere(Propagator): @@ -69,39 +76,76 @@ class VlasovAmpere(Propagator): * For :class:`~struphy.models.hybrid.ColdPlasmaVlasov`: :math:`c_1 = \nu\alpha^2/\varepsilon_\textrm{c} \,, \, c_2 = 1/\varepsilon_\textrm{h}` """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) - - return dct - - def __init__( - self, - e: BlockVector, - particles: Particles6D, - *, - c1: float = 1.0, - c2: float = 1.0, - solver=options(default=True)["solver"], - ): - super().__init__(e, particles) - - self._c1 = c1 - self._c2 = c2 - self._info = solver["info"] + class Variables: + def __init__(self): + self._e = FEECVariable = None + self._ions: PICVariable = None + + @property + def e(self) -> FEECVariable: + return self._e + + @e.setter + def e(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._e = new + + @property + def ions(self) -> PICVariable: + return self._ions + + @ions.setter + def ions(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles6D" + self._ions = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsCouplingCoeffs = Literal["VlasovAmpere", "VlasovMaxwell", "ColdPlasma"] + # propagator options + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + coeffs: OptsCouplingCoeffs = "VlasovAmpere" + + def __post_init__(self): + # checks + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + check_option(self.coeffs, self.OptsCouplingCoeffs) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new + + def allocate(self): + # equation coeffs + if self.options.coeffs == "VlasovAmpere": + self._c1 = c1 + self._c2 = c2 + + self._info = self.options.solver_params.info # get accumulation kernel accum_kernel = accum_kernels.vlasov_maxwell From aaa58fd6c95f9a4280bde85bb240def1512099c4 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 11 Sep 2025 15:30:49 +0200 Subject: [PATCH 112/292] VlasovAmpereOneSecies is running --- src/struphy/main.py | 4 +- src/struphy/models/fluid.py | 7 +- src/struphy/models/hybrid.py | 17 +- src/struphy/models/kinetic.py | 282 +++++++----------- src/struphy/models/species.py | 50 +++- src/struphy/models/toy.py | 6 +- src/struphy/propagators/base.py | 14 +- .../propagators/propagators_coupling.py | 80 +++-- src/struphy/propagators/propagators_fields.py | 243 ++++++++------- .../propagators/propagators_markers.py | 19 +- 10 files changed, 358 insertions(+), 364 deletions(-) diff --git a/src/struphy/main.py b/src/struphy/main.py index bdfdfccd2..55d8e79f7 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -190,9 +190,7 @@ def run( grid = grids.TensorProductGrid(Nel=Nel) # allocate derham-related objects - if derham_opts is not None: - model.allocate_feec(grid, derham_opts) - else: + if derham_opts is None: p = (3, 3, 3) spl_kind = (False, False, False) print(f"\nNo Derham options specified - creating Derham with {p = } and {spl_kind = } for projecting equilibrium.") diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index c307e49f9..304116ab9 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -1,5 +1,4 @@ import numpy as np -from dataclasses import dataclass from mpi4py import MPI from struphy.models.base import StruphyModel @@ -40,8 +39,6 @@ class LinearMHD(StruphyModel): 1. :class:`~struphy.propagators.propagators_fields.ShearAlfven` 2. :class:`~struphy.propagators.propagators_fields.Magnetosonic` - - :ref:`Model info `: """ ## species @@ -59,7 +56,7 @@ def __init__(self): self.init_variables() ## propagators - + class Propagators: def __init__(self): self.shear_alf = propagators_fields.ShearAlfven() @@ -71,7 +68,7 @@ def __init__(self): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # 1. instantiate all species, variables + # 1. instantiate all species self.em_fields = self.EMFields() self.mhd = self.MHD() diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index 76afea648..ee93662ca 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -974,27 +974,22 @@ class ColdPlasmaVlasov(StruphyModel): &\frac{\partial \mathbf B}{\partial t} + \nabla\times\mathbf E = 0\,, \\[2mm] -&\frac{\partial \mathbf E}{\partial t} + \nabla\times\mathbf B = - \frac{\alpha^2}{\varepsilon_\textnormal{c}} \left( \mathbf j_\textnormal{c} + \nu \int_{\mathbb{R}^3} \mathbf{v} f \, \text{d}^3 \mathbf{v} \right) \,, + \frac{\alpha^2}{\varepsilon_\textnormal{h}} \left( \mathbf j_\textnormal{c} + \int_{\mathbb{R}^3} \mathbf{v} f \, \text{d}^3 \mathbf{v} \right) \,, where :math:`(n_0,\mathbf B_0)` denotes a (inhomogeneous) background and .. math:: - \alpha = \frac{\hat \Omega_\textnormal{p,cold}}{\hat \Omega_\textnormal{c,cold}}\,, \qquad \varepsilon_\textnormal{c} = \frac{1}{\hat \Omega_\textnormal{c,cold} \hat t}\,, \qquad \varepsilon_\textnormal{h} = \frac{1}{\hat \Omega_\textnormal{c,hot} \hat t} \,, \qquad \nu = \frac{Z_\textnormal{h}}{Z_\textnormal{c}}\,. + \alpha = \frac{\hat \Omega_\textnormal{p,cold}}{\hat \Omega_\textnormal{c,cold}}\,, \qquad \varepsilon_\textnormal{c} = \frac{1}{\hat \Omega_\textnormal{c,cold} \hat t}\,, \qquad \varepsilon_\textnormal{h} = \frac{1}{\hat \Omega_\textnormal{c,hot} \hat t} \,. At initial time the Poisson equation is solved once to weakly satisfy the Gauss law: .. math:: \begin{align} - \nabla \cdot \mathbf{E} & = \nu \frac{\alpha^2}{\varepsilon_\textnormal{c}} \int_{\mathbb{R}^3} f \, \text{d}^3 \mathbf{v}\,. + \nabla \cdot \mathbf{E} & = \nu \frac{\alpha^2}{\varepsilon_\textnormal{h}} \int_{\mathbb{R}^3} f \, \text{d}^3 \mathbf{v}\,. \end{align} - Note - ---------- - If hot and cold particles are of the same species (:math:`Z_\textnormal{c} = Z_\textnormal{h} \,, A_\textnormal{c} = A_\textnormal{h}`) then :math:`\varepsilon_\textnormal{c} = \varepsilon_\textnormal{h}` and :math:`\nu = 1`. - - :ref:`propagators` (called in sequence): 1. :class:`~struphy.propagators.propagators_fields.Maxwell` @@ -1003,8 +998,6 @@ class ColdPlasmaVlasov(StruphyModel): 4. :class:`~struphy.propagators.propagators_markers.PushVxB` 5. :class:`~struphy.propagators.propagators_markers.PushEta` 6. :class:`~struphy.propagators.propagators_coupling.VlasovAmpere` - - :ref:`Model info `: """ @staticmethod @@ -1111,13 +1104,13 @@ def __init__(self, params, comm, clone_config=None): self._kwargs[propagators_markers.PushVxB] = { "algo": algo_vxb, - "kappa": 1.0 / self._epsilon_cold, + "kappa": 1.0 / self._epsilon_hot, "b2": self.pointer["b_field"], "b2_add": self._b_background, } self._kwargs[propagators_coupling.VlasovAmpere] = { - "c1": self._nu * self._alpha**2 / self._epsilon_cold, + "c1": self._alpha**2 / self._epsilon_hot, "c2": 1.0 / self._epsilon_hot, "solver": params_coupling, } diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 9964e72ad..1f7bc98ee 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -1,9 +1,15 @@ import numpy as np +from mpi4py import MPI from struphy.kinetic_background.base import KineticBackground from struphy.models.base import StruphyModel from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers +from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies +from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable +from struphy.pic.accumulation.particles_to_grid import AccumulatorVector + +rank = MPI.COMM_WORLD.Get_rank() class VlasovAmpereOneSpecies(StruphyModel): @@ -73,148 +79,107 @@ class VlasovAmpereOneSpecies(StruphyModel): 1. :class:`~struphy.propagators.propagators_markers.PushEta` 2. :class:`~struphy.propagators.propagators_coupling.VlasovAmpere` 3. :class:`~struphy.propagators.propagators_markers.PushVxB` - - :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["em_fields"]["e_field"] = "Hcurl" - dct["kinetic"]["species1"] = "Particles6D" - return dct + ## species + + class EMFields(FieldSpecies): + def __init__(self): + self.e_field = FEECVariable(space="Hcurl") + self.phi = FEECVariable(space="H1") + self.init_variables() + + class KineticIons(KineticSpecies): + def __init__(self): + self.var = PICVariable(space="Particles6D") + self.init_variables() + + ## propagators + + class Propagators: + def __init__(self): + self.push_eta = propagators_markers.PushEta() + self.push_vxb = propagators_markers.PushVxB() + self.coupling_va = propagators_coupling.VlasovAmpere() + + ## abstract methods + + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.kinetic_ions = self.KineticIons() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.push_eta.variables.var = self.kinetic_ions.var + self.propagators.push_vxb.variables.ions = self.kinetic_ions.var + self.propagators.coupling_va.variables.e = self.em_fields.e_field + self.propagators.coupling_va.variables.ions = self.kinetic_ions.var + + # define scalars for update_scalar_quantities + self.add_scalar("en_E") + self.add_scalar("en_f", compute="from_particles", variable=self.kinetic_ions.var) + self.add_scalar("en_tot") + + # initial Poisson (not a propagator used in time stepping) + self.initial_poisson = propagators_fields.Poisson() + self.initial_poisson.variables.phi = self.em_fields.phi - @staticmethod - def bulk_species(): - return "species1" + @property + def bulk_species(self): + return self.kinetic_ions - @staticmethod - def velocity_scale(): + @property + def velocity_scale(self): return "light" - @staticmethod - def propagators_dct(): - return { - propagators_markers.PushEta: ["species1"], - propagators_markers.PushVxB: ["species1"], - propagators_coupling.VlasovAmpere: ["e_field", "species1"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - # add special options - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["em_fields"], - option=propagators_fields.ImplicitDiffusion, - dct=dct, - ) - cls.add_option( - species=["kinetic", "species1"], - key="override_eq_params", - option=[False, {"alpha": 1.0, "epsilon": -1.0}], - dct=dct, - ) - return dct - - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - from mpi4py.MPI import IN_PLACE, SUM + def allocate_helpers(self): + self._tmp = np.empty(1, dtype=float) - # get species paramaters - species1_params = params["kinetic"]["species1"] + def update_scalar_quantities(self): + # e*M1*e/2 + e = self.em_fields.e_field.spline.vector + en_E = 0.5 * self.mass_ops.M1.dot_inner(e, e) + self.update_scalar("en_E", en_E) - # Get coupling strength - if species1_params["options"]["override_eq_params"]: - self._alpha = species1_params["options"]["override_eq_params"]["alpha"] - self._epsilon = species1_params["options"]["override_eq_params"]["epsilon"] - print( - f"\n!!! Override equation parameters: {self._alpha = } and {self._epsilon = }.", + # alpha^2 / 2 / N * sum_p w_p v_p^2 + particles = self.kinetic_ions.var.particles + alpha = self.kinetic_ions.equation_params.alpha + self._tmp[0] = ( + alpha**2 + / (2 * particles.Np) + * np.dot( + particles.markers_wo_holes[:, 3] ** 2 + + particles.markers_wo_holes[:, 4] ** 2 + + particles.markers_wo_holes[:, 5] ** 2, + particles.markers_wo_holes[:, 6], ) - else: - self._alpha = self.equation_params["species1"]["alpha"] - self._epsilon = self.equation_params["species1"]["epsilon"] - - # Check if it is control-variate method - self._control_variate = species1_params["markers"]["control_variate"] - - # check mean velocity - # TODO: assert f0.params[] == 0. - - # Initialize background magnetic field from MHD equilibrium - if self.projected_equil: - self._b_background = self.projected_equil.b2 - else: - self._b_background = None - - # propagator parameters - self._poisson_params = params["em_fields"]["options"]["ImplicitDiffusion"]["solver"] - algo_eta = params["kinetic"]["species1"]["options"]["PushEta"]["algo"] - if self._b_background is not None: - algo_vxb = params["kinetic"]["species1"]["options"]["PushVxB"]["algo"] - params_coupling = params["em_fields"]["options"]["VlasovAmpere"]["solver"] - - # set keyword arguments for propagators - self._kwargs[propagators_markers.PushEta] = { - "algo": algo_eta, - } - - # Only add PushVxB if magnetic field is not zero - self._kwargs[propagators_markers.PushVxB] = None - if self._b_background is not None: - self._kwargs[propagators_markers.PushVxB] = { - "algo": algo_vxb, - "b2": self._b_background, - "kappa": 1.0 / self._epsilon, - } - - self._kwargs[propagators_coupling.VlasovAmpere] = { - "c1": self._alpha**2 / self._epsilon, - "c2": 1.0 / self._epsilon, - "solver": params_coupling, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during the simulation - self.add_scalar("en_E") - self.add_scalar("en_f") - self.add_scalar("en_tot") - - # MPI operations needed for scalar variables - self._mpi_sum = SUM - self._mpi_in_place = IN_PLACE + ) + self.update_scalar("en_f", self._tmp[0]) - # temporaries - self._tmp = np.empty(1, dtype=float) + # en_tot = en_w + en_e + self.update_scalar("en_tot", en_E + self._tmp[0]) - def initialize_from_params(self): + def allocate_propagators(self): """Solve initial Poisson equation. :meta private: """ - from struphy.pic.accumulation.particles_to_grid import AccumulatorVector - # initialize fields and particles - super().initialize_from_params() + super().allocate_propagators() - if self.rank_world == 0: + if MPI.COMM_WORLD.Get_rank() == 0: print("\nINITIAL POISSON SOLVE:") # use control variate method - self.pointer["species1"].update_weights() + particles = self.kinetic_ions.var.particles + particles.update_weights() # sanity check # self.pointer['species1'].show_distribution_function( @@ -222,65 +187,50 @@ def initialize_from_params(self): # accumulate charge density charge_accum = AccumulatorVector( - self.pointer["species1"], + particles, "H1", accum_kernels.charge_density_0form, self.mass_ops, self.domain.args_domain, ) - charge_accum(self.pointer["species1"].vdim) + charge_accum(particles.vdim) # another sanity check: compute FE coeffs of density # charge_accum.show_accumulated_spline_field(self.mass_ops) - - # Instantiate Poisson solver - _phi = self.derham.Vh["0"].zeros() - poisson_solver = propagators_fields.ImplicitDiffusion( - _phi, - sigma_1=0.0, - sigma_2=0.0, - sigma_3=1.0, - rho=self._alpha**2 / self._epsilon * charge_accum.vectors[0], - solver=self._poisson_params, - ) - + + alpha = self.kinetic_ions.equation_params.alpha + epsilon = self.kinetic_ions.equation_params.epsilon + + self.initial_poisson.options.rho = alpha**2 / epsilon * charge_accum.vectors[0] + # self.initial_poisson.variables.phi.allocate(self.derham, domain=self.domain) + self.initial_poisson.allocate() + # Solve with dt=1. and compute electric field - if self.rank_world == 0: + if MPI.COMM_WORLD.Get_rank() == 0: print("\nSolving initial Poisson problem...") - poisson_solver(1.0) + self.initial_poisson(1.0) - self.derham.grad.dot(-_phi, out=self.pointer["e_field"]) - if self.rank_world == 0: + phi = self.initial_poisson.variables.phi.spline.vector + self.derham.grad.dot(-phi, out=self.em_fields.e_field.spline.vector) + if MPI.COMM_WORLD.Get_rank() == 0: print("Done.") - def update_scalar_quantities(self): - # e*M1*e/2 - en_E = 0.5 * self.mass_ops.M1.dot_inner(self.pointer["e_field"], self.pointer["e_field"]) - self.update_scalar("en_E", en_E) - - # alpha^2 / 2 / N * sum_p w_p v_p^2 - self._tmp[0] = ( - self._alpha**2 - / (2 * self.pointer["species1"].Np) - * np.dot( - self.pointer["species1"].markers_wo_holes[:, 3] ** 2 - + self.pointer["species1"].markers_wo_holes[:, 4] ** 2 - + self.pointer["species1"].markers_wo_holes[:, 5] ** 2, - self.pointer["species1"].markers_wo_holes[:, 6], - ) - ) - if self.comm_world is not None: - self.comm_world.Allreduce( - self._mpi_in_place, - self._tmp, - op=self._mpi_sum, - ) - self.update_scalar("en_f", self._tmp[0]) - - # en_tot = en_w + en_e - self.update_scalar("en_tot", en_E + self._tmp[0]) - + ## default parameters + def generate_default_parameter_file(self, path = None, prompt = True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "coupling_va.Options" in line: + new_file += [line] + new_file += ["model.initial_poisson.options = model.initial_poisson.Options()\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) class VlasovMaxwellOneSpecies(StruphyModel): r"""Vlasov-Maxwell equations for one species. diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 4cf94e097..5bf80b179 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -4,6 +4,7 @@ from typing import Callable import numpy as np from mpi4py import MPI +import warnings from struphy.fields_background.base import FluidEquilibrium from struphy.kinetic_background.base import KineticBackground @@ -45,15 +46,31 @@ def mass_number(self) -> int: """Mass number in units of proton mass.""" return self._mass_number - def set_phys_params(self, charge_number: int = 1, mass_number: int = 1): - """Set charge- and mass number.""" + def set_phys_params(self, + charge_number: int = 1, + mass_number: int = 1, + alpha: float = None, + epsilon: float = None, + kappa: float = None, + ): + """Set charge- and mass number. Set equation parameters (alpha, epsilon, ...) to override units.""" self._charge_number = charge_number self._mass_number = mass_number + self.alpha = alpha + self.epsilon = epsilon + self.kappa = kappa class EquationParameters: """Normalization parameters of one species, appearing in scaled equations.""" - def __init__(self, species, units: Units = None, verbose: bool = False): + def __init__(self, + species, + units: Units = None, + alpha: float = None, + epsilon: float = None, + kappa: float = None, + verbose: bool = False, + ): if units is None: units = Units() @@ -67,9 +84,23 @@ def __init__(self, species, units: Units = None, verbose: bool = False): om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters - self.alpha = om_p / om_c - self.epsilon = 1.0 / (om_c * units.t) - self.kappa = om_p * units.t + if alpha is None: + self.alpha = om_p / om_c + else: + self.alpha = alpha + warnings.warn(f"Override equation parameter {self.alpha = }") + + if epsilon is None: + self.epsilon = 1.0 / (om_c * units.t) + else: + self.epsilon = epsilon + warnings.warn(f"Override equation parameter {self.epsilon = }") + + if kappa is None: + self.kappa = om_p * units.t + else: + self.kappa = kappa + warnings.warn(f"Override equation parameter {self.kappa = }") if verbose and MPI.COMM_WORLD.Get_rank() == 0: print(f'\nSet normalization parameters for species {species.__class__.__name__}:') @@ -87,7 +118,12 @@ def setup_equation_params(self, units: Units, verbose=False): * epsilon = 1 / (cyclotron frequency * time unit) * kappa = plasma frequency * time unit """ - self._equation_params = self.EquationParameters(species=self, units=units, verbose=verbose) + self._equation_params = self.EquationParameters(species=self, + units=units, + alpha=self.alpha, + epsilon=self.epsilon, + kappa=self.kappa, + verbose=verbose,) diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 902882a31..0a346b4e3 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -54,7 +54,7 @@ def __init__(self): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # 1. instantiate all species, variables + # 1. instantiate all species self.em_fields = self.EMFields() # 2. instantiate all propagators @@ -130,7 +130,7 @@ def __init__(self): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}' ***") - # 1. instantiate all species, variables + # 1. instantiate all species self.kinetic_ions = self.KineticIons() # 2. instantiate all propagators @@ -218,7 +218,7 @@ def __init__(self): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}' ***") - # 1. instantiate all species, variables + # 1. instantiate all species self.kinetic_ions = self.KineticIons() # 2. instantiate all propagators diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index b1cfcc0a8..2f1e5181a 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -98,18 +98,18 @@ def update_feec_variables(self, **new_coeffs): ------- diffs : dict max_diff for all feec variables. - """ + """ diffs = {} - assert len(new_coeffs) == len(self.variables.__dict__), f"Coefficients must be passed in the following order: {self.variables}, but is {new_coeffs}." - for ((k, new), (kp, vp)) in zip(new_coeffs.items(), self.variables.__dict__.items()): - assert k in kp, f"Variable name '{k}' not in '{kp}'; variables must be passed in the order {self.variables.__dict__}." + for var, new in new_coeffs.items(): + assert "_" + var in self.variables.__dict__, f"{var} not in {self.variables.__dict__}." assert isinstance(new, (StencilVector, BlockVector)) - assert isinstance(vp, FEECVariable) - old = vp.spline.vector + old_var = getattr(self.variables, var) + assert isinstance(old_var, FEECVariable) + old = old_var.spline.vector assert new.space == old.space # calculate maximum of difference abs(new - old) - diffs[k] = np.max(np.abs(new.toarray() - old.toarray())) + diffs[var] = np.max(np.abs(new.toarray() - old.toarray())) # copy new coeffs into old new.copy(out=old) diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index 1aa70a9d2..af0ea32a3 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -33,9 +33,9 @@ class VlasovAmpere(Propagator): .. math:: -& \int_\Omega \frac{\partial \mathbf E}{\partial t} \cdot \mathbf F\,\textrm d \mathbf x = - c_1 \int_\Omega \int_{\mathbb{R}^3} f \mathbf{v} \cdot \mathbf F \, \text{d}^3 \mathbf{v} \,\textrm d \mathbf x \qquad \forall \, \mathbf F \in H(\textnormal{curl}) \,, + \frac{\alpha^2}{\varepsilon} \int_\Omega \int_{\mathbb{R}^3} f \mathbf{v} \cdot \mathbf F \, \text{d}^3 \mathbf{v} \,\textrm d \mathbf x \qquad \forall \, \mathbf F \in H(\textnormal{curl}) \,, \\[2mm] - &\frac{\partial f}{\partial t} + c_2\, \mathbf{E} + &\frac{\partial f}{\partial t} + \frac{1}{\varepsilon}\, \mathbf{E} \cdot \frac{\partial f}{\partial \mathbf{v}} = 0 \,. :ref:`time_discret`: Crank-Nicolson (implicit mid-point). System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`, such that @@ -49,8 +49,8 @@ class VlasovAmpere(Propagator): = \frac{\Delta t}{2} \begin{bmatrix} - 0 & - c_1 \mathbb L^1 \bar{DF^{-1}} \bar{\mathbf w} \\ - c_2 \bar{DF^{-\top}} \left(\mathbb L^1\right)^\top & 0 + 0 & - \frac{\alpha^2}{\varepsilon} \mathbb L^1 \bar{DF^{-1}} \bar{\mathbf w} \\ + \frac{1}{\varepsilon} \bar{DF^{-\top}} \left(\mathbb L^1\right)^\top & 0 \end{bmatrix} \begin{bmatrix} \mathbf{e}^{n+1} + \mathbf{e}^n \\ @@ -61,24 +61,18 @@ class VlasovAmpere(Propagator): .. math:: - A = \mathbb M^1\,,\qquad B = \frac{c_1}{2} \mathbb L^1 \bar{DF^{-1}} \bar{\mathbf w}\,,\qquad C = - \frac{c_2}{2} \bar{DF^{-\top}} \left(\mathbb L^1\right)^\top \,. + A = \mathbb M^1\,,\qquad B = \frac{\alpha^2}{2\varepsilon} \mathbb L^1 \bar{DF^{-1}} \bar{\mathbf w}\,,\qquad C = - \frac{1}{2\varepsilon} \bar{DF^{-\top}} \left(\mathbb L^1\right)^\top \,. The accumulation matrix and vector assembled in :class:`~struphy.pic.accumulation.particles_to_grid.Accumulator` are .. math:: M = BC \,,\qquad V = B \mathbf V \,. - - Note - ---------- - * For :class:`~struphy.models.kinetic.VlasovAmpereOneSpecies`: :math:`c_1 = \kappa^2 \,, \, c_2 = 1` - * For :class:`~struphy.models.kinetic.VlasovMaxwellOneSpecies`: :math:`c_1 = \alpha^2/\varepsilon \,, \, c_2 = 1/\varepsilon` - * For :class:`~struphy.models.hybrid.ColdPlasmaVlasov`: :math:`c_1 = \nu\alpha^2/\varepsilon_\textrm{c} \,, \, c_2 = 1/\varepsilon_\textrm{h}` """ class Variables: def __init__(self): - self._e = FEECVariable = None + self._e: FEECVariable = None self._ions: PICVariable = None @property @@ -106,19 +100,14 @@ def __init__(self): @dataclass class Options: - # specific literals - OptsCouplingCoeffs = Literal["VlasovAmpere", "VlasovMaxwell", "ColdPlasma"] - # propagator options solver: OptsSymmSolver = "pcg" precond: OptsMassPrecond = "MassMatrixPreconditioner" solver_params: SolverParameters = None - coeffs: OptsCouplingCoeffs = "VlasovAmpere" def __post_init__(self): # checks check_option(self.solver, OptsSymmSolver) check_option(self.precond, OptsMassPrecond) - check_option(self.coeffs, self.OptsCouplingCoeffs) # defaults if self.solver_params is None: @@ -140,10 +129,12 @@ def options(self, new): self._options = new def allocate(self): - # equation coeffs - if self.options.coeffs == "VlasovAmpere": - self._c1 = c1 - self._c2 = c2 + # scaling factors + alpha = self.variables.ions.species.equation_params.alpha + epsilon = self.variables.ions.species.equation_params.epsilon + + self._c1 = alpha**2 / epsilon + self._c2 = 1.0 / epsilon self._info = self.options.solver_params.info @@ -151,6 +142,8 @@ def allocate(self): accum_kernel = accum_kernels.vlasov_maxwell # Initialize Accumulator object + particles = self.variables.ions.particles + self._accum = Accumulator( particles, "Hcurl", @@ -162,19 +155,19 @@ def allocate(self): ) # Create buffers to store temporarily e and its sum with old e - self._e_tmp = e.space.zeros() - self._e_scale = e.space.zeros() - self._e_sum = e.space.zeros() + self._e_tmp = self.derham.Vh["1"].zeros() + self._e_scale = self.derham.Vh["1"].zeros() + self._e_sum = self.derham.Vh["1"].zeros() # ================================ # ========= Schur Solver ========= # ================================ # Preconditioner - if solver["type"][1] == None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(self.mass_ops.M1) # Define block matrix [[A B], [C I]] (without time step size dt in the diagonals) @@ -185,11 +178,9 @@ def allocate(self): self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, ) # Instantiate particle pusher @@ -217,14 +208,14 @@ def __call__(self, dt): self._schur_solver.BC = self._accum.operators[0] self._schur_solver.BC *= -self._c1 * self._c2 / 4.0 - # Vector for schur solver + # Vector for Schur solver self._e_scale *= 0.0 self._e_scale += self._accum.vectors[0] self._e_scale *= self._c1 / 2.0 # new e coeffs self._e_tmp, info = self._schur_solver( - self.feec_vars[0], + self.variables.e.spline.vector, self._e_scale, dt, out=self._e_tmp, @@ -232,7 +223,7 @@ def __call__(self, dt): # mid-point e-field (no tmps created here) self._e_sum *= 0.0 - self._e_sum += self.feec_vars[0] + self._e_sum += self.variables.e.spline.vector self._e_sum += self._e_tmp self._e_sum *= 0.5 @@ -240,29 +231,30 @@ def __call__(self, dt): self._pusher(dt) # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() + if self.variables.ions.species.weights_params.control_variate: + self.variables.ions.particles.update_weights() # write new coeffs into self.variables - (max_de,) = self.feec_vars_update(self._e_tmp) + (max_de,) = self.update_feec_variables(e=self._e_tmp) # Print out max differences for weights and e-field if self._info: print("Status for VlasovMaxwell:", info["success"]) print("Iterations for VlasovMaxwell:", info["niter"]) print("Maxdiff e1 for VlasovMaxwell:", max_de) - buffer_idx = self.particles[0].bufferindex + particles = self.variables.ions.particles + buffer_idx = particles.bufferindex max_diff = np.max( np.abs( np.sqrt( - self.particles[0].markers_wo_holes[:, 3] ** 2 - + self.particles[0].markers_wo_holes[:, 4] ** 2 - + self.particles[0].markers_wo_holes[:, 5] ** 2, + particles.markers_wo_holes[:, 3] ** 2 + + particles.markers_wo_holes[:, 4] ** 2 + + particles.markers_wo_holes[:, 5] ** 2, ) - np.sqrt( - self.particles[0].markers_wo_holes[:, buffer_idx + 3] ** 2 - + self.particles[0].markers_wo_holes[:, buffer_idx + 4] ** 2 - + self.particles[0].markers_wo_holes[:, buffer_idx + 5] ** 2, + particles.markers_wo_holes[:, buffer_idx + 3] ** 2 + + particles.markers_wo_holes[:, buffer_idx + 4] ** 2 + + particles.markers_wo_holes[:, buffer_idx + 5] ** 2, ), ), ) diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 8ebaf991e..e514429e0 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -2662,63 +2662,84 @@ class ImplicitDiffusion(Propagator): solver : dict Parameters for the iterative solver (see ``__init__`` for details). """ + class Variables: + def __init__(self): + self._phi: FEECVariable = None + + @property + def phi(self) -> FEECVariable: + return self._phi + + @phi.setter + def phi(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._phi = new - @staticmethod - def options(default=False): - dct = {} - dct["model"] = { - "sigma_1": 1.0, - "sigma_2": 0.0, - "sigma_3": 1.0, - "stab_mat": ["M0", "M0ad", "Id"], - "diffusion_mat": ["M1", "M1perp"], - } - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) - - return dct - - def __init__( - self, - phi: StencilVector, - *, - sigma_1: float = options()["model"]["sigma_1"], - sigma_2: float = options()["model"]["sigma_2"], - sigma_3: float = options()["model"]["sigma_3"], - divide_by_dt: bool = False, - stab_mat: str = options(default=True)["model"]["stab_mat"], - diffusion_mat: str = options(default=True)["model"]["diffusion_mat"], - rho: StencilVector | tuple | list | Callable = None, - x0: StencilVector = None, - solver: dict = options(default=True)["solver"], - ): - assert phi.space == self.derham.Vh["0"] - - super().__init__(phi) + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsStabMat = Literal["M0", "M0ad", "Id"] + OptsDiffusionMat = Literal["M1", "M1perp"] + # propagator options + sigma_1: float = 1.0 + sigma_2: float = 0.0 + sigma_3: float = 1.0 + divide_by_dt: bool = False + stab_mat: OptsStabMat = "M0" + diffusion_mat: OptsDiffusionMat = "M1" + rho: StencilVector | tuple | list | Callable = None + x0: StencilVector = None + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + + def __post_init__(self): + # checks + check_option(self.stab_mat, self.OptsStabMat) + check_option(self.diffusion_mat, self.OptsDiffusionMat) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f' {k}: {v}') + self._options = new + def allocate(self): # always stabilize - if np.abs(sigma_1) < 1e-14: - sigma_1 = 1e-14 - print(f"Stabilizing Poisson solve with {sigma_1 = }") + if np.abs(self.options.sigma_1) < 1e-14: + self.options.sigma_1 = 1e-14 + print(f"Stabilizing Poisson solve with {self.options.sigma_1 = }") # model parameters - self._sigma_1 = sigma_1 - self._sigma_2 = sigma_2 - self._sigma_3 = sigma_3 - self._divide_by_dt = divide_by_dt + self._sigma_1 = self.options.sigma_1 + self._sigma_2 = self.options.sigma_2 + self._sigma_3 = self.options.sigma_3 + self._divide_by_dt = self.options.divide_by_dt + + phi = self.variables.phi.spline.vector # collect rhs + rho = self.options.rho + if rho is None: self._rho = [phi.space.zeros()] else: @@ -2743,18 +2764,19 @@ def __init__( self._rho = rho # initial guess and solver params - self._x0 = x0 - self._info = solver["info"] + self._x0 = self.options.x0 + self._info = self.options.solver_params.info - if stab_mat == "Id": + if self.options.stab_mat == "Id": stab_mat = IdentityOperator(phi.space) else: - stab_mat = getattr(self.mass_ops, stab_mat) + stab_mat = getattr(self.mass_ops, self.options.stab_mat) - print(f"{diffusion_mat = }") - if isinstance(diffusion_mat, str): - diffusion_mat = getattr(self.mass_ops, diffusion_mat) + print(f"{self.options.diffusion_mat = }") + if isinstance(self.options.diffusion_mat, str): + diffusion_mat = getattr(self.mass_ops, self.options.diffusion_mat) else: + diffusion_mat = self.options.diffusion_mat assert isinstance(diffusion_mat, WeightedMassOperator) assert diffusion_mat.domain == self.derham.grad.codomain assert diffusion_mat.codomain == self.derham.grad.codomain @@ -2764,7 +2786,7 @@ def __init__( self._diffusion_op = self.derham.grad.T @ diffusion_mat @ self.derham.grad # preconditioner and solver for Ax=b - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: # TODO: waiting for multigrid preconditioner @@ -2773,13 +2795,13 @@ def __init__( # solver just with A_2, but will be set during call with dt self._solver = inverse( self._diffusion_op, - solver["type"][0], + self.options.solver, pc=pc, x0=self.x0, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + recycle=self.options.solver_params.recycle, ) # allocate memory for solution @@ -2858,7 +2880,7 @@ def __call__(self, dt): sig_3 = self._sigma_3 # compute rhs - phin = self.feec_vars[0] + phin = self.variables.phi.spline.vector rhs = self._stab_mat.dot(phin, out=self._rhs) rhs *= sig_2 @@ -2882,7 +2904,7 @@ def __call__(self, dt): if self._info: print(info) - self.feec_vars_update(out) + self.update_feec_variables(phi=out) class Poisson(ImplicitDiffusion): @@ -2929,52 +2951,51 @@ class Poisson(ImplicitDiffusion): solver : dict Parameters for the iterative solver (see ``__init__`` for details). """ - - @staticmethod - def options(default=False): - dct = {} - dct["stabilization"] = { - "stab_eps": 0.0, - "stab_mat": ["Id", "M0", "M0ad"], - } - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) - - return dct - - def __init__( - self, - phi: StencilVector, - *, - stab_eps: float = 0.0, - stab_mat: str = options(default=True)["stabilization"]["stab_mat"], - rho: StencilVector | tuple | list | Callable = None, - x0: StencilVector = None, - solver: dict = options(default=True)["solver"], - ): - super().__init__( - phi, - sigma_1=stab_eps, - sigma_2=0.0, - sigma_3=1.0, - divide_by_dt=False, - stab_mat=stab_mat, - diffusion_mat="M1", - rho=rho, - x0=x0, - solver=solver, - ) + @dataclass + class Options: + # specific literals + OptsStabMat = Literal["M0", "M0ad", "Id"] + # propagator options + stab_eps: float = 0.0 + stab_mat: OptsStabMat = "M0" + rho: StencilVector | tuple | list | Callable = None + x0: StencilVector = None + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + + def __post_init__(self): + # checks + check_option(self.stab_mat, self.OptsStabMat) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + # Poisson solve (-> set some params of parent class) + self.sigma_1 = self.stab_eps + self.sigma_2 = 0.0 + self.sigma_3 = 1.0 + self.divide_by_dt = False + self.diffusion_mat = "M1" + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + if "sigma" not in k and k not in ("divide_by_dt", "diffusion_mat"): + print(f' {k}: {v}') + self._options = new class VariationalMomentumAdvection(Propagator): diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index b87493005..74f70c5ef 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -125,9 +125,9 @@ class PushVxB(Propagator): .. math:: - \frac{\textnormal d \mathbf v_p(t)}{\textnormal d t} = \kappa \, \mathbf v_p(t) \times (\mathbf B + \mathbf B_{\text{add}}) \,, + \frac{\textnormal d \mathbf v_p(t)}{\textnormal d t} = \frac{1}{\varepsilon} \, \mathbf v_p(t) \times (\mathbf B + \mathbf B_{\text{add}}) \,, - where :math:`\kappa \in \mathbb R` is a constant scaling factor, and for rotation vector :math:`\mathbf B` and optional, additional fixed rotation + where :math:`\varepsilon = 1/(\hat\Omega_c \hat t)` is a constant scaling factor, and for rotation vector :math:`\mathbf B` and optional, additional fixed rotation vector :math:`\mathbf B_{\text{add}}`, both given as a 2-form: .. math:: @@ -159,7 +159,6 @@ class Options: OptsAlgo = Literal["analytic", "implicit"] # propagator options algo: OptsAlgo = "analytic" - kappa: float = 1.0 b2_var: FEECVariable = None def __post_init__(self): @@ -182,6 +181,9 @@ def options(self, new): self._options = new def allocate(self): + # scaling factor + self._epsilon = self.variables.ions.species.equation_params.epsilon + # TODO: treat PolarVector as well, but polar splines are being reworked at the moment if self.projected_equil is not None: self._b2 = self.projected_equil.b2 @@ -235,9 +237,10 @@ def __call__(self, dt): # extract coefficients to tensor product space b_full: BlockVector = self._E2T.dot(tmp, out=self._b_full) b_full.update_ghost_regions() + b_full /= self._epsilon # call pusher kernel - self._pusher(self.options.kappa * dt) + self._pusher(dt) # update_weights if self.variables.ions.particles.control_variate: @@ -503,6 +506,9 @@ def options(self, new): self._options = new def allocate(self): + # scaling factor + self._epsilon = self.variables.ions.species.equation_params.epsilon + # magnetic equilibrium field unit_b1 = self.projected_equil.unit_b1 self._gradB1 = self.projected_equil.gradB1 @@ -533,7 +539,6 @@ def allocate(self): self._phi = self.options.phi.spline.vector self._evaluate_e_field = self.options.evaluate_e_field self._e_field = self.derham.Vh["1"].zeros() - self._epsilon = self.variables.ions.species.equation_params.epsilon # choose method particles = self.variables.ions.particles @@ -937,6 +942,9 @@ def options(self, new): self._options = new def allocate(self): + # scaling factor + self._epsilon = self.variables.ions.species.equation_params.epsilon + # magnetic equilibrium field self._gradB1 = self.projected_equil.gradB1 b2 = self.projected_equil.b2 @@ -968,7 +976,6 @@ def allocate(self): self._phi = self.options.phi.spline.vector self._evaluate_e_field = self.options.evaluate_e_field self._e_field = self.derham.Vh["1"].zeros() - self._epsilon = self.variables.ions.species.equation_params.epsilon # choose method particles = self.variables.ions.particles From 32baaf2af652b5f4dd449da23529990818ff95c7 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 12 Sep 2025 08:09:07 +0200 Subject: [PATCH 113/292] new class BinningPlot --- src/struphy/models/base.py | 14 ++-- src/struphy/models/species.py | 8 +- src/struphy/models/variables.py | 59 +++++++-------- src/struphy/pic/utilities.py | 29 +++++++ .../post_processing/post_processing_tools.py | 75 ++++++++++--------- 5 files changed, 109 insertions(+), 76 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 914fd2acb..e7b9c6da0 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -626,7 +626,6 @@ def allocate_variables(self, verbose: bool = False): # self._pointer[key] = val["obj"].vector - def integrate(self, dt, split_algo="LieTrotter"): """ Advance the model by a time step ``dt`` by sequentially calling its Propagators. @@ -735,7 +734,7 @@ def update_distr_functions(self): str_dn = f"d{i + 1}" dim_to_int[str_dn] = 3 + obj.vdim + 3 + i - if species.f_binned is not None: + if species.binning_plots: for slice_i, edges in var.kinetic_data["bin_edges"].items(): comps = slice_i.split("_") components = [False] * (3 + obj.vdim + 3 + obj.n_cols_diagnostics) @@ -1093,11 +1092,14 @@ def initialize_data_output(self, data: DataContainer, size): data.add_data({key_spec_restart: obj._markers}) + # TODO: kinetic_data should be a KineticData object, not a dict for key1, val1 in var.kinetic_data.items(): key_dat = os.path.join(key_spec, key1) - # case of "f" and "df" - if isinstance(val1, dict): + if key1 == "bin_edges": + continue + elif key1 == "f" or key1 == "df": + assert isinstance(val1, dict) for key2, val2 in val1.items(): key_f = os.path.join(key_dat, key2) data.add_data({key_f: val2}) @@ -1327,7 +1329,7 @@ def generate_default_parameter_file( exclude = f"# model.{sn}.{vn}.save_data = False\n" elif isinstance(var, PICVariable): has_pic = True - init_pert_pic = f"perturbation = perturbations.TorusModesCos()\n" + init_pert_pic = f"\nperturbation = perturbations.TorusModesCos()" if "6D" in var.space: init_bckgr_pic = f"\nmaxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n" init_bckgr_pic += f"maxwellian_2 = maxwellians.Maxwellian3D(n=(0.1, None))\n" @@ -1347,7 +1349,7 @@ def generate_default_parameter_file( file.write("from struphy.initial import perturbations\n") file.write("from struphy.kinetic_background import maxwellians\n") - file.write("from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters\n") + file.write("from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters, BinningPlot\n") file.write("from struphy import main\n") file.write("\n# import model, set verbosity\n") diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 5bf80b179..ca8ef48c4 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -13,7 +13,9 @@ from struphy.models.variables import Variable from struphy.pic.utilities import (LoadingParameters, WeightsParameters, - BoundaryParameters,) + BoundaryParameters, + BinningPlot, + ) class Species(metaclass=ABCMeta): @@ -191,12 +193,12 @@ def set_sorting_boxes(self, def set_save_data(self, n_markers: int | float = 3, - f_binned: dict = None, + binning_plots: tuple[BinningPlot] = (), n_sph: dict = None, ): """Saving marker orits, binned data and kernel density reconstructions.""" self.n_markers = n_markers - self.f_binned = f_binned + self.binning_plots = binning_plots self.n_sph = n_sph diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 863bca007..58b6f6885 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -214,38 +214,31 @@ def allocate(self, self.particles.initialize_weights() # for storing the binned distribution function - if self.species.f_binned is not None: - slices = self.species.f_binned["slices"] - n_bins = self.species.f_binned["n_bins"] - ranges = self.species.f_binned["ranges"] - - self.kinetic_data["f"] = {} - self.kinetic_data["df"] = {} - self.kinetic_data["bin_edges"] = {} - if len(slices) > 0: - for i, sli in enumerate(slices): - assert ((len(sli) - 2) / 3).is_integer() - assert len(slices[i].split("_")) == len(ranges[i]) == len(n_bins[i]), ( - f"Number of slices names ({len(slices[i].split('_'))}), number of bins ({len(n_bins[i])}), and number of ranges ({len(ranges[i])}) are inconsistent with each other!\n\n" - ) - self.kinetic_data["bin_edges"][sli] = [] - dims = (len(sli) - 2) // 3 + 1 - for j in range(dims): - self.kinetic_data["bin_edges"][sli] += [ - np.linspace( - ranges[i][j][0], - ranges[i][j][1], - n_bins[i][j] + 1, - ), - ] - self.kinetic_data["f"][sli] = np.zeros( - n_bins[i], - dtype=float, - ) - self.kinetic_data["df"][sli] = np.zeros( - n_bins[i], - dtype=float, - ) + self.kinetic_data["bin_edges"] = {} + self.kinetic_data["f"] = {} + self.kinetic_data["df"] = {} + + for bin_plot in self.species.binning_plots: + sli = bin_plot.slice + n_bins = bin_plot.n_bins + ranges = bin_plot.ranges + + assert ((len(sli) - 2) / 3).is_integer(), f"Binning coordinates must be separated by '_', but reads {sli}." + assert len(sli.split("_")) == len(ranges) == len(n_bins), ( + f"Number of slices names ({len(sli.split('_'))}), number of bins ({len(n_bins)}), and number of ranges ({len(ranges)}) are inconsistent with each other!\n\n" + ) + self.kinetic_data["bin_edges"][sli] = [] + dims = (len(sli) - 2) // 3 + 1 + for j in range(dims): + self.kinetic_data["bin_edges"][sli] += [ + np.linspace( + ranges[j][0], + ranges[j][1], + n_bins[j] + 1, + ), + ] + self.kinetic_data["f"][sli] = np.zeros(n_bins, dtype=float) + self.kinetic_data["df"][sli] = np.zeros(n_bins, dtype=float) # for storing an sph evaluation of the density n if self.species.n_sph is not None: @@ -268,7 +261,7 @@ def allocate(self, self.kinetic_data["n_sph"] += [np.zeros(ee1.shape, dtype=float)] # other data (wave-particle power exchange, etc.) - # TODO + # TODO class SPHVariable(Variable): diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index c0047bf92..db9599a5a 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -132,6 +132,35 @@ def __init__(self, self.bc_sph = bc_sph +class BinningPlot: + """Binning plot of marker distribution in phase space. + + Parameters + ---------- + slice : str + Coordinate-slice in phase space to bin. A combination of "e1", "e2", "e3", "v1", etc., separated by an underscore "_". + For example, "e1" showas a 1D binning plot over eta1, whereas "e1_v1" shows a 2D binning plot over eta1 and v1. + + n_bins : int | tuple[int] + Number of bins for each coordinate. + + ranges : tuple[int] | tuple[tuple[int]]= (0.0, 1.0) + Binning range (as an interval in R) for each coordinate. + """ + def __init__(self, + slice: str = "e1", + n_bins: int | tuple[int] = 128, + ranges: tuple[float] | tuple[tuple[float]]= (0.0, 1.0),): + self.slice = slice + + if isinstance(n_bins, int): + n_bins = (n_bins,) + self.n_bins = n_bins + + if not isinstance(ranges[0], tuple): + ranges = (ranges,) + self.ranges = ranges + def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): """ diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index f12be939f..98d2f1e3a 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -10,12 +10,15 @@ from struphy.fields_background.base import FluidEquilibrium from struphy.fields_background import equils from struphy.io.setup import import_parameters_py -from struphy.models.base import setup_derham +from struphy.models.base import setup_derham, StruphyModel +from struphy.models.species import KineticSpecies +from struphy.models.variables import PICVariable from struphy.feec.psydac_derham import SplineFunction from struphy.io.options import EnvironmentOptions, Units, Time from struphy.topology.grids import TensorProductGrid from struphy.geometry import domains from struphy.geometry.base import Domain +from struphy.kinetic_background.base import KineticBackground class ParamsIn: @@ -27,7 +30,8 @@ def __init__(self, domain = None, equil = None, grid: TensorProductGrid = None, - derham_opts = None): + derham_opts = None, + model: StruphyModel = None,): self.env = env self.units = units self.time_opts = time_opts @@ -35,6 +39,7 @@ def __init__(self, self.equil = equil self.grid = grid self.derham_opts = derham_opts + self.model = model def get_params_of_run(path: str) -> ParamsIn: @@ -60,6 +65,7 @@ def get_params_of_run(path: str) -> ParamsIn: equil = params_in.equil grid = params_in.grid derham_opts = params_in.derham_opts + model = params_in.model elif os.path.exists(bin_path): with open(os.path.join(path, "env.bin"), "rb") as f: @@ -83,6 +89,8 @@ def get_params_of_run(path: str) -> ParamsIn: grid = pickle.load(f) with open(os.path.join(path, "derham_opts.bin"), "rb") as f: derham_opts = pickle.load(f) + with open(os.path.join(path, "model.bin"), "rb") as f: + model = pickle.load(f) else: raise FileNotFoundError(f"Neither of the paths {params_path} or {bin_path} exists.") @@ -96,6 +104,7 @@ def get_params_of_run(path: str) -> ParamsIn: equil=equil, grid=grid, derham_opts=derham_opts, + model=model, ) @@ -497,9 +506,6 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma step : int, optional Whether to do post-processing at every time step (step=1, default), every second time step (step=2), etc. """ - - print(f"{domain = }") - # get # of MPI processes from meta.txt file with open(os.path.join(path_in, "meta.yml"), "r") as f: meta = yaml.load(f, Loader=yaml.FullLoader) @@ -614,21 +620,15 @@ def post_process_f(path_in, path_out, species, step=1, compute_bckgr=False): Whether to do post-processing at every time step (step=1, default), every second time step (step=2), etc. compute_bckgr : bool - Whehter to compute the kinetic background values and add them to the binning data. + Whether to compute the kinetic background values and add them to the binning data. This is used if non-standard weights are binned. """ + # get # of MPI processes from meta.txt file + with open(os.path.join(path_in, "meta.yml"), "r") as f: + meta = yaml.load(f, Loader=yaml.FullLoader) + nproc = meta["MPI processes"] - # get model name and # of MPI processes from meta.txt file - with open(os.path.join(path_in, "meta.txt"), "r") as f: - lines = f.readlines() - - nproc = lines[4].split()[-1] - - # load parameters - with open(os.path.join(path_in, "parameters.yml"), "r") as f: - params = yaml.load(f, Loader=yaml.FullLoader) - - # open hdf5 files + # open hdf5 files and get names and number of saved markers of kinetic species files = [ h5py.File( os.path.join( @@ -641,6 +641,9 @@ def post_process_f(path_in, path_out, species, step=1, compute_bckgr=False): for i in range(int(nproc)) ] + # import parameters + params = get_params_of_run(path_in) + # directory for .npy files path_distr = os.path.join(path_out, "distribution_function") @@ -692,23 +695,27 @@ def post_process_f(path_in, path_out, species, step=1, compute_bckgr=False): np.save(os.path.join(path_slice, "delta_f_binned.npy"), data_df) if compute_bckgr: - bckgr_params = params["kinetic"][species]["background"] - - f_bckgr = None - for fi, maxw_params in bckgr_params.items(): - if fi[-2] == "_": - fi_type = fi[:-2] - else: - fi_type = fi - - if f_bckgr is None: - f_bckgr = getattr(maxwellians, fi_type)( - maxw_params=maxw_params, - ) - else: - f_bckgr = f_bckgr + getattr(maxwellians, fi_type)( - maxw_params=maxw_params, - ) + # bckgr_params = params["kinetic"][species]["background"] + + # f_bckgr = None + # for fi, maxw_params in bckgr_params.items(): + # if fi[-2] == "_": + # fi_type = fi[:-2] + # else: + # fi_type = fi + + # if f_bckgr is None: + # f_bckgr = getattr(maxwellians, fi_type)( + # maxw_params=maxw_params, + # ) + # else: + # f_bckgr = f_bckgr + getattr(maxwellians, fi_type)( + # maxw_params=maxw_params, + # ) + + spec: KineticSpecies = getattr(params.model, species) + var: PICVariable = spec.var + f_bckgr: KineticBackground = var.backgrounds # load all grids of the variables of f grid_tot = [] From f270e5a997852ea3343bc2a510b7a9a1f8956df0 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 12 Sep 2025 08:48:28 +0200 Subject: [PATCH 114/292] added Simdata.f dict for distribution function --- src/struphy/main.py | 46 +++++++++++++++++++++++++++++++++++---------- 1 file changed, 36 insertions(+), 10 deletions(-) diff --git a/src/struphy/main.py b/src/struphy/main.py index 55d8e79f7..c12b21024 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -610,6 +610,7 @@ class SimData: def __init__(self, path: str): self.path = path self._orbits = {} + self._f = {} self._spline_values = {} self.sph_species = {} self.grids_log: list[np.ndarray] = None @@ -621,6 +622,11 @@ def orbits(self) -> dict[str, np.ndarray]: """Keys: species name. Values: 3d arrays indexed by (n, p, a), where 'n' is the time index, 'p' the particle index and 'a' the attribute index.""" return self._orbits + @property + def f(self) -> dict[str, dict[str, dict[str, np.ndarray]]]: + """Keys: species name. Values: dicts of slice names ('e1_v1' etc.) holding dicts of corresponding np.arrays for plotting.""" + return self._f + @property def spline_values(self) -> dict[str, dict[str, np.ndarray]]: """Keys: species name. Values: dicts of variable names with values being 3d arrays on the grid.""" @@ -701,7 +707,7 @@ def load_data(path: str) -> SimData: # species folders species = next(os.walk(path_fields))[1] for spec in species: - simdata.spline_values[spec] = {} + simdata._spline_values[spec] = {} # simdata.arrays[spec] = {} path_spec = os.path.join(path_fields, spec) wlk = os.walk(path_spec) @@ -712,7 +718,7 @@ def load_data(path: str) -> SimData: var = file.split(".")[0] with open(os.path.join(path_spec, file), "rb") as f: # try: - simdata.spline_values[spec][var] = pickle.load(f) + simdata._spline_values[spec][var] = pickle.load(f) # simdata.arrays[spec][var] = pickle.load(f) if os.path.exists(path_kinetic): @@ -727,8 +733,8 @@ def load_data(path: str) -> SimData: for folder in sub_folders: path_dat = os.path.join(path_spec, folder) sub_wlk = os.walk(path_dat) - files = next(sub_wlk)[2] if "orbits" in folder: + files = next(sub_wlk)[2] Nt = len(files) // 2 n = 0 for file in files: @@ -737,10 +743,25 @@ def load_data(path: str) -> SimData: step = int(file.split(".")[0].split("_")[-1]) tmp = np.load(os.path.join(path_dat, file)) if n == 0: - simdata.orbits[spec] = np.zeros((Nt, *tmp.shape), dtype=float) - simdata.orbits[spec][step] = tmp + simdata._orbits[spec] = np.zeros((Nt, *tmp.shape), dtype=float) + simdata._orbits[spec][step] = tmp n += 1 + elif "distribution_function" in folder: + simdata._f[spec] = {} + slices = next(sub_wlk)[1] + # print(f"{slices = }") + for sli in slices: + simdata._f[spec][sli] = {} + # print(f"{sli = }") + files = next(sub_wlk)[2] + # print(f"{files = }") + for file in files: + name = file.split(".")[0] + tmp = np.load(os.path.join(path_dat, sli, file)) + # print(f"{name = }") + simdata._f[spec][sli][name] = tmp else: + print(f"{folder = }") raise NotImplementedError # # simdata.pic_species[spec][folder] = {} # tmp = {} @@ -755,16 +776,21 @@ def load_data(path: str) -> SimData: print("\nThe following data has been loaded:") print(f"{simdata.time_grid_size = }") print(f"{simdata.spline_grid_resolution = }") - print(f"simdata.spline_values:") + print(f"\nsimdata.spline_values:") for k, v in simdata.spline_values.items(): - print(f" {k}:") + print(f" {k}") for kk, vv in v.items(): print(f" {kk}") - print(f"simdata.orbits:") + print(f"\nsimdata.orbits:") for k, v in simdata.orbits.items(): print(f" {k}") - # for kk, vv in v.items(): - # print(f" {kk}") + print(f"\nsimdata.f:") + for k, v in simdata.f.items(): + print(f" {k}") + for kk, vv in v.items(): + print(f" {kk}") + for kkk, vvv in vv.items(): + print(f" {kkk}") print(f"simdata.sph_species:") return simdata From bc1cdc352e9d65a48f58103e51960b7030768970 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 12 Sep 2025 09:36:39 +0200 Subject: [PATCH 115/292] new method Maxwellian.add_perturbation enables to extract the background easily --- src/struphy/kinetic_background/base.py | 21 +++++++++++++++++++-- src/struphy/pic/base.py | 7 ++++--- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index b9a1d5459..7cb4344f6 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -470,7 +470,7 @@ def __call__(self, *args): return res - def _evaluate_moment(self, eta1, eta2, eta3, *, name="n"): + def _evaluate_moment(self, eta1, eta2, eta3, *, name: str = "n", add_perturbation: bool = None): """Scalar moment evaluation as background + perturbation. Parameters @@ -480,6 +480,9 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name="n"): name : str Which moment to evaluate (see varaible "dct" below). + + add_perturbation : bool | None + Whether to add the perturbation defined in maxw_params. If None, is taken from self.add_perturbation. Returns ------- @@ -544,8 +547,11 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name="n"): out += background(*etas) # add perturbation + if add_perturbation is None: + add_perturbation = self.add_perturbation + perturbation = params[1] - if perturbation is not None: + if perturbation is not None and add_perturbation: assert isinstance(perturbation, Perturbation) if eta1.ndim == 1: out += perturbation(eta1, eta2, eta3) @@ -553,6 +559,17 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name="n"): out += perturbation(*etas) return out + + @property + def add_perturbation(self) -> bool: + if not hasattr(self, "_add_perturbation"): + self._add_perturbation = True + return self._add_perturbation + + @add_perturbation.setter + def add_perturbation(self, new): + assert isinstance(new, bool) + self._add_perturbation = new class CanonicalMaxwellian(metaclass=ABCMeta): diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index c826e20ea..f6f20cf5b 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -19,7 +19,7 @@ from struphy.geometry.utilities import TransformedPformComponent from struphy.initial.base import Perturbation from struphy.io.output_handling import DataContainer -from struphy.kinetic_background.base import KineticBackground +from struphy.kinetic_background.base import KineticBackground, Maxwellian from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments from struphy.pic import sampling_kernels, sobol_seq from struphy.pic.pushing.pusher_utilities_kernels import reflect @@ -553,7 +553,7 @@ def u_init(self): return self._u_init @property - def f0(self): + def f0(self) -> Maxwellian: assert hasattr(self, "_f0"), AttributeError( "No background distribution available, please run self._set_background_function()", ) @@ -948,7 +948,8 @@ def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): return dom_arr, tuple(nprocs) def _set_background_function(self): - self._f0 = self.background + self._f0 = copy.deepcopy(self.background) + self.f0.add_perturbation = False # self._f0 = None # if isinstance(self.bckgr_params, FluidEquilibrium): # self._f0 = self.bckgr_params From 72c59ce86a77f9dee7c48f1f16f373ddd2c8a72b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 12 Sep 2025 11:46:33 +0200 Subject: [PATCH 116/292] fix bug in initial Poisson solve of VlasovAmpereOneSpecies; add flaf with_B0 to model __init__ --- src/struphy/models/kinetic.py | 29 ++++++++++++++----- src/struphy/pic/utilities.py | 4 +-- src/struphy/propagators/propagators_fields.py | 3 +- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 1f7bc98ee..0c9adc8dd 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -8,6 +8,7 @@ from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable from struphy.pic.accumulation.particles_to_grid import AccumulatorVector +from struphy.feec.projectors import L2Projector rank = MPI.COMM_WORLD.Get_rank() @@ -97,27 +98,31 @@ def __init__(self): ## propagators class Propagators: - def __init__(self): + def __init__(self, with_B0: bool = True): self.push_eta = propagators_markers.PushEta() - self.push_vxb = propagators_markers.PushVxB() + if with_B0: + self.push_vxb = propagators_markers.PushVxB() self.coupling_va = propagators_coupling.VlasovAmpere() ## abstract methods - def __init__(self): + def __init__(self, with_B0: bool = True): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + self.with_B0 = with_B0 # 1. instantiate all species self.em_fields = self.EMFields() self.kinetic_ions = self.KineticIons() # 2. instantiate all propagators - self.propagators = self.Propagators() + self.propagators = self.Propagators(with_B0=with_B0) # 3. assign variables to propagators self.propagators.push_eta.variables.var = self.kinetic_ions.var - self.propagators.push_vxb.variables.ions = self.kinetic_ions.var + if with_B0: + self.propagators.push_vxb.variables.ions = self.kinetic_ions.var self.propagators.coupling_va.variables.e = self.em_fields.e_field self.propagators.coupling_va.variables.ions = self.kinetic_ions.var @@ -198,12 +203,14 @@ def allocate_propagators(self): # another sanity check: compute FE coeffs of density # charge_accum.show_accumulated_spline_field(self.mass_ops) - + alpha = self.kinetic_ions.equation_params.alpha epsilon = self.kinetic_ions.equation_params.epsilon - self.initial_poisson.options.rho = alpha**2 / epsilon * charge_accum.vectors[0] - # self.initial_poisson.variables.phi.allocate(self.derham, domain=self.domain) + l2_proj = L2Projector(space_id='H1', mass_ops=self.mass_ops) + rho_coeffs = l2_proj.solve(charge_accum.vectors[0]) + + self.initial_poisson.options.rho = alpha**2 / epsilon * rho_coeffs self.initial_poisson.allocate() # Solve with dt=1. and compute electric field @@ -225,6 +232,12 @@ def generate_default_parameter_file(self, path = None, prompt = True): if "coupling_va.Options" in line: new_file += [line] new_file += ["model.initial_poisson.options = model.initial_poisson.Options()\n"] + elif "push_vxb.Options" in line: + new_file += ["if model.with_B0:\n"] + new_file += [" " + line] + elif "set_save_data" in line: + new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] + new_file += ["model.kinetic_ions.set_save_data(binning_plots=(binplot,))\n"] else: new_file += [line] diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index db9599a5a..773281422 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -53,9 +53,9 @@ class LoadingParameters: Key in .hdf5 file's restart/ folder where marker array is stored. """ def __init__(self, - Np: int = 100, + Np: int = None, ppc: int = None, - ppb: int = None, + ppb: int = 10, loading: OptsLoading = "pseudo_random", seed: int = None, moments: tuple = None, diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index e514429e0..4e7b22553 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -2772,7 +2772,6 @@ def allocate(self): else: stab_mat = getattr(self.mass_ops, self.options.stab_mat) - print(f"{self.options.diffusion_mat = }") if isinstance(self.options.diffusion_mat, str): diffusion_mat = getattr(self.mass_ops, self.options.diffusion_mat) else: @@ -2957,7 +2956,7 @@ class Options: OptsStabMat = Literal["M0", "M0ad", "Id"] # propagator options stab_eps: float = 0.0 - stab_mat: OptsStabMat = "M0" + stab_mat: OptsStabMat = "Id" rho: StencilVector | tuple | list | Callable = None x0: StencilVector = None solver: OptsSymmSolver = "pcg" From ab855250037561038ca30f25d0a186c9d962aa30 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 12 Sep 2025 12:56:18 +0200 Subject: [PATCH 117/292] add line_profiler to some functions --- src/struphy/main.py | 2 ++ src/struphy/pic/base.py | 7 +++++++ src/struphy/propagators/propagators_coupling.py | 3 +++ src/struphy/propagators/propagators_fields.py | 9 +++++++++ src/struphy/propagators/propagators_markers.py | 9 +++++++++ 5 files changed, 30 insertions(+) diff --git a/src/struphy/main.py b/src/struphy/main.py index c12b21024..b7df20cbf 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -11,6 +11,7 @@ import pickle import h5py import copy +from line_profiler import profile from struphy.fields_background.base import FluidEquilibriumWithB from struphy.io.output_handling import DataContainer @@ -39,6 +40,7 @@ from struphy.post_processing.orbits import orbits_tools +@profile def run( model: StruphyModel, *, diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index f6f20cf5b..6565f2c09 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,6 +2,7 @@ import os import warnings from abc import ABCMeta, abstractmethod +from line_profiler import profile import h5py import numpy as np @@ -1546,6 +1547,7 @@ def draw_markers( self.mpi_sort_markers() self.do_sort() + @profile def mpi_sort_markers( self, apply_bc: bool = True, @@ -1706,6 +1708,7 @@ def initialize_weights( else: self.weights = self.weights0 + @profile def update_weights(self): """ Applies the control variate method, i.e. updates the time-dependent marker weights @@ -1869,6 +1872,7 @@ def _find_outside_particles(self, axis): return outside_inds + @profile def apply_kinetic_bc(self, newton=False): """ Apply boundary conditions to markers that are outside of the logical unit cube. @@ -2381,6 +2385,7 @@ def sort_boxed_particles_numpy(self): self._argsort_array[:] = self._markers[:, sorting_axis].argsort() self._markers[:, :] = self._markers[self._argsort_array] + @profile def put_particles_in_boxes(self): """Assign the right box to the particles and the list of the particles to each box. If sorting_boxes was instantiated with an MPI comm, then the particles in the @@ -2408,6 +2413,7 @@ def put_particles_in_boxes(self): ) self.update_ghost_particles() + @profile def do_sort(self): """Assign the particles to boxes and then sort them.""" nx = self._sorting_boxes.nx @@ -2789,6 +2795,7 @@ def self_communication_boxes(self): self.markers[holes_inds[np.arange(self._send_info_box[self.mpi_rank])]] = self._send_list_box[self.mpi_rank] + @profile def communicate_boxes(self, verbose=False): if verbose: n_valid = np.count_nonzero(self.valid_mks) diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index af0ea32a3..6119d556a 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from mpi4py import MPI from typing import Literal +from line_profiler import profile from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -128,6 +129,7 @@ def options(self, new): print(f' {k}: {v}') self._options = new + @profile def allocate(self): # scaling factors alpha = self.variables.ions.species.equation_params.alpha @@ -200,6 +202,7 @@ def allocate(self): alpha_in_kernel=1.0, ) + @profile def __call__(self, dt): # accumulate self._accum() diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 4e7b22553..a4b0dfe40 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -5,6 +5,7 @@ from dataclasses import dataclass from typing import Literal, get_args import copy +from line_profiler import profile import numpy as np import scipy as sc @@ -130,6 +131,7 @@ def options(self, new): print(f' {k}: {v}') self._options = new + @profile def allocate(self): # obtain needed matrices M1 = self.mass_ops.M1 @@ -202,6 +204,7 @@ def f2(t, y1, y2, out: BlockVector = out2): self._e_tmp2 = self.variables.e.spline.vector.space.zeros() self._b_tmp1 = self.variables.b.spline.vector.space.zeros() + @profile def __call__(self, dt): # current FE coeffs en = self.variables.e.spline.vector @@ -542,6 +545,7 @@ def options(self, new): print(f' {k}: {v}') self._options = new + @profile def allocate(self): u_space = self.options.u_space @@ -620,6 +624,7 @@ def f2(t, y1, y2, out: BlockVector = out2): self._u_tmp2 = self.variables.u.spline.vector.space.zeros() self._b_tmp1 = self.variables.b.spline.vector.space.zeros() + @profile def __call__(self, dt): # current FE coeffs un = self.variables.u.spline.vector @@ -1007,6 +1012,7 @@ def options(self, new): print(f' {k}: {v}') self._options = new + @profile def allocate(self): u_space = self.options.u_space @@ -1072,6 +1078,7 @@ def allocate(self): self._byn1 = self._B.codomain.zeros() self._byn2 = self._B.codomain.zeros() + @profile def __call__(self, dt): # current FE coeffs nn = self.variables.n.spline.vector @@ -2723,6 +2730,7 @@ def options(self, new): print(f' {k}: {v}') self._options = new + @profile def allocate(self): # always stabilize if np.abs(self.options.sigma_1) < 1e-14: @@ -2867,6 +2875,7 @@ def x0(self, value): else: self._x0[:] = value[:] + @profile def __call__(self, dt): # set parameters if self._divide_by_dt: diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 74f70c5ef..17d786080 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -6,6 +6,7 @@ import copy from dataclasses import dataclass from mpi4py import MPI +from line_profiler import profile from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -84,6 +85,7 @@ def options(self, new): print(f' {k}: {v}') self._options = new + @profile def allocate(self): # get kernel kernel = pusher_kernels.push_eta_stage @@ -112,6 +114,7 @@ def allocate(self): mpi_sort="each", ) + @profile def __call__(self, dt): self._pusher(dt) @@ -180,6 +183,7 @@ def options(self, new): print(f' {k}: {v}') self._options = new + @profile def allocate(self): # scaling factor self._epsilon = self.variables.ions.species.equation_params.epsilon @@ -228,6 +232,7 @@ def allocate(self): # transposed extraction operator PolarVector --> BlockVector (identity map in case of no polar splines) self._E2T: LinearOperator = self.derham.extraction_ops["2"].transpose() + @profile def __call__(self, dt): # sum up total magnetic field tmp = self._b2.copy(out=self._tmp) @@ -505,6 +510,7 @@ def options(self, new): print(f' {k}: {v}') self._options = new + @profile def allocate(self): # scaling factor self._epsilon = self.variables.ions.species.equation_params.epsilon @@ -812,6 +818,7 @@ def allocate(self): verbose=self.options.verbose, ) + @profile def __call__(self, dt): # electric field # TODO: add out to __neg__ of StencilVector @@ -941,6 +948,7 @@ def options(self, new): print(f' {k}: {v}') self._options = new + @profile def allocate(self): # scaling factor self._epsilon = self.variables.ions.species.equation_params.epsilon @@ -1258,6 +1266,7 @@ def allocate(self): verbose=self.options.verbose, ) + @profile def __call__(self, dt): # electric field # TODO: add out to __neg__ of StencilVector From ef3ba9d0c8f5696c67a3c17d14086ae1d915ee70 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 12 Sep 2025 13:50:01 +0200 Subject: [PATCH 118/292] fix base_units loading; allow comm=None in Particles --- src/struphy/main.py | 4 ++-- src/struphy/models/base.py | 16 ++++++++-------- src/struphy/models/variables.py | 6 +++++- src/struphy/pic/base.py | 1 + src/struphy/pic/pushing/pusher.py | 2 ++ .../post_processing/post_processing_tools.py | 14 +++++++------- 6 files changed, 25 insertions(+), 18 deletions(-) diff --git a/src/struphy/main.py b/src/struphy/main.py index b7df20cbf..0fe7d4db6 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -126,8 +126,8 @@ def run( else: with open(os.path.join(path_out, "env.bin"), 'wb') as f: pickle.dump(env, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "units.bin"), 'wb') as f: - pickle.dump(units, f, pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_out, "base_units.bin"), 'wb') as f: + pickle.dump(base_units, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "time_opts.bin"), 'wb') as f: pickle.dump(time_opts, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "domain.bin"), 'wb') as f: diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index e7b9c6da0..1336ce599 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -6,6 +6,7 @@ import os import yaml import struphy +from line_profiler import profile import numpy as np import yaml @@ -164,18 +165,13 @@ def species(self): ## allocate methods - def allocate_feec(self, - grid: TensorProductGrid, - derham_opts: DerhamOptions, - comm: MPI.Intracomm = None, - clone_config: CloneConfig = None, - ): + def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): # create discrete derham sequence - if clone_config is None: + if self.clone_config is None: derham_comm = MPI.COMM_WORLD else: - derham_comm = clone_config.sub_comm + derham_comm = self.clone_config.sub_comm self._derham = setup_derham( grid, @@ -578,6 +574,7 @@ def add_time_state(self, time_state): if isinstance(prop, Propagator): prop.add_time_state(time_state) + @profile def allocate_variables(self, verbose: bool = False): """ Allocate memory for model variables and set initial conditions. @@ -626,6 +623,7 @@ def allocate_variables(self, verbose: bool = False): # self._pointer[key] = val["obj"].vector + @profile def integrate(self, dt, split_algo="LieTrotter"): """ Advance the model by a time step ``dt`` by sequentially calling its Propagators. @@ -671,6 +669,7 @@ def integrate(self, dt, split_algo="LieTrotter"): f"Splitting scheme {split_algo} not available.", ) + @profile def update_markers_to_be_saved(self): """ Writes markers with IDs that are supposed to be saved into corresponding array. @@ -714,6 +713,7 @@ def update_markers_to_be_saved(self): var.kinetic_data["markers"][:] = -1.0 var.kinetic_data["markers"][:n_markers_on_proc] = obj.markers[markers_on_proc] + @profile def update_distr_functions(self): """ Writes distribution functions slices that are supposed to be saved into corresponding array. diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 58b6f6885..1931115c0 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -184,8 +184,12 @@ def allocate(self, kinetic_class = getattr(particles, self.space) + comm_world = MPI.COMM_WORLD + if comm_world.Get_size() == 1: + comm_world = None + self._particles: Particles = kinetic_class( - comm_world=MPI.COMM_WORLD, + comm_world=comm_world, clone_config=clone_config, domain_decomp=domain_decomp, mpi_dims_mask=self.species.dims_mask, diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 6565f2c09..76dec4d01 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1743,6 +1743,7 @@ def reset_marker_ids(self): )[self.mpi_rank] self.marker_ids = first_marker_id + np.arange(self.n_mks_loc, dtype=int) + @profile def binning(self, components, bin_edges, divide_by_jac=True): r"""Computes full-f and delta-f distribution functions via marker binning in logical space. Numpy's histogramdd is used, following the algorithm outlined in :ref:`binning`. diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 88c796ec4..b1a5b47ff 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -2,6 +2,7 @@ import numpy as np from mpi4py.MPI import IN_PLACE, SUM +from line_profiler import profile from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments from struphy.pic.base import Particles @@ -162,6 +163,7 @@ def __init__( else: self._box_comm = False + @profile def __call__(self, dt: float): """ Applies the chosen pusher kernel by a time step dt, diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 98d2f1e3a..c79b21ca1 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -14,7 +14,7 @@ from struphy.models.species import KineticSpecies from struphy.models.variables import PICVariable from struphy.feec.psydac_derham import SplineFunction -from struphy.io.options import EnvironmentOptions, Units, Time +from struphy.io.options import EnvironmentOptions, BaseUnits, Time from struphy.topology.grids import TensorProductGrid from struphy.geometry import domains from struphy.geometry.base import Domain @@ -25,7 +25,7 @@ class ParamsIn: """Holds the input parameters of a Struphy simulation as attributes.""" def __init__(self, env: EnvironmentOptions = None, - units: Units = None, + base_units: BaseUnits = None, time_opts: Time = None, domain = None, equil = None, @@ -33,7 +33,7 @@ def __init__(self, derham_opts = None, model: StruphyModel = None,): self.env = env - self.units = units + self.units = base_units self.time_opts = time_opts self.domain = domain self.equil = equil @@ -59,7 +59,7 @@ def get_params_of_run(path: str) -> ParamsIn: if os.path.exists(params_path): params_in = import_parameters_py(params_path) env = params_in.env - units = params_in.units + base_units = params_in.base_units time_opts = params_in.time_opts domain = params_in.domain equil = params_in.equil @@ -70,8 +70,8 @@ def get_params_of_run(path: str) -> ParamsIn: elif os.path.exists(bin_path): with open(os.path.join(path, "env.bin"), "rb") as f: env = pickle.load(f) - with open(os.path.join(path, "units.bin"), "rb") as f: - units = pickle.load(f) + with open(os.path.join(path, "base_units.bin"), "rb") as f: + base_units = pickle.load(f) with open(os.path.join(path, "time_opts.bin"), "rb") as f: time_opts = pickle.load(f) with open(os.path.join(path, "domain.bin"), "rb") as f: @@ -98,7 +98,7 @@ def get_params_of_run(path: str) -> ParamsIn: print("done.") return ParamsIn(env=env, - units=units, + base_units=base_units, time_opts=time_opts, domain=domain, equil=equil, From f2a403f62b3cf4bfeebc41d3676a2442f3f90661 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 12 Sep 2025 14:04:07 +0200 Subject: [PATCH 119/292] added tutorial 5 --- .../tutorial_05_vlasov_maxwell.ipynb | 378 ++++++++++++++++++ 1 file changed, 378 insertions(+) create mode 100644 doc/tutorials/tutorial_05_vlasov_maxwell.ipynb diff --git a/doc/tutorials/tutorial_05_vlasov_maxwell.ipynb b/doc/tutorials/tutorial_05_vlasov_maxwell.ipynb new file mode 100644 index 000000000..149d60427 --- /dev/null +++ b/doc/tutorials/tutorial_05_vlasov_maxwell.ipynb @@ -0,0 +1,378 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 9 - Vlasov-Ampère equations\n", + "\n", + "The equations we will solve are described in the model [VlasovAmpereOneSpecies](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_kinetic.html#struphy.models.kinetic.VlasovAmpereOneSpecies).\n", + "To create the default parameter file from the console:\n", + "\n", + "```\n", + "struphy params VlasovAmpereOneSpecies\n", + "```\n", + "\n", + "Adapt the parameters and run the model with\n", + "\n", + "```\n", + "python3 params_VlasovAmpereOneSpecies.py\n", + "```\n", + "\n", + "In this notebook we shall re-create the parameter file and perform some tests.\n", + "\n", + "## Weak Landau damping\n", + "\n", + "1. Imports:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.io.options import EnvironmentOptions, BaseUnits, Time\n", + "from struphy.geometry import domains\n", + "from struphy.fields_background import equils\n", + "from struphy.topology import grids\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.io.options import FieldsBackground\n", + "from struphy.initial import perturbations\n", + "from struphy.kinetic_background import maxwellians\n", + "from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters, BinningPlot\n", + "from struphy import main\n", + "\n", + "from struphy.models.kinetic import VlasovAmpereOneSpecies" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "2. Generic options:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# environment options\n", + "env = EnvironmentOptions()\n", + "\n", + "# units\n", + "base_units = BaseUnits()\n", + "\n", + "# time stepping\n", + "time_opts = Time(dt = 0.05, Tend = 0.5)#, Tend = 3.5\n", + "\n", + "# geometry\n", + "r1 = 12.56\n", + "domain = domains.Cuboid(r1=r1)\n", + "\n", + "# fluid equilibrium (can be used as part of initial conditions)\n", + "equil = None\n", + "\n", + "# grid\n", + "grid = grids.TensorProductGrid(Nel=(32, 1, 1))\n", + "\n", + "# derham options\n", + "derham_opts = DerhamOptions()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "3. Model instance and physics parameters:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model = VlasovAmpereOneSpecies()\n", + "\n", + "model.kinetic_ions.set_phys_params(alpha=1.0, epsilon=1.0)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "4. Kinetic species parameters:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "loading_params = LoadingParameters(ppc=10000)\n", + "weights_params = WeightsParameters(control_variate=True)\n", + "boundary_params = BoundaryParameters()\n", + "model.kinetic_ions.set_markers(loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params)\n", + "model.kinetic_ions.set_sorting_boxes()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# particle binning\n", + "binplot_1 = BinningPlot(slice=\"e1\", n_bins=128, ranges=(0.0, 1.0))\n", + "binplot_2 = BinningPlot(slice=\"v1\", n_bins=128, ranges=(-5.0, 5.0))\n", + "binplot_3 = BinningPlot(slice=\"e1_v1\", n_bins=(128, 128), ranges=((0.0, 1.0), (-5.0, 5.0)))\n", + "\n", + "binning_plots = (binplot_1, binplot_2, binplot_3)\n", + "\n", + "model.kinetic_ions.set_save_data(binning_plots=binning_plots)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "5. Propagator options:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model.propagators.push_eta.options = model.propagators.push_eta.Options()\n", + "model.propagators.coupling_va.options = model.propagators.coupling_va.Options()\n", + "model.initial_poisson.options = model.initial_poisson.Options(stab_eps=1e-12)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "6. Initial conditions:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "perturbation = perturbations.ModesCos(ls=[1], amps=[0.001])\n", + "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", + "model.kinetic_ions.var.add_background(background)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Let us run the model. However, depending on the confguration, running in a notebook might be very slow. In order to get fast execution, run from the console. First, create the default parameter file and rename it \n", + "\n", + "```\n", + "struphy params VlasovAmpereOneSpecies\n", + "mv params_VlasovAmpereOneSpecies.py landau.py\n", + "```\n", + "\n", + "Adapt it with the parameters from this notebook. Start the run with\n", + "\n", + "```\n", + "python landau.py\n", + "```\n", + "\n", + "or \n", + "\n", + "```\n", + "mpirun -n 2 python landau.py\n", + "```\n", + "\n", + "for a run on two threads. Line profiling can be enabled with \n", + "\n", + "```\n", + "LINE_PROFILE=1 mpirun -n 2 landau.py\n", + "```\n", + "\n", + "Let us look a the slower run in the notebook:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "verbose = True\n", + "\n", + "main.run(model, \n", + " params_path=None, \n", + " env=env, \n", + " base_units=base_units, \n", + " time_opts=time_opts, \n", + " domain=domain, \n", + " equil=equil, \n", + " grid=grid, \n", + " derham_opts=derham_opts, \n", + " verbose=verbose, \n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "main.pproc(path, celldivide=8)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "simdata = main.load_data(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# plot in v1\n", + "from matplotlib import pyplot as plt\n", + "\n", + "v1_bins = simdata.f[\"kinetic_ions\"][\"v1\"][\"grid_v1\"]\n", + "f_v1_init = simdata.f[\"kinetic_ions\"][\"v1\"][\"f_binned\"][0]\n", + "\n", + "plt.plot(v1_bins, f_v1_init)\n", + "plt.xlabel('vx')\n", + "plt.title('Initial Maxwellian');" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# plot in e1\n", + "\n", + "e1_bins = simdata.f[\"kinetic_ions\"][\"e1\"][\"grid_e1\"]\n", + "df_e1_init = simdata.f[\"kinetic_ions\"][\"e1\"][\"delta_f_binned\"][0]\n", + "\n", + "plt.plot(e1_bins, df_e1_init)\n", + "plt.xlabel('$\\eta_1$')\n", + "plt.title('Initial spatial perturbation');" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# plot in e1-v1\n", + "\n", + "e1_bins = simdata.f[\"kinetic_ions\"][\"e1_v1\"][\"grid_e1\"]\n", + "v1_bins = simdata.f[\"kinetic_ions\"][\"e1_v1\"][\"grid_v1\"]\n", + "f_init = simdata.f[\"kinetic_ions\"][\"e1_v1\"][\"f_binned\"][0]\n", + "df_init = simdata.f[\"kinetic_ions\"][\"e1_v1\"][\"delta_f_binned\"][0]\n", + "f_end = simdata.f[\"kinetic_ions\"][\"e1_v1\"][\"f_binned\"][-1]\n", + "df_end = simdata.f[\"kinetic_ions\"][\"e1_v1\"][\"delta_f_binned\"][-1]\n", + "\n", + "plt.figure(figsize=(14, 10))\n", + "\n", + "plt.subplot(2, 2, 1)\n", + "plt.pcolor(e1_bins, v1_bins, f_init.T)\n", + "plt.xlabel('$\\eta_1$')\n", + "plt.ylabel('$v_x$')\n", + "plt.title('Initial Maxwellian')\n", + "plt.colorbar()\n", + "\n", + "plt.subplot(2, 2, 2)\n", + "plt.pcolor(e1_bins, v1_bins, df_init.T)\n", + "plt.xlabel('$\\eta_1$')\n", + "plt.ylabel('$v_x$')\n", + "plt.title('Initial perturbation')\n", + "plt.colorbar()\n", + "\n", + "plt.subplot(2, 2, 3)\n", + "plt.pcolor(e1_bins, v1_bins, f_end.T)\n", + "plt.xlabel('$\\eta_1$')\n", + "plt.ylabel('$v_x$')\n", + "plt.title('Final Maxwellian')\n", + "plt.colorbar()\n", + "\n", + "plt.subplot(2, 2, 4)\n", + "plt.pcolor(e1_bins, v1_bins, df_end.T)\n", + "plt.xlabel('$\\eta_1$')\n", + "plt.ylabel('$v_x$')\n", + "plt.title('Final perturbation')\n", + "plt.colorbar();" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# electric field\n", + "\n", + "e1, e2, e3 = simdata.grids_log \n", + "e_vals = simdata.spline_values[\"em_fields\"][\"e_field_log\"][0][0]\n", + "\n", + "plt.plot(e1, e_vals[:, 0, 0], label='E')\n", + "plt.xlabel('$\\eta_1$')\n", + "plt.title('Initial electric field')\n", + "plt.legend();" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# electric energy\n", + "\n", + "# plt.plot(time_vec, np.log(energy_E))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "env", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} From 508445f7f54b7091a7cb22543c798ce7b3a38707 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 12 Sep 2025 14:34:18 +0000 Subject: [PATCH 120/292] Format all source files and and temporary linting check to CI --- .gitlab-ci.yml | 14 ++ src/struphy/console/params.py | 6 +- src/struphy/diagnostics/diagn_tools.py | 28 ++- src/struphy/feec/psydac_derham.py | 56 +++-- src/struphy/feec/utilities.py | 2 +- .../fields_background/projected_equils.py | 5 +- src/struphy/geometry/utilities.py | 31 +-- src/struphy/initial/base.py | 15 +- src/struphy/initial/perturbations.py | 236 ++++++++++++------ .../initial/tests/test_init_perturbations.py | 7 +- src/struphy/io/inp/params_Maxwell.py | 40 +-- src/struphy/io/inp/params_Maxwell_lw.py | 8 +- .../io/inp/verification/Maxwell_coaxial.py | 15 +- src/struphy/io/options.py | 97 +++---- src/struphy/io/setup.py | 13 +- src/struphy/kinetic_background/base.py | 19 +- src/struphy/kinetic_background/maxwellians.py | 22 +- src/struphy/linear_algebra/schur_solver.py | 15 +- src/struphy/linear_algebra/solver.py | 3 +- .../tests/test_saddle_point_propagator.py | 10 +- src/struphy/main.py | 199 ++++++++------- src/struphy/models/base.py | 213 +++++++++------- src/struphy/models/fluid.py | 51 ++-- src/struphy/models/kinetic.py | 43 ++-- src/struphy/models/species.py | 171 +++++++------ src/struphy/models/tests/test_LinearMHD.py | 135 +++++----- src/struphy/models/tests/test_Maxwell.py | 135 +++++----- src/struphy/models/tests/test_models.py | 71 +++--- src/struphy/models/tests/util.py | 2 +- src/struphy/models/tests/verification.py | 1 - src/struphy/models/toy.py | 72 +++--- src/struphy/models/variables.py | 146 ++++++----- src/struphy/ode/utils.py | 18 +- src/struphy/physics/physics.py | 3 +- src/struphy/pic/base.py | 45 ++-- src/struphy/pic/particles.py | 11 +- src/struphy/pic/pushing/pusher.py | 2 +- src/struphy/pic/utilities.py | 128 +++++----- .../post_processing/post_processing_tools.py | 126 ++++++---- src/struphy/post_processing/pproc_struphy.py | 4 +- src/struphy/propagators/base.py | 49 ++-- .../propagators/propagators_coupling.py | 46 ++-- src/struphy/propagators/propagators_fields.py | 158 ++++++------ .../propagators/propagators_markers.py | 132 +++++----- src/struphy/topology/grids.py | 4 +- src/struphy/tutorials/tests/test_tutorials.py | 2 +- 46 files changed, 1425 insertions(+), 1184 deletions(-) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 3c051d103..db65d9bc0 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -942,6 +942,20 @@ lint_repo: # Lint all files in current branch, FAIL the CI pipeline if any files are incorrectly formatted - struphy lint all --output-format plain --verbose +lint_repo_temporary: + stage: lint + needs: [] + extends: + - .rules_startup + - .image_ubuntu_latest + script: + - !reference [.scripts, inspect_directory] + - !reference [.scripts, create_venv] + - pip install -e .[dev] # We have to install struphy in editable mode for struphy lint to work + - !reference [.scripts, inspect_struphy] + # Lint all files in current branch, FAIL the CI pipeline if any files are incorrectly formatted + - struphy lint all --output-format plain --verbose + # Lint struphy with `struphy lint` lint_branch_report: stage: lint diff --git a/src/struphy/console/params.py b/src/struphy/console/params.py index d09883405..1648c7383 100644 --- a/src/struphy/console/params.py +++ b/src/struphy/console/params.py @@ -1,5 +1,5 @@ -from struphy.models.base import StruphyModel from struphy.models import fluid, hybrid, kinetic, toy +from struphy.models.base import StruphyModel def struphy_params(model_name: str, params_path: str, yes: bool = False): @@ -12,7 +12,7 @@ def struphy_params(model_name: str, params_path: str, yes: bool = False): params_path : str An alternative file name to the default params_.yml. - + yes : bool If true, say yes on prompt to overwrite .yml FILE """ @@ -23,6 +23,6 @@ def struphy_params(model_name: str, params_path: str, yes: bool = False): model: StruphyModel = model_class() except AttributeError: pass - + prompt = not yes model.generate_default_parameter_file(path=params_path, prompt=prompt) diff --git a/src/struphy/diagnostics/diagn_tools.py b/src/struphy/diagnostics/diagn_tools.py index 841bff70c..dc50161fb 100644 --- a/src/struphy/diagnostics/diagn_tools.py +++ b/src/struphy/diagnostics/diagn_tools.py @@ -69,14 +69,14 @@ def power_spectrum_2d( fit_branches: int How many branches to fit in the dispersion relation. Default=0 means no fits are made. - + noise_level: float Sets the threshold above which local maxima in the power spectrum are taken into account. - Computed as threshold = max(spectrum) * noise_level. - + Computed as threshold = max(spectrum) * noise_level. + extr_oder: int Order given to argrelextrema. - + fit_degree: tuple[int] Degree of fitting polynomial for each branch (fit_branches) of power spectrum. @@ -99,7 +99,7 @@ def power_spectrum_2d( dispersion : np.array 2d array of shape (omega.size, kvec.size) holding the fft. - + coeffs : list[list] List of fitting coefficients (lenght is fit_branches). """ @@ -155,8 +155,8 @@ def power_spectrum_2d( if fit_branches > 0: assert len(fit_degree) == fit_branches # determine maxima for each k - k_start = kvec.size // 8 # take only first half of k-vector - k_end = kvec.size // 2 # take only first half of k-vector + k_start = kvec.size // 8 # take only first half of k-vector + k_end = kvec.size // 2 # take only first half of k-vector k_fit = [] omega_fit = {} for n in range(fit_branches): @@ -172,16 +172,18 @@ def power_spectrum_2d( intersec.sort() # print(f"{intersec = }") # print(f"{[omega[intersec[n]] for n in range(fit_branches)]}") - assert len(intersec) == fit_branches, f"Number of found branches {len(intersec)} is not {fit_branches = }! \ + assert len(intersec) == fit_branches, ( + f"Number of found branches {len(intersec)} is not {fit_branches = }! \ Try to lower 'noise_level' or increase 'extr_order'." + ) k_fit += [k] for n in range(fit_branches): omega_fit[n] += [omega[intersec[n]]] - + # fit coeffs = [] for m, om in omega_fit.items(): - coeffs += [np.polyfit(k_fit, om, deg=fit_degree[n])] + coeffs += [np.polyfit(k_fit, om, deg=fit_degree[n])] print(f"\nFitted {coeffs = }") if do_plot: @@ -206,14 +208,16 @@ def power_spectrum_2d( ax.set_title(title) ax.set_xlabel("$k$ [a.u.]") ax.set_ylabel(r"$\omega$ [a.u.]") - + if fit_branches > 0: for n, cs in enumerate(coeffs): + def fun(k): - out = k*0.0 + out = k * 0.0 for i, c in enumerate(np.flip(cs)): out += c * k**i return out + ax.plot(kvec, fun(kvec), "r:", label=f"fit_{n + 1}") # analytic solution: diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 0158dcffc..4fb115a0c 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -22,19 +22,18 @@ from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.local_projectors_kernels import get_local_problem_size, select_quasi_points from struphy.feec.projectors import CommutingProjector, CommutingProjectorLocal -from struphy.fields_background.base import MHDequilibrium +from struphy.fields_background.base import FluidEquilibrium, MHDequilibrium from struphy.fields_background.equils import set_defaults from struphy.geometry.base import Domain from struphy.geometry.utilities import TransformedPformComponent from struphy.initial import perturbations, utilities +from struphy.initial.base import Perturbation +from struphy.initial.perturbations import Noise +from struphy.io.options import FieldsBackground, GivenInBasis, NoiseDirections from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator -from struphy.io.options import FieldsBackground, NoiseDirections, GivenInBasis -from struphy.initial.perturbations import Noise -from struphy.initial.base import Perturbation -from struphy.fields_background.base import FluidEquilibrium class Derham: @@ -899,10 +898,10 @@ def create_spline_function( perturbations : Perturbation | list For the initial condition. - + domain : Domain Mapping for pullback/transform of initial condition. - + equil : FLuidEquilibrium Fluid background used for inital condition. """ @@ -1413,10 +1412,10 @@ class SplineFunction: perturbations : Perturbation | list For the initial condition. - + domain : Domain Mapping for pullback/transform of initial condition. - + equil : FluidEquilibrium Fluid background used for inital condition. """ @@ -1478,8 +1477,8 @@ def __init__( ) else: self._nbasis = [tuple([space.nbasis for space in vec_space.spaces]) for vec_space in self.fem_space.spaces] - - if verbose and MPI.COMM_WORLD.Get_rank() == 0: + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: print(f"\nAllocated SplineFuntion '{self.name}' in space '{self.space_id}'.") if self.backgrounds is not None or self.perturbations is not None: @@ -1504,12 +1503,12 @@ def space_key(self): def derham(self): """3d Derham complex struphy.feec.psydac_derham.Derham.""" return self._derham - + @property def domain(self): """Mapping for pullback/transform of initial condition.""" return self._domain - + @property def equil(self): """Fluid equilibirum used for initial condition.""" @@ -1676,7 +1675,7 @@ def initialize_coeffs( # if self.perturbations is not None: # print(f"Attention: overwriting perturbation parameters for {self.name}") self._perturbations = perturbations - + # set domain if domain is not None: # if self.domain is not None: @@ -1685,7 +1684,7 @@ def initialize_coeffs( if isinstance(self.backgrounds, FieldsBackground): self._backgrounds = [self.backgrounds] - + if isinstance(self.perturbations, Perturbation): self._perturbations = [self.perturbations] @@ -1701,13 +1700,14 @@ def initialize_coeffs( assert isinstance(fb, FieldsBackground) if MPI.COMM_WORLD.Get_rank() == 0: print(f"Adding background {fb} ...") - + # special case of const if fb.type == "LogicalConst": vals = fb.values assert isinstance(vals, (list, tuple)) if self.space_id in {"H1", "L2"}: + def f_tmp(e1, e2, e3): return vals[0] + 0.0 * e1 @@ -1759,10 +1759,12 @@ def f_tmp(e1, e2, e3): # special case of white noise in logical space for different components if isinstance(ptb, Noise): # set white noise FE coefficients - self._add_noise(direction=ptb.direction, - amp=ptb.amp, - seed=ptb.seed, - n=ptb.comp,) + self._add_noise( + direction=ptb.direction, + amp=ptb.amp, + seed=ptb.seed, + n=ptb.comp, + ) # perturbation class elif isinstance(ptb, Perturbation): if self.space_id in {"H1", "L2"}: @@ -1773,9 +1775,9 @@ def f_tmp(e1, e2, e3): domain=domain, ) elif self.space_id in {"Hcurl", "Hdiv", "H1vec"}: - fun_vec = [None]*3 + fun_vec = [None] * 3 fun_vec[ptb.comp] = ptb - + # pullback callable for each component fun = [] for comp in range(3): @@ -1793,7 +1795,7 @@ def f_tmp(e1, e2, e3): self.vector += self.derham.P[self.space_key](fun) # TODO: re-add Eigfun and InitFromOutput in new framework - + # loading of MHD eigenfunction (legacy code, might not be up to date) # elif "EigFun" in _type: # print("Warning: Eigfun is not regularly tested ...") @@ -2219,7 +2221,13 @@ def _flag_pts_not_on_proc(self, *etas): E2[~E2_on_proc] = -1.0 E3[~E3_on_proc] = -1.0 - def _add_noise(self, direction: NoiseDirections = "e3", amp: float = 0.0001, seed: int = None, n: int = None,): + def _add_noise( + self, + direction: NoiseDirections = "e3", + amp: float = 0.0001, + seed: int = None, + n: int = None, + ): """Add noise to a vector component where init_comps==True, otherwise leave at zero. Parameters diff --git a/src/struphy/feec/utilities.py b/src/struphy/feec/utilities.py index d0c89108b..2541f9a63 100644 --- a/src/struphy/feec/utilities.py +++ b/src/struphy/feec/utilities.py @@ -2,9 +2,9 @@ from psydac.api.essential_bc import apply_essential_bc_stencil from psydac.fem.tensor import TensorFemSpace from psydac.fem.vector import VectorFemSpace +from psydac.linalg.basic import Vector from psydac.linalg.block import BlockLinearOperator, BlockVector from psydac.linalg.stencil import StencilMatrix, StencilVector -from psydac.linalg.basic import Vector import struphy.feec.utilities_kernels as kernels from struphy.feec import banded_to_stencil_kernels as bts diff --git a/src/struphy/fields_background/projected_equils.py b/src/struphy/fields_background/projected_equils.py index 427398ee8..26fa4f9c8 100644 --- a/src/struphy/fields_background/projected_equils.py +++ b/src/struphy/fields_background/projected_equils.py @@ -1,11 +1,12 @@ +from psydac.linalg.block import BlockVector +from psydac.linalg.stencil import StencilVector + from struphy.feec.psydac_derham import Derham from struphy.fields_background.base import ( FluidEquilibrium, FluidEquilibriumWithB, MHDequilibrium, ) -from psydac.linalg.stencil import StencilVector -from psydac.linalg.block import BlockVector class ProjectedFluidEquilibrium: diff --git a/src/struphy/geometry/utilities.py b/src/struphy/geometry/utilities.py index f321a1c8b..6449174df 100644 --- a/src/struphy/geometry/utilities.py +++ b/src/struphy/geometry/utilities.py @@ -1,21 +1,22 @@ # from __future__ import annotations "Domain-related utility functions." +from typing import Callable + import numpy as np + # from typing import TYPE_CHECKING from scipy.optimize import newton, root, root_scalar from scipy.sparse import csc_matrix from scipy.sparse.linalg import splu -from typing import Callable from struphy.bsplines import bsplines as bsp -from struphy.geometry.base import PoloidalSplineTorus -from struphy.geometry.utilities_kernels import weighted_arc_lengths_flux_surface -from struphy.linear_algebra.linalg_kron import kron_lusolve_2d # if TYPE_CHECKING: -from struphy.geometry.base import Domain +from struphy.geometry.base import Domain, PoloidalSplineTorus +from struphy.geometry.utilities_kernels import weighted_arc_lengths_flux_surface from struphy.io.options import GivenInBasis +from struphy.linear_algebra.linalg_kron import kron_lusolve_2d def field_line_tracing( @@ -353,7 +354,7 @@ class TransformedPformComponent: out_form : str The p-form representation of the output: '0', '1', '2' '3' or 'v'. - + comp : int Which component of the vector-valued function to return (=0 for scalars). @@ -361,14 +362,14 @@ class TransformedPformComponent: All things mapping. If None, the input fun is just evaluated and not transformed at __call__. """ - def __init__(self, - fun: Callable | list, - given_in_basis: GivenInBasis, - out_form: str, - comp: int = 0, - domain: Domain=None, - ): - + def __init__( + self, + fun: Callable | list, + given_in_basis: GivenInBasis, + out_form: str, + comp: int = 0, + domain: Domain = None, + ): if isinstance(fun, list): assert len(fun) == 1 or len(fun) == 3 else: @@ -377,8 +378,10 @@ def __init__(self, self._fun = [] for f in fun: if f is None: + def f_zero(x, y, z): return 0 * x + self._fun += [f_zero] else: assert callable(f) diff --git a/src/struphy/initial/base.py b/src/struphy/initial/base.py index 899053937..cafca9dbe 100644 --- a/src/struphy/initial/base.py +++ b/src/struphy/initial/base.py @@ -1,24 +1,24 @@ -from typing import Callable from abc import ABCMeta, abstractmethod +from typing import Callable from struphy.io.options import GivenInBasis, check_option class Perturbation(metaclass=ABCMeta): """Base class for perturbations that can be chosen as initial conditions.""" - + @abstractmethod def __call__(self, eta1, eta2, eta3, flat_eval=False): pass - + def prepare_eval_pts(self): # TODO: we could prepare the arguments via a method in this base class (flat_eval, sparse meshgrid, etc.). pass - + @property def given_in_basis(self) -> str: r"""In which basis the perturbation is represented, must be set in child class (use the setter below). - + Either * '0', '1', '2' or '3' for a p-form basis * 'v' for a vector-field basis @@ -32,7 +32,7 @@ def given_in_basis(self) -> str: def given_in_basis(self, new: str): check_option(new, GivenInBasis) self._given_in_basis = new - + @property def comp(self) -> int: """Which component of vector is perturbed (=0 for scalar-valued functions). @@ -40,9 +40,8 @@ def comp(self) -> int: if not hasattr(self, "_comp"): self._comp = 0 return self._comp - + @comp.setter def comp(self, new: int): assert new in (0, 1, 2) self._comp = new - \ No newline at end of file diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 23a211c64..45c7d7780 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -1,18 +1,20 @@ #!/usr/bin/env python3 "Analytical perturbations." +from dataclasses import dataclass + import numpy as np import scipy import scipy.special -from dataclasses import dataclass + from struphy.initial.base import Perturbation -from struphy.io.options import NoiseDirections, GivenInBasis, check_option +from struphy.io.options import GivenInBasis, NoiseDirections, check_option @dataclass class Noise(Perturbation): """White noise for FEEC coefficients. - + Parameters ---------- direction: str @@ -24,15 +26,18 @@ class Noise(Perturbation): seed: int Seed for the random number generator. """ + direction: NoiseDirections = "e3" amp: float = 0.0001 seed: int = None - comp: int = 0 + comp: int = 0 given_in_basis: GivenInBasis = "0" - - def __post_init__(self,): + + def __post_init__( + self, + ): check_option(self.direction, NoiseDirections) - + def __call__(self): pass @@ -86,10 +91,10 @@ class ModesSin(Perturbation): Lx, Ly, Lz : float Domain lengths. - + given_in_basis : str In which basis the perturbation is represented (see base class). - + comp : int Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ @@ -167,7 +172,7 @@ def __init__( ] else: raise ValueError(f"Profile function {pfun} is not defined..") - + self._ls = ls self._ms = ms self._ns = ns @@ -176,7 +181,7 @@ def __init__( self._Ly = Ly self._Lz = Lz self._theta = theta - + # use the setters self.given_in_basis = given_in_basis self.comp = comp @@ -228,15 +233,23 @@ class ModesCos(Perturbation): given_in_basis : str In which basis the perturbation is represented (see base class). - + comp : int Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, ls=None, ms=None, ns=None, amps=(1e-4,), Lx=1.0, Ly=1.0, Lz=1.0, - given_in_basis: GivenInBasis = "0", - comp: int = 0,): - + def __init__( + self, + ls=None, + ms=None, + ns=None, + amps=(1e-4,), + Lx=1.0, + Ly=1.0, + Lz=1.0, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): if ls is not None: n_modes = len(ls) elif ms is not None: @@ -274,7 +287,7 @@ def __init__(self, ls=None, ms=None, ns=None, amps=(1e-4,), Lx=1.0, Ly=1.0, Lz=1 self._Lx = Lx self._Ly = Ly self._Lz = Lz - + # use the setters self.given_in_basis = given_in_basis self.comp = comp @@ -312,7 +325,7 @@ def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28): self._r2 = a2 self._a = a self._b = b - + # use the setters self.given_in_basis = "norm" self.comp = 0 @@ -354,7 +367,7 @@ def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28): self._r2 = a2 self._a = a self._b = b - + # use the setters self.given_in_basis = "norm" self.comp = 1 @@ -393,7 +406,7 @@ def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28): self._r2 = a2 self._a = a self._b = b - + # use the setters self.given_in_basis = "norm" self.comp = 2 @@ -452,17 +465,23 @@ class TorusModesSin(Perturbation): given_in_basis : str In which basis the perturbation is represented (see base class). - + comp : int Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, ms=None, ns=None, amps=(1e-4,), pfuns=("sin",), pfun_params=None, - given_in_basis: GivenInBasis = "0", - comp: int = 0,): - + def __init__( + self, + ms=None, + ns=None, + amps=(1e-4,), + pfuns=("sin",), + pfun_params=None, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): assert "physical" not in given_in_basis - + if ms is not None: n_modes = len(ms) elif ns is not None: @@ -517,7 +536,7 @@ def __init__(self, ms=None, ns=None, amps=(1e-4,), pfuns=("sin",), pfun_params=N ] else: raise ValueError(f"Profile function {pfun} is not defined..") - + # use the setters self.given_in_basis = given_in_basis self.comp = comp @@ -578,18 +597,23 @@ class TorusModesCos(Perturbation): given_in_basis : str In which basis the perturbation is represented (see base class). - + comp : int Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, ms: tuple = (2,), ns: tuple = (1,), amps: tuple = (0.1,), - pfuns: tuple = ("sin",), pfun_params=None, - given_in_basis: GivenInBasis = "0", - comp: int = 0,): - + def __init__( + self, + ms: tuple = (2,), + ns: tuple = (1,), + amps: tuple = (0.1,), + pfuns: tuple = ("sin",), + pfun_params=None, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): assert "physical" not in given_in_basis - + if ms is not None: n_modes = len(ms) elif ns is not None: @@ -648,7 +672,7 @@ def __init__(self, ms: tuple = (2,), ns: tuple = (1,), amps: tuple = (0.1,), raise ValueError( 'Profile function must be "sin" or "cos" or "exp".', ) - + # use the setters self.given_in_basis = given_in_basis self.comp = comp @@ -656,7 +680,6 @@ def __init__(self, ms: tuple = (2,), ns: tuple = (1,), amps: tuple = (0.1,), def __call__(self, eta1, eta2, eta3): val = 0.0 for mi, ni, pfun, amp in zip(self._ms, self._ns, self._pfuns, self._amps): - val += ( amp * pfun(eta1) @@ -684,22 +707,26 @@ class Shear_x(Perturbation): delta : float Characteristic size of the shear layer - + given_in_basis : str In which basis the perturbation is represented (see base class). - + comp : int Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, amp=1e-4, delta=1 / 15, given_in_basis: GivenInBasis = "0", - comp: int = 0,): - - assert "physical" not in given_in_basis, f'Perturbation {self.__name__} can only be used in logical space.' - + def __init__( + self, + amp=1e-4, + delta=1 / 15, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): + assert "physical" not in given_in_basis, f"Perturbation {self.__name__} can only be used in logical space." + self._amp = amp self._delta = delta - + # use the setters self.given_in_basis = given_in_basis self.comp = comp @@ -726,22 +753,26 @@ class Shear_y(Perturbation): delta : float Characteristic size of the shear layer - + given_in_basis : str In which basis the perturbation is represented (see base class). - + comp : int Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, amp=1e-4, delta=1 / 15, given_in_basis: GivenInBasis = "0", - comp: int = 0,): - - assert "physical" not in given_in_basis, f'Perturbation {self.__name__} can only be used in logical space.' - + def __init__( + self, + amp=1e-4, + delta=1 / 15, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): + assert "physical" not in given_in_basis, f"Perturbation {self.__name__} can only be used in logical space." + self._amp = amp self._delta = delta - + # use the setters self.given_in_basis = given_in_basis self.comp = comp @@ -771,19 +802,23 @@ class Shear_z(Perturbation): given_in_basis : str In which basis the perturbation is represented (see base class). - + comp : int Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, amp=1e-4, delta=1 / 15, given_in_basis: GivenInBasis = "0", - comp: int = 0,): - - assert "physical" not in given_in_basis, f'Perturbation {self.__name__} can only be used in logical space.' - + def __init__( + self, + amp=1e-4, + delta=1 / 15, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): + assert "physical" not in given_in_basis, f"Perturbation {self.__name__} can only be used in logical space." + self._amp = amp self._delta = delta - + # use the setters self.given_in_basis = given_in_basis self.comp = comp @@ -810,22 +845,26 @@ class Erf_z(Perturbation): delta : float Characteristic size of the shear layer - + given_in_basis : str In which basis the perturbation is represented (see base class). - + comp : int Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, amp=1e-4, delta=1 / 15, given_in_basis: GivenInBasis = "0", - comp: int = 0,): - - assert "physical" not in given_in_basis, f'Perturbation {self.__name__} can only be used in logical space.' - + def __init__( + self, + amp=1e-4, + delta=1 / 15, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): + assert "physical" not in given_in_basis, f"Perturbation {self.__name__} can only be used in logical space." + self._amp = amp self._delta = delta - + # use the setters self.given_in_basis = given_in_basis self.comp = comp @@ -888,14 +927,23 @@ class RestelliAnalyticSolutionVelocity(Perturbation): in plasma physics, Journal of Computational Physics 2018. """ - def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0, comp: int = 0,): + def __init__( + self, + a=1.0, + R0=2.0, + B0=10.0, + Bp=12.5, + alpha=0.1, + beta=1.0, + comp: int = 0, + ): self._a = a self._R0 = R0 self._B0 = B0 self._Bp = Bp self._alpha = alpha self._beta = beta - + # use the setters self.given_in_basis = "physical" self.comp = comp @@ -980,14 +1028,23 @@ class RestelliAnalyticSolutionVelocity_2(Perturbation): in plasma physics, Journal of Computational Physics 2018. """ - def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0, comp: int = 0,): + def __init__( + self, + a=1.0, + R0=2.0, + B0=10.0, + Bp=12.5, + alpha=0.1, + beta=1.0, + comp: int = 0, + ): self._a = a self._R0 = R0 self._B0 = B0 self._Bp = Bp self._alpha = alpha self._beta = beta - + # use the setter self.given_in_basis = "physical" self.comp = comp @@ -1072,14 +1129,23 @@ class RestelliAnalyticSolutionVelocity_3(Perturbation): in plasma physics, Journal of Computational Physics 2018. """ - def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0, comp: int = 0,): + def __init__( + self, + a=1.0, + R0=2.0, + B0=10.0, + Bp=12.5, + alpha=0.1, + beta=1.0, + comp: int = 0, + ): self._a = a self._R0 = R0 self._B0 = B0 self._Bp = Bp self._alpha = alpha self._beta = beta - + # use the setters self.given_in_basis = "physical" self.comp = comp @@ -1169,7 +1235,7 @@ def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0): self._Bp = Bp self._alpha = alpha self._beta = beta - + # use the setter self.given_in_basis = "physical" @@ -1218,11 +1284,17 @@ class ManufacturedSolutionVelocity(Perturbation): Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, species="Ions", dimension="1D", b0=1.0, comp: int = 0,): + def __init__( + self, + species="Ions", + dimension="1D", + b0=1.0, + comp: int = 0, + ): self._b = b0 self._species = species self._dimension = dimension - + # use the setters self.given_in_basis = "physical" self.comp = comp @@ -1324,7 +1396,7 @@ class ManufacturedSolutionPotential(Perturbation): def __init__(self, dimension="1D", b0=1.0): self._ab = b0 self._dimension = dimension - + # use the setter self.given_in_basis = "physical" @@ -1373,11 +1445,17 @@ class ManufacturedSolutionVelocity_2(Perturbation): Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions) """ - def __init__(self, species="Ions", dimension="1D", b0=1.0, comp: int = 0,): + def __init__( + self, + species="Ions", + dimension="1D", + b0=1.0, + comp: int = 0, + ): self._b = b0 self._species = species self._dimension = dimension - + # use the setters self.given_in_basis = "physical" self.comp = comp diff --git a/src/struphy/initial/tests/test_init_perturbations.py b/src/struphy/initial/tests/test_init_perturbations.py index 9087e0219..3b46728c1 100644 --- a/src/struphy/initial/tests/test_init_perturbations.py +++ b/src/struphy/initial/tests/test_init_perturbations.py @@ -1,5 +1,6 @@ import inspect from copy import deepcopy + import pytest @@ -196,7 +197,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False perturbation_0 = perturbation perturbation_1 = deepcopy(perturbation) perturbation_2 = deepcopy(perturbation) - + params = { key: { "given_in_basis": [fun_form] * 3, @@ -249,7 +250,9 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False ) fun1_xyz, fun2_xyz, fun3_xyz = domain.push([tmp1, tmp2, tmp3], eee1, eee2, eee3, kind="v") else: - fun1_xyz, fun2_xyz, fun3_xyz = domain.push([perturbation, perturbation, perturbation], eee1, eee2, eee3, kind=fun_form) + fun1_xyz, fun2_xyz, fun3_xyz = domain.push( + [perturbation, perturbation, perturbation], eee1, eee2, eee3, kind=fun_form + ) fun_xyz_vec = [fun1_xyz, fun2_xyz, fun3_xyz] diff --git a/src/struphy/io/inp/params_Maxwell.py b/src/struphy/io/inp/params_Maxwell.py index cd477269f..984b81620 100644 --- a/src/struphy/io/inp/params_Maxwell.py +++ b/src/struphy/io/inp/params_Maxwell.py @@ -1,15 +1,14 @@ -from struphy.io.options import EnvironmentOptions, Units, Time -from struphy.geometry import domains +from struphy import main from struphy.fields_background import equils -from struphy.topology import grids -from struphy.io.options import DerhamOptions -from struphy.io.options import FieldsBackground +from struphy.geometry import domains from struphy.initial import perturbations +from struphy.io.options import DerhamOptions, EnvironmentOptions, FieldsBackground, Time, Units from struphy.kinetic_background import maxwellians -from struphy import main # import model, set verbosity from struphy.models.toy import Maxwell as Model +from struphy.topology import grids + verbose = True # environment options @@ -41,23 +40,24 @@ # initial conditions (background + perturbation) model.em_fields.b_field.add_background(FieldsBackground()) -model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=0)) -model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=1)) -model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=2)) +model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos(given_in_basis="v", comp=0)) +model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos(given_in_basis="v", comp=1)) +model.em_fields.b_field.add_perturbation(perturbations.TorusModesCos(given_in_basis="v", comp=2)) # optional: exclude variables from saving # model.em_fields.b_field.save_data = False if __name__ == "__main__": # start run - main.run(model, - params_path=__file__, - env=env, - units=units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=verbose, - ) \ No newline at end of file + main.run( + model, + params_path=__file__, + env=env, + units=units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) diff --git a/src/struphy/io/inp/params_Maxwell_lw.py b/src/struphy/io/inp/params_Maxwell_lw.py index 5ac1d1ef4..f39582146 100644 --- a/src/struphy/io/inp/params_Maxwell_lw.py +++ b/src/struphy/io/inp/params_Maxwell_lw.py @@ -1,12 +1,12 @@ -from struphy.io.options import Units, Time, DerhamOptions, FieldsBackground from struphy.fields_background import equils from struphy.geometry import domains from struphy.initial import perturbations +from struphy.io.options import DerhamOptions, FieldsBackground, Time, Units from struphy.kinetic_background import maxwellians -from struphy.topology import grids -# import model +# import model from struphy.models.toy import Maxwell as Model +from struphy.topology import grids # light-weight model instance model = Model() @@ -63,4 +63,4 @@ # exclude variables from saving # model.em_fields.e_field.save_data = False -# model.em_fields.b_field.save_data = False \ No newline at end of file +# model.em_fields.b_field.save_data = False diff --git a/src/struphy/io/inp/verification/Maxwell_coaxial.py b/src/struphy/io/inp/verification/Maxwell_coaxial.py index 9bf9b836c..394242234 100644 --- a/src/struphy/io/inp/verification/Maxwell_coaxial.py +++ b/src/struphy/io/inp/verification/Maxwell_coaxial.py @@ -1,20 +1,21 @@ -from struphy.io.options import Units, Time, DerhamOptions, FieldsBackground from struphy.fields_background import equils from struphy.geometry import domains from struphy.initial import perturbations +from struphy.io.options import DerhamOptions, FieldsBackground, Time, Units from struphy.kinetic_background import maxwellians -from struphy.topology import grids -# import model +# import model from struphy.models.toy import Maxwell as Model +from struphy.topology import grids + verbose = True # units units = Units() # geometry -a1=2.326744 -a2=3.686839 +a1 = 2.326744 +a2 = 3.686839 domain = domains.HollowCylinder(a1=a1, a2=a2, Lz=2.0) # fluid equilibrium (can be used as part of initial conditions) @@ -28,7 +29,7 @@ Nel=(32, 64, 1), p=(3, 3, 1), spl_kind=(False, True, True), - dirichlet_bc=[[True, True], [False, False], [False, False]] + dirichlet_bc=[[True, True], [False, False], [False, False]], ) # derham options @@ -56,4 +57,4 @@ model.em_fields.b_field.add_perturbation( perturbations.CoaxialWaveguideMagnetic(m=3, a1=a1, a2=a2), verbose=verbose, -) \ No newline at end of file +) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index fb91d4d4d..601f228ef 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -1,11 +1,11 @@ +import os from dataclasses import dataclass from typing import Literal, get_args + import numpy as np -import os from struphy.physics.physics import ConstantsOfNature - ## Literal options # time @@ -21,7 +21,7 @@ # perturbations NoiseDirections = Literal["e1", "e2", "e3", "e1e2", "e1e3", "e2e3", "e1e2e3"] -GivenInBasis = Literal['0', '1', '2', '3', 'v', 'physical', 'physical_at_eta', 'norm', None] +GivenInBasis = Literal["0", "1", "2", "3", "v", "physical", "physical_at_eta", "norm", None] # solvers OptsSymmSolver = Literal["pcg", "cg"] @@ -32,13 +32,14 @@ OptsPICSpace = Literal["Particles6D", "DeltaFParticles6D", "Particles5D", "Particles3D"] OptsMarkerBC = Literal["periodic", "reflect"] OptsRecontructBC = Literal["periodic", "mirror", "fixed"] -OptsLoading = Literal["pseudo_random", 'sobol_standard', 'sobol_antithetic', 'external', 'restart', "tesselation"] +OptsLoading = Literal["pseudo_random", "sobol_standard", "sobol_antithetic", "external", "restart", "tesselation"] OptsSpatialLoading = Literal["uniform", "disc"] OptsMPIsort = Literal["each", "last", None] ## Option classes + @dataclass class Time: """Time stepping options. @@ -47,10 +48,10 @@ class Time: ---------- dt : float Time step. - + Tend : float End time. - + split_algo : SplitAlgos Splitting algorithm (the order of the propagators is defined in the model). """ @@ -80,9 +81,10 @@ class BaseUnits: Unit of particle number density in 1e20/m^3. kBT : float, optional - Unit of internal energy in keV. + Unit of internal energy in keV. Only in effect if the velocity scale is set to 'thermal'. """ + x: float = 1.0 B: float = 1.0 n: float = 1.0 @@ -97,56 +99,55 @@ class Units: def __init__(self, base: BaseUnits = None): if base is None: base = BaseUnits() - + self._x = base.x self._B = base.B self._n = base.n * 1e20 self._kBT = base.kBT - + @property def x(self): return self._x - + @property def B(self): return self._B - + @property def n(self): """Unit of particle number density in 1/m^3.""" return self._n - + @property def kBT(self): return self._kBT - + @property def v(self): """Unit of velocity in m/s.""" return self._v - + @property def t(self): """Unit of time in s.""" return self._t - + @property def p(self): """Unit of pressure in Pa, equal to B^2/mu0 if velocity_scale='alfvén'.""" return self._p - + @property def rho(self): """Unit of mass density in kg/m^3.""" return self._rho - + @property def j(self): """Unit of current density in A/m^2.""" return self._j - - def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk: int = None, - verbose=False): + + 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.""" from mpi4py import MPI @@ -176,7 +177,7 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk # time (s) self._t = self.x / self.v - + # return if no bulk is present if A_bulk is None: self._p = None @@ -184,17 +185,27 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk self._j = None else: # pressure (Pa), equal to B^2/mu0 if velocity_scale='alfvén' - self._p = A_bulk * con.mH * self.n * self.v ** 2 + self._p = A_bulk * con.mH * self.n * self.v**2 # mass density (kg/m^3) self._rho = A_bulk * con.mH * self.n # current density (A/m^2) self._j = con.e * self.n * self.v - + # print to screen if verbose and MPI.COMM_WORLD.Get_rank() == 0: - units_used = (" m", " T", " m⁻³", "keV", " m/s", " s", " bar", " kg/m³", " A/m²",) + units_used = ( + " m", + " T", + " m⁻³", + "keV", + " m/s", + " s", + " bar", + " kg/m³", + " A/m²", + ) print("") for (k, v), u in zip(self.__dict__.items(), units_used): if v is None: @@ -220,19 +231,20 @@ class DerhamOptions: dirichlet_bc : tuple[tuple[bool]] Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). - + nquads : tuple[int] Number of Gauss-Legendre quadrature points in each direction (default = p, leads to exact integration of degree 2p-1 polynomials). nq_pr : tuple[int] Number of Gauss-Legendre quadrature points in each direction for geometric projectors (default = p+1, leads to exact integration of degree 2p+1 polynomials). - + polar_ck : PolarRegularity Smoothness at a polar singularity at eta_1=0 (default -1 : standard tensor product splines, OR 1 : C1 polar splines) local_projectors : bool Whether to build the local commuting projectors based on quasi-inter-/histopolation. """ + p: tuple = (1, 1, 1) spl_kind: tuple = (True, True, True) dirichlet_bc: tuple = ((False, False), (False, False), (False, False)) @@ -248,16 +260,16 @@ def __post_init__(self): @dataclass class FieldsBackground: """Options for backgrounds in configuration (=position) space. - + Parameters ---------- type : BackgroundTypes Type of background. - + values : tuple[float] - Values for LogicalConst on the unit cube. + Values for LogicalConst on the unit cube. Can be length 1 for scalar functions; must be length 3 for vector-valued functions. - + variable : str Name of the function in FluidEquilibrium that should be the background. """ @@ -268,22 +280,22 @@ class FieldsBackground: def __post_init__(self): check_option(self.type, BackgroundTypes) - - + + @dataclass class EnvironmentOptions: - """Environment options for launching run on current architecture - (these options do not influence the simulation result). + """Environment options for launching run on current architecture + (these options do not influence the simulation result). Parameters ---------- out_folders : str - The directory where all sim_folders are stored. - + The directory where all sim_folders are stored. + sim_folder : str Folder in 'out_folders/' for the current simulation (default='sim_1'). Will create the folder if it does not exist OR cleans the folder for new runs. - + restart : bool Whether to restart a run (default=False). @@ -307,20 +319,19 @@ class EnvironmentOptions: save_step: int = 1 sort_step: int = 0 num_clones: int = 1 - + def __post_init__(self): self.path_out: str = os.path.join(self.out_folders, self.sim_folder) - + def __repr__(self): for k, v in self.__dict__.items(): print(f"{k}:".ljust(20), v) - - + + def check_option(opt, options): """Check if opt is contained in options; if opt is a list, checks for each element.""" opts = get_args(options) if not isinstance(opt, list): opt = [opt] for o in opt: - assert o in opts, f"Option '{o}' is not in {opts}." - + assert o in opts, f"Option '{o}' is not in {opts}." diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index fdb4a7699..ed23926df 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -1,18 +1,17 @@ -import importlib.util -import sys import glob +import importlib.util import os import shutil -import yaml +import sys from types import ModuleType +import yaml from mpi4py import MPI -from struphy.io.options import DerhamOptions -from struphy.topology.grids import TensorProductGrid from struphy.geometry.base import Domain from struphy.geometry.domains import Cuboid -from struphy.io.options import Units, Time +from struphy.io.options import DerhamOptions, Time, Units +from struphy.topology.grids import TensorProductGrid def import_parameters_py(params_path: str) -> ModuleType: @@ -121,7 +120,7 @@ def setup_derham( Nel = grid.Nel # mpi mpi_dims_mask = grid.mpi_dims_mask - + # spline degrees p = options.p # spline types (clamped vs. periodic) diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 7cb4344f6..353127c21 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -1,9 +1,10 @@ "Base classes for kinetic backgrounds." from abc import ABCMeta, abstractmethod -import numpy as np from typing import Callable +import numpy as np + from struphy.fields_background.base import FluidEquilibrium from struphy.fields_background.equils import set_defaults from struphy.initial.base import Perturbation @@ -349,7 +350,7 @@ def vth(self, *etas): @abstractmethod def maxw_params(self) -> dict: """Parameters dictionary defining moments of the Maxwellian.""" - + def check_maxw_params(self): for k, v in self.maxw_params.items(): assert isinstance(k, str) @@ -480,7 +481,7 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name: str = "n", add_perturbatio name : str Which moment to evaluate (see varaible "dct" below). - + add_perturbation : bool | None Whether to add the perturbation defined in maxw_params. If None, is taken from self.add_perturbation. @@ -494,7 +495,7 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name: str = "n", add_perturbatio assert isinstance(eta2, np.ndarray) assert isinstance(eta3, np.ndarray) assert eta1.shape == eta2.shape == eta3.shape - + params = self.maxw_params[name] assert isinstance(params, tuple) assert len(params) == 2 @@ -545,11 +546,11 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name: str = "n", add_perturbatio out += background(eta1, eta2, eta3) else: out += background(*etas) - + # add perturbation if add_perturbation is None: add_perturbation = self.add_perturbation - + perturbation = params[1] if perturbation is not None and add_perturbation: assert isinstance(perturbation, Perturbation) @@ -557,15 +558,15 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name: str = "n", add_perturbatio out += perturbation(eta1, eta2, eta3) else: out += perturbation(*etas) - + return out - + @property def add_perturbation(self) -> bool: if not hasattr(self, "_add_perturbation"): self._add_perturbation = True return self._add_perturbation - + @add_perturbation.setter def add_perturbation(self, new): assert isinstance(new, bool) diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index dcd1ff3ee..eb60e09b6 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -1,13 +1,14 @@ "Maxwellian (Gaussian) distributions in velocity space." -import numpy as np from typing import Callable +import numpy as np + from struphy.fields_background.base import FluidEquilibriumWithB from struphy.fields_background.equils import set_defaults +from struphy.initial.base import Perturbation from struphy.kinetic_background import moment_functions from struphy.kinetic_background.base import CanonicalMaxwellian, Maxwellian -from struphy.initial.base import Perturbation class Maxwellian3D(Maxwellian): @@ -16,7 +17,7 @@ class Maxwellian3D(Maxwellian): Parameters ---------- n, ui, vthi : tuple - Moments of the Maxwellian as tuples. The first entry defines the background + Moments of the Maxwellian as tuples. The first entry defines the background (float for constant background or callable), the second entry defines a Perturbation (can be None). """ @@ -30,7 +31,6 @@ def __init__( vth2: tuple[float | Callable[..., float], Perturbation] = (1.0, None), vth3: tuple[float | Callable[..., float], Perturbation] = (1.0, None), ): - self._maxw_params = {} self._maxw_params["n"] = n self._maxw_params["u1"] = u1 @@ -39,7 +39,7 @@ def __init__( self._maxw_params["vth1"] = vth1 self._maxw_params["vth2"] = vth2 self._maxw_params["vth3"] = vth3 - + self.check_maxw_params() # factors multiplied onto the defined moments n, u and vth (can be set via setter) @@ -52,7 +52,7 @@ def __init__( @property def maxw_params(self): return self._maxw_params - + @property def coords(self): """Coordinates of the Maxwellian6D, :math:`(v_1, v_2, v_3)`.""" @@ -144,9 +144,9 @@ class GyroMaxwellian2D(Maxwellian): Parameters ---------- n, u_para, u_perp, vth_para, vth_perp : tuple - Moments of the Maxwellian as tuples. The first entry defines the background + Moments of the Maxwellian as tuples. The first entry defines the background (float for constant background or callable), the second entry defines a Perturbation (can be None). - + maxw_params : dict Parameters for the kinetic background. @@ -161,6 +161,7 @@ class GyroMaxwellian2D(Maxwellian): if True it is multiplied by the Jacobian determinant |v_perp| of the polar coordinate transofrmation (default = False). """ + def __init__( self, n: tuple[float | Callable[..., float], Perturbation] = (1.0, None), @@ -171,14 +172,13 @@ def __init__( equil: FluidEquilibriumWithB = None, volume_form: bool = True, ): - self._maxw_params = {} self._maxw_params["n"] = n self._maxw_params["u_para"] = u_para self._maxw_params["u_perp"] = u_perp self._maxw_params["vth_para"] = vth_para self._maxw_params["vth_perp"] = vth_perp - + self.check_maxw_params() # volume form represenation @@ -264,7 +264,7 @@ def velocity_jacobian_det(self, eta1, eta2, eta3, *v): def volume_form(self) -> bool: """Boolean. True if the background is represented as a volume form (thus including the velocity Jacobian |v_perp|).""" return self._volume_form - + @property def equil(self) -> FluidEquilibriumWithB: """Fluid background with B-field.""" diff --git a/src/struphy/linear_algebra/schur_solver.py b/src/struphy/linear_algebra/schur_solver.py index fc154f01c..dd41af54b 100644 --- a/src/struphy/linear_algebra/schur_solver.py +++ b/src/struphy/linear_algebra/schur_solver.py @@ -1,6 +1,7 @@ from psydac.linalg.basic import IdentityOperator, LinearOperator, Vector from psydac.linalg.block import BlockLinearOperator, BlockVector from psydac.linalg.solvers import inverse + from struphy.linear_algebra.solver import SolverParameters @@ -47,16 +48,20 @@ class SchurSolver: Must correspond to the chosen solver. """ - def __init__(self, A: LinearOperator, BC: LinearOperator, - solver_name: str, - precond = None, # TODO: add Preconditioner base class - solver_params: SolverParameters = None,): + def __init__( + self, + A: LinearOperator, + BC: LinearOperator, + solver_name: str, + precond=None, # TODO: add Preconditioner base class + solver_params: SolverParameters = None, + ): assert isinstance(A, LinearOperator) assert isinstance(BC, LinearOperator) assert A.domain == BC.domain assert A.codomain == BC.codomain - + if solver_params is None: solver_params = SolverParameters() diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py index a47f1d592..0e36cfcaf 100644 --- a/src/struphy/linear_algebra/solver.py +++ b/src/struphy/linear_algebra/solver.py @@ -1,5 +1,6 @@ from dataclasses import dataclass + @dataclass class SolverParameters: """Parameters for psydac solvers.""" @@ -8,4 +9,4 @@ class SolverParameters: maxiter: int = 3000 info: bool = False verbose: bool = False - recycle: bool = True \ No newline at end of file + recycle: bool = True diff --git a/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py b/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py index 3561c49cc..e0b4143d7 100644 --- a/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py +++ b/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py @@ -20,9 +20,9 @@ def test_propagator1D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): from struphy.feec.utilities import compare_arrays from struphy.fields_background.equils import HomogenSlab from struphy.geometry import domains - from struphy.propagators.propagators_fields import TwoFluidQuasiNeutralFull - from struphy.models.variables import FEECVariable from struphy.initial import perturbations + from struphy.models.variables import FEECVariable + from struphy.propagators.propagators_fields import TwoFluidQuasiNeutralFull mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -93,13 +93,13 @@ def test_propagator1D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): uvec.add_perturbation(pp_u) uvec.allocate(derham, domain, eq_mhd) - + u_evec.add_perturbation(pp_ue) u_evec.allocate(derham, domain, eq_mhd) - + potentialvec.add_perturbation(pp_potential) potentialvec.allocate(derham, domain, eq_mhd) - + uinitial.allocate(derham, domain, eq_mhd) # uvec.initialize_coeffs(domain=domain, pert_params=pp_u) diff --git a/src/struphy/main.py b/src/struphy/main.py index 0fe7d4db6..bbca0c526 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -1,44 +1,45 @@ -from typing import Optional, TypedDict +import copy +import datetime +import glob import os +import pickle +import shutil import sysconfig import time -import datetime -import glob +from typing import Optional, TypedDict + +import h5py import numpy as np +from line_profiler import profile from mpi4py import MPI from pyevtk.hl import gridToVTK -import shutil -import pickle -import h5py -import copy -from line_profiler import profile -from struphy.fields_background.base import FluidEquilibriumWithB +from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB +from struphy.fields_background.equils import HomogenSlab +from struphy.geometry import domains +from struphy.geometry.base import Domain +from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, Time, Units from struphy.io.output_handling import DataContainer -from struphy.io.setup import setup_folders +from struphy.io.setup import import_parameters_py, setup_folders from struphy.models.base import StruphyModel -from struphy.profiling.profiling import ProfileManager -from struphy.utils.clone_config import CloneConfig -from struphy.utils.utils import dict_to_yaml -from struphy.pic.base import Particles from struphy.models.species import Species from struphy.models.variables import FEECVariable -from struphy.io.options import BaseUnits, Units, Time, EnvironmentOptions -from struphy.io.setup import import_parameters_py -from struphy.geometry.base import Domain -from struphy.geometry import domains -from struphy.fields_background.base import FluidEquilibrium -from struphy.fields_background.equils import HomogenSlab -from struphy.topology.grids import TensorProductGrid -from struphy.io.options import DerhamOptions -from struphy.post_processing.post_processing_tools import create_femfields, eval_femfields, create_vtk -from struphy.topology import grids -from struphy.io.options import DerhamOptions -from struphy.post_processing.post_processing_tools import (post_process_markers, - post_process_f, - post_process_n_sph) +from struphy.pic.base import Particles from struphy.post_processing.orbits import orbits_tools - +from struphy.post_processing.post_processing_tools import ( + create_femfields, + create_vtk, + eval_femfields, + post_process_f, + post_process_markers, + post_process_n_sph, +) +from struphy.profiling.profiling import ProfileManager +from struphy.topology import grids +from struphy.topology.grids import TensorProductGrid +from struphy.utils.clone_config import CloneConfig +from struphy.utils.utils import dict_to_yaml + @profile def run( @@ -71,15 +72,15 @@ def run( size = comm.Get_size() start_simulation = time.time() - + # check model assert hasattr(model, "propagators"), "Attribute 'self.propagators' must be set in model __init__!" model_name = model.__class__.__name__ model.verbose = verbose - + if rank == 0: print(f"\n*** Starting run for model '{model_name}':") - + # meta-data path_out = env.path_out restart = env.restart @@ -87,7 +88,7 @@ def run( save_step = env.save_step sort_step = env.sort_step num_clones = env.num_clones - + meta = {} meta["platform"] = sysconfig.get_platform() meta["python version"] = sysconfig.get_python_version() @@ -99,20 +100,22 @@ def run( meta["restart"] = restart meta["max wall-clock [min]"] = max_runtime meta["save interval [steps]"] = save_step - + if rank == 0: print("\nMETADATA:") for k, v in meta.items(): - print(f'{k}:'.ljust(25), v) - + print(f"{k}:".ljust(25), v) + # creating output folders - setup_folders(path_out=path_out, - restart=restart, - verbose=verbose,) - + setup_folders( + path_out=path_out, + restart=restart, + verbose=verbose, + ) + # add derived units units = Units(base_units) - + # save parameter file if rank == 0: # save python param file @@ -124,30 +127,30 @@ def run( ) # pickle struphy objects else: - with open(os.path.join(path_out, "env.bin"), 'wb') as f: + with open(os.path.join(path_out, "env.bin"), "wb") as f: pickle.dump(env, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "base_units.bin"), 'wb') as f: + with open(os.path.join(path_out, "base_units.bin"), "wb") as f: pickle.dump(base_units, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "time_opts.bin"), 'wb') as f: + with open(os.path.join(path_out, "time_opts.bin"), "wb") as f: pickle.dump(time_opts, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "domain.bin"), 'wb') as f: + with open(os.path.join(path_out, "domain.bin"), "wb") as f: # WORKAROUND: cannot pickle pyccelized classes at the moment tmp_dct = {"name": domain.__class__.__name__, "params": domain.params} pickle.dump(tmp_dct, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "equil.bin"), 'wb') as f: + with open(os.path.join(path_out, "equil.bin"), "wb") as f: # WORKAROUND: cannot pickle pyccelized classes at the moment if equil is not None: tmp_dct = {"name": equil.__class__.__name__, "params": equil.params} else: tmp_dct = {} pickle.dump(tmp_dct, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "grid.bin"), 'wb') as f: + with open(os.path.join(path_out, "grid.bin"), "wb") as f: pickle.dump(grid, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "derham_opts.bin"), 'wb') as f: + with open(os.path.join(path_out, "derham_opts.bin"), "wb") as f: pickle.dump(derham_opts, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "model.bin"), 'wb') as f: + with open(os.path.join(path_out, "model.bin"), "wb") as f: pickle.dump(model, f, pickle.HIGHEST_PROTOCOL) - + # config clones if comm is None: clone_config = None @@ -163,10 +166,10 @@ def run( clone_config.print_clone_config() if model.kinetic_species: clone_config.print_particle_config() - + model.clone_config = clone_config comm.Barrier() - + ## configure model instance # units @@ -177,41 +180,45 @@ def run( else: A_bulk = model.bulk_species.mass_number Z_bulk = model.bulk_species.charge_number - model.units.derive_units(velocity_scale=model.velocity_scale, - A_bulk=A_bulk, - Z_bulk=Z_bulk, - verbose=verbose,) - + model.units.derive_units( + velocity_scale=model.velocity_scale, + A_bulk=A_bulk, + Z_bulk=Z_bulk, + verbose=verbose, + ) + # domain and fluid background model.setup_domain_and_equil(domain, equil) - + # default grid if grid is None: Nel = (16, 16, 16) print(f"\nNo grid specified - using TensorProductGrid with {Nel = }.") grid = grids.TensorProductGrid(Nel=Nel) - + # allocate derham-related objects if derham_opts is None: p = (3, 3, 3) spl_kind = (False, False, False) - print(f"\nNo Derham options specified - creating Derham with {p = } and {spl_kind = } for projecting equilibrium.") + print( + f"\nNo Derham options specified - creating Derham with {p = } and {spl_kind = } for projecting equilibrium." + ) derham_opts = DerhamOptions(p=p, spl_kind=spl_kind) model.allocate_feec(grid, derham_opts) - + # equation paramters model.setup_equation_params(units=model.units, verbose=verbose) - + # allocate variables model.allocate_variables(verbose=verbose) model.allocate_helpers() - + # pass info to propagators model.allocate_propagators() - + # plasma parameters model.compute_plasma_params(verbose=verbose) - + if rank < 32: if rank == 0: print("") @@ -311,7 +318,7 @@ def run( data.file.close() end_simulation = time.time() if rank == 0: - print(f"\nTime steps done: {time_state["index"][0]}") + print(f"\nTime steps done: {time_state['index'][0]}") print( "wall-clock time of simulation [sec]: ", end_simulation - start_simulation, @@ -386,7 +393,7 @@ def run( meta["wall-clock time[min]"] = (end_simulation - start_simulation) / 60 comm.Barrier() - + if rank == 0: # save meta-data dict_to_yaml(meta, os.path.join(path_out, "meta.yml")) @@ -438,16 +445,16 @@ def pproc( if MPI.COMM_WORLD.Get_rank() == 0: print(f"\n*** Start post-processing of {path}:") - + # load light-weight model instance from simulation try: params_in = import_parameters_py(os.path.join(path, "parameters.py")) model: StruphyModel = params_in.model domain: Domain = params_in.domain except FileNotFoundError: - with open(os.path.join(path, "model.bin"), 'rb') as f: + with open(os.path.join(path, "model.bin"), "rb") as f: model: StruphyModel = pickle.load(f) - with open(os.path.join(path, "domain.bin"), 'rb') as f: + with open(os.path.join(path, "domain.bin"), "rb") as f: # domain: Domain = pickle.load(f) # print(f"{domain = }") # print(f"{domain.params = }") @@ -489,7 +496,7 @@ def pproc( for name in file["kinetic"].keys(): kinetic_species += [name] kinetic_kinds += [next(iter(model.species[name].variables.values())).space] - + # check for saved markers if "markers" in file["kinetic"][name]: exist_kinetic["markers"] = True @@ -501,16 +508,14 @@ def pproc( exist_kinetic["n_sph"] = True else: exist_kinetic = None - + file.close() # field post-processing if exist_fields: fields, t_grid = create_femfields(path, step=step) - point_data, grids_log, grids_phy = eval_femfields( - path, fields, celldivide=[celldivide, celldivide, celldivide] - ) + point_data, grids_log, grids_phy = eval_femfields(path, fields, celldivide=[celldivide, celldivide, celldivide]) if physical: point_data_phy, grids_log, grids_phy = eval_femfields( @@ -578,7 +583,14 @@ def pproc( # markers if exist_kinetic["markers"]: - post_process_markers(path, path_kinetics_species, species, domain, kinetic_kinds[n], step,) + post_process_markers( + path, + path_kinetics_species, + species, + domain, + kinetic_kinds[n], + step, + ) if guiding_center: assert kinetic_kinds[n] == "Particles6D" @@ -603,12 +615,13 @@ def pproc( class SimData: """Holds post-processed Struphy data as attributes. - + Parameters ---------- path : str Absolute path of simulation output folder to post-process. """ + def __init__(self, path: str): self.path = path self._orbits = {} @@ -618,22 +631,22 @@ def __init__(self, path: str): self.grids_log: list[np.ndarray] = None self.grids_phy: list[np.ndarray] = None self.t_grid: np.ndarray = None - + @property def orbits(self) -> dict[str, np.ndarray]: """Keys: species name. Values: 3d arrays indexed by (n, p, a), where 'n' is the time index, 'p' the particle index and 'a' the attribute index.""" return self._orbits - + @property def f(self) -> dict[str, dict[str, dict[str, np.ndarray]]]: """Keys: species name. Values: dicts of slice names ('e1_v1' etc.) holding dicts of corresponding np.arrays for plotting.""" return self._f - + @property def spline_values(self) -> dict[str, dict[str, np.ndarray]]: """Keys: species name. Values: dicts of variable names with values being 3d arrays on the grid.""" return self._spline_values - + @property def Nt(self) -> dict[str, int]: """Number of available time points (snap shots) for each species.""" @@ -642,7 +655,7 @@ def Nt(self) -> dict[str, int]: for spec, orbs in self.orbits.items(): self._Nt[spec] = orbs.shape[0] return self._Nt - + @property def Np(self) -> dict[str, int]: """Number of particle orbits for each species.""" @@ -651,7 +664,7 @@ def Np(self) -> dict[str, int]: for spec, orbs in self.orbits.items(): self._Np[spec] = orbs.shape[1] return self._Np - + @property def Nattr(self) -> dict[str, int]: """Number of particle attributes for each species.""" @@ -660,7 +673,7 @@ def Nattr(self) -> dict[str, int]: for spec, orbs in self.orbits.items(): self._Nattr[spec] = orbs.shape[2] return self._Nattr - + @property def spline_grid_resolution(self): if self.grids_log is not None: @@ -668,7 +681,7 @@ def spline_grid_resolution(self): else: res = None return res - + @property def time_grid_size(self): return self.t_grid.size @@ -676,7 +689,7 @@ def time_grid_size(self): def load_data(path: str) -> SimData: """Load data generated during post-processing. - + Parameters ---------- path : str @@ -689,23 +702,22 @@ def load_data(path: str) -> SimData: print(f"{path = }") simdata = SimData(path) - + # load time grid simdata.t_grid = np.load(os.path.join(path_pproc, "t_grid.npy")) # data paths path_fields = os.path.join(path_pproc, "fields_data") path_kinetic = os.path.join(path_pproc, "kinetic_data") - + # load point data if os.path.exists(path_fields): - # grids with open(os.path.join(path_fields, "grids_log.bin"), "rb") as f: simdata.grids_log = pickle.load(f) with open(os.path.join(path_fields, "grids_phy.bin"), "rb") as f: simdata.grids_phy = pickle.load(f) - + # species folders species = next(os.walk(path_fields))[1] for spec in species: @@ -722,9 +734,8 @@ def load_data(path: str) -> SimData: # try: simdata._spline_values[spec][var] = pickle.load(f) # simdata.arrays[spec][var] = pickle.load(f) - + if os.path.exists(path_kinetic): - # species folders species = next(os.walk(path_kinetic))[1] print(f"{species = }") @@ -774,7 +785,7 @@ def load_data(path: str) -> SimData: # tmp[var] = np.load(os.path.join(path_dat, file)) # # sort dict # simdata.pic_species[spec][folder] = dict(sorted(tmp.items())) - + print("\nThe following data has been loaded:") print(f"{simdata.time_grid_size = }") print(f"{simdata.spline_grid_resolution = }") @@ -794,7 +805,7 @@ def load_data(path: str) -> SimData: for kkk, vvv in vv.items(): print(f" {kkk}") print(f"simdata.sph_species:") - + return simdata diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 1336ce599..289c5c4b8 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1,43 +1,42 @@ import inspect import operator +import os from abc import ABCMeta, abstractmethod from functools import reduce from typing import Callable -import os -import yaml -import struphy -from line_profiler import profile import numpy as np import yaml +from line_profiler import profile from mpi4py import MPI +from psydac.linalg.stencil import StencilVector +import struphy from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import WeightedMassOperators from struphy.feec.psydac_derham import SplineFunction from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB, MHDequilibrium +from struphy.fields_background.equils import HomogenSlab from struphy.fields_background.projected_equils import ( ProjectedFluidEquilibrium, ProjectedFluidEquilibriumWithB, ProjectedMHDequilibrium, ) -from struphy.io.setup import setup_derham, descend_options_dict -from struphy.profiling.profiling import ProfileManager -from struphy.utils.clone_config import CloneConfig -from struphy.utils.utils import dict_to_yaml, read_state -from struphy.models.species import Species, FieldSpecies, FluidSpecies, KineticSpecies, DiagnosticSpecies -from struphy.models.variables import FEECVariable, PICVariable, SPHVariable -from struphy.io.options import BaseUnits, Units, Time, DerhamOptions -from struphy.topology.grids import TensorProductGrid from struphy.geometry.base import Domain from struphy.geometry.domains import Cuboid -from struphy.fields_background.equils import HomogenSlab +from struphy.io.options import BaseUnits, DerhamOptions, Time, Units +from struphy.io.output_handling import DataContainer +from struphy.io.setup import descend_options_dict, setup_derham +from struphy.kinetic_background import maxwellians +from struphy.models.species import DiagnosticSpecies, FieldSpecies, FluidSpecies, KineticSpecies, Species +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.pic import particles from struphy.pic.base import Particles +from struphy.profiling.profiling import ProfileManager from struphy.propagators.base import Propagator -from struphy.kinetic_background import maxwellians -from psydac.linalg.stencil import StencilVector -from struphy.io.output_handling import DataContainer +from struphy.topology.grids import TensorProductGrid +from struphy.utils.clone_config import CloneConfig +from struphy.utils.utils import dict_to_yaml, read_state class StruphyModel(metaclass=ABCMeta): @@ -55,12 +54,12 @@ class StruphyModel(metaclass=ABCMeta): @abstractmethod class Propagators: pass - + @abstractmethod def __init__(self): """Light-weight init of model.""" - - @property + + @property @abstractmethod def bulk_species() -> Species: """Bulk species of the plasma. Must be an attribute of species_static().""" @@ -90,9 +89,9 @@ def setup_equation_params(self, units: Units, verbose=False): for _, species in self.kinetic_species.items(): assert isinstance(species, KineticSpecies) species.setup_equation_params(units=units, verbose=verbose) - + def setup_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): - """If a numerical equilibirum is used, the domain is taken from this equilibirum. """ + """If a numerical equilibirum is used, the domain is taken from this equilibirum.""" if equil is not None: self._equil = equil if "Numerical" in self.equil.__class__.__name__: @@ -102,8 +101,8 @@ def setup_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): self._equil.domain = domain else: self._domain = domain - self._equil = None - + self._equil = None + if MPI.COMM_WORLD.Get_rank() == 0 and self.verbose: print("\nDOMAIN:") print(f"type:".ljust(25), self.domain.__class__.__name__) @@ -118,9 +117,9 @@ def setup_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): print((key + ":").ljust(25), val) else: print("None.") - - ## species - + + ## species + @property def field_species(self) -> dict: if not hasattr(self, "_field_species"): @@ -129,7 +128,7 @@ def field_species(self) -> dict: if isinstance(v, FieldSpecies): self._field_species[k] = v return self._field_species - + @property def fluid_species(self) -> dict: if not hasattr(self, "_fluid_species"): @@ -138,7 +137,7 @@ def fluid_species(self) -> dict: if isinstance(v, FluidSpecies): self._fluid_species[k] = v return self._fluid_species - + @property def kinetic_species(self) -> dict: if not hasattr(self, "_kinetic_species"): @@ -147,7 +146,7 @@ def kinetic_species(self) -> dict: if isinstance(v, KineticSpecies): self._kinetic_species[k] = v return self._kinetic_species - + @property def diagnostic_species(self) -> dict: if not hasattr(self, "_diagnostic_species"): @@ -156,17 +155,16 @@ def diagnostic_species(self) -> dict: if isinstance(v, DiagnosticSpecies): self._diagnostic_species[k] = v return self._diagnostic_species - + @property def species(self): if not hasattr(self, "_species"): self._species = self.field_species | self.fluid_species | self.kinetic_species return self._species - + ## allocate methods - - def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): + def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): # create discrete derham sequence if self.clone_config is None: derham_comm = MPI.COMM_WORLD @@ -180,7 +178,7 @@ def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): domain=self.domain, verbose=self.verbose, ) - + # create weighted mass operators self._mass_ops = WeightedMassOperators( self.derham, @@ -188,7 +186,7 @@ def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): verbose=self.verbose, eq_mhd=self.equil, ) - + # create projected equilibrium if isinstance(self.equil, MHDequilibrium): self._projected_equil = ProjectedMHDequilibrium( @@ -204,10 +202,10 @@ def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): self._projected_equil = ProjectedFluidEquilibrium( self.equil, self.derham, - ) + ) else: self._projected_equil = None - + def allocate_propagators(self): # set propagators base class attributes (then available to all propagators) Propagator.derham = self.derham @@ -221,7 +219,7 @@ def allocate_propagators(self): eq_mhd=self.equil, ) Propagator.projected_equil = self.projected_equil - + assert len(self.prop_list) > 0, "No propagators in this model, check the model class." for prop in self.prop_list: assert isinstance(prop, Propagator) @@ -256,7 +254,7 @@ def equation_params(self): def clone_config(self): """Config in case domain clones are used.""" return self._clone_config - + @clone_config.setter def clone_config(self, new): assert isinstance(new, CloneConfig) or new is None @@ -291,9 +289,9 @@ def projected_equil(self): def units(self) -> Units: """All Struphy units.""" return self._units - + @units.setter - def units(self, new) : + def units(self, new): assert isinstance(new, Units) self._units = new @@ -301,7 +299,7 @@ def units(self, new) : def mass_ops(self): """WeighteMassOperators object, see :ref:`mass_ops`.""" return self._mass_ops - + @property def prop_list(self): """List of Propagator objects.""" @@ -437,8 +435,7 @@ def setInDict(dataDict, mapList, value): assert key is not None, "Must provide key if option is not a class." setInDict(dct, species + ["options"] + key, option) - def add_scalar(self, - name: str, variable: PICVariable | SPHVariable = None, compute=None, summands=None): + def add_scalar(self, name: str, variable: PICVariable | SPHVariable = None, compute=None, summands=None): """ Add a scalar to be saved during the simulation. @@ -585,7 +582,11 @@ def allocate_variables(self, verbose: bool = False): assert isinstance(spec, FieldSpecies) for k, v in spec.variables.items(): assert isinstance(v, FEECVariable) - v.allocate(derham=self.derham, domain=self.domain, equil=self.equil,) + v.allocate( + derham=self.derham, + domain=self.domain, + equil=self.equil, + ) # allocate memory for FE coeffs of fluid variables if self.fluid_species: @@ -593,21 +594,27 @@ def allocate_variables(self, verbose: bool = False): assert isinstance(spec, FluidSpecies) for k, v in spec.variables.items(): assert isinstance(v, FEECVariable) - v.allocate(derham=self.derham, domain=self.domain, equil=self.equil,) - + v.allocate( + derham=self.derham, + domain=self.domain, + equil=self.equil, + ) + # allocate memory for marker arrays of kinetic variables if self.kinetic_species: for species, spec in self.kinetic_species.items(): assert isinstance(spec, KineticSpecies) for k, v in spec.variables.items(): assert isinstance(v, (PICVariable, SPHVariable)) - v.allocate(clone_config=self.clone_config, - derham=self.derham, - domain=self.domain, - equil=self.equil, - projected_equil=self.projected_equil, - verbose=verbose,) - + v.allocate( + clone_config=self.clone_config, + derham=self.derham, + domain=self.domain, + equil=self.equil, + projected_equil=self.projected_equil, + verbose=verbose, + ) + # TODO: allocate memory for FE coeffs of diagnostics # if self.params.diagnostic_fields is not None: # for key, val in self.diagnostics.items(): @@ -1031,7 +1038,7 @@ def initialize_data_output(self, data: DataContainer, size): feec_species = self.field_species | self.fluid_species | self.diagnostic_species for species, val in feec_species.items(): assert isinstance(val, Species) - + species_path = os.path.join("feec", species) species_path_restart = os.path.join("restart", species) @@ -1108,7 +1115,11 @@ def initialize_data_output(self, data: DataContainer, size): for dim in range(dims): data.file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = ( var.kinetic_data["bin_edges"][key2][dim][:-1] - + (var.kinetic_data["bin_edges"][key2][dim][1] - var.kinetic_data["bin_edges"][key2][dim][0]) / 2 + + ( + var.kinetic_data["bin_edges"][key2][dim][1] + - var.kinetic_data["bin_edges"][key2][dim][0] + ) + / 2 ) # case of "n_sph" elif isinstance(val1, list): @@ -1254,19 +1265,19 @@ def generate_default_parameter_file( prompt : bool Whether to prompt for overwriting the specified .yml file. - + Returns ------- params_path : str The path of the parameter file. """ - + if path is None: path = os.path.join(os.getcwd(), f"params_{self.__class__.__name__}.py") # create new default file try: - file = open(path, "x") + file = open(path, "x") except FileExistsError: if not prompt: yn = "Y" @@ -1288,21 +1299,21 @@ def generate_default_parameter_file( file = open(path, "x") else: print("exiting ...") - return - + return + file.write("from struphy.io.options import EnvironmentOptions, BaseUnits, Time\n") file.write("from struphy.geometry import domains\n") file.write("from struphy.fields_background import equils\n") - + species_params = "\n# species parameters\n" - kinetic_params = "" + kinetic_params = "" has_plasma = False has_feec = False has_pic = False has_sph = False for sn, species in self.species.items(): assert isinstance(species, Species) - + if isinstance(species, (FluidSpecies, KineticSpecies)): has_plasma = True species_params += f"model.{sn}.set_phys_params()\n" @@ -1313,7 +1324,7 @@ def generate_default_parameter_file( kinetic_params += f"model.{sn}.set_markers(loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params)\n" kinetic_params += f"model.{sn}.set_sorting_boxes()\n" kinetic_params += f"model.{sn}.set_save_data()\n" - + for vn, var in species.variables.items(): if isinstance(var, FEECVariable): has_feec = True @@ -1322,10 +1333,12 @@ def generate_default_parameter_file( init_pert_feec = f"model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos())\n" else: init_bckgr_feec = f"model.{sn}.{vn}.add_background(FieldsBackground())\n" - init_pert_feec = f"model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=0))\n\ + init_pert_feec = ( + f"model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=0))\n\ model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=1))\n\ model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=2))\n" - + ) + exclude = f"# model.{sn}.{vn}.save_data = False\n" elif isinstance(var, PICVariable): has_pic = True @@ -1338,65 +1351,67 @@ def generate_default_parameter_file( init_bckgr_pic += f"maxwellian_2 = maxwellians.GyroMaxwellian2D(n=(0.1, None))\n" init_bckgr_pic += f"background = maxwellian_1 + maxwellian_2\n" init_bckgr_pic += f"model.{sn}.{vn}.add_background(background)\n" - + exclude = f"# model.....save_data = False\n" elif isinstance(var, SPHVariable): has_sph = True - - file.write("from struphy.topology import grids\n") + + file.write("from struphy.topology import grids\n") file.write("from struphy.io.options import DerhamOptions\n") file.write("from struphy.io.options import FieldsBackground\n") file.write("from struphy.initial import perturbations\n") - + file.write("from struphy.kinetic_background import maxwellians\n") - file.write("from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters, BinningPlot\n") + file.write( + "from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters, BinningPlot\n" + ) file.write("from struphy import main\n") - + file.write("\n# import model, set verbosity\n") file.write(f"from {self.__module__} import {self.__class__.__name__}\n") file.write("verbose = True\n") - + file.write("\n# environment options\n") file.write("env = EnvironmentOptions()\n") - + file.write("\n# units\n") file.write("base_units = BaseUnits()\n") - + file.write("\n# time stepping\n") file.write("time_opts = Time()\n") - + file.write("\n# geometry\n") file.write("domain = domains.Cuboid()\n") - + file.write("\n# fluid equilibrium (can be used as part of initial conditions)\n") file.write("equil = equils.HomogenSlab()\n") - + if has_feec: grid = "grid = grids.TensorProductGrid()\n" derham = "derham_opts = DerhamOptions()\n" else: grid = "grid = None\n" derham = "derham_opts = None\n" - + file.write("\n# grid\n") file.write(grid) - + file.write("\n# derham options\n") file.write(derham) - + file.write("\n# light-weight model instance\n") file.write(f"model = {self.__class__.__name__}()\n") - + if has_plasma: file.write(species_params) - + if has_pic: file.write(kinetic_params) - + file.write("\n# propagator options\n") for prop in self.propagators.__dict__: file.write(f"model.propagators.{prop}.options = model.propagators.{prop}.Options()\n") - + file.write("\n# initial conditions (background + perturbation)\n") if has_feec: file.write(init_bckgr_feec) @@ -1404,13 +1419,14 @@ def generate_default_parameter_file( if has_pic: file.write(init_pert_pic) file.write(init_bckgr_pic) - + file.write("\n# optional: exclude variables from saving\n") file.write(exclude) - + file.write('\nif __name__ == "__main__":\n') file.write(" # start run\n") - file.write(" main.run(model, \n\ + file.write( + " main.run(model, \n\ params_path=__file__, \n\ env=env, \n\ base_units=base_units, \n\ @@ -1420,13 +1436,16 @@ def generate_default_parameter_file( grid=grid, \n\ derham_opts=derham_opts, \n\ verbose=verbose, \n\ - )") - + )" + ) + file.close() - - print(f"\nDefault parameter file for '{self.__class__.__name__}' has been created in {path}.\n\ -You can now launch with 'struphy run {self.__class__.__name__}' or with 'struphy run -i params_{self.__class__.__name__}.py'") - + + print( + f"\nDefault parameter file for '{self.__class__.__name__}' has been created in {path}.\n\ +You can now launch with 'struphy run {self.__class__.__name__}' or with 'struphy run -i params_{self.__class__.__name__}.py'" + ) + return path ################### @@ -1617,12 +1636,12 @@ def compute_plasma_params(self, verbose=True): eta3 = np.linspace(h / 2.0, 1.0 - h / 2.0, 20) ## global parameters - + # plasma volume (hat x^3) det_tmp = self.domain.jacobian_det(eta1, eta2, eta3) vol1 = np.mean(np.abs(det_tmp)) # plasma volume (m⁻³) - plasma_volume = vol1 * self.units.x ** 3 + plasma_volume = vol1 * self.units.x**3 # transit length (m) transit_length = plasma_volume ** (1 / 3) # magnetic field (T) diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 304116ab9..5b88a1fab 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -1,12 +1,12 @@ import numpy as np from mpi4py import MPI +from psydac.linalg.block import BlockVector from struphy.models.base import StruphyModel -from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies -from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable +from struphy.models.species import FieldSpecies, FluidSpecies, KineticSpecies +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.polar.basic import PolarVector -from psydac.linalg.block import BlockVector +from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers rank = MPI.COMM_WORLD.Get_rank() @@ -40,21 +40,21 @@ class LinearMHD(StruphyModel): 1. :class:`~struphy.propagators.propagators_fields.ShearAlfven` 2. :class:`~struphy.propagators.propagators_fields.Magnetosonic` """ - + ## species - + class EMFields(FieldSpecies): def __init__(self): self.b_field = FEECVariable(space="Hdiv") self.init_variables() - + class MHD(FluidSpecies): def __init__(self): - self.density = FEECVariable(space="L2") + self.density = FEECVariable(space="L2") self.velocity = FEECVariable(space="Hdiv") self.pressure = FEECVariable(space="L2") self.init_variables() - + ## propagators class Propagators: @@ -67,14 +67,14 @@ def __init__(self): def __init__(self): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - + # 1. instantiate all species self.em_fields = self.EMFields() self.mhd = self.MHD() # 2. instantiate all propagators self.propagators = self.Propagators() - + # 3. assign variables to propagators self.propagators.shear_alf.variables.u = self.mhd.velocity self.propagators.shear_alf.variables.b = self.em_fields.b_field @@ -82,7 +82,7 @@ def __init__(self): self.propagators.mag_sonic.variables.n = self.mhd.density self.propagators.mag_sonic.variables.u = self.mhd.velocity self.propagators.mag_sonic.variables.p = self.mhd.pressure - + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_p") @@ -106,14 +106,20 @@ def allocate_helpers(self): self._ones.tp[:] = 1.0 else: self._ones[:] = 1.0 - - self._tmp_b1: BlockVector = self.derham.Vh["2"].zeros() # TODO: replace derham.Vh dict by class + + self._tmp_b1: BlockVector = self.derham.Vh["2"].zeros() # TODO: replace derham.Vh dict by class self._tmp_b2: BlockVector = self.derham.Vh["2"].zeros() - def update_scalar_quantities(self): + def update_scalar_quantities(self): # perturbed fields - en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.mhd.velocity.spline.vector, self.mhd.velocity.spline.vector,) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.em_fields.b_field.spline.vector, self.em_fields.b_field.spline.vector,) + en_U = 0.5 * self.mass_ops.M2n.dot_inner( + self.mhd.velocity.spline.vector, + self.mhd.velocity.spline.vector, + ) + en_B = 0.5 * self.mass_ops.M2.dot_inner( + self.em_fields.b_field.spline.vector, + self.em_fields.b_field.spline.vector, + ) en_p = self.mhd.pressure.spline.vector.inner(self._ones) / (5 / 3 - 1) self.update_scalar("en_U", en_U) @@ -139,22 +145,23 @@ def update_scalar_quantities(self): en_Btot = self._tmp_b1.inner(self._tmp_b2) / 2 self.update_scalar("en_B_tot", en_Btot) - + ## default parameters - def generate_default_parameter_file(self, path = None, prompt = True): + def generate_default_parameter_file(self, path=None, prompt=True): params_path = super().generate_default_parameter_file(path=path, prompt=prompt) new_file = [] with open(params_path, "r") as f: for line in f: if "mag_sonic.Options" in line: - new_file += ["model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n"] + new_file += [ + "model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n" + ] else: new_file += [line] - + with open(params_path, "w") as f: for line in new_file: f.write(line) - class LinearExtendedMHDuniform(StruphyModel): diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 0c9adc8dd..0c19aadef 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -1,14 +1,14 @@ import numpy as np from mpi4py import MPI +from struphy.feec.projectors import L2Projector from struphy.kinetic_background.base import KineticBackground from struphy.models.base import StruphyModel +from struphy.models.species import FieldSpecies, FluidSpecies, KineticSpecies +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.pic.accumulation import accum_kernels, accum_kernels_gc -from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies -from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable from struphy.pic.accumulation.particles_to_grid import AccumulatorVector -from struphy.feec.projectors import L2Projector +from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers rank = MPI.COMM_WORLD.Get_rank() @@ -83,20 +83,20 @@ class VlasovAmpereOneSpecies(StruphyModel): """ ## species - + class EMFields(FieldSpecies): def __init__(self): self.e_field = FEECVariable(space="Hcurl") self.phi = FEECVariable(space="H1") self.init_variables() - + class KineticIons(KineticSpecies): def __init__(self): self.var = PICVariable(space="Particles6D") self.init_variables() - + ## propagators - + class Propagators: def __init__(self, with_B0: bool = True): self.push_eta = propagators_markers.PushEta() @@ -109,28 +109,28 @@ def __init__(self, with_B0: bool = True): def __init__(self, with_B0: bool = True): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - + self.with_B0 = with_B0 - + # 1. instantiate all species self.em_fields = self.EMFields() self.kinetic_ions = self.KineticIons() # 2. instantiate all propagators self.propagators = self.Propagators(with_B0=with_B0) - + # 3. assign variables to propagators self.propagators.push_eta.variables.var = self.kinetic_ions.var if with_B0: self.propagators.push_vxb.variables.ions = self.kinetic_ions.var self.propagators.coupling_va.variables.e = self.em_fields.e_field self.propagators.coupling_va.variables.ions = self.kinetic_ions.var - + # define scalars for update_scalar_quantities self.add_scalar("en_E") self.add_scalar("en_f", compute="from_particles", variable=self.kinetic_ions.var) self.add_scalar("en_tot") - + # initial Poisson (not a propagator used in time stepping) self.initial_poisson = propagators_fields.Poisson() self.initial_poisson.variables.phi = self.em_fields.phi @@ -160,8 +160,8 @@ def update_scalar_quantities(self): / (2 * particles.Np) * np.dot( particles.markers_wo_holes[:, 3] ** 2 - + particles.markers_wo_holes[:, 4] ** 2 - + particles.markers_wo_holes[:, 5] ** 2, + + particles.markers_wo_holes[:, 4] ** 2 + + particles.markers_wo_holes[:, 5] ** 2, particles.markers_wo_holes[:, 6], ) ) @@ -206,13 +206,13 @@ def allocate_propagators(self): alpha = self.kinetic_ions.equation_params.alpha epsilon = self.kinetic_ions.equation_params.epsilon - - l2_proj = L2Projector(space_id='H1', mass_ops=self.mass_ops) + + l2_proj = L2Projector(space_id="H1", mass_ops=self.mass_ops) rho_coeffs = l2_proj.solve(charge_accum.vectors[0]) - + self.initial_poisson.options.rho = alpha**2 / epsilon * rho_coeffs self.initial_poisson.allocate() - + # Solve with dt=1. and compute electric field if MPI.COMM_WORLD.Get_rank() == 0: print("\nSolving initial Poisson problem...") @@ -224,7 +224,7 @@ def allocate_propagators(self): print("Done.") ## default parameters - def generate_default_parameter_file(self, path = None, prompt = True): + def generate_default_parameter_file(self, path=None, prompt=True): params_path = super().generate_default_parameter_file(path=path, prompt=prompt) new_file = [] with open(params_path, "r") as f: @@ -240,11 +240,12 @@ def generate_default_parameter_file(self, path = None, prompt = True): new_file += ["model.kinetic_ions.set_save_data(binning_plots=(binplot,))\n"] else: new_file += [line] - + with open(params_path, "w") as f: for line in new_file: f.write(line) + class VlasovMaxwellOneSpecies(StruphyModel): r"""Vlasov-Maxwell equations for one species. diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index ca8ef48c4..814554465 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,30 +1,32 @@ +import warnings from abc import ABCMeta, abstractmethod -from dataclasses import dataclass from copy import deepcopy +from dataclasses import dataclass from typing import Callable + import numpy as np from mpi4py import MPI -import warnings from struphy.fields_background.base import FluidEquilibrium -from struphy.kinetic_background.base import KineticBackground from struphy.io.options import Units -from struphy.physics.physics import ConstantsOfNature +from struphy.kinetic_background.base import KineticBackground from struphy.models.variables import Variable -from struphy.pic.utilities import (LoadingParameters, - WeightsParameters, - BoundaryParameters, - BinningPlot, - ) +from struphy.physics.physics import ConstantsOfNature +from struphy.pic.utilities import ( + BinningPlot, + BoundaryParameters, + LoadingParameters, + WeightsParameters, +) class Species(metaclass=ABCMeta): """Single species of a StruphyModel.""" - + @abstractmethod def __init__(self): self.init_variables() - + # set species attribute for each variable def init_variables(self): self._variables = {} @@ -32,72 +34,74 @@ def init_variables(self): if isinstance(v, Variable): v._name = k v._species = self - self._variables[k] = v - + self._variables[k] = v + @property def variables(self) -> dict: return self._variables - + @property def charge_number(self) -> int: """Charge number in units of elementary charge.""" return self._charge_number - + @property def mass_number(self) -> int: """Mass number in units of proton mass.""" return self._mass_number - def set_phys_params(self, - charge_number: int = 1, - mass_number: int = 1, - alpha: float = None, - epsilon: float = None, - kappa: float = None, - ): + def set_phys_params( + self, + charge_number: int = 1, + mass_number: int = 1, + alpha: float = None, + epsilon: float = None, + kappa: float = None, + ): """Set charge- and mass number. Set equation parameters (alpha, epsilon, ...) to override units.""" self._charge_number = charge_number self._mass_number = mass_number self.alpha = alpha self.epsilon = epsilon - self.kappa = kappa - + self.kappa = kappa + class EquationParameters: """Normalization parameters of one species, appearing in scaled equations.""" - - def __init__(self, - species, - units: Units = None, - alpha: float = None, - epsilon: float = None, - kappa: float = None, - verbose: bool = False, - ): + + def __init__( + self, + species, + units: Units = None, + alpha: float = None, + epsilon: float = None, + kappa: float = None, + verbose: bool = False, + ): if units is None: units = Units() - + Z = species.charge_number A = species.mass_number - + con = ConstantsOfNature() # relevant frequencies om_p = np.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) - + # compute equation parameters if alpha is None: self.alpha = om_p / om_c else: self.alpha = alpha warnings.warn(f"Override equation parameter {self.alpha = }") - + if epsilon is None: self.epsilon = 1.0 / (om_c * units.t) else: self.epsilon = epsilon warnings.warn(f"Override equation parameter {self.epsilon = }") - + if kappa is None: self.kappa = om_p * units.t else: @@ -105,29 +109,29 @@ def __init__(self, warnings.warn(f"Override equation parameter {self.kappa = }") if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f'\nSet normalization parameters for species {species.__class__.__name__}:') + print(f"\nSet normalization parameters for species {species.__class__.__name__}:") for key, val in self.__dict__.items(): print((key + ":").ljust(25), "{:4.3e}".format(val)) - + @property def equation_params(self) -> EquationParameters: return self._equation_params - + def setup_equation_params(self, units: Units, verbose=False): """Set the following equation parameters: - + * alpha = plasma-frequenca / cyclotron frequency * epsilon = 1 / (cyclotron frequency * time unit) * kappa = plasma frequency * time unit """ - self._equation_params = self.EquationParameters(species=self, - units=units, - alpha=self.alpha, - epsilon=self.epsilon, - kappa=self.kappa, - verbose=verbose,) - - + self._equation_params = self.EquationParameters( + species=self, + units=units, + alpha=self.alpha, + epsilon=self.epsilon, + kappa=self.kappa, + verbose=verbose, + ) class FieldSpecies(Species): @@ -136,74 +140,77 @@ class FieldSpecies(Species): class FluidSpecies(Species): """Single fluid species in 3d configuration space.""" + pass class KineticSpecies(Species): """Single kinetic species in 3d + vdim phase space.""" - - def set_markers(self, - loading_params: LoadingParameters = None, - weights_params: WeightsParameters = None, - boundary_params: BoundaryParameters = None, - bufsize: float = 1.0, - ): + + def set_markers( + self, + loading_params: LoadingParameters = None, + weights_params: WeightsParameters = None, + boundary_params: BoundaryParameters = None, + bufsize: float = 1.0, + ): """Set marker parameters for loading, weight calculation, kernel density reconstruction and boundary conditions. - + Parameters ---------- loading_params : LoadingParameters - + weights_params : WeightsParameters - + boundary_params : BoundaryParameters - + bufsize : float Size of buffer (as multiple of total size, default=.25) in markers array.""" - + # defaults if loading_params is None: loading_params = LoadingParameters() - + if weights_params is None: weights_params = WeightsParameters() - + if boundary_params is None: boundary_params = BoundaryParameters() - + self.loading_params = loading_params self.weights_params = weights_params self.boundary_params = boundary_params self.bufsize = bufsize - - def set_sorting_boxes(self, - do_sort: bool = False, - sorting_frequency: int = 0, - boxes_per_dim: tuple = (16, 1, 1), - box_bufsize: float = 2.0, - dims_maks: tuple = (True, True, True), - ): + + def set_sorting_boxes( + self, + do_sort: bool = False, + sorting_frequency: int = 0, + boxes_per_dim: tuple = (16, 1, 1), + box_bufsize: float = 2.0, + dims_maks: tuple = (True, True, True), + ): """For sorting markers in memory.""" self.do_sort = do_sort self.sorting_fequency = sorting_frequency self.boxes_per_dim = boxes_per_dim self.box_bufsize = box_bufsize self.dims_mask = dims_maks - - def set_save_data(self, - n_markers: int | float = 3, - binning_plots: tuple[BinningPlot] = (), - n_sph: dict = None, - ): + + def set_save_data( + self, + n_markers: int | float = 3, + binning_plots: tuple[BinningPlot] = (), + n_sph: dict = None, + ): """Saving marker orits, binned data and kernel density reconstructions.""" self.n_markers = n_markers self.binning_plots = binning_plots - self.n_sph = n_sph + self.n_sph = n_sph class DiagnosticSpecies(Species): """Diagnostic species (fields) without mass and charge.""" - pass - + pass diff --git a/src/struphy/models/tests/test_LinearMHD.py b/src/struphy/models/tests/test_LinearMHD.py index 5b994fd64..8ff0d28b8 100644 --- a/src/struphy/models/tests/test_LinearMHD.py +++ b/src/struphy/models/tests/test_LinearMHD.py @@ -1,26 +1,27 @@ -from mpi4py import MPI import os + import numpy as np import pytest +from mpi4py import MPI -from struphy.io.options import EnvironmentOptions, BaseUnits, Time -from struphy.geometry import domains +from struphy import main +from struphy.diagnostics.diagn_tools import power_spectrum_2d from struphy.fields_background import equils -from struphy.topology import grids -from struphy.io.options import DerhamOptions, FieldsBackground +from struphy.geometry import domains from struphy.initial import perturbations +from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time from struphy.kinetic_background import maxwellians -from struphy import main -from struphy.diagnostics.diagn_tools import power_spectrum_2d +from struphy.topology import grids test_folder = os.path.join(os.getcwd(), "verification_tests") @pytest.mark.mpi(min_size=3) -@pytest.mark.parametrize('algo', ["implicit", "explicit"]) +@pytest.mark.parametrize("algo", ["implicit", "explicit"]) def test_slab_waves_1d(algo: str, do_plot: bool = False): # import model, set verbosity from struphy.models.fluid import LinearMHD + verbose = True # environment options @@ -52,7 +53,7 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): # light-weight model instance model = LinearMHD() - + # species parameters model.mhd.set_phys_params() @@ -66,90 +67,88 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): model.mhd.velocity.add_perturbation(perturbations.Noise(amp=0.1, comp=2, seed=123)) # start run - main.run(model, - params_path=None, - env=env, - base_units=base_units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=verbose, - ) - + main.run( + model, + params_path=None, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) + # post processing if MPI.COMM_WORLD.Get_rank() == 0: main.pproc(env.path_out) - + # diagnostics if MPI.COMM_WORLD.Get_rank() == 0: simdata = main.load_data(env.path_out) # first fft u_of_t = simdata.spline_values["mhd"]["velocity_log"] - - Bsquare = (B0x**2 + B0y**2 + B0z**2) + + Bsquare = B0x**2 + B0y**2 + B0z**2 p0 = beta * Bsquare / 2 - disp_params = {'B0x': B0x, - 'B0y': B0y, - 'B0z': B0z, - 'p0': p0, - 'n0': n0, - 'gamma': 5/3} - - _1, _2, _3, coeffs = power_spectrum_2d(u_of_t, - "velocity_log", - grids=simdata.grids_log, - grids_mapped=simdata.grids_phy, - component=0, - slice_at=[0, 0, None], - do_plot=do_plot, - disp_name='MHDhomogenSlab', - disp_params=disp_params, - fit_branches=1, - noise_level=0.5, - extr_order=10, - fit_degree=(1,), + disp_params = {"B0x": B0x, "B0y": B0y, "B0z": B0z, "p0": p0, "n0": n0, "gamma": 5 / 3} + + _1, _2, _3, coeffs = power_spectrum_2d( + u_of_t, + "velocity_log", + grids=simdata.grids_log, + grids_mapped=simdata.grids_phy, + component=0, + slice_at=[0, 0, None], + do_plot=do_plot, + disp_name="MHDhomogenSlab", + disp_params=disp_params, + fit_branches=1, + noise_level=0.5, + extr_order=10, + fit_degree=(1,), ) - + # assert vA = np.sqrt(Bsquare / n0) v_alfven = vA * B0z / np.sqrt(Bsquare) print(f"{v_alfven = }") assert np.abs(coeffs[0][0] - v_alfven) < 0.07 - + # second fft p_of_t = simdata.spline_values["mhd"]["pressure_log"] - - _1, _2, _3, coeffs = power_spectrum_2d(p_of_t, - "pressure_log", - grids=simdata.grids_log, - grids_mapped=simdata.grids_phy, - component=0, - slice_at=[0, 0, None], - do_plot=do_plot, - disp_name='MHDhomogenSlab', - disp_params=disp_params, - fit_branches=2, - noise_level=0.4, - extr_order=10, - fit_degree=(1, 1), + + _1, _2, _3, coeffs = power_spectrum_2d( + p_of_t, + "pressure_log", + grids=simdata.grids_log, + grids_mapped=simdata.grids_phy, + component=0, + slice_at=[0, 0, None], + do_plot=do_plot, + disp_name="MHDhomogenSlab", + disp_params=disp_params, + fit_branches=2, + noise_level=0.4, + extr_order=10, + fit_degree=(1, 1), ) - + # assert - gamma = 5/3 + gamma = 5 / 3 cS = np.sqrt(gamma * p0 / n0) - - delta = (4 * B0z ** 2 * cS**2 * vA**2) / ((cS**2 + vA**2) ** 2 * Bsquare) + + delta = (4 * B0z**2 * cS**2 * vA**2) / ((cS**2 + vA**2) ** 2 * Bsquare) v_slow = np.sqrt(1 / 2 * (cS**2 + vA**2) * (1 - np.sqrt(1 - delta))) v_fast = np.sqrt(1 / 2 * (cS**2 + vA**2) * (1 + np.sqrt(1 - delta))) print(f"{v_slow = }") print(f"{v_fast = }") assert np.abs(coeffs[0][0] - v_slow) < 0.05 assert np.abs(coeffs[1][0] - v_fast) < 0.19 - - -if __name__ == '__main__': - test_slab_waves_1d(algo="implicit", do_plot=True) \ No newline at end of file + + +if __name__ == "__main__": + test_slab_waves_1d(algo="implicit", do_plot=True) diff --git a/src/struphy/models/tests/test_Maxwell.py b/src/struphy/models/tests/test_Maxwell.py index 297092f4f..f191a93c4 100644 --- a/src/struphy/models/tests/test_Maxwell.py +++ b/src/struphy/models/tests/test_Maxwell.py @@ -1,28 +1,29 @@ -from mpi4py import MPI import os + import numpy as np import pytest -from scipy.special import jv, yn from matplotlib import pyplot as plt +from mpi4py import MPI +from scipy.special import jv, yn -from struphy.io.options import EnvironmentOptions, BaseUnits, Time -from struphy.geometry import domains +from struphy import main +from struphy.diagnostics.diagn_tools import power_spectrum_2d from struphy.fields_background import equils -from struphy.topology import grids -from struphy.io.options import DerhamOptions, FieldsBackground +from struphy.geometry import domains from struphy.initial import perturbations +from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time from struphy.kinetic_background import maxwellians -from struphy import main -from struphy.diagnostics.diagn_tools import power_spectrum_2d +from struphy.topology import grids test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") @pytest.mark.mpi(min_size=3) -@pytest.mark.parametrize('algo', ["implicit", "explicit"]) +@pytest.mark.parametrize("algo", ["implicit", "explicit"]) def test_light_wave_1d(algo: str, do_plot: bool = False): # import model, set verbosity from struphy.models.toy import Maxwell + verbose = True # environment options @@ -58,51 +59,54 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1, seed=123)) # # start run - main.run(model, - params_path=None, - env=env, - base_units=base_units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=verbose, - ) - + main.run( + model, + params_path=None, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) + # post processing if MPI.COMM_WORLD.Get_rank() == 0: main.pproc(env.path_out) - + # diagnostics if MPI.COMM_WORLD.Get_rank() == 0: simdata = main.load_data(env.path_out) - # fft + # fft E_of_t = simdata.spline_values["em_fields"]["e_field_log"] - _1, _2, _3, coeffs = power_spectrum_2d(E_of_t, - "e_field_log", - grids=simdata.grids_log, - grids_mapped=simdata.grids_phy, - component=0, - slice_at=[0, 0, None], - do_plot=do_plot, - disp_name='Maxwell1D', - fit_branches=1, - noise_level=0.5, - extr_order=10, - fit_degree=(1,), + _1, _2, _3, coeffs = power_spectrum_2d( + E_of_t, + "e_field_log", + grids=simdata.grids_log, + grids_mapped=simdata.grids_phy, + component=0, + slice_at=[0, 0, None], + do_plot=do_plot, + disp_name="Maxwell1D", + fit_branches=1, + noise_level=0.5, + extr_order=10, + fit_degree=(1,), ) - + # assert c_light_speed = 1.0 assert np.abs(coeffs[0][0] - c_light_speed) < 0.02 - - + + @pytest.mark.mpi(min_size=4) def test_coaxial(do_plot: bool = False): - # import model, set verbosity + # import model, set verbosity from struphy.models.toy import Maxwell + verbose = True # environment options @@ -114,7 +118,7 @@ def test_coaxial(do_plot: bool = False): # time time_opts = Time(dt=0.05, Tend=10.0) - + # geometry a1 = 2.326744 a2 = 3.686839 @@ -128,7 +132,8 @@ def test_coaxial(do_plot: bool = False): grid = grids.TensorProductGrid(Nel=(32, 64, 1)) # derham options - derham_opts = DerhamOptions(p=(3, 3, 1), + derham_opts = DerhamOptions( + p=(3, 3, 1), spl_kind=(False, True, True), dirichlet_bc=((True, True), (False, False), (False, False)), ) @@ -146,22 +151,23 @@ def test_coaxial(do_plot: bool = False): model.em_fields.b_field.add_perturbation(perturbations.CoaxialWaveguideMagnetic(m=m, a1=a1, a2=a2)) # start run - main.run(model, - params_path=None, - env=env, - base_units=base_units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=verbose, - ) - + main.run( + model, + params_path=None, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) + # post processing if MPI.COMM_WORLD.Get_rank() == 0: main.pproc(env.path_out, physical=True) - + # diagnostics if MPI.COMM_WORLD.Get_rank() == 0: # get parameters @@ -169,10 +175,10 @@ def test_coaxial(do_plot: bool = False): split_algo = time_opts.split_algo Nel = grid.Nel modes = m - + # load data simdata = main.load_data(env.path_out) - + t_grid = simdata.t_grid grids_phy = simdata.grids_phy e_field_phy = simdata.spline_values["em_fields"]["e_field_phy"] @@ -198,7 +204,9 @@ def E_theta(X, Y, Z, m, t): """Electrical field in azimuthal direction of coaxial cabel""" r = (X**2 + Y**2) ** 0.5 theta = np.arctan2(Y, X) - return ((m / r * jv(m, r) - jv(m + 1, r)) - 0.28 * (m / r * yn(m, r) - yn(m + 1, r))) * np.sin(m * theta - t) + return ((m / r * jv(m, r) - jv(m + 1, r)) - 0.28 * (m / r * yn(m, r) - yn(m + 1, r))) * np.sin( + m * theta - t + ) def to_E_r(X, Y, E_x, E_y): r = (X**2 + Y**2) ** 0.5 @@ -209,7 +217,7 @@ def to_E_theta(X, Y, E_x, E_y): r = (X**2 + Y**2) ** 0.5 theta = np.arctan2(Y, X) return -np.sin(theta) * E_x + np.cos(theta) * E_y - + # plot if do_plot: vmin = E_theta(X, Y, grids_phy[0], modes, 0).min() @@ -232,7 +240,7 @@ def to_E_theta(X, Y, E_x, E_y): ax2.set_xlabel("Numerical") fig.suptitle(f"Exact and Simulated $E_\\theta$ Field {dt=}, {split_algo=}, {Nel=}", fontsize=14) plt.show() - + # assert Ex_tend = e_field_phy[t_grid[-1]][0][:, :, 0] Ey_tend = e_field_phy[t_grid[-1]][1][:, :, 0] @@ -252,15 +260,12 @@ def to_E_theta(X, Y, E_x, E_y): print("") assert rel_err_Bz < 0.0021, f"Assertion for magnetic field Maxwell failed: {rel_err_Bz = }" print(f"Assertion for magnetic field Maxwell passed ({rel_err_Bz = }).") - assert rel_err_Etheta < 0.0021, ( - f"Assertion for electric (E_theta) field Maxwell failed: {rel_err_Etheta = }" - ) + assert rel_err_Etheta < 0.0021, f"Assertion for electric (E_theta) field Maxwell failed: {rel_err_Etheta = }" print(f"Assertion for electric field Maxwell passed ({rel_err_Etheta = }).") assert rel_err_Er < 0.0021, f"Assertion for electric (E_r) field Maxwell failed: {rel_err_Er = }" print(f"Assertion for electric field Maxwell passed ({rel_err_Er = }).") - - + if __name__ == "__main__": # test_light_wave_1d(algo="explicit", do_plot=True) - test_coaxial(do_plot=True) \ No newline at end of file + test_coaxial(do_plot=True) diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index 1487679ee..c0c455799 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -1,14 +1,14 @@ -import pytest import inspect import os from types import ModuleType -from struphy.models import toy, fluid, kinetic, hybrid -from struphy.models.base import StruphyModel -from struphy.io.setup import import_parameters_py -from struphy.io.options import EnvironmentOptions -from struphy import main +import pytest +from struphy import main +from struphy.io.options import EnvironmentOptions +from struphy.io.setup import import_parameters_py +from struphy.models import fluid, hybrid, kinetic, toy +from struphy.models.base import StruphyModel # available models toy_models = [] @@ -41,20 +41,20 @@ # generic function for calling model tests -def call_test(model_name: str, module: ModuleType, verbose=True): +def call_test(model_name: str, module: ModuleType, verbose=True): print(f"\n*** Testing '{model_name}':") model = getattr(module, model_name)() assert isinstance(model, StruphyModel) - + # generate paramater file for testing path = os.path.join(test_folder, f"params_{model_name}.py") model.generate_default_parameter_file(path=path, prompt=False) del model print("\nDeleting light-weight instance ...") - + # set environment options env = EnvironmentOptions(out_folders=test_folder, sim_folder=f"{model_name}_test") - + # read parameters params_in = import_parameters_py(path) units = params_in.units @@ -63,45 +63,44 @@ def call_test(model_name: str, module: ModuleType, verbose=True): equil = params_in.equil grid = params_in.grid derham_opts = params_in.derham_opts - + # test model = params_in.model - main.run(model, - params_path=path, - env=env, - units=units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - verbose=verbose, - ) + main.run( + model, + params_path=path, + env=env, + units=units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) # specific tests -@pytest.mark.parametrize('model_name', toy_models) +@pytest.mark.parametrize("model_name", toy_models) def test_toy(model_name: str, verbose=True): call_test(model_name=model_name, module=toy, verbose=verbose) - -@pytest.mark.parametrize('model_name', fluid_models) + + +@pytest.mark.parametrize("model_name", fluid_models) def test_fluid(model_name: str, verbose=True): call_test(model_name=model_name, module=fluid, verbose=verbose) - -@pytest.mark.parametrize('model_name', kinetic_models) + + +@pytest.mark.parametrize("model_name", kinetic_models) def test_kinetic(model_name: str, verbose=True): call_test(model_name=model_name, module=kinetic, verbose=verbose) - -@pytest.mark.parametrize('model_name', hybrid_models) + + +@pytest.mark.parametrize("model_name", hybrid_models) def test_hybrid(model_name: str, verbose=True): call_test(model_name=model_name, module=hybrid, verbose=verbose) - - -if __name__ == '__main__': + +if __name__ == "__main__": test_toy("Maxwell") test_fluid("LinearMHD") - - - - \ No newline at end of file diff --git a/src/struphy/models/tests/util.py b/src/struphy/models/tests/util.py index bbd330088..18dcd8086 100644 --- a/src/struphy/models/tests/util.py +++ b/src/struphy/models/tests/util.py @@ -9,8 +9,8 @@ import struphy from struphy.console.main import recursive_get_files from struphy.io.setup import descend_options_dict -from struphy.struphy import run from struphy.models.base import StruphyModel +from struphy.struphy import run libpath = struphy.__path__[0] diff --git a/src/struphy/models/tests/verification.py b/src/struphy/models/tests/verification.py index f70ef7402..c88e2ff32 100644 --- a/src/struphy/models/tests/verification.py +++ b/src/struphy/models/tests/verification.py @@ -1,7 +1,6 @@ import os import pickle from pathlib import Path -from mpi4py import MPI import h5py import numpy as np diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 0a346b4e3..ce1bf9035 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,13 +1,13 @@ +from dataclasses import dataclass import numpy as np -from dataclasses import dataclass from mpi4py import MPI from struphy.models.base import StruphyModel -from struphy.propagators.base import Propagator +from struphy.models.species import FieldSpecies, FluidSpecies, KineticSpecies +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.models.species import KineticSpecies, FluidSpecies, FieldSpecies -from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable +from struphy.propagators.base import Propagator rank = MPI.COMM_WORLD.Get_rank() @@ -33,43 +33,43 @@ class Maxwell(StruphyModel): 1. :class:`~struphy.propagators.propagators_fields.Maxwell` """ - + ## species - + class EMFields(FieldSpecies): def __init__(self): self.e_field = FEECVariable(space="Hcurl") self.b_field = FEECVariable(space="Hdiv") self.init_variables() - + ## propagators - + class Propagators: def __init__(self): self.maxwell = propagators_fields.Maxwell() ## abstract methods - def __init__(self): + def __init__(self): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - + # 1. instantiate all species self.em_fields = self.EMFields() # 2. instantiate all propagators self.propagators = self.Propagators() - + # 3. assign variables to propagators self.propagators.maxwell.variables.e = self.em_fields.e_field self.propagators.maxwell.variables.b = self.em_fields.b_field - + # define scalars for update_scalar_quantities self.add_scalar("electric energy") self.add_scalar("magnetic energy") self.add_scalar("total energy") - - @property + + @property def bulk_species(self): return None @@ -81,13 +81,17 @@ def allocate_helpers(self): pass def update_scalar_quantities(self): - en_E = 0.5 * self.mass_ops.M1.dot_inner(self.em_fields.e_field.spline.vector, self.em_fields.e_field.spline.vector) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.em_fields.b_field.spline.vector, self.em_fields.b_field.spline.vector) + en_E = 0.5 * self.mass_ops.M1.dot_inner( + self.em_fields.e_field.spline.vector, self.em_fields.e_field.spline.vector + ) + en_B = 0.5 * self.mass_ops.M2.dot_inner( + self.em_fields.b_field.spline.vector, self.em_fields.b_field.spline.vector + ) self.update_scalar("electric energy", en_E) self.update_scalar("magnetic energy", en_B) self.update_scalar("total energy", en_E + en_B) - + class Vlasov(StruphyModel): r"""Vlasov equation in static background magnetic field. @@ -109,16 +113,16 @@ class Vlasov(StruphyModel): 1. :class:`~struphy.propagators.propagators_markers.PushVxB` 2. :class:`~struphy.propagators.propagators_markers.PushEta` """ - + ## species - + class KineticIons(KineticSpecies): def __init__(self): self.var = PICVariable(space="Particles6D") self.init_variables() - + ## propagators - + class Propagators: def __init__(self): self.push_vxb = propagators_markers.PushVxB() @@ -129,17 +133,17 @@ def __init__(self): def __init__(self): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}' ***") - + # 1. instantiate all species self.kinetic_ions = self.KineticIons() # 2. instantiate all propagators self.propagators = self.Propagators() - + # 3. assign variables to propagators self.propagators.push_vxb.variables.ions = self.kinetic_ions.var self.propagators.push_eta.variables.var = self.kinetic_ions.var - + # define scalars for update_scalar_quantities self.add_scalar("en_f", compute="from_particles", variable=self.kinetic_ions.var) @@ -150,9 +154,9 @@ def bulk_species(self): @property def velocity_scale(self): return "cyclotron" - + def allocate_helpers(self): - self._tmp = np.empty(1, dtype=float) + self._tmp = np.empty(1, dtype=float) def update_scalar_quantities(self): particles = self.kinetic_ions.var.particles @@ -197,16 +201,16 @@ class GuidingCenter(StruphyModel): 1. :class:`~struphy.propagators.propagators_markers.PushGuidingCenterBxEstar` 2. :class:`~struphy.propagators.propagators_markers.PushGuidingCenterParallel` """ - + ## species - + class KineticIons(KineticSpecies): def __init__(self): self.var = PICVariable(space="Particles5D") self.init_variables() - + ## propagators - + class Propagators: def __init__(self): self.push_bxe = propagators_markers.PushGuidingCenterBxEstar() @@ -217,17 +221,17 @@ def __init__(self): def __init__(self): if rank == 0: print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}' ***") - + # 1. instantiate all species self.kinetic_ions = self.KineticIons() # 2. instantiate all propagators self.propagators = self.Propagators() - + # 3. assign variables to propagators self.propagators.push_bxe.variables.ions = self.kinetic_ions.var self.propagators.push_parallel.variables.ions = self.kinetic_ions.var - + # define scalars for update_scalar_quantities self.add_scalar("en_fv", compute="from_particles", variable=self.kinetic_ions.var) self.add_scalar("en_fB", compute="from_particles", variable=self.kinetic_ions.var) @@ -240,7 +244,7 @@ def bulk_species(self): @property def velocity_scale(self): return "alfvén" - + def allocate_helpers(self): self._en_fv = np.empty(1, dtype=float) self._en_fB = np.empty(1, dtype=float) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 1931115c0..ef4cdb6b2 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -1,63 +1,69 @@ - # for type checking (cyclic imports) from __future__ import annotations -from typing import TYPE_CHECKING from abc import ABCMeta, abstractmethod -from mpi4py import MPI +from typing import TYPE_CHECKING + import numpy as np +from mpi4py import MPI from struphy.feec.psydac_derham import Derham, SplineFunction -from struphy.io.options import (FieldsBackground, OptsFEECSpace, OptsPICSpace, check_option,) -from struphy.initial.perturbations import Perturbation -from struphy.geometry.base import Domain from struphy.fields_background.base import FluidEquilibrium from struphy.fields_background.projected_equils import ProjectedFluidEquilibrium -from struphy.pic.base import Particles +from struphy.geometry.base import Domain +from struphy.initial.perturbations import Perturbation +from struphy.io.options import ( + FieldsBackground, + OptsFEECSpace, + OptsPICSpace, + check_option, +) from struphy.kinetic_background.base import KineticBackground from struphy.pic import particles +from struphy.pic.base import Particles from struphy.utils.clone_config import CloneConfig if TYPE_CHECKING: - from struphy.models.species import Species, KineticSpecies, FieldSpecies, FluidSpecies + from struphy.models.species import FieldSpecies, FluidSpecies, KineticSpecies, Species + class Variable(metaclass=ABCMeta): """Single variable (unknown) of a Species.""" - + @abstractmethod def allocate(self): """Alocate object and memory for variable.""" - + @property def backgrounds(self): if not hasattr(self, "_backgrounds"): self._backgrounds = None return self._backgrounds - + @property def perturbations(self): if not hasattr(self, "_perturbations"): self._perturbations = None return self._perturbations - + @property def save_data(self): """Store variable data during simulation (default=True).""" if not hasattr(self, "_save_data"): self._save_data = True return self._save_data - + @save_data.setter def save_data(self, new): assert isinstance(new, bool) self._save_data = new - + @property def species(self) -> Species: if not hasattr(self, "_species"): self._species = None return self._species - + @property def __name__(self): if not hasattr(self, "_name"): @@ -72,11 +78,13 @@ def add_background(self, background, verbose=True): if not isinstance(self.backgrounds, list): self._backgrounds = [self.backgrounds] self._backgrounds += [background] - + if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added background '{background.__class__.__name__}' with:") + print( + f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added background '{background.__class__.__name__}' with:" + ) for k, v in background.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") def add_perturbation(self, perturbation: Perturbation, verbose=True): if not hasattr(self, "_perturbations") or self.perturbations is None: @@ -85,95 +93,101 @@ def add_perturbation(self, perturbation: Perturbation, verbose=True): if not isinstance(self.perturbations, list): self._perturbations = [self.perturbations] self._perturbations += [perturbation] - + if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added perturbation '{perturbation.__class__.__name__}' with:") + print( + f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added perturbation '{perturbation.__class__.__name__}' with:" + ) for k, v in perturbation.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") + - class FEECVariable(Variable): def __init__(self, space: OptsFEECSpace = "H1"): check_option(space, OptsFEECSpace) self._space = space - + @property def space(self): return self._space - + @property def spline(self) -> SplineFunction: return self._spline - + @property def species(self) -> FieldSpecies | FluidSpecies: if not hasattr(self, "_species"): self._species = None return self._species - + def add_background(self, background: FieldsBackground, verbose=True): super().add_background(background, verbose=verbose) - - def allocate(self, derham: Derham, domain: Domain = None, equil: FluidEquilibrium = None,): + + def allocate( + self, + derham: Derham, + domain: Domain = None, + equil: FluidEquilibrium = None, + ): self._spline = derham.create_spline_function( - name=self.__name__, - space_id=self.space, - backgrounds=self.backgrounds, - perturbations=self.perturbations, - domain=domain, - equil=equil, - ) - - + name=self.__name__, + space_id=self.space, + backgrounds=self.backgrounds, + perturbations=self.perturbations, + domain=domain, + equil=equil, + ) + + class PICVariable(Variable): def __init__(self, space: OptsPICSpace = "Particles6D"): check_option(space, OptsPICSpace) self._space = space self._kinetic_data = {} - + @property def space(self): return self._space - + @property def particles(self) -> Particles: return self._particles - + @property def kinetic_data(self): return self._kinetic_data - + @property def species(self) -> KineticSpecies: if not hasattr(self, "_species"): self._species = None return self._species - + @property def n_as_volume_form(self) -> bool: """Whether the number density n is given as a volume form or scalar function (=default).""" if not hasattr(self, "_n_as_volume_form"): self._n_as_volume_form = False return self._n_as_volume_form - - def add_background(self, - background: KineticBackground, - n_as_volume_form: bool = False, - verbose=True): + + def add_background(self, background: KineticBackground, n_as_volume_form: bool = False, verbose=True): self._n_as_volume_form = n_as_volume_form super().add_background(background, verbose=verbose) - - def allocate(self, - clone_config: CloneConfig = None, - derham: Derham = None, - domain: Domain = None, - equil: FluidEquilibrium = None, - projected_equil: ProjectedFluidEquilibrium = None, - verbose: bool = False, - ): - - #assert isinstance(self.species, KineticSpecies) - assert isinstance(self.backgrounds, KineticBackground), f"List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." + + def allocate( + self, + clone_config: CloneConfig = None, + derham: Derham = None, + domain: Domain = None, + equil: FluidEquilibrium = None, + projected_equil: ProjectedFluidEquilibrium = None, + verbose: bool = False, + ): + # assert isinstance(self.species, KineticSpecies) + assert isinstance(self.backgrounds, KineticBackground), ( + f"List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." + ) if derham is None: domain_decomp = None @@ -209,7 +223,7 @@ def allocate(self, equation_params=self.species.equation_params, verbose=verbose, ) - + if self.species.do_sort: sort = True else: @@ -221,12 +235,12 @@ def allocate(self, self.kinetic_data["bin_edges"] = {} self.kinetic_data["f"] = {} self.kinetic_data["df"] = {} - + for bin_plot in self.species.binning_plots: sli = bin_plot.slice n_bins = bin_plot.n_bins ranges = bin_plot.ranges - + assert ((len(sli) - 2) / 3).is_integer(), f"Binning coordinates must be separated by '_', but reads {sli}." assert len(sli.split("_")) == len(ranges) == len(n_bins), ( f"Number of slices names ({len(sli.split('_'))}), number of bins ({len(n_bins)}), and number of ranges ({len(ranges)}) are inconsistent with each other!\n\n" @@ -265,8 +279,8 @@ def allocate(self, self.kinetic_data["n_sph"] += [np.zeros(ee1.shape, dtype=float)] # other data (wave-particle power exchange, etc.) - # TODO - - + # TODO + + class SPHVariable(Variable): - pass \ No newline at end of file + pass diff --git a/src/struphy/ode/utils.py b/src/struphy/ode/utils.py index 14af9b848..83300ae58 100644 --- a/src/struphy/ode/utils.py +++ b/src/struphy/ode/utils.py @@ -1,16 +1,16 @@ -import numpy as np from dataclasses import dataclass from typing import Literal, get_args +import numpy as np OptsButcher = Literal[ - "rk4", - "forward_euler", - "heun2", - "rk2", - "heun3", - "3/8 rule", - ] + "rk4", + "forward_euler", + "heun2", + "rk2", + "heun3", + "3/8 rule", +] @dataclass @@ -29,7 +29,7 @@ class ButcherTableau: algo : OptsButcher Name of the RK method. """ - + algo: OptsButcher = "rk4" def __post_init__(self): diff --git a/src/struphy/physics/physics.py b/src/struphy/physics/physics.py index ce2c4eced..898e9a45e 100644 --- a/src/struphy/physics/physics.py +++ b/src/struphy/physics/physics.py @@ -1,5 +1,6 @@ from dataclasses import dataclass + @dataclass class ConstantsOfNature: e = 1.602176634e-19 # elementary charge (C) @@ -7,4 +8,4 @@ class ConstantsOfNature: mu0 = 1.25663706212e-6 # magnetic constant (N/A^2) eps0 = 8.8541878128e-12 # vacuum permittivity (F/m) kB = 1.380649e-23 # Boltzmann constant (J/K) - c = 299792458 # speed of light (m/s) \ No newline at end of file + c = 299792458 # speed of light (m/s) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 76dec4d01..6bcacaad3 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,11 +2,11 @@ import os import warnings from abc import ABCMeta, abstractmethod -from line_profiler import profile import h5py import numpy as np import scipy.special as sp +from line_profiler import profile from mpi4py import MPI from mpi4py.MPI import Intracomm from sympy.ntheory import factorint @@ -19,9 +19,10 @@ from struphy.geometry.base import Domain from struphy.geometry.utilities import TransformedPformComponent from struphy.initial.base import Perturbation +from struphy.io.options import OptsLoading from struphy.io.output_handling import DataContainer -from struphy.kinetic_background.base import KineticBackground, Maxwellian from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments +from struphy.kinetic_background.base import KineticBackground, Maxwellian from struphy.pic import sampling_kernels, sobol_seq from struphy.pic.pushing.pusher_utilities_kernels import reflect from struphy.pic.sorting_kernels import ( @@ -38,13 +39,13 @@ naive_evaluation_flat, naive_evaluation_meshgrid, ) +from struphy.pic.utilities import ( + BoundaryParameters, + LoadingParameters, + WeightsParameters, +) from struphy.utils import utils from struphy.utils.clone_config import CloneConfig -from struphy.pic.utilities import (LoadingParameters, - WeightsParameters, - BoundaryParameters, - ) -from struphy.io.options import OptsLoading class Particles(metaclass=ABCMeta): @@ -100,7 +101,7 @@ class Particles(metaclass=ABCMeta): weights_params : WeightsParameters Parameters for particle weights. - + boundary_params : BoundaryParameters Parameters for particle boundary conditions. @@ -118,7 +119,7 @@ class Particles(metaclass=ABCMeta): background : KineticBackground Kinetic background parameters. - + n_as_volume_form: bool Whether the number density n is given as a volume form or scalar function (=default). @@ -173,14 +174,14 @@ def __init__( self._equil = equil self._projected_equil = projected_equil self._equation_params = equation_params - + # defaults if loading_params is None: loading_params = LoadingParameters() - + if weights_params is None: weights_params = WeightsParameters() - + if boundary_params is None: boundary_params = BoundaryParameters() @@ -292,8 +293,10 @@ def __init__( if isinstance(background, FluidEquilibrium): self._pforms = (False, False) else: - self._pforms = (n_as_volume_form, - self.background.volume_form,) + self._pforms = ( + n_as_volume_form, + self.background.volume_form, + ) # set background function self._set_background_function() @@ -499,11 +502,11 @@ def background(self) -> KineticBackground: @property def loading_params(self) -> LoadingParameters: return self._loading_params - + @property def weights_params(self) -> WeightsParameters: return self._weights_params - + @property def boundary_params(self) -> BoundaryParameters: """Parameters for marker loading.""" @@ -513,7 +516,7 @@ def boundary_params(self) -> BoundaryParameters: def reject_weights(self): """Whether to reect weights below threshold.""" return self._reject_weights - + @property def threshold(self): """Threshold for rejecting weights.""" @@ -1106,12 +1109,12 @@ def _initialize_sorting_boxes(self): def _generate_sampling_moments(self): """Automatically determine moments for sampling distribution (Gaussian) from the given background.""" - + if self.loading_params.moments is None: - self.loading_params.moments = tuple([0.0]*self.vdim + [1.0]*self.vdim) - + self.loading_params.moments = tuple([0.0] * self.vdim + [1.0] * self.vdim) + # TODO: reformulate this function with KineticBackground methods - + # ns = [] # us = [] # vths = [] diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 7a7da2054..1a8560375 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -299,6 +299,7 @@ class Particles5D(Particles): **kwargs : dict Parameters for markers, see :class:`~struphy.pic.base.Particles`. """ + @classmethod def default_background(cls): return maxwellians.GyroMaxwellian2D() @@ -408,11 +409,11 @@ def svol(self, eta1, eta2, eta3, *v): } self._svol = maxwellians.GyroMaxwellian2D( - n = (1.0, None), - u_para = (self.loading_params.moments[0], None), - u_perp = (self.loading_params.moments[1], None), - vth_para = (self.loading_params.moments[2], None), - vth_perp = (self.loading_params.moments[3], None), + n=(1.0, None), + u_para=(self.loading_params.moments[0], None), + u_perp=(self.loading_params.moments[1], None), + vth_para=(self.loading_params.moments[2], None), + vth_perp=(self.loading_params.moments[3], None), volume_form=True, equil=self._magn_bckgr, ) diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index b1a5b47ff..bb48607df 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -1,8 +1,8 @@ "Accelerated particle pushing." import numpy as np -from mpi4py.MPI import IN_PLACE, SUM from line_profiler import profile +from mpi4py.MPI import IN_PLACE, SUM from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments from struphy.pic.base import Particles diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index 773281422..e66ab8753 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -1,15 +1,17 @@ import numpy as np import struphy.pic.utilities_kernels as utils -from struphy.io.options import (OptsLoading, - OptsSpatialLoading, - OptsMarkerBC, - OptsRecontructBC,) +from struphy.io.options import ( + OptsLoading, + OptsMarkerBC, + OptsRecontructBC, + OptsSpatialLoading, +) class LoadingParameters: """Parameters for particle loading. - + Parameters ---------- Np : int @@ -17,57 +19,58 @@ class LoadingParameters: ppc : int Particles to load per cell if a grid is defined. Cells are defined from ``domain_array``. - + ppb : int Particles to load per sorting box. Sorting boxes are defined from ``boxes_per_dim``. - + loading : OptsLoading How to load markers: multiple options for Monte-Carlo, or "tesselation" for positioning them on a regular grid. - + seed : int Seed for random generator. If None, no seed is taken. - + moments : tuple Mean velocities and temperatures for the Gaussian sampling distribution. If None, these are auto-calculated form the given background. - + spatial : OptsSpatialLoading Draw uniformly in eta, or draw uniformly on the "disc" image of (eta1, eta2). - + specific_markers : tuple[tuple] Each entry is a tuple of phase space coordinates (floats) of a specific marker to be initialized. - + n_quad : int Number of quadrature points for tesselation. - + dir_external : str Load markers from external .hdf5 file (absolute path). - + dir_particles_abs : str Load markers from restart .hdf5 file (absolute path). - + dir_particles : str Load markers from restart .hdf5 file (relative path to output folder). - + restart_key : str Key in .hdf5 file's restart/ folder where marker array is stored. """ - def __init__(self, - Np: int = None, - ppc: int = None, - ppb: int = 10, - loading: OptsLoading = "pseudo_random", - seed: int = None, - moments: tuple = None, - spatial: OptsSpatialLoading = "uniform", - specific_markers: tuple[tuple] = None, - n_quad: int = 1, - dir_exrernal: str = None, - dir_particles: str = None, - dir_particles_abs: str = None, - restart_key: str = None, - ): + def __init__( + self, + Np: int = None, + ppc: int = None, + ppb: int = 10, + loading: OptsLoading = "pseudo_random", + seed: int = None, + moments: tuple = None, + spatial: OptsSpatialLoading = "uniform", + specific_markers: tuple[tuple] = None, + n_quad: int = 1, + dir_exrernal: str = None, + dir_particles: str = None, + dir_particles_abs: str = None, + restart_key: str = None, + ): self.Np = Np self.ppc = ppc self.ppb = ppb @@ -81,27 +84,29 @@ def __init__(self, self.dir_particles = dir_particles self.dir_particles_abs = dir_particles_abs self.restart_key = restart_key - - + + class WeightsParameters: """Paramters for particle weights. - + Parameters ---------- control_variate : bool Whether to use a control variate for noise reduction. - + reject_weights : bool Whether to reject weights below threshold. - + threshold : float Threshold for rejecting weights. """ - def __init__(self, - control_variate: bool = False, - reject_weights: bool = False, - threshold: float = 0.0,): - + + def __init__( + self, + control_variate: bool = False, + reject_weights: bool = False, + threshold: float = 0.0, + ): self.control_variate = control_variate self.reject_weights = reject_weights self.threshold = threshold @@ -109,7 +114,7 @@ def __init__(self, class BoundaryParameters: """Parameters for particle boundary and sph reconstruction boundary conditions. - + Parameters ---------- bc : tuple[OptsMarkerBC] @@ -118,45 +123,50 @@ class BoundaryParameters: bc_refill : list Either 'inner' or 'outer'. - + bc_sph : tuple[OptsRecontructBC] Boundary conditions for sph kernel reconstruction. """ - def __init__(self, - bc: tuple[OptsMarkerBC] = ("periodic", "periodic", "periodic"), - bc_refill = None, - bc_sph: tuple[OptsRecontructBC] = ("periodic", "periodic", "periodic"), - ): + + def __init__( + self, + bc: tuple[OptsMarkerBC] = ("periodic", "periodic", "periodic"), + bc_refill=None, + bc_sph: tuple[OptsRecontructBC] = ("periodic", "periodic", "periodic"), + ): self.bc = bc self.bc_refill = bc_refill self.bc_sph = bc_sph - + class BinningPlot: """Binning plot of marker distribution in phase space. - + Parameters ---------- slice : str Coordinate-slice in phase space to bin. A combination of "e1", "e2", "e3", "v1", etc., separated by an underscore "_". For example, "e1" showas a 1D binning plot over eta1, whereas "e1_v1" shows a 2D binning plot over eta1 and v1. - + n_bins : int | tuple[int] Number of bins for each coordinate. - + ranges : tuple[int] | tuple[tuple[int]]= (0.0, 1.0) Binning range (as an interval in R) for each coordinate. """ - def __init__(self, - slice: str = "e1", - n_bins: int | tuple[int] = 128, - ranges: tuple[float] | tuple[tuple[float]]= (0.0, 1.0),): + + def __init__( + self, + slice: str = "e1", + n_bins: int | tuple[int] = 128, + ranges: tuple[float] | tuple[tuple[float]] = (0.0, 1.0), + ): self.slice = slice - + if isinstance(n_bins, int): - n_bins = (n_bins,) + n_bins = (n_bins,) self.n_bins = n_bins - + if not isinstance(ranges[0], tuple): ranges = (ranges,) self.ranges = ranges diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index c79b21ca1..14c9bb8dc 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -1,37 +1,41 @@ import os +import pickle import shutil + import h5py import numpy as np import yaml from tqdm import tqdm -import pickle -from struphy.kinetic_background import maxwellians -from struphy.fields_background.base import FluidEquilibrium +from struphy.feec.psydac_derham import SplineFunction from struphy.fields_background import equils +from struphy.fields_background.base import FluidEquilibrium +from struphy.geometry import domains +from struphy.geometry.base import Domain +from struphy.io.options import BaseUnits, EnvironmentOptions, Time from struphy.io.setup import import_parameters_py -from struphy.models.base import setup_derham, StruphyModel +from struphy.kinetic_background import maxwellians +from struphy.kinetic_background.base import KineticBackground +from struphy.models.base import StruphyModel, setup_derham from struphy.models.species import KineticSpecies from struphy.models.variables import PICVariable -from struphy.feec.psydac_derham import SplineFunction -from struphy.io.options import EnvironmentOptions, BaseUnits, Time from struphy.topology.grids import TensorProductGrid -from struphy.geometry import domains -from struphy.geometry.base import Domain -from struphy.kinetic_background.base import KineticBackground class ParamsIn: """Holds the input parameters of a Struphy simulation as attributes.""" - def __init__(self, - env: EnvironmentOptions = None, - base_units: BaseUnits = None, - time_opts: Time = None, - domain = None, - equil = None, - grid: TensorProductGrid = None, - derham_opts = None, - model: StruphyModel = None,): + + def __init__( + self, + env: EnvironmentOptions = None, + base_units: BaseUnits = None, + time_opts: Time = None, + domain=None, + equil=None, + grid: TensorProductGrid = None, + derham_opts=None, + model: StruphyModel = None, + ): self.env = env self.units = base_units self.time_opts = time_opts @@ -44,15 +48,15 @@ def __init__(self, def get_params_of_run(path: str) -> ParamsIn: """Retrieve parameters of finished Struphy run. - + Parameters ---------- path : str Absolute path of simulation output folder. """ - + print(f"\nReading in paramters from {path} ... ") - + params_path = os.path.join(path, "parameters.py") bin_path = os.path.join(path, "env.bin") @@ -66,7 +70,7 @@ def get_params_of_run(path: str) -> ParamsIn: grid = params_in.grid derham_opts = params_in.derham_opts model = params_in.model - + elif os.path.exists(bin_path): with open(os.path.join(path, "env.bin"), "rb") as f: env = pickle.load(f) @@ -83,7 +87,7 @@ def get_params_of_run(path: str) -> ParamsIn: equil_dct = pickle.load(f) if equil_dct: equil: FluidEquilibrium = getattr(equils, equil_dct["name"])(**equil_dct["params"]) - else: + else: equil = None with open(os.path.join(path, "grid.bin"), "rb") as f: grid = pickle.load(f) @@ -91,21 +95,22 @@ def get_params_of_run(path: str) -> ParamsIn: derham_opts = pickle.load(f) with open(os.path.join(path, "model.bin"), "rb") as f: model = pickle.load(f) - + else: raise FileNotFoundError(f"Neither of the paths {params_path} or {bin_path} exists.") - + print("done.") - - return ParamsIn(env=env, - base_units=base_units, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, - model=model, - ) + + return ParamsIn( + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + model=model, + ) def create_femfields( @@ -127,7 +132,7 @@ def create_femfields( ------- fields : dict Nested dictionary holding :class:`~struphy.feec.psydac_derham.SplineFunction`: fields[t][name] contains the Field with the name "name" in the hdf5 file at time t. - + t_grid : np.ndarray Time grid. """ @@ -135,15 +140,16 @@ def create_femfields( with open(os.path.join(path, "meta.yml"), "r") as f: meta = yaml.load(f, Loader=yaml.FullLoader) nproc = meta["MPI processes"] - + # import parameters params_in = get_params_of_run(path) - derham = setup_derham(params_in.grid, - params_in.derham_opts, - comm=None, - domain=params_in.domain, - ) + derham = setup_derham( + params_in.grid, + params_in.derham_opts, + comm=None, + domain=params_in.domain, + ) # get fields names, space IDs and time grid from 0-th rank hdf5 file file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") @@ -166,7 +172,11 @@ def create_femfields( for species, vars in space_ids.items(): fields[t][species] = {} for var, id in vars.items(): - fields[t][species][var] = derham.create_spline_function(var, id, verbose=False,) + fields[t][species][var] = derham.create_spline_function( + var, + id, + verbose=False, + ) # get hdf5 data print("") @@ -303,12 +313,11 @@ def eval_femfields( point_data[species] = {} for name, field in vars.items(): point_data[species][name] = {} - + print("\nEvaluating fields ...") for t in tqdm(fields): for species, vars in fields[t].items(): for name, field in vars.items(): - assert isinstance(field, SplineFunction) space_id = field.space_id @@ -373,7 +382,7 @@ def eval_femfields( else: point_data[species][name][t].append(temp_val[j]) - + return point_data, grids_log, grids_phy @@ -406,7 +415,7 @@ def create_vtk( """ from pyevtk.hl import gridToVTK - + for species, vars in point_data.items(): species_path = os.path.join(path, species, "vtk" + physical * "_phy") try: @@ -445,7 +454,14 @@ def create_vtk( ) -def post_process_markers(path_in: str, path_out: str, species: str, domain: Domain, kind: str = "Particles6D", step: int = 1,): +def post_process_markers( + path_in: str, + path_out: str, + species: str, + domain: Domain, + kind: str = "Particles6D", + step: int = 1, +): """Computes the Cartesian (x, y, z) coordinates of saved markers during a simulation and writes them to a .npy files and to .txt files. Also saves the weights. @@ -496,7 +512,7 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma species : str Name of the species for which the post processing should be performed. - + domain : Domain Domain object. @@ -531,7 +547,7 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma # directory for .txt files and marker index which will be saved path_orbits = os.path.join(path_out, "orbits") - + if "5D" in kind: save_index = list(range(0, 6)) + [10] + [-1] elif "6D" in kind or "SPH" in kind: @@ -554,7 +570,7 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma # loop over time grid for n in tqdm(range(int((nt - 1) / step) + 1)): # clear buffer - temp[:, :] = 0. + temp[:, :] = 0.0 # create text file for this time step and this species file_npy = os.path.join( @@ -575,7 +591,7 @@ def post_process_markers(path_in: str, path_out: str, species: str, domain: Doma # sorting out lost particles ids = temp[:, -1].astype("int") ids_lost_particles = np.setdiff1d(np.arange(n_markers), ids) - ids_removed_particles = np.nonzero(temp[:, 0] == -1.)[0] + ids_removed_particles = np.nonzero(temp[:, 0] == -1.0)[0] ids_lost_particles = np.array(list(set(ids_lost_particles) | set(ids_removed_particles)), dtype=int) lost_particles_mask[:] = False lost_particles_mask[ids_lost_particles] = True @@ -712,9 +728,9 @@ def post_process_f(path_in, path_out, species, step=1, compute_bckgr=False): # f_bckgr = f_bckgr + getattr(maxwellians, fi_type)( # maxw_params=maxw_params, # ) - - spec: KineticSpecies = getattr(params.model, species) - var: PICVariable = spec.var + + spec: KineticSpecies = getattr(params.model, species) + var: PICVariable = spec.var f_bckgr: KineticBackground = var.backgrounds # load all grids of the variables of f diff --git a/src/struphy/post_processing/pproc_struphy.py b/src/struphy/post_processing/pproc_struphy.py index 9524abe60..06f64c37f 100644 --- a/src/struphy/post_processing/pproc_struphy.py +++ b/src/struphy/post_processing/pproc_struphy.py @@ -94,9 +94,9 @@ def main( exist_kinetic["n_sph"] = True else: exist_kinetic = None - + file.close() - + # import parameters params_in = import_parameters_py(os.path.join(path, "parameters.py")) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 2f1e5181a..7a8878eb5 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -1,20 +1,21 @@ "Propagator base class." from abc import ABCMeta, abstractmethod -import numpy as np -from mpi4py import MPI from dataclasses import dataclass from typing import Literal +import numpy as np +from mpi4py import MPI +from psydac.linalg.block import BlockVector +from psydac.linalg.stencil import StencilVector + from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import WeightedMassOperators from struphy.feec.psydac_derham import Derham -from struphy.geometry.base import Domain -from struphy.models.variables import Variable, FEECVariable, PICVariable, SPHVariable -from psydac.linalg.stencil import StencilVector -from psydac.linalg.block import BlockVector from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB +from struphy.geometry.base import Domain from struphy.io.options import check_option +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable class Propagator(metaclass=ABCMeta): @@ -26,45 +27,47 @@ class Propagator(metaclass=ABCMeta): in one of the modules ``propagators_fields.py``, ``propagators_markers.py`` or ``propagators_coupling.py``. Only propagators that update both a FEEC and a PIC species go into ``propagators_coupling.py``. """ + @abstractmethod class Variables: """Define variable names and types to be updated by the propagator.""" + def __init__(self): self._var1 = None - + @property def var1(self): return self._var1 - + @var1.setter def var1(self, new): assert isinstance(new, PICVariable) assert new.space == "Particles6D" self._var1 = new - + @abstractmethod def __init__(self): self.variables = self.Variables() - + @abstractmethod @dataclass class Options: # specific literals OptsTemplate = Literal["implicit", "explicit"] # propagator options - opt1: str = "implicit", - + opt1: str = ("implicit",) + def __post_init__(self): # checks check_option(self.opt1, self.OptsTemplate) - + @property @abstractmethod def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter @abstractmethod def options(self, new): @@ -72,13 +75,13 @@ def options(self, new): if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new - + @abstractmethod def allocate(self): """Allocate all data/objects of the instance.""" - + @abstractmethod def __call__(self, dt: float): """Update variables from t -> t + dt. @@ -89,16 +92,16 @@ def __call__(self, dt: float): dt : float Time step size. """ - + def update_feec_variables(self, **new_coeffs): r"""Return max_diff = max(abs(new - old)) for each new_coeffs, update feec coefficients and update ghost regions. - + Returns ------- diffs : dict max_diff for all feec variables. - """ + """ diffs = {} for var, new in new_coeffs.items(): assert "_" + var in self.variables.__dict__, f"{var} not in {self.variables.__dict__}." @@ -107,7 +110,7 @@ def update_feec_variables(self, **new_coeffs): assert isinstance(old_var, FEECVariable) old = old_var.spline.vector assert new.space == old.space - + # calculate maximum of difference abs(new - old) diffs[var] = np.max(np.abs(new.toarray() - old.toarray())) @@ -205,7 +208,7 @@ def projected_equil(self, new): @property def time_state(self): """A pointer to the time variable of the dynamics ('t').""" - return self._time_state + return self._time_state def add_time_state(self, time_state): """Add a pointer to the time variable of the dynamics ('t'). @@ -313,4 +316,4 @@ def add_eval_kernel( comps, args_eval, ) - ] \ No newline at end of file + ] diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index 6119d556a..cd629ffe4 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -1,20 +1,23 @@ "Particle and FEEC variables are updated." -import numpy as np from dataclasses import dataclass -from mpi4py import MPI from typing import Literal -from line_profiler import profile +import numpy as np +from line_profiler import profile +from mpi4py import MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector from struphy.feec import preconditioner from struphy.feec.linear_operators import LinOpWithTransp +from struphy.io.options import OptsGenSolver, OptsMassPrecond, OptsSymmSolver, OptsVecSpace, check_option from struphy.io.setup import descend_options_dict from struphy.kinetic_background.base import Maxwellian from struphy.kinetic_background.maxwellians import Maxwellian3D from struphy.linear_algebra.schur_solver import SchurSolver +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable, PICVariable from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.pic.accumulation.particles_to_grid import Accumulator from struphy.pic.particles import Particles5D, Particles6D @@ -22,9 +25,6 @@ from struphy.pic.pushing.pusher import Pusher from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.models.variables import FEECVariable, PICVariable -from struphy.linear_algebra.solver import SolverParameters -from struphy.io.options import (check_option, OptsSymmSolver, OptsMassPrecond, OptsGenSolver, OptsVecSpace) class VlasovAmpere(Propagator): @@ -75,58 +75,58 @@ class Variables: def __init__(self): self._e: FEECVariable = None self._ions: PICVariable = None - - @property + + @property def e(self) -> FEECVariable: return self._e - + @e.setter def e(self, new): assert isinstance(new, FEECVariable) assert new.space == "Hcurl" self._e = new - - @property + + @property def ions(self) -> PICVariable: return self._ions - + @ions.setter def ions(self, new): assert isinstance(new, PICVariable) assert new.space == "Particles6D" self._ions = new - + def __init__(self): self.variables = self.Variables() - + @dataclass class Options: - solver: OptsSymmSolver = "pcg" + solver: OptsSymmSolver = "pcg" precond: OptsMassPrecond = "MassMatrixPreconditioner" solver_params: SolverParameters = None - + def __post_init__(self): # checks check_option(self.solver, OptsSymmSolver) check_option(self.precond, OptsMassPrecond) - + # defaults if self.solver_params is None: self.solver_params = SolverParameters() - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @profile @@ -134,10 +134,10 @@ def allocate(self): # scaling factors alpha = self.variables.ions.species.equation_params.alpha epsilon = self.variables.ions.species.equation_params.epsilon - + self._c1 = alpha**2 / epsilon self._c2 = 1.0 / epsilon - + self._info = self.options.solver_params.info # get accumulation kernel @@ -145,7 +145,7 @@ def allocate(self): # Initialize Accumulator object particles = self.variables.ions.particles - + self._accum = Accumulator( particles, "Hcurl", diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index a4b0dfe40..6c12fb5a3 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -1,14 +1,14 @@ "Only FEEC variables are updated." +import copy from collections.abc import Callable from copy import deepcopy from dataclasses import dataclass from typing import Literal, get_args -import copy -from line_profiler import profile import numpy as np import scipy as sc +from line_profiler import profile from mpi4py import MPI from numpy import zeros from psydac.linalg.basic import IdentityOperator, ZeroOperator @@ -33,11 +33,14 @@ from struphy.fields_background.equils import set_defaults from struphy.geometry.utilities import TransformedPformComponent from struphy.initial import perturbations +from struphy.io.options import OptsGenSolver, OptsMassPrecond, OptsSymmSolver, OptsVecSpace, check_option from struphy.io.setup import descend_options_dict from struphy.kinetic_background.base import Maxwellian from struphy.kinetic_background.maxwellians import GyroMaxwellian2D, Maxwellian3D from struphy.linear_algebra.saddle_point import SaddlePointSolver from struphy.linear_algebra.schur_solver import SchurSolver +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.ode.solvers import ODEsolverFEEC from struphy.ode.utils import ButcherTableau, OptsButcher from struphy.pic.accumulation import accum_kernels, accum_kernels_gc @@ -46,10 +49,6 @@ from struphy.pic.particles import Particles5D, Particles6D from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.models.variables import Variable -from struphy.linear_algebra.solver import SolverParameters -from struphy.io.options import (check_option, OptsSymmSolver, OptsMassPrecond, OptsGenSolver, OptsVecSpace) -from struphy.models.variables import FEECVariable, PICVariable, SPHVariable class Maxwell(Propagator): @@ -64,25 +63,26 @@ class Maxwell(Propagator): :ref:`time_discret`: Crank-Nicolson (implicit mid-point). System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`. """ + class Variables: def __init__(self): self._e: FEECVariable = None self._b: FEECVariable = None - - @property + + @property def e(self) -> FEECVariable: return self._e - + @e.setter def e(self, new): assert isinstance(new, FEECVariable) assert new.space == "Hcurl" self._e = new - - @property + + @property def b(self) -> FEECVariable: return self._b - + @b.setter def b(self, new): assert isinstance(new, FEECVariable) @@ -91,44 +91,44 @@ def b(self, new): def __init__(self): self.variables = self.Variables() - + @dataclass class Options: # specific literals OptsAlgo = Literal["implicit", "explicit"] # propagator options algo: OptsAlgo = "implicit" - solver: OptsSymmSolver = "pcg" + solver: OptsSymmSolver = "pcg" precond: OptsMassPrecond = "MassMatrixPreconditioner" solver_params: SolverParameters = None butcher: ButcherTableau = None - + def __post_init__(self): # checks check_option(self.algo, self.OptsAlgo) check_option(self.solver, OptsSymmSolver) - check_option(self.precond, OptsMassPrecond) - + check_option(self.precond, OptsMassPrecond) + # defaults if self.solver_params is None: self.solver_params = SolverParameters() - + if self.algo == "explicit" and self.butcher is None: self.butcher = ButcherTableau() - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @profile @@ -476,25 +476,26 @@ class ShearAlfven(Propagator): where :math:`\alpha \in \{1, 2, v\}` and :math:`\mathbb M^\rho_\alpha` is a weighted mass matrix in :math:`\alpha`-space, the weight being :math:`\rho_0`, the MHD equilibirum density. The solution of the above system is based on the :ref:`Schur complement `. """ + class Variables: def __init__(self): self._u: FEECVariable = None self._b: FEECVariable = None - - @property + + @property def u(self) -> FEECVariable: return self._u - + @u.setter def u(self, new): assert isinstance(new, FEECVariable) assert new.space in ("Hcurl", "Hdiv", "H1vec") self._u = new - - @property + + @property def b(self) -> FEECVariable: return self._b - + @b.setter def b(self, new): assert isinstance(new, FEECVariable) @@ -503,52 +504,52 @@ def b(self, new): def __init__(self): self.variables = self.Variables() - + @dataclass class Options: # specific literals OptsAlgo = Literal["implicit", "explicit"] # propagator options u_space: OptsVecSpace = "Hdiv" - algo: OptsAlgo = "implicit" - solver: OptsSymmSolver = "pcg" - precond: OptsMassPrecond = "MassMatrixPreconditioner" + algo: OptsAlgo = "implicit" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" solver_params: SolverParameters = None butcher: ButcherTableau = None - + def __post_init__(self): # checks check_option(self.u_space, OptsVecSpace) check_option(self.algo, self.OptsAlgo) check_option(self.solver, OptsSymmSolver) - check_option(self.precond, OptsMassPrecond) - + check_option(self.precond, OptsMassPrecond) + # defaults if self.solver_params is None: self.solver_params = SolverParameters() - + if self.algo == "explicit" and self.butcher is None: self.butcher = ButcherTableau() - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @profile def allocate(self): u_space = self.options.u_space - + # define block matrix [[A B], [C I]] (without time step size dt in the diagonals) id_M = "M" + self.derham.space_to_form[u_space] + "n" id_T = "T" + self.derham.space_to_form[u_space] @@ -938,36 +939,37 @@ class Magnetosonic(Propagator): \boldsymbol{\rho}^{n+1} = \boldsymbol{\rho}^n - \frac{\Delta t}{2} \mathbb D \mathcal Q^\alpha (\mathbf u^{n+1} + \mathbf u^n) \,. """ + class Variables: def __init__(self): self._n: FEECVariable = None self._u: FEECVariable = None self._p: FEECVariable = None - - @property + + @property def n(self) -> FEECVariable: return self._n - + @n.setter def n(self, new): assert isinstance(new, FEECVariable) assert new.space == "L2" self._n = new - - @property + + @property def u(self) -> FEECVariable: return self._u - + @u.setter def u(self, new): assert isinstance(new, FEECVariable) assert new.space in ("Hcurl", "Hdiv", "H1vec") self._u = new - - @property + + @property def p(self) -> FEECVariable: return self._p - + @p.setter def p(self, new): assert isinstance(new, FEECVariable) @@ -976,40 +978,40 @@ def p(self, new): def __init__(self): self.variables = self.Variables() - + @dataclass class Options: b_field: FEECVariable = None - u_space: OptsVecSpace = "Hdiv" + u_space: OptsVecSpace = "Hdiv" solver: OptsGenSolver = "pbicgstab" - precond: OptsMassPrecond = "MassMatrixPreconditioner" + precond: OptsMassPrecond = "MassMatrixPreconditioner" solver_params: SolverParameters = None - + def __post_init__(self): # checks check_option(self.u_space, OptsVecSpace) check_option(self.solver, OptsGenSolver) - check_option(self.precond, OptsMassPrecond) - + check_option(self.precond, OptsMassPrecond) + # defaults if self.b_field is None: self.b_field = FEECVariable(space="Hdiv") if self.solver_params is None: self.solver_params = SolverParameters() - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @profile @@ -1105,7 +1107,7 @@ def __call__(self, dt): nn1 += nn diffs = self.update_feec_variables(n=nn1, u=un1, p=pn1) - + if self._info and MPI.COMM_WORLD.Get_rank() == 0: print("Status for Magnetosonic:", info["success"]) print("Iterations for Magnetosonic:", info["niter"]) @@ -2669,14 +2671,15 @@ class ImplicitDiffusion(Propagator): solver : dict Parameters for the iterative solver (see ``__init__`` for details). """ + class Variables: def __init__(self): self._phi: FEECVariable = None - - @property + + @property def phi(self) -> FEECVariable: return self._phi - + @phi.setter def phi(self, new): assert isinstance(new, FEECVariable) @@ -2685,7 +2688,7 @@ def phi(self, new): def __init__(self): self.variables = self.Variables() - + @dataclass class Options: # specific literals @@ -2703,31 +2706,31 @@ class Options: solver: OptsSymmSolver = "pcg" precond: OptsMassPrecond = "MassMatrixPreconditioner" solver_params: SolverParameters = None - + def __post_init__(self): # checks check_option(self.stab_mat, self.OptsStabMat) check_option(self.diffusion_mat, self.OptsDiffusionMat) check_option(self.solver, OptsSymmSolver) - check_option(self.precond, OptsMassPrecond) - + check_option(self.precond, OptsMassPrecond) + # defaults if self.solver_params is None: self.solver_params = SolverParameters() - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @profile @@ -2747,7 +2750,7 @@ def allocate(self): # collect rhs rho = self.options.rho - + if rho is None: self._rho = [phi.space.zeros()] else: @@ -2959,6 +2962,7 @@ class Poisson(ImplicitDiffusion): solver : dict Parameters for the iterative solver (see ``__init__`` for details). """ + @dataclass class Options: # specific literals @@ -2971,30 +2975,30 @@ class Options: solver: OptsSymmSolver = "pcg" precond: OptsMassPrecond = "MassMatrixPreconditioner" solver_params: SolverParameters = None - + def __post_init__(self): # checks check_option(self.stab_mat, self.OptsStabMat) check_option(self.solver, OptsSymmSolver) - check_option(self.precond, OptsMassPrecond) - + check_option(self.precond, OptsMassPrecond) + # defaults if self.solver_params is None: self.solver_params = SolverParameters() - + # Poisson solve (-> set some params of parent class) self.sigma_1 = self.stab_eps self.sigma_2 = 0.0 self.sigma_3 = 1.0 self.divide_by_dt = False self.diffusion_mat = "M1" - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) @@ -3002,7 +3006,7 @@ def options(self, new): print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): if "sigma" not in k and k not in ("divide_by_dt", "diffusion_mat"): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @@ -7535,7 +7539,7 @@ class TwoFluidQuasiNeutralFull(Propagator): def allocate(self): pass - + def set_options(self, **kwargs): pass diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 17d786080..fc2aa2de9 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -1,20 +1,23 @@ "Only particle variables are updated." -import numpy as np -from numpy import array, polynomial, random -from typing import Literal, get_args import copy from dataclasses import dataclass -from mpi4py import MPI -from line_profiler import profile +from typing import Literal, get_args +import numpy as np +from line_profiler import profile +from mpi4py import MPI +from numpy import array, polynomial, random +from psydac.linalg.basic import LinearOperator from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector from struphy.feec.mass import WeightedMassOperators from struphy.fields_background.base import MHDequilibrium from struphy.fields_background.equils import set_defaults +from struphy.io.options import OptsMPIsort, check_option from struphy.io.setup import descend_options_dict +from struphy.models.variables import FEECVariable, PICVariable from struphy.ode.utils import ButcherTableau from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.pic.base import Particles @@ -23,9 +26,6 @@ from struphy.pic.pushing.pusher import Pusher from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.models.variables import FEECVariable, PICVariable -from struphy.io.options import check_option, OptsMPIsort -from psydac.linalg.basic import LinearOperator class PushEta(Propagator): @@ -45,14 +45,15 @@ class PushEta(Propagator): * Explicit from :class:`~struphy.ode.utils.ButcherTableau` """ + class Variables: def __init__(self): self._var: PICVariable = None - - @property + + @property def var(self) -> PICVariable: return self._var - + @var.setter def var(self, new): assert isinstance(new, PICVariable) @@ -60,29 +61,29 @@ def var(self, new): def __init__(self): self.variables = self.Variables() - + @dataclass class Options: butcher: ButcherTableau = None - + def __post_init__(self): # defaults if self.butcher is None: self.butcher = ButcherTableau() - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @profile @@ -139,23 +140,24 @@ class PushVxB(Propagator): Available algorithms: ``analytic``, ``implicit``. """ + class Variables: def __init__(self): self._ions: PICVariable = None - - @property + + @property def ions(self) -> PICVariable: return self._ions - + @ions.setter def ions(self, new): assert isinstance(new, PICVariable) assert new.space == "Particles6D" self._ions = new - + def __init__(self): self.variables = self.Variables() - + @dataclass class Options: # specific literals @@ -163,44 +165,44 @@ class Options: # propagator options algo: OptsAlgo = "analytic" b2_var: FEECVariable = None - + def __post_init__(self): # checks check_option(self.algo, self.OptsAlgo) - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @profile def allocate(self): # scaling factor self._epsilon = self.variables.ions.species.equation_params.epsilon - + # TODO: treat PolarVector as well, but polar splines are being reworked at the moment if self.projected_equil is not None: self._b2 = self.projected_equil.b2 assert self._b2.space == self.derham.Vh["2"] else: self._b2 = self.derham.Vh["2"].zeros() - + if self.options.b2_var is None: self._b2_var = None else: assert self.options.b2_var.spline.vector.space == self.derham.Vh["2"] self._b2_var = self.options.b2_var.spline.vector - + # allocate dummy vectors to avoid temporary array allocations self._tmp = self.derham.Vh["2"].zeros() self._b_full = self.derham.Vh["2"].zeros() @@ -448,30 +450,33 @@ class PushGuidingCenterBxEstar(Propagator): * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_1st_order_newton` * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_bxEstar_discrete_gradient_2nd_order` """ + class Variables: def __init__(self): self._ions: PICVariable = None - - @property + + @property def ions(self) -> PICVariable: return self._ions - + @ions.setter def ions(self, new): assert isinstance(new, PICVariable) assert new.space == "Particles5D" self._ions = new - + def __init__(self): self.variables = self.Variables() - + @dataclass class Options: # specific literals - OptsAlgo = Literal["discrete_gradient_2nd_order", - "discrete_gradient_1st_order", - "discrete_gradient_1st_order_newton", - "explicit",] + OptsAlgo = Literal[ + "discrete_gradient_2nd_order", + "discrete_gradient_1st_order", + "discrete_gradient_1st_order_newton", + "explicit", + ] # propagator options phi: FEECVariable = None evaluate_e_field: bool = False @@ -482,39 +487,39 @@ class Options: tol: float = 1e-7 mpi_sort: OptsMPIsort = "each" verbose: bool = False - + def __post_init__(self): # checks check_option(self.algo, self.OptsAlgo) check_option(self.mpi_sort, OptsMPIsort) - + # defaults if self.phi is None: self.phi = FEECVariable(space="H1") - + if self.algo == "explicit" and self.butcher is None: self.butcher = ButcherTableau() - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @profile def allocate(self): # scaling factor self._epsilon = self.variables.ions.species.equation_params.epsilon - + # magnetic equilibrium field unit_b1 = self.projected_equil.unit_b1 self._gradB1 = self.projected_equil.gradB1 @@ -548,7 +553,7 @@ def allocate(self): # choose method particles = self.variables.ions.particles - + if "discrete_gradient" in self.options.algo: # place for storing data during iteration first_free_idx = particles.args_markers.first_free_idx @@ -886,30 +891,33 @@ class PushGuidingCenterParallel(Propagator): * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_Bstar_discrete_gradient_1st_order_newton` * :func:`~struphy.pic.pushing.pusher_kernels_gc.push_gc_Bstar_discrete_gradient_2nd_order` """ + class Variables: def __init__(self): self._ions: PICVariable = None - - @property + + @property def ions(self) -> PICVariable: return self._ions - + @ions.setter def ions(self, new): assert isinstance(new, PICVariable) assert new.space == "Particles5D" self._ions = new - + def __init__(self): self.variables = self.Variables() - + @dataclass class Options: # specific literals - OptsAlgo = Literal["discrete_gradient_2nd_order", - "discrete_gradient_1st_order", - "discrete_gradient_1st_order_newton", - "explicit",] + OptsAlgo = Literal[ + "discrete_gradient_2nd_order", + "discrete_gradient_1st_order", + "discrete_gradient_1st_order_newton", + "explicit", + ] # propagator options phi: FEECVariable = None evaluate_e_field: bool = False @@ -920,39 +928,39 @@ class Options: tol: float = 1e-7 mpi_sort: OptsMPIsort = "each" verbose: bool = False - + def __post_init__(self): # checks check_option(self.algo, self.OptsAlgo) check_option(self.mpi_sort, OptsMPIsort) - + # defaults if self.phi is None: self.phi = FEECVariable(space="H1") - + if self.algo == "explicit" and self.butcher is None: self.butcher = ButcherTableau() - + @property def options(self) -> Options: if not hasattr(self, "_options"): self._options = self.Options() return self._options - + @options.setter def options(self, new): assert isinstance(new, self.Options) if MPI.COMM_WORLD.Get_rank() == 0: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): - print(f' {k}: {v}') + print(f" {k}: {v}") self._options = new @profile def allocate(self): # scaling factor self._epsilon = self.variables.ions.species.equation_params.epsilon - + # magnetic equilibrium field self._gradB1 = self.projected_equil.gradB1 b2 = self.projected_equil.b2 @@ -987,7 +995,7 @@ def allocate(self): # choose method particles = self.variables.ions.particles - + if "discrete_gradient" in self.options.algo: # place for storing data during iteration first_free_idx = particles.args_markers.first_free_idx diff --git a/src/struphy/topology/grids.py b/src/struphy/topology/grids.py index e66dae1ea..0593bce05 100644 --- a/src/struphy/topology/grids.py +++ b/src/struphy/topology/grids.py @@ -1,6 +1,7 @@ -import numpy as np from dataclasses import dataclass +import numpy as np + @dataclass class TensorProductGrid: @@ -15,5 +16,6 @@ class TensorProductGrid: True if the dimension is to be used in the domain decomposition (=default for each dimension). If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. """ + Nel: tuple = (16, 1, 1) mpi_dims_mask: tuple = (True, True, True) diff --git a/src/struphy/tutorials/tests/test_tutorials.py b/src/struphy/tutorials/tests/test_tutorials.py index cfa726aee..17cf2069d 100644 --- a/src/struphy/tutorials/tests/test_tutorials.py +++ b/src/struphy/tutorials/tests/test_tutorials.py @@ -5,8 +5,8 @@ from mpi4py import MPI import struphy -from struphy.struphy import run from struphy.post_processing import pproc_struphy +from struphy.struphy import run comm = MPI.COMM_WORLD rank = comm.Get_rank() From 8302d5d11e1b6f70808143709871dd14e0d1a91f Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Sat, 13 Sep 2025 07:33:14 +0000 Subject: [PATCH 121/292] Move tutorials to base folder --- .gitignore | 2 + doc/conf.py | 16 +- doc/sections/tutorials.rst | 6 +- .../tutorial_01_parameter_files.ipynb | 46 +++--- .../tutorial_02_test_particles.ipynb | 146 +++++++++--------- .../tutorial_05_vlasov_maxwell.ipynb | 0 6 files changed, 116 insertions(+), 100 deletions(-) rename {doc/tutorials => tutorials}/tutorial_01_parameter_files.ipynb (96%) rename {doc/tutorials => tutorials}/tutorial_02_test_particles.ipynb (96%) rename {doc/tutorials => tutorials}/tutorial_05_vlasov_maxwell.ipynb (100%) diff --git a/.gitignore b/.gitignore index f529a8f8a..e01ccf71d 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,8 @@ share/python-wheels/ *.egg MANIFEST doc/_build/ +doc/source/* +tutorials/sim* *STUBDIR* .VSCodeCounter/ .pymon diff --git a/doc/conf.py b/doc/conf.py index 9f98eba4e..c5a605edb 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -10,10 +10,24 @@ # add these directories to sys.path here. If the directory is relative to the # documentation root, use os.path.abspath to make it absolute, like shown here. # -# import os +import os +import shutil # import sys # sys.path.insert(0, os.path.abspath('.')) +def copy_tutorials(app): + src = os.path.abspath("../tutorials") + dst = os.path.abspath("source/tutorials") + + # Remove existing target directory if it exists + if os.path.exists(dst): + shutil.rmtree(dst) + + shutil.copytree(src, dst) + +def setup(app): + app.connect("builder-inited", copy_tutorials) + with open("../src/struphy/console/main.py") as f: exec(f.read()) diff --git a/doc/sections/tutorials.rst b/doc/sections/tutorials.rst index c1b4a414b..d0052ad2a 100644 --- a/doc/sections/tutorials.rst +++ b/doc/sections/tutorials.rst @@ -3,7 +3,7 @@ Tutorials ========= -All notebooks are available at https://gitlab.mpcdf.mpg.de/struphy/struphy/-/blob/devel/doc/tutorials/. +All notebooks are available at https://gitlab.mpcdf.mpg.de/struphy/struphy/-/blob/devel/tutorials/. The objects used in these notebooks are the same as in the available :ref:`models`. They can thus be used for MPI parallel runs in HPC applications. @@ -11,6 +11,6 @@ They can thus be used for MPI parallel runs in HPC applications. .. toctree:: :maxdepth: 1 :caption: Notebook tutorials: + :glob: - ../tutorials/tutorial_01_parameter_files - ../tutorials/tutorial_02_test_particles + ../source/tutorials/* diff --git a/doc/tutorials/tutorial_01_parameter_files.ipynb b/tutorials/tutorial_01_parameter_files.ipynb similarity index 96% rename from doc/tutorials/tutorial_01_parameter_files.ipynb rename to tutorials/tutorial_01_parameter_files.ipynb index e4c5006a1..13a6c1839 100644 --- a/doc/tutorials/tutorial_01_parameter_files.ipynb +++ b/tutorials/tutorial_01_parameter_files.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "d34c79c5", + "id": "0", "metadata": {}, "source": [ "# 1 - Parameters and `struphy.main`\n", @@ -45,7 +45,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3ecab659", + "id": "1", "metadata": {}, "outputs": [], "source": [ @@ -67,7 +67,7 @@ }, { "cell_type": "markdown", - "id": "5cf6d9c7", + "id": "2", "metadata": {}, "source": [ "All parameter files import the modules listed above, even though some of them might not be needed in a specific model. \n", @@ -91,7 +91,7 @@ { "cell_type": "code", "execution_count": null, - "id": "cc43d2fc", + "id": "3", "metadata": {}, "outputs": [], "source": [ @@ -125,7 +125,7 @@ }, { "cell_type": "markdown", - "id": "74e6f739", + "id": "4", "metadata": {}, "source": [ "## Part 3: Model instance\n", @@ -140,7 +140,7 @@ { "cell_type": "code", "execution_count": null, - "id": "83dc7f7f", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -153,7 +153,7 @@ }, { "cell_type": "markdown", - "id": "aa34840a", + "id": "6", "metadata": {}, "source": [ "In case of a kinetic species, one can also set parameters regarding marker drawing, box sorting and data saving: " @@ -162,7 +162,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6c498fb3", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -179,7 +179,7 @@ }, { "cell_type": "markdown", - "id": "b0f65b0a", + "id": "8", "metadata": {}, "source": [ "## Part 4: Propagator options\n", @@ -192,7 +192,7 @@ { "cell_type": "code", "execution_count": null, - "id": "be6875e7", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -203,7 +203,7 @@ }, { "cell_type": "markdown", - "id": "b1ef8b97", + "id": "10", "metadata": {}, "source": [ "## Part 5: Initial conditions\n", @@ -214,7 +214,7 @@ { "cell_type": "code", "execution_count": null, - "id": "cc0ac424", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -227,7 +227,7 @@ }, { "cell_type": "markdown", - "id": "879978af", + "id": "12", "metadata": {}, "source": [ "## Part 6: `main.run`\n", @@ -238,7 +238,7 @@ { "cell_type": "code", "execution_count": null, - "id": "03764138", + "id": "13", "metadata": {}, "outputs": [], "source": [ @@ -257,7 +257,7 @@ }, { "cell_type": "markdown", - "id": "f9ce5099", + "id": "14", "metadata": {}, "source": [ "## Post processing: `main.pproc`\n", @@ -268,7 +268,7 @@ { "cell_type": "code", "execution_count": null, - "id": "dd7cfe62", + "id": "15", "metadata": {}, "outputs": [], "source": [ @@ -280,7 +280,7 @@ }, { "cell_type": "markdown", - "id": "0076c65c", + "id": "16", "metadata": {}, "source": [ "One can also post-process directly from the console:\n", @@ -300,7 +300,7 @@ }, { "cell_type": "markdown", - "id": "e317f88d", + "id": "17", "metadata": {}, "source": [ "## Loading data: `main.load_data`\n", @@ -311,7 +311,7 @@ { "cell_type": "code", "execution_count": null, - "id": "78b8cbab", + "id": "18", "metadata": {}, "outputs": [], "source": [ @@ -320,7 +320,7 @@ }, { "cell_type": "markdown", - "id": "08544a91", + "id": "19", "metadata": {}, "source": [ "## Plotting particle orbits\n", @@ -331,7 +331,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0ea08d24", + "id": "20", "metadata": {}, "outputs": [], "source": [ @@ -344,7 +344,7 @@ }, { "cell_type": "markdown", - "id": "001437bc", + "id": "21", "metadata": {}, "source": [ "Let us plot the orbits:" @@ -353,7 +353,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7af3facd", + "id": "22", "metadata": {}, "outputs": [], "source": [ diff --git a/doc/tutorials/tutorial_02_test_particles.ipynb b/tutorials/tutorial_02_test_particles.ipynb similarity index 96% rename from doc/tutorials/tutorial_02_test_particles.ipynb rename to tutorials/tutorial_02_test_particles.ipynb index 3d89d1ec0..d2711ffbd 100644 --- a/doc/tutorials/tutorial_02_test_particles.ipynb +++ b/tutorials/tutorial_02_test_particles.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "724e47a7", + "id": "0", "metadata": {}, "source": [ "# 2 - Test particles\n", @@ -33,7 +33,7 @@ { "cell_type": "code", "execution_count": null, - "id": "541d473c", + "id": "1", "metadata": {}, "outputs": [], "source": [ @@ -55,7 +55,7 @@ }, { "cell_type": "markdown", - "id": "8861c9ef", + "id": "2", "metadata": {}, "source": [ "Note that we set `verbose = False` which will be passed to `main.run` to supress screen output during the simulation.\n", @@ -66,7 +66,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b6629b5e", + "id": "3", "metadata": {}, "outputs": [], "source": [ @@ -77,7 +77,7 @@ }, { "cell_type": "markdown", - "id": "5d1ba880", + "id": "4", "metadata": {}, "source": [ "The other generic options will be the same for both simulations. Here, we just perform one time step und load a cylindrical geometry:" @@ -86,7 +86,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4dbe46b5", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -105,7 +105,7 @@ }, { "cell_type": "markdown", - "id": "ee94df63", + "id": "6", "metadata": {}, "source": [ "We can already look a t the simulation domain:" @@ -114,7 +114,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d4174c86", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -123,7 +123,7 @@ }, { "cell_type": "markdown", - "id": "361147f5", + "id": "8", "metadata": {}, "source": [ "We can leave the equilibrium, grid and Derham complex empty:" @@ -132,7 +132,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4a03b02e", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -148,7 +148,7 @@ }, { "cell_type": "markdown", - "id": "32615bbd", + "id": "10", "metadata": {}, "source": [ "For each simulation, we must create the light-weight model instance and set parameters:" @@ -157,7 +157,7 @@ { "cell_type": "code", "execution_count": null, - "id": "853ac716", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -172,7 +172,7 @@ }, { "cell_type": "markdown", - "id": "71ad7471", + "id": "12", "metadata": {}, "source": [ "For the second simulation, in the parameters for particle loading we choose `spatial=\"disc\"` in order to draw uniformly on the cross section of the cylinder: " @@ -181,7 +181,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c13b1c6c", + "id": "13", "metadata": {}, "outputs": [], "source": [ @@ -207,7 +207,7 @@ }, { "cell_type": "markdown", - "id": "06b41735", + "id": "14", "metadata": {}, "source": [ "Propagator options and initial conditions shall be the same in both simulations:" @@ -216,7 +216,7 @@ { "cell_type": "code", "execution_count": null, - "id": "585c3808", + "id": "15", "metadata": {}, "outputs": [], "source": [ @@ -231,7 +231,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9394a942", + "id": "16", "metadata": {}, "outputs": [], "source": [ @@ -245,7 +245,7 @@ }, { "cell_type": "markdown", - "id": "6ec98182", + "id": "17", "metadata": {}, "source": [ "Let us now run the first simulation:" @@ -254,7 +254,7 @@ { "cell_type": "code", "execution_count": null, - "id": "aaa18727", + "id": "18", "metadata": {}, "outputs": [], "source": [ @@ -273,7 +273,7 @@ }, { "cell_type": "markdown", - "id": "6099f2e8", + "id": "19", "metadata": {}, "source": [ "And now the second simulation:" @@ -282,7 +282,7 @@ { "cell_type": "code", "execution_count": null, - "id": "24b65c51", + "id": "20", "metadata": {}, "outputs": [], "source": [ @@ -301,7 +301,7 @@ }, { "cell_type": "markdown", - "id": "5834fb2f", + "id": "21", "metadata": {}, "source": [ "We now post-process both runs, load the generated data and plot the initial particle positions on a cross section of the cylinder:" @@ -310,7 +310,7 @@ { "cell_type": "code", "execution_count": null, - "id": "679e992b", + "id": "22", "metadata": {}, "outputs": [], "source": [ @@ -325,7 +325,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a98937e5", + "id": "23", "metadata": {}, "outputs": [], "source": [ @@ -336,7 +336,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0aaffb2e", + "id": "24", "metadata": {}, "outputs": [], "source": [ @@ -373,7 +373,7 @@ }, { "cell_type": "markdown", - "id": "cbb0dfd5", + "id": "25", "metadata": {}, "source": [ "## Reflecting boundary conditions \n", @@ -385,7 +385,7 @@ { "cell_type": "code", "execution_count": null, - "id": "aef8e67a", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -397,7 +397,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b23d8ff8", + "id": "27", "metadata": {}, "outputs": [], "source": [ @@ -416,7 +416,7 @@ }, { "cell_type": "markdown", - "id": "31de5e6d", + "id": "28", "metadata": {}, "source": [ "We still have to set the propagator options and the initial conditions:" @@ -425,7 +425,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f9f1874f", + "id": "29", "metadata": {}, "outputs": [], "source": [ @@ -442,7 +442,7 @@ }, { "cell_type": "markdown", - "id": "fdffdb7c", + "id": "30", "metadata": {}, "source": [ "We can now run the simulation, then post-process the data and plot the resulting orbits:" @@ -451,7 +451,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0ec1aaeb", + "id": "31", "metadata": {}, "outputs": [], "source": [ @@ -471,7 +471,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d90b4117", + "id": "32", "metadata": {}, "outputs": [], "source": [ @@ -482,7 +482,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3e8b5c2b", + "id": "33", "metadata": {}, "outputs": [], "source": [ @@ -491,7 +491,7 @@ }, { "cell_type": "markdown", - "id": "2b1736a1", + "id": "34", "metadata": {}, "source": [ "Under `simdata.orbits[]` one finds a three-dimensional numpy array; the first index refers to the time step, the second index to the particle and the third index to the particel attribute. The first three attributes are the partciel positions, followed by the velocities and the (initial and time-dependent) weights." @@ -500,7 +500,7 @@ { "cell_type": "code", "execution_count": null, - "id": "80bb8873", + "id": "35", "metadata": {}, "outputs": [], "source": [ @@ -514,7 +514,7 @@ { "cell_type": "code", "execution_count": null, - "id": "81afac14", + "id": "36", "metadata": {}, "outputs": [], "source": [ @@ -544,7 +544,7 @@ }, { "cell_type": "markdown", - "id": "bbc72273", + "id": "37", "metadata": {}, "source": [ "## Particles in a cylinder with a magnetic field\n", @@ -568,7 +568,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ac87beba", + "id": "38", "metadata": {}, "outputs": [], "source": [ @@ -580,7 +580,7 @@ }, { "cell_type": "markdown", - "id": "706486a1", + "id": "39", "metadata": {}, "source": [ "In order to project the equilibrium on the spline basis for fast evaluation in the particle kernels, we need a Derham complex:" @@ -589,7 +589,7 @@ { "cell_type": "code", "execution_count": null, - "id": "48fc6e08", + "id": "40", "metadata": {}, "outputs": [], "source": [ @@ -599,7 +599,7 @@ }, { "cell_type": "markdown", - "id": "f55b6f60", + "id": "41", "metadata": {}, "source": [ "Now we create the light-weight instance of the model and set the species options. We shall `remove` particles that hit the boundary in $\\eta_1$ (radial) direction:" @@ -608,7 +608,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d8b94d0e", + "id": "42", "metadata": {}, "outputs": [], "source": [ @@ -639,7 +639,7 @@ }, { "cell_type": "markdown", - "id": "5026b318", + "id": "43", "metadata": {}, "source": [ "Now the usual procedure: run, post-process, load data and finally plot the orbits:" @@ -648,7 +648,7 @@ { "cell_type": "code", "execution_count": null, - "id": "34893626", + "id": "44", "metadata": {}, "outputs": [], "source": [ @@ -669,7 +669,7 @@ { "cell_type": "code", "execution_count": null, - "id": "966bb91e", + "id": "45", "metadata": {}, "outputs": [], "source": [ @@ -682,7 +682,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4cfe0ee8", + "id": "46", "metadata": {}, "outputs": [], "source": [ @@ -695,7 +695,7 @@ { "cell_type": "code", "execution_count": null, - "id": "fb7eaf92", + "id": "47", "metadata": {}, "outputs": [], "source": [ @@ -723,7 +723,7 @@ }, { "cell_type": "markdown", - "id": "f10f537b", + "id": "48", "metadata": {}, "source": [ "## Particles in a Tokamak equilibrium\n", @@ -734,7 +734,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f1c0a4f2", + "id": "49", "metadata": {}, "outputs": [], "source": [ @@ -746,7 +746,7 @@ }, { "cell_type": "markdown", - "id": "60c488fa", + "id": "50", "metadata": {}, "source": [ "Since [EQDSKequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils_sub.html#struphy.fields_background.mhd_equil.equils.EQDSKequilibrium) is an [AxisymmMHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils_sub.html#struphy.fields_background.mhd_equil.base.AxisymmMHDequilibrium), which in turn is a [CartesianMHDequilibrium](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils_sub.html#struphy.fields_background.mhd_equil.base.CartesianMHDequilibrium), we are free to choose any mapping for the simulation (e.g. a Cuboid for Cartesian coordinates). In order to be conforming to the boundary of the equilibrium, we shall choose the [Tokamak](https://struphy.pages.mpcdf.de/struphy/sections/subsections/domains_avail.html#struphy.geometry.domains.Tokamak) mapping:" @@ -755,7 +755,7 @@ { "cell_type": "code", "execution_count": null, - "id": "f3e56de4", + "id": "51", "metadata": {}, "outputs": [], "source": [ @@ -772,7 +772,7 @@ }, { "cell_type": "markdown", - "id": "c3135576", + "id": "52", "metadata": {}, "source": [ "The [Tokamak](https://struphy.pages.mpcdf.de/struphy/sections/subsections/domains_avail.html#struphy.geometry.domains.Tokamak) domain is a [PoloidalSplineTorus](https://struphy.pages.mpcdf.de/struphy/sections/subsections/domains_base.html#struphy.geometry.base.PoloidalSplineTorus), hence\n", @@ -823,7 +823,7 @@ { "cell_type": "code", "execution_count": null, - "id": "cf741b53", + "id": "53", "metadata": {}, "outputs": [], "source": [ @@ -841,7 +841,7 @@ { "cell_type": "code", "execution_count": null, - "id": "85f37793", + "id": "54", "metadata": {}, "outputs": [], "source": [ @@ -856,7 +856,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c36dc211", + "id": "55", "metadata": {}, "outputs": [], "source": [ @@ -873,7 +873,7 @@ { "cell_type": "code", "execution_count": null, - "id": "078039c9", + "id": "56", "metadata": {}, "outputs": [], "source": [ @@ -931,7 +931,7 @@ }, { "cell_type": "markdown", - "id": "f7dc2c1c", + "id": "57", "metadata": {}, "source": [ "We now set up a simualtion of 4 specific particle orbits in this equilibrium:" @@ -940,7 +940,7 @@ { "cell_type": "code", "execution_count": null, - "id": "65feff34", + "id": "58", "metadata": {}, "outputs": [], "source": [ @@ -977,7 +977,7 @@ }, { "cell_type": "markdown", - "id": "9de13919", + "id": "59", "metadata": {}, "source": [ "We again need a Derham complex for the projection of the equilibirum onto the spline basis:" @@ -986,7 +986,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4579b2af", + "id": "60", "metadata": {}, "outputs": [], "source": [ @@ -1000,7 +1000,7 @@ }, { "cell_type": "markdown", - "id": "9e50f541", + "id": "61", "metadata": {}, "source": [ "We aim to simulate 15000 time steps with a second-order splitting algorithm:" @@ -1009,7 +1009,7 @@ { "cell_type": "code", "execution_count": null, - "id": "86583b41", + "id": "62", "metadata": {}, "outputs": [], "source": [ @@ -1031,7 +1031,7 @@ { "cell_type": "code", "execution_count": null, - "id": "aae7877e", + "id": "63", "metadata": {}, "outputs": [], "source": [ @@ -1047,7 +1047,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3be0d0a3", + "id": "64", "metadata": {}, "outputs": [], "source": [ @@ -1060,7 +1060,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ae248ef1", + "id": "65", "metadata": {}, "outputs": [], "source": [ @@ -1086,7 +1086,7 @@ }, { "cell_type": "markdown", - "id": "602f4408", + "id": "66", "metadata": {}, "source": [ "## Guiding-centers in a Tokamak equilibrium\n", @@ -1097,7 +1097,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a5182604", + "id": "67", "metadata": {}, "outputs": [], "source": [ @@ -1137,7 +1137,7 @@ { "cell_type": "code", "execution_count": null, - "id": "041005d2", + "id": "68", "metadata": {}, "outputs": [], "source": [ @@ -1196,7 +1196,7 @@ { "cell_type": "code", "execution_count": null, - "id": "e792ad7e", + "id": "69", "metadata": {}, "outputs": [], "source": [ @@ -1218,7 +1218,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ef3b687f", + "id": "70", "metadata": {}, "outputs": [], "source": [ @@ -1234,7 +1234,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9819e15f", + "id": "71", "metadata": {}, "outputs": [], "source": [ @@ -1247,7 +1247,7 @@ { "cell_type": "code", "execution_count": null, - "id": "05cc5d7d", + "id": "72", "metadata": {}, "outputs": [], "source": [ diff --git a/doc/tutorials/tutorial_05_vlasov_maxwell.ipynb b/tutorials/tutorial_05_vlasov_maxwell.ipynb similarity index 100% rename from doc/tutorials/tutorial_05_vlasov_maxwell.ipynb rename to tutorials/tutorial_05_vlasov_maxwell.ipynb From 3faca93ec288cdf4d726fcfed706d73c291d9336 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 18 Sep 2025 12:22:11 +0000 Subject: [PATCH 122/292] Fix unit testing --- .gitlab-ci.yml | 64 ++- pyproject.toml | 11 + src/struphy/conftest.py | 16 +- src/struphy/console/main.py | 80 ++-- src/struphy/console/test.py | 99 +--- src/struphy/console/tests/test_console.py | 9 +- src/struphy/feec/tests/test_basis_ops.py | 8 +- src/struphy/feec/tests/test_eval_field.py | 72 +-- src/struphy/feec/tests/test_field_init.py | 146 +++--- src/struphy/feec/tests/test_mass_matrices.py | 31 +- src/struphy/initial/perturbations.py | 87 +++- .../initial/tests/test_init_perturbations.py | 39 +- src/struphy/kinetic_background/base.py | 185 +------- src/struphy/kinetic_background/maxwellians.py | 249 +++++++--- .../kinetic_background/moment_functions.py | 54 --- .../kinetic_background/tests/test_base.py | 4 +- .../tests/test_maxwellians.py | 286 ++++++------ .../tests/test_saddle_point_propagator.py | 21 +- .../tests/test_saddlepoint_massmatrices.py | 2 +- src/struphy/main.py | 56 ++- src/struphy/models/base.py | 25 +- src/struphy/models/species.py | 2 +- src/struphy/models/tests/test_fluid_models.py | 28 -- .../models/tests/test_hybrid_models.py | 28 -- .../models/tests/test_kinetic_models.py | 28 -- src/struphy/models/tests/test_models.py | 172 +++++-- src/struphy/models/tests/test_toy_models.py | 28 -- ...t_LinearMHD.py => test_verif_LinearMHD.py} | 0 ...{test_Maxwell.py => test_verif_Maxwell.py} | 0 src/struphy/models/tests/test_xxpproc.py | 69 --- src/struphy/models/tests/util.py | 426 ------------------ src/struphy/models/variables.py | 16 + src/struphy/ode/tests/test_ode_feec.py | 10 +- src/struphy/pic/base.py | 68 ++- src/struphy/pic/particles.py | 148 +++--- src/struphy/pic/tests/test_accum_vec_H1.py | 14 +- src/struphy/pic/tests/test_accumulation.py | 5 +- src/struphy/pic/tests/test_binning.py | 359 +++++++-------- src/struphy/pic/tests/test_draw_parallel.py | 14 +- src/struphy/pic/tests/test_pushers.py | 42 +- src/struphy/pic/tests/test_sorting.py | 5 +- src/struphy/pic/tests/test_sph.py | 57 ++- src/struphy/pic/tests/test_tesselation.py | 31 +- .../post_processing/post_processing_tools.py | 38 +- src/struphy/propagators/propagators_fields.py | 3 +- .../tests/test_gyrokinetic_poisson.py | 194 +++++--- src/struphy/propagators/tests/test_poisson.py | 122 +++-- src/struphy/tutorials/tests/test_tutorials.py | 175 ------- src/struphy/utils/utils.py | 14 +- tutorials/tutorial_02_test_particles.ipynb | 2 +- tutorials/tutorial_05_vlasov_maxwell.ipynb | 7 +- 51 files changed, 1503 insertions(+), 2146 deletions(-) delete mode 100644 src/struphy/kinetic_background/moment_functions.py delete mode 100644 src/struphy/models/tests/test_fluid_models.py delete mode 100644 src/struphy/models/tests/test_hybrid_models.py delete mode 100644 src/struphy/models/tests/test_kinetic_models.py delete mode 100644 src/struphy/models/tests/test_toy_models.py rename src/struphy/models/tests/{test_LinearMHD.py => test_verif_LinearMHD.py} (100%) rename src/struphy/models/tests/{test_Maxwell.py => test_verif_Maxwell.py} (100%) delete mode 100644 src/struphy/models/tests/test_xxpproc.py delete mode 100644 src/struphy/models/tests/util.py delete mode 100644 src/struphy/tutorials/tests/test_tutorials.py diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index db65d9bc0..580ef9eee 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -97,6 +97,7 @@ stages: rules: - !reference [.rules_common, skip_pages_and_scheduled] - if: $CI_MERGE_REQUEST_TARGET_BRANCH_NAME == "devel" && $CI_PIPELINE_SOURCE == "merge_request_event" + - if: $CI_MERGE_REQUEST_TARGET_BRANCH_NAME == "318-parameter-file-as-py" && $CI_PIPELINE_SOURCE == "merge_request_event" .rules_mr_to_master: rules: @@ -398,6 +399,7 @@ stages: inspect_struphy: - struphy -p - struphy -h + - struphy --refresh-models - struphy run -h unit_tests: - struphy compile --status @@ -405,30 +407,39 @@ stages: - struphy test unit model_tests: - struphy compile --status - - struphy test models --fast - - struphy test models --fast --verification --mpi 1 - - struphy test models --fast --verification --mpi 4 - - struphy test models --fast --verification --mpi 4 --nclones 2 - - struphy test DriftKineticElectrostaticAdiabatic --mpi 2 --nclones 2 + - struphy test LinearMHD + - struphy test toy + - struphy test models --mpi 1 + - struphy test models --mpi 2 + - struphy test models --mpi 3 + - struphy test models --mpi 4 + - struphy test verification --mpi 4 + # - struphy test models --fast --verification --mpi 4 + # - struphy test models --fast --verification --mpi 4 --nclones 2 + # - struphy test DriftKineticElectrostaticAdiabatic --mpi 2 --nclones 2 quickstart_tests: - - struphy --set-i . - - struphy --set-o . - struphy -p - struphy -h - - struphy run -h - - struphy params VlasovMaxwellOneSpecies -y + - struphy params VlasovAmpereOneSpecies + - ls -1a + - mv params_VlasovAmpereOneSpecies.py test.py + - python3 test.py + - ls sim_1/ + - mpirun -n 2 python3 test.py + - LINE_PROFILE=1 mpirun -n 2 python3 test.py + # - struphy pproc my_first_sim + # - ls my_first_sim/post_processing/fields_data/ + # - ls my_first_sim/post_processing/kinetic_data/ + # - struphy params VlasovMaxwellOneSpecies --options + # - struphy run -i test.yml -o my_first_sim_1clone --mpi 4 --nclones 1 + # - struphy run -i test.yml -o my_first_sim_2clone --mpi 4 --nclones 2 + # - struphy run -i test.yml -o my_first_sim_4clone --mpi 4 --nclones 4 + # - struphy pproc my_first_sim_1clone my_first_sim_2clone my_first_sim_4clone + tutorial_tests: + - pwd - ls -1a - - mv params_VlasovMaxwellOneSpecies.yml test.yml - - struphy run -i test.yml -o my_first_sim - - ls my_first_sim/ - - struphy pproc my_first_sim - - ls my_first_sim/post_processing/fields_data/ - - ls my_first_sim/post_processing/kinetic_data/ - - struphy params VlasovMaxwellOneSpecies --options - - struphy run -i test.yml -o my_first_sim_1clone --mpi 4 --nclones 1 - - struphy run -i test.yml -o my_first_sim_2clone --mpi 4 --nclones 2 - - struphy run -i test.yml -o my_first_sim_4clone --mpi 4 --nclones 4 - - struphy pproc my_first_sim_1clone my_first_sim_2clone my_first_sim_4clone + - which python + - jupyter nbconvert --to notebook --execute tutorials/*.ipynb build_images: - buildah images - buildah build -t gitlab-registry.mpcdf.mpg.de/struphy/struphy/almalinux-latest -f docker/almalinux-latest.dockerfile . @@ -554,10 +565,21 @@ quickstart_tests: - !reference [.scripts, install_on_push] - !reference [.scripts, quickstart_tests] +# tutorial_tests: +# stage: test +# extends: +# - .rules_mr_to_devel +# - .image_gitlab_mpcdf_struphy +# - .before_script_load_modules +# - .variables_push +# script: +# - !reference [.scripts, install_on_push] +# - !reference [.scripts, tutorial_tests] + pages_tests: stage: test extends: - - .rules_mr_to_devel + - .rules_scheduled - .image_gitlab_mpcdf_struphy - .before_script_load_modules - .variables_push diff --git a/pyproject.toml b/pyproject.toml index 53c485685..e3c087296 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ 'argcomplete', 'pytest', 'pytest-mpi', + 'line_profiler', ] [project.license] @@ -163,3 +164,13 @@ ignore = [ "D211", "D213", ] + +[tool.pytest.ini_options] +markers = [ + "models", + "toy", + "fluid", + "kinetic", + "hybrid", + "single", +] diff --git a/src/struphy/conftest.py b/src/struphy/conftest.py index 5a936455c..05e10b55e 100644 --- a/src/struphy/conftest.py +++ b/src/struphy/conftest.py @@ -1,18 +1,14 @@ def pytest_addoption(parser): - parser.addoption("--fast", action="store_true") parser.addoption("--with-desc", action="store_true") parser.addoption("--vrbose", action="store_true") - parser.addoption("--verification", action="store_true") parser.addoption("--show-plots", action="store_true") parser.addoption("--nclones", type=int, default=1) + parser.addoption("--model-name", type=str, default="Maxwell") def pytest_generate_tests(metafunc): # This is called for every test. Only get/set command line arguments - # if the argument is specified in the list of test "fixturenames". - option_value = metafunc.config.option.fast - if "fast" in metafunc.fixturenames and option_value is not None: - metafunc.parametrize("fast", [option_value]) + # if the argument is specified in the list of test "fixturenames".]) option_value = metafunc.config.option.with_desc if "with_desc" in metafunc.fixturenames and option_value is not None: @@ -22,10 +18,6 @@ def pytest_generate_tests(metafunc): if "vrbose" in metafunc.fixturenames and option_value is not None: metafunc.parametrize("vrbose", [option_value]) - option_value = metafunc.config.option.verification - if "verification" in metafunc.fixturenames and option_value is not None: - metafunc.parametrize("verification", [option_value]) - option_value = metafunc.config.option.nclones if "nclones" in metafunc.fixturenames and option_value is not None: metafunc.parametrize("nclones", [option_value]) @@ -33,3 +25,7 @@ def pytest_generate_tests(metafunc): option_value = metafunc.config.option.show_plots if "show_plots" in metafunc.fixturenames and option_value is not None: metafunc.parametrize("show_plots", [option_value]) + + option_value = metafunc.config.option.model_name + if "model_name" in metafunc.fixturenames and option_value is not None: + metafunc.parametrize("model_name", [option_value]) diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index ed5bd895c..b12fc9ad8 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -15,10 +15,10 @@ import yaml # struphy path -import struphy as _ +import struphy from struphy.utils import utils -libpath = _.__path__[0] +libpath = struphy.__path__[0] __version__ = importlib.metadata.version("struphy") # version message @@ -63,15 +63,16 @@ def struphy(): # Load the models and messages model_message = "All models are listed on https://struphy.pages.mpcdf.de/struphy/sections/models.html" list_models = [] - try: - with open(os.path.join(libpath, "models", "models_list"), "rb") as fp: - list_models = pickle.load(fp) - # with open(os.path.join(libpath, "models", "models_message"), "rb") as fp: - # model_message, fluid_message, kinetic_message, hybrid_message, toy_message = pickle.load( - # fp, - # ) - except: - print("run: struphy --refresh-models") + ml_path = os.path.join(libpath, "models", "models_list") + if not os.path.isfile(ml_path): + utils.refresh_models() + + with open(ml_path, "rb") as fp: + list_models = pickle.load(fp) + with open(os.path.join(libpath, "models", "models_message"), "rb") as fp: + model_message, fluid_message, kinetic_message, hybrid_message, toy_message = pickle.load( + fp, + ) # 0. basic options add_parser_basic_options(parser, i_path, o_path, b_path) @@ -128,18 +129,18 @@ def struphy(): sys.exit(0) # display subset of models - # model_flags = [ - # (args.fluid, fluid_message), - # (args.kinetic, kinetic_message), - # (args.hybrid, hybrid_message), - # (args.toy, toy_message), - # ] - - # for flag, message in model_flags: - # if flag: - # print(message) - # print("For more info on Struphy models, visit https://struphy.pages.mpcdf.de/struphy/sections/models.html") - # sys.exit(0) + model_flags = [ + (args.fluid, fluid_message), + (args.kinetic, kinetic_message), + (args.hybrid, hybrid_message), + (args.toy, toy_message), + ] + + for flag, message in model_flags: + if flag: + print(message) + print("For more info on Struphy models, visit https://struphy.pages.mpcdf.de/struphy/sections/models.html") + sys.exit(0) # Set default input path if args.set_i: @@ -924,11 +925,19 @@ def add_parser_test(subparsers, list_models): parser_test.add_argument( "group", type=str, - choices=list_models + ["models"] + ["unit"] + ["fluid"] + ["kinetic"] + ["hybrid"] + ["toy"], + choices=list_models + + ["models"] + + ["unit"] + + ["fluid"] + + ["kinetic"] + + ["hybrid"] + + ["toy"] + + ["verification"], metavar="GROUP", help='can be either:\na) a model name \ \nb) "models" for testing of all models (or "fluid", "kinetic", "hybrid", "toy" for testing just a sub-group) \ - \nc) "unit" for performing unit tests', + \nc) "verification" for running all verification tests \ + \nd) "unit" for performing unit tests', ) parser_test.add_argument( @@ -939,27 +948,12 @@ def add_parser_test(subparsers, list_models): default=2, ) - parser_test.add_argument( - "-f", - "--fast", - help="test model(s) just in slab geometry (Cuboid)", - action="store_true", - ) - parser_test.add_argument( "--with-desc", help="include DESC equilibrium in tests (mem consuming)", action="store_true", ) - parser_test.add_argument( - "-T", - "--Tend", - type=float, - help="if GROUP=a), simulation end time in units of the model (default=0.015 with dt=0.005), data is only saved at TEND if set", - default=None, - ) - parser_test.add_argument( "-v", "--vrbose", @@ -967,12 +961,6 @@ def add_parser_test(subparsers, list_models): action="store_true", ) - parser_test.add_argument( - "--verification", - help="perform verification runs specified in io/inp/verification/", - action="store_true", - ) - parser_test.add_argument( "--nclones", type=int, diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 8fec47ca8..09a0872ac 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -4,12 +4,9 @@ def struphy_test( group: str, *, - mpi: int = 4, - fast: bool = False, + mpi: int = 2, with_desc: bool = False, - Tend: float = None, vrbose: bool = False, - verification: bool = False, show_plots: bool = False, nclones: int = 1, ): @@ -19,13 +16,10 @@ def struphy_test( Parameters ---------- group : str - Test identifier: "unit", "models", "fluid", "kinetic", "hybrid", "toy" or a model name. + Test identifier: "unit", "models", "fluid", "kinetic", "hybrid", "toy", "verification" or a model name. mpi : int - Number of MPI processes used in tests (must be >1, default=4). - - fast : bool - Whether to test models just in slab geometry. + Number of MPI processes used in tests (must be >1, default=2). with_desc : bool Whether to include DESC equilibrium in unit tests (mem consuming). @@ -36,9 +30,6 @@ def struphy_test( vrbose : bool Show full screen output. - verification : bool - Whether to run verification tests specified in io/inp/tests. - show_plots : bool Show plots of tests. """ @@ -76,7 +67,7 @@ def struphy_test( cmd += ["--show-plots"] subp_run(cmd) - elif "models" in group: + elif group in {"models", "fluid", "kinetic", "hybrid", "toy"}: cmd = [ "mpirun", "-n", @@ -84,103 +75,57 @@ def struphy_test( "pytest", "-k", "_models", + "-m", + group, "-s", "--with-mpi", ] - if fast: - cmd += ["--fast"] if vrbose: cmd += ["--vrbose"] - if verification: - cmd += ["--verification"] if nclones > 1: cmd += ["--nclones", f"{nclones}"] if show_plots: cmd += ["--show-plots"] subp_run(cmd) - # test post processing of models - if not verification: - cmd = [ - "pytest", - "-k", - "pproc", - "-s", - ] - subp_run(cmd) - - elif group in {"fluid", "kinetic", "hybrid", "toy"}: + elif "verification" in group: cmd = [ "mpirun", "-n", str(mpi), "pytest", "-k", - group + "_models", + "_verif_", "-s", "--with-mpi", ] - if fast: - cmd += ["--fast"] if vrbose: cmd += ["--vrbose"] - if verification: - cmd += ["--verification"] if nclones > 1: cmd += ["--nclones", f"{nclones}"] if show_plots: cmd += ["--show-plots"] subp_run(cmd) - if not verification: - from struphy.models.tests.test_xxpproc import test_pproc_codes - - test_pproc_codes(group=group) - else: - import os - import pickle - - import struphy - - libpath = struphy.__path__[0] - - with open(os.path.join(libpath, "models", "models_message"), "rb") as fp: - model_message, fluid_message, kinetic_message, hybrid_message, toy_message = pickle.load( - fp, - ) - - if group in toy_message: - mtype = "toy" - elif group in fluid_message: - mtype = "fluid" - elif group in kinetic_message: - mtype = "kinetic" - elif group in hybrid_message: - mtype = "hybrid" - else: - raise ValueError(f"{group} is not a valid model name.") - - py_file = os.path.join(libpath, "models", "tests", "util.py") - cmd = [ "mpirun", "-n", str(mpi), - "python3", - py_file, - mtype, + "pytest", + "-k", + "_models", + "-m", + "single", + "-s", + "--with-mpi", + "--model-name", group, - str(Tend), - str(fast), - str(vrbose), - str(verification), - str(nclones), - str(show_plots), ] + if vrbose: + cmd += ["--vrbose"] + if nclones > 1: + cmd += ["--nclones", f"{nclones}"] + if show_plots: + cmd += ["--show-plots"] subp_run(cmd) - - if not verification: - from struphy.models.tests.test_xxpproc import test_pproc_codes - - test_pproc_codes(group) diff --git a/src/struphy/console/tests/test_console.py b/src/struphy/console/tests/test_console.py index ce5facbb4..781f25426 100644 --- a/src/struphy/console/tests/test_console.py +++ b/src/struphy/console/tests/test_console.py @@ -78,7 +78,7 @@ def split_command(command): # ["units", "Maxwell", "--input-abs", "/params.yml"], # Test cases for 'params' sub-command ["params", "Maxwell"], - ["params", "Vlasov", "--options"], + ["params", "Vlasov"], # ["params", "Maxwell", "-f", "params_Maxwell.yml"], # Test cases for 'profile' sub-command ["profile", "sim_1"], @@ -93,7 +93,7 @@ def split_command(command): # Test cases for 'test' sub-command ["test", "models"], ["test", "unit"], - ["test", "Maxwell", "--Tend", "1.0"], + ["test", "Maxwell"], ["test", "hybrid", "--mpi", "8"], ], ) @@ -360,10 +360,9 @@ def mock_remove(path): @pytest.mark.parametrize("model", ["Maxwell"]) @pytest.mark.parametrize("file", ["params_Maxwell.yml", "params_Maxwel2.yml"]) @pytest.mark.parametrize("yes", [True]) -@pytest.mark.parametrize("options", [True, False]) -def test_struphy_params(tmp_path, model, file, yes, options): +def test_struphy_params(tmp_path, model, file, yes): file_path = os.path.join(tmp_path, file) - struphy_params(model, str(file_path), yes=yes, options=options) + struphy_params(model, str(file_path), yes=yes) @pytest.mark.mpi_skip diff --git a/src/struphy/feec/tests/test_basis_ops.py b/src/struphy/feec/tests/test_basis_ops.py index 67b810faa..3c7e52301 100644 --- a/src/struphy/feec/tests/test_basis_ops.py +++ b/src/struphy/feec/tests/test_basis_ops.py @@ -461,7 +461,7 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): @pytest.mark.parametrize("spl_kind", [[False, True, True], [False, True, False]]) @pytest.mark.parametrize( "dirichlet_bc", - [None, [[False, True], [False, False], [False, True]], [[False, False], [False, False], [True, False]]], + [None, [(False, True), (False, False), (False, True)], [(False, False), (False, False), (True, False)]], ) @pytest.mark.parametrize("mapping", [["IGAPolarCylinder", {"a": 1.0, "Lz": 3.0}]]) def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): @@ -516,9 +516,11 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal if dirichlet_bc is not None: for i, knd in enumerate(spl_kind): if knd: - dirichlet_bc[i] = [False, False] + dirichlet_bc[i] = (False, False) else: - dirichlet_bc = [[False, False]] * 3 + dirichlet_bc = [(False, False)] * 3 + + dirichlet_bc = tuple(dirichlet_bc) # derham object nq_el = [p[0] + 1, p[1] + 1, p[2] + 1] diff --git a/src/struphy/feec/tests/test_eval_field.py b/src/struphy/feec/tests/test_eval_field.py index de547c56b..9fb829942 100644 --- a/src/struphy/feec/tests/test_eval_field.py +++ b/src/struphy/feec/tests/test_eval_field.py @@ -14,6 +14,7 @@ def test_eval_field(Nel, p, spl_kind): from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import compare_arrays from struphy.geometry.base import Domain + from struphy.initial import perturbations comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -28,55 +29,28 @@ def test_eval_field(Nel, p, spl_kind): n3 = derham.create_spline_function("density", "L2") uv = derham.create_spline_function("velocity", "H1vec") - # initialize fields as forms - comps = { - "pressure": "0", - "e_field": ["1", "1", "1"], - "b_field": ["2", "2", "2"], - "density": "3", - "velocity": ["v", "v", "v"], - } - # initialize with sin/cos perturbations - pert_params_p0 = {"ModesCos": {"given_in_basis": "0", "ls": [0], "ms": [0], "ns": [1], "amps": [5.0]}} - - pert_params_E1 = { - "ModesCos": { - "given_in_basis": ["1", "1", "1"], - "ls": [[0], [0], [0]], - "ms": [[0], [0], [0]], - "ns": [[1], [1], [1]], - "amps": [[5.0], [5.0], [5.0]], - } - } - - pert_params_B2 = { - "ModesCos": { - "given_in_basis": ["2", "2", "2"], - "ls": [[0], [0], [0]], - "ms": [[0], [0], [0]], - "ns": [[1], [1], [1]], - "amps": [[5.0], [5.0], [5.0]], - } - } - - pert_params_n3 = {"ModesCos": {"given_in_basis": "3", "ls": [0], "ms": [0], "ns": [1], "amps": [5.0]}} - - pert_params_uv = { - "ModesCos": { - "given_in_basis": ["v", "v", "v"], - "ls": [[0], [0], [0]], - "ms": [[0], [0], [0]], - "ns": [[1], [1], [1]], - "amps": [[5.0], [5.0], [5.0]], - } - } - - p0.initialize_coeffs(pert_params=pert_params_p0) - E1.initialize_coeffs(pert_params=pert_params_E1) - B2.initialize_coeffs(pert_params=pert_params_B2) - n3.initialize_coeffs(pert_params=pert_params_n3) - uv.initialize_coeffs(pert_params=pert_params_uv) + pert_p0 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,)) + + pert_E1_1 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,), given_in_basis="1", comp=0) + pert_E1_2 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,), given_in_basis="1", comp=1) + pert_E1_3 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,), given_in_basis="1", comp=2) + + pert_B2_1 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,), given_in_basis="2", comp=0) + pert_B2_2 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,), given_in_basis="2", comp=1) + pert_B2_3 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,), given_in_basis="2", comp=2) + + pert_n3 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,)) + + pert_uv_1 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,), given_in_basis="v", comp=0) + pert_uv_2 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,), given_in_basis="v", comp=1) + pert_uv_3 = perturbations.ModesCos(ls=(0,), ms=(0,), ns=(1,), amps=(5.0,), given_in_basis="v", comp=2) + + p0.initialize_coeffs(perturbations=pert_p0) + E1.initialize_coeffs(perturbations=[pert_E1_1, pert_E1_2, pert_E1_3]) + B2.initialize_coeffs(perturbations=[pert_B2_1, pert_B2_2, pert_B2_3]) + n3.initialize_coeffs(perturbations=pert_n3) + uv.initialize_coeffs(perturbations=[pert_uv_1, pert_uv_2, pert_uv_3]) # evaluation points for meshgrid eta1 = np.linspace(0, 1, 11) @@ -528,6 +502,8 @@ def test_eval_field(Nel, p, spl_kind): [np.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] ) + print("\nAll assertions passed.") + if __name__ == "__main__": test_eval_field([8, 9, 10], [3, 2, 4], [False, False, True]) diff --git a/src/struphy/feec/tests/test_field_init.py b/src/struphy/feec/tests/test_field_init.py index 237b8979e..0e8b38e48 100644 --- a/src/struphy/feec/tests/test_field_init.py +++ b/src/struphy/feec/tests/test_field_init.py @@ -14,6 +14,7 @@ def test_bckgr_init_const(Nel, p, spl_kind, spaces, vec_comps): from mpi4py import MPI from struphy.feec.psydac_derham import Derham + from struphy.io.options import FieldsBackground comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -37,18 +38,18 @@ def test_bckgr_init_const(Nel, p, spl_kind, spaces, vec_comps): for i, space in enumerate(spaces): field = derham.create_spline_function("name_" + str(i), space) if space in ("H1", "L2"): - bckgr_params = {"LogicalConst": {"values": val}} - field.initialize_coeffs(bckgr_params=bckgr_params) + background = FieldsBackground(type="LogicalConst", values=(val,)) + field.initialize_coeffs(backgrounds=background) print( f"\n{rank = }, {space = }, after init:\n {np.max(np.abs(field(*meshgrids) - val)) = }", ) # print(f'{field(*meshgrids) = }') assert np.allclose(field(*meshgrids), val) else: - bckgr_params = {"LogicalConst": {"values": [val, None, val]}} - field.initialize_coeffs(bckgr_params=bckgr_params) - for j in range(3): - if bckgr_params["LogicalConst"]["values"][j]: + background = FieldsBackground(type="LogicalConst", values=(val, None, val)) + field.initialize_coeffs(backgrounds=background) + for j, val in enumerate(background.values): + if val is not None: print( f"\n{rank = }, {space = }, after init:\n {j = }, {np.max(np.abs(field(*meshgrids)[j] - val)) = }", ) @@ -71,8 +72,9 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show from struphy.feec.psydac_derham import Derham from struphy.fields_background import equils - from struphy.fields_background.base import FluidEquilibriumWithB + from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB from struphy.geometry import domains + from struphy.io.options import FieldsBackground comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -81,11 +83,11 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show derham = Derham(Nel, p, spl_kind, comm=comm) # background parameters - bckgr_params_0 = {"MHD": {"variable": "absB0"}} - bckgr_params_1 = {"MHD": {"variable": "u1"}} - bckgr_params_2 = {"MHD": {"variable": "u2"}} - bckgr_params_3 = {"MHD": {"variable": "p3"}} - bckgr_params_4 = {"MHD": {"variable": "uv"}} + bckgr_0 = FieldsBackground(type="FluidEquilibrium", variable="absB0") + bckgr_1 = FieldsBackground(type="FluidEquilibrium", variable="u1") + bckgr_2 = FieldsBackground(type="FluidEquilibrium", variable="u2") + bckgr_3 = FieldsBackground(type="FluidEquilibrium", variable="p3") + bckgr_4 = FieldsBackground(type="FluidEquilibrium", variable="uv") # evaluation grids for comparisons e1 = np.linspace(0.0, 1.0, Nel[0]) @@ -106,6 +108,9 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show continue mhd_equil = val() + if not isinstance(mhd_equil, FluidEquilibriumWithB): + continue + print(f"{mhd_equil.params = }") if "AdhocTorus" in key: @@ -148,38 +153,34 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show field_0 = derham.create_spline_function( "name_0", "H1", - bckgr_params=bckgr_params_0, + backgrounds=bckgr_0, + equil=mhd_equil, ) field_1 = derham.create_spline_function( "name_1", "Hcurl", - bckgr_params=bckgr_params_1, + backgrounds=bckgr_1, + equil=mhd_equil, ) field_2 = derham.create_spline_function( "name_2", "Hdiv", - bckgr_params=bckgr_params_2, + backgrounds=bckgr_2, + equil=mhd_equil, ) field_3 = derham.create_spline_function( "name_3", "L2", - bckgr_params=bckgr_params_3, + backgrounds=bckgr_3, + equil=mhd_equil, ) field_4 = derham.create_spline_function( "name_4", "H1vec", - bckgr_params=bckgr_params_4, + backgrounds=bckgr_4, + equil=mhd_equil, ) - field_1.initialize_coeffs(bckgr_obj=mhd_equil) - print("field_1 initialized.") - field_2.initialize_coeffs(bckgr_obj=mhd_equil) - print("field_2 initialized.") - field_3.initialize_coeffs(bckgr_obj=mhd_equil) - print("field_3 initialized.") - field_4.initialize_coeffs(bckgr_obj=mhd_equil) - print("field_4 initialized.") - # scalar spaces print( f"{np.max(np.abs(field_3(*meshgrids) - mhd_equil.p3(*meshgrids))) / np.max(np.abs(mhd_equil.p3(*meshgrids)))}" @@ -193,8 +194,6 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show ) if isinstance(mhd_equil, FluidEquilibriumWithB): - field_0.initialize_coeffs(bckgr_obj=mhd_equil) - print("field_0 initialized.") print( f"{np.max(np.abs(field_0(*meshgrids) - mhd_equil.absB0(*meshgrids))) / np.max(np.abs(mhd_equil.absB0(*meshgrids)))}" ) @@ -1088,28 +1087,34 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): from struphy.feec.psydac_derham import Derham from struphy.initial.perturbations import ModesCos, ModesSin + from struphy.io.options import FieldsBackground comm = MPI.COMM_WORLD rank = comm.Get_rank() # background parameters - avg_0 = 1.2 - avg_1 = [None, 2.6, 3.7] - avg_2 = [2, 3, 4.2] + avg_0 = (1.2,) + avg_1 = (0.0, 2.6, 3.7) + avg_2 = (2, 3, 4.2) - bckgr_params_0 = {"LogicalConst": {"values": avg_0}} - bckgr_params_1 = {"LogicalConst": {"values": avg_1}} - bckgr_params_2 = {"LogicalConst": {"values": avg_2}} + bckgr_0 = FieldsBackground(type="LogicalConst", values=avg_0) + bckgr_1 = FieldsBackground(type="LogicalConst", values=avg_1) + bckgr_2 = FieldsBackground(type="LogicalConst", values=avg_2) # perturbations ms_s = [0, 2] ns_s = [1, 1] amps = [0.2] - f_sin = ModesSin(ms=ms_s, ns=ns_s, amps=amps) + f_sin_0 = ModesSin(ms=ms_s, ns=ns_s, amps=amps) + f_sin_11 = ModesSin(ms=ms_s, ns=ns_s, amps=amps, given_in_basis="1", comp=0) + f_sin_13 = ModesSin(ms=ms_s, ns=ns_s, amps=amps, given_in_basis="1", comp=2) ms_c = [1] ns_c = [0] - f_cos = ModesCos(ms=ms_c, ns=ns_c, amps=amps) + f_cos_0 = ModesCos(ms=ms_c, ns=ns_c, amps=amps) + f_cos_11 = ModesCos(ms=ms_c, ns=ns_c, amps=amps, given_in_basis="1", comp=0) + f_cos_12 = ModesCos(ms=ms_c, ns=ns_c, amps=amps, given_in_basis="1", comp=1) + f_cos_22 = ModesCos(ms=ms_c, ns=ns_c, amps=amps, given_in_basis="2", comp=1) pert_params_0 = { "ModesSin": { @@ -1153,13 +1158,11 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): # Psydac discrete Derham sequence and fields derham = Derham(Nel, p, spl_kind, comm=comm) - field_0 = derham.create_spline_function("name_0", "H1") - field_1 = derham.create_spline_function("name_1", "Hcurl") - field_2 = derham.create_spline_function("name_2", "Hdiv") - - field_0.initialize_coeffs(bckgr_params=bckgr_params_0, pert_params=pert_params_0) - field_1.initialize_coeffs(bckgr_params=bckgr_params_1, pert_params=pert_params_1) - field_2.initialize_coeffs(bckgr_params=bckgr_params_2, pert_params=pert_params_2) + field_0 = derham.create_spline_function("name_0", "H1", backgrounds=bckgr_0, perturbations=[f_sin_0, f_cos_0]) + field_1 = derham.create_spline_function( + "name_1", "Hcurl", backgrounds=bckgr_1, perturbations=[f_sin_11, f_sin_13, f_cos_11, f_cos_12] + ) + field_2 = derham.create_spline_function("name_2", "Hdiv", backgrounds=bckgr_2, perturbations=[f_cos_22]) # evaluation grids for comparisons e1 = np.linspace(0.0, 1.0, Nel[0]) @@ -1167,24 +1170,16 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): e3 = np.linspace(0.0, 1.0, Nel[2]) meshgrids = np.meshgrid(e1, e2, e3, indexing="ij") - fun_0 = avg_0 + f_sin(*meshgrids) + f_cos(*meshgrids) - - for i, a in enumerate(avg_1): - if a is None: - avg_1[i] = 0.0 - - for i, a in enumerate(avg_2): - if a is None: - avg_2[i] = 0.0 + fun_0 = avg_0 + f_sin_0(*meshgrids) + f_cos_0(*meshgrids) fun_1 = [ - avg_1[0] + f_sin(*meshgrids) + +f_cos(*meshgrids), - avg_1[1] + f_cos(*meshgrids), - avg_1[2] + f_sin(*meshgrids), + avg_1[0] + f_sin_11(*meshgrids) + f_cos_11(*meshgrids), + avg_1[1] + f_cos_12(*meshgrids), + avg_1[2] + f_sin_13(*meshgrids), ] fun_2 = [ avg_2[0] + 0.0 * meshgrids[0], - avg_2[1] + f_cos(*meshgrids), + avg_2[1] + f_cos_22(*meshgrids), avg_2[2] + 0.0 * meshgrids[0], ] @@ -1321,6 +1316,7 @@ def test_noise_init(Nel, p, spl_kind, space, direction): from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import compare_arrays + from struphy.initial.perturbations import Noise comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -1333,16 +1329,10 @@ def test_noise_init(Nel, p, spl_kind, space, direction): field_np = derham_np.create_spline_function("field", space) # initial conditions - pert_params = { - "noise": { - "comps": [True, False, False], - "direction": direction, - "amp": 0.0001, - "seed": 1234, - }, - } - field.initialize_coeffs(pert_params=pert_params) - field_np.initialize_coeffs(pert_params=pert_params) + pert = Noise(direction=direction, amp=0.0001, seed=1234, comp=0) + + field.initialize_coeffs(perturbations=pert) + field_np.initialize_coeffs(perturbations=pert) # print('#'*80) # print(f'npts={field.vector[0].space.npts}, npts_np={field_np.vector[0].space.npts}') @@ -1360,15 +1350,15 @@ def test_noise_init(Nel, p, spl_kind, space, direction): if __name__ == "__main__": # test_bckgr_init_const([8, 10, 12], [1, 2, 3], [False, False, True], [ # 'H1', 'Hcurl', 'Hdiv'], [True, True, False]) - test_bckgr_init_mhd( - [18, 24, 12], - [1, 2, 1], - [ - False, - True, - True, - ], - show_plot=True, - ) - # test_sincos_init_const([1, 32, 32], [1, 3, 3], [True]*3, show_plot=True) - # test_noise_init([4, 8, 6], [1, 1, 1], [True, True, True], "Hcurl", "e1") + # test_bckgr_init_mhd( + # [18, 24, 12], + # [1, 2, 1], + # [ + # False, + # True, + # True, + # ], + # show_plot=False, + # ) + test_sincos_init_const([1, 32, 32], [1, 3, 3], [True] * 3, show_plot=True) + test_noise_init([4, 8, 6], [1, 1, 1], [True, True, True], "Hcurl", "e1") diff --git a/src/struphy/feec/tests/test_mass_matrices.py b/src/struphy/feec/tests/test_mass_matrices.py index d1031c32e..5f05c3228 100644 --- a/src/struphy/feec/tests/test_mass_matrices.py +++ b/src/struphy/feec/tests/test_mass_matrices.py @@ -7,7 +7,7 @@ @pytest.mark.parametrize("spl_kind", [[False, True, True], [True, False, True]]) @pytest.mark.parametrize( "dirichlet_bc", - [None, [[False, True], [True, False], [False, False]], [[True, False], [False, True], [False, False]]], + [None, [(False, True), (True, False), (False, False)], [(True, False), (False, True), (False, False)]], ) @pytest.mark.parametrize("mapping", [["Colella", {"Lx": 1.0, "Ly": 6.0, "alpha": 0.1, "Lz": 10.0}]]) def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): @@ -102,10 +102,11 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): if dirichlet_bc is not None: for i, knd in enumerate(spl_kind): if knd: - dirichlet_bc[i] = [False, False] + dirichlet_bc[i] = (False, False) else: - dirichlet_bc = [[False, False]] * 3 + dirichlet_bc = [(False, False)] * 3 + dirichlet_bc = tuple(dirichlet_bc) print(f"{dirichlet_bc = }") # derham object @@ -371,7 +372,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): @pytest.mark.parametrize("spl_kind", [[False, True, True], [False, True, False]]) @pytest.mark.parametrize( "dirichlet_bc", - [None, [[False, True], [False, False], [False, True]], [[False, False], [False, False], [True, False]]], + [None, [(False, True), (False, False), (False, True)], [(False, False), (False, False), (True, False)]], ) @pytest.mark.parametrize("mapping", [["IGAPolarCylinder", {"a": 1.0, "Lz": 3.0}]]) def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): @@ -433,9 +434,11 @@ def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): if dirichlet_bc is not None: for i, knd in enumerate(spl_kind): if knd: - dirichlet_bc[i] = [False, False] + dirichlet_bc[i] = (False, False) else: - dirichlet_bc = [[False, False]] * 3 + dirichlet_bc = [(False, False)] * 3 + + dirichlet_bc = tuple(dirichlet_bc) # derham object derham = Derham( @@ -566,7 +569,7 @@ def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): @pytest.mark.parametrize("spl_kind", [[False, True, True], [False, True, False]]) @pytest.mark.parametrize( "dirichlet_bc", - [None, [[False, True], [False, False], [False, True]], [[False, False], [False, False], [True, False]]], + [None, [(False, True), (False, False), (False, True)], [(False, False), (False, False), (True, False)]], ) @pytest.mark.parametrize("mapping", [["HollowCylinder", {"a1": 0.1, "a2": 1.0, "Lz": 18.84955592153876}]]) def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): @@ -664,9 +667,11 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots if dirichlet_bc is not None: for i, knd in enumerate(spl_kind): if knd: - dirichlet_bc[i] = [False, False] + dirichlet_bc[i] = (False, False) else: - dirichlet_bc = [[False, False]] * 3 + dirichlet_bc = [(False, False)] * 3 + + dirichlet_bc = tuple(dirichlet_bc) # derham object derham = Derham(Nel, p, spl_kind, comm=mpi_comm, dirichlet_bc=dirichlet_bc) @@ -872,7 +877,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots @pytest.mark.parametrize("spl_kind", [[False, True, True], [False, True, False]]) @pytest.mark.parametrize( "dirichlet_bc", - [None, [[False, True], [False, False], [False, True]], [[False, False], [False, False], [True, False]]], + [None, [(False, True), (False, False), (False, True)], [(False, False), (False, False), (True, False)]], ) @pytest.mark.parametrize("mapping", [["IGAPolarCylinder", {"a": 1.0, "Lz": 3.0}]]) def test_mass_preconditioner_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): @@ -937,9 +942,11 @@ def test_mass_preconditioner_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show if dirichlet_bc is not None: for i, knd in enumerate(spl_kind): if knd: - dirichlet_bc[i] = [False, False] + dirichlet_bc[i] = (False, False) else: - dirichlet_bc = [[False, False]] * 3 + dirichlet_bc = [(False, False)] * 3 + + dirichlet_bc = tuple(dirichlet_bc) # derham object derham = Derham( diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 45c7d7780..1ff365c8b 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -66,20 +66,20 @@ class ModesSin(Perturbation): Parameters ---------- - ls : tuple | list + ls : tuple[int] Mode numbers in x-direction (kx = l*2*pi/Lx). - ms : tuple | list + ms : tuple[int] Mode numbers in y-direction (ky = m*2*pi/Ly). - ns : tuple | list + ns : tuple[int] Mode numbers in z-direction (kz = n*2*pi/Lz). - amps : tuple | list + amps : tuple[float] Amplitude of each mode. theta : tuple | list - Phase of each mode + Phase of each mode. pfuns : tuple | list[str] "Id" or "localize" define the profile functions. @@ -101,11 +101,11 @@ class ModesSin(Perturbation): def __init__( self, - ls=None, - ms=None, - ns=None, - amps=(1e-4,), - theta=None, + ls: tuple[int] = None, + ms: tuple[int] = None, + ns: tuple[int] = None, + amps: tuple[float] = (1e-4,), + theta: tuple[float] = None, pfuns=("Id",), pfuns_params=(0.0,), Lx=1.0, @@ -216,16 +216,16 @@ class ModesCos(Perturbation): Parameters ---------- - ls : tuple | list + ls : tuple[int] Mode numbers in x-direction (kx = l*2*pi/Lx). - ms : tuple | list + ms : tuple[int] Mode numbers in y-direction (ky = m*2*pi/Ly). - ns : tuple | list + ns : tuple[int] Mode numbers in z-direction (kz = n*2*pi/Lz). - amps : tuple | list + amps : tuple[float] Amplitude of each mode. Lx, Ly, Lz : float @@ -240,10 +240,10 @@ class ModesCos(Perturbation): def __init__( self, - ls=None, - ms=None, - ns=None, - amps=(1e-4,), + ls: tuple[int] = None, + ms: tuple[int] = None, + ns: tuple[int] = None, + amps: tuple[float] = (1e-4,), Lx=1.0, Ly=1.0, Lz=1.0, @@ -1516,3 +1516,54 @@ def __call__(self, x, y, z): else: raise ValueError(f"Invalid species '{self._species}'. Must be 'Ions' or 'Electrons'.") + + +class ITPA_density(Perturbation): + r"""ITPA radial density profile in `A. Könies et al. 2018 `_ + + .. math:: + + n(\eta_1) = n_0*c_3\exp\left[-\frac{c_2}{c_1}\tanh\left(\frac{\eta_1 - c_0}{c_2}\right)\right]\,. + """ + + def __init__( + self, + n0: float = 0.00720655, + c: tuple = (0.491230, 0.298228, 0.198739, 0.521298), + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): + """ + Parameters + ---------- + n0 : float + ITPA profile density + + c : tuple | list + 4 ITPA profile coefficients + """ + + assert len(c) == 4 + + self._n0 = n0 + self._c = c + + # use the setters + self.given_in_basis = "physical" + self.comp = comp + + def __call__(self, eta1, eta2=None, eta3=None): + val = 0.0 + + if self._c[2] == 0.0: + val = self._c[3] - 0 * eta1 + else: + val = ( + self._n0 + * self._c[3] + * np.exp( + -self._c[2] / self._c[1] * np.tanh((eta1 - self._c[0]) / self._c[2]), + ) + ) + + return val diff --git a/src/struphy/initial/tests/test_init_perturbations.py b/src/struphy/initial/tests/test_init_perturbations.py index 3b46728c1..105043c01 100644 --- a/src/struphy/initial/tests/test_init_perturbations.py +++ b/src/struphy/initial/tests/test_init_perturbations.py @@ -80,7 +80,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False form_vector += ["physical"] for key, val in inspect.getmembers(perturbations): - if inspect.isclass(val): + if inspect.isclass(val) and val.__module__ == perturbations.__name__: print(key, val) if "Modes" not in key: @@ -109,10 +109,21 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False if form in ("0", "3"): for n, fun_form in enumerate(form_scalar): + if "Torus" in key and fun_form == "physical": + continue + + if "Modes" in key and fun_form == "physical": + perturbation._Lx = Lx + perturbation._Ly = Ly + perturbation._Lz = Lz + else: + perturbation._Lx = 1.0 + perturbation._Ly = 1.0 + perturbation._Lz = 1.0 # use the setter perturbation.given_in_basis = fun_form - var = FEECVariable(name=form, space=space) + var = FEECVariable(space=space) var.add_perturbation(perturbation) var.allocate(derham, domain) field = var.spline @@ -194,6 +205,17 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False else: for n, fun_form in enumerate(form_vector): + if "Torus" in key and fun_form == "physical": + continue + + if "Modes" in key and fun_form == "physical": + perturbation._Lx = Lx + perturbation._Ly = Ly + perturbation._Lz = Lz + else: + perturbation._Lx = 1.0 + perturbation._Ly = 1.0 + perturbation._Lz = 1.0 perturbation_0 = perturbation perturbation_1 = deepcopy(perturbation) perturbation_2 = deepcopy(perturbation) @@ -222,7 +244,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False perturbation_2.given_in_basis = fun_form perturbation_2.comp = 2 - var = FEECVariable(name=form, space=space) + var = FEECVariable(space=space) var.add_perturbation(perturbation_0) var.add_perturbation(perturbation_1) var.add_perturbation(perturbation_2) @@ -306,11 +328,8 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False if __name__ == "__main__": # mapping = ['Colella', {'Lx': 4., 'Ly': 5., 'alpha': .07, 'Lz': 6.}] - # mapping = ['HollowCylinder', {'a1': 0.1}] + mapping = ["HollowCylinder", {"a1": 0.1}] # mapping = ['Cuboid', {'l1': 0., 'r1': 4., 'l2': 0., 'r2': 5., 'l3': 0., 'r3': 6.}] - # test_init_modes([16, 16, 16], [2, 3, 4], [False, True, True], - # mapping, - # combine_comps=None, - # do_plot=False) - mapping = ["HollowTorus", {"tor_period": 1}] - test_init_modes([16, 14, 14], [2, 3, 4], [False, True, True], mapping, combine_comps=None, do_plot=True) + test_init_modes([16, 16, 16], [2, 3, 4], [False, True, True], mapping, combine_comps=None, do_plot=False) + # mapping = ["HollowTorus", {"tor_period": 1}] + # test_init_modes([16, 14, 14], [2, 3, 4], [False, True, True], mapping, combine_comps=None, do_plot=True) diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 353127c21..ca84a8ec5 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -5,11 +5,10 @@ import numpy as np -from struphy.fields_background.base import FluidEquilibrium +from struphy.fields_background.base import FluidEquilibriumWithB from struphy.fields_background.equils import set_defaults from struphy.initial.base import Perturbation from struphy.initial.utilities import Noise -from struphy.kinetic_background import moment_functions class KineticBackground(metaclass=ABCMeta): @@ -146,6 +145,10 @@ def __init__(self, f1, f2): self._f1 = f1 self._f2 = f2 + if hasattr(f1, "_equil"): + assert f1.equil is f2.equil + self._equil = f1.equil + @property def coords(self): """Coordinates of the distribution.""" @@ -166,6 +169,13 @@ def volume_form(self): """Boolean. True if the background is represented as a volume form (thus including the velocity Jacobian).""" return self._f1.volume_form + @property + def equil(self) -> FluidEquilibriumWithB: + """Fluid background with B-field.""" + if not hasattr(self, "_equil"): + self._equil = None + return self._equil + def velocity_jacobian_det(self, eta1, eta2, eta3, *v): """Jacobian determinant of the velocity coordinate transformation.""" return self._f1.velocity_jacobian_det(eta1, eta2, eta3, *v) @@ -356,7 +366,8 @@ def check_maxw_params(self): assert isinstance(k, str) assert isinstance(v, tuple), f"Maxwallian parameter {k} must be tuple, but is {v}" assert len(v) == 2 - assert isinstance(v[0], (float, int, Callable[..., float])) + + assert isinstance(v[0], (float, int, Callable)) assert isinstance(v[1], Perturbation) or v[1] is None @classmethod @@ -542,10 +553,10 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name: str = "n", add_perturbatio out += background else: assert callable(background) - if eta1.ndim == 1: - out += background(eta1, eta2, eta3) - else: - out += background(*etas) + # if eta1.ndim == 1: + # out += background(eta1, eta2, eta3) + # else: + out += background(*etas) # add perturbation if add_perturbation is None: @@ -571,163 +582,3 @@ def add_perturbation(self) -> bool: def add_perturbation(self, new): assert isinstance(new, bool) self._add_perturbation = new - - -class CanonicalMaxwellian(metaclass=ABCMeta): - r"""Base class for a canonical Maxwellian distribution function. - It is defined by three constants of motion in the axissymmetric toroidal system: - - - Shifted canonical toroidal momentum - - .. math:: - - \psi_c = \psi + \frac{m_s F}{q_s B}v_\parallel - \text{sign}(v_\parallel)\sqrt{2(\epsilon - \mu B)}\frac{m_sF}{q_sB} \mathcal{H}(\epsilon - \mu B), - - - Energy - - .. math:: - - \epsilon = \frac{1}{2}m_sv_\parallel² + \mu B, - - - Magnetic moment - - .. math:: - - \mu = \frac{m_s v_\perp²}{2B}, - - where :math:`\psi` is the poloidal magnetic flux function, :math:`F=F(\psi)` is the poloidal current function and :math:`\mathcal{H}` is the Heaviside function. - - With the three constants of motion, a canonical Maxwellian distribution function is defined as - - .. math:: - - F(\psi_c, \epsilon, \mu) = \frac{n(\psi_c)}{(2\pi)^{3/2}v_\text{th}³(\psi_c)} \text{exp}\left[ - \frac{\epsilon}{v_\text{th}²(\psi_c)}\right]. - - """ - - @property - @abstractmethod - def coords(self): - """Coordinates of the distribution.""" - pass - - @abstractmethod - def velocity_jacobian_det(self, eta1, eta2, eta3, *v): - """Jacobian determinant of the velocity coordinate transformation.""" - pass - - @abstractmethod - def n(self, psic): - """Number density (0-form). - - Parameters - ---------- - psic : numpy.arrays - Shifted canonical toroidal momentum. - - Returns - ------- - A numpy.array with the density evaluated at evaluation points (same shape as etas). - """ - pass - - @abstractmethod - def vth(self, psic): - """Thermal velocities (0-forms). - - Parameters - ---------- - psic : numpy.arrays - Shifted canonical toroidal momentum. - - Returns - ------- - A numpy.array with the thermal velocity evaluated at evaluation points (one dimension more than etas). - The additional dimension is in the first index. - """ - pass - - def gaussian(self, e, vth=1.0): - """3-dim. normal distribution, to which array-valued thermal velocities can be passed. - - Parameters - ---------- - e : float | array-like - Energy. - - vth : float | array-like - Thermal velocity evaluated at psic. - - Returns - ------- - An array of size(e). - """ - - if isinstance(vth, np.ndarray): - assert e.shape == vth.shape, f"{e.shape = } but {vth.shape = }" - - return 2.0 * np.sqrt(e / np.pi) / vth**3 * np.exp(-e / vth**2) - - def __call__(self, *args): - """Evaluates the canonical Maxwellian distribution function. - - There are two use-cases for this function in the code: - - 1. Evaluating for particles ("flat evaluation", inputs are all 1D of length N_p) - 2. Evaluating the function on a meshgrid (in phase space). - - Hence all arguments must always have - - 1. the same shape - 2. either ndim = 1 or ndim = 3. - - Parameters - ---------- - *args : array_like - Position-velocity arguments in the order energy, magnetic moment, canonical toroidal momentum. - - Returns - ------- - f : np.ndarray - The evaluated Maxwellian. - """ - - # Check that all args have the same shape - shape0 = np.shape(args[0]) - for i, arg in enumerate(args): - assert np.shape(arg) == shape0, f"Argument {i} has {np.shape(arg) = }, but must be {shape0 = }." - assert np.ndim(arg) == 1 or np.ndim(arg) == 3, ( - f"{np.ndim(arg) = } not allowed for canonical Maxwellian evaluation." - ) # flat or meshgrid evaluation - - # Get result evaluated with each particles' psic - res = self.n(args[2]) - vths = self.vth(args[2]) - - # take care of correct broadcasting, assuming args come from phase space meshgrid - if np.ndim(args[0]) == 3: - # move eta axes to the back - arg_t = np.moveaxis(args[0], 0, -1) - arg_t = np.moveaxis(arg_t, 0, -1) - arg_t = np.moveaxis(arg_t, 0, -1) - - # broadcast - res_broad = res + 0.0 * arg_t - - # move eta axes to the front - res = np.moveaxis(res_broad, -1, 0) - res = np.moveaxis(res, -1, 0) - res = np.moveaxis(res, -1, 0) - - # Multiply result with gaussian in energy - if np.ndim(args[0]) == 3: - vth_broad = vths + 0.0 * arg_t - vth = np.moveaxis(vth_broad, -1, 0) - vth = np.moveaxis(vth, -1, 0) - vth = np.moveaxis(vth, -1, 0) - else: - vth = vths - - res *= self.gaussian(args[0], vth=vth) - - return res diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index eb60e09b6..5e240a1b7 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -7,8 +7,7 @@ from struphy.fields_background.base import FluidEquilibriumWithB from struphy.fields_background.equils import set_defaults from struphy.initial.base import Perturbation -from struphy.kinetic_background import moment_functions -from struphy.kinetic_background.base import CanonicalMaxwellian, Maxwellian +from struphy.kinetic_background.base import Maxwellian class Maxwellian3D(Maxwellian): @@ -23,13 +22,13 @@ class Maxwellian3D(Maxwellian): def __init__( self, - n: tuple[float | Callable[..., float], Perturbation] = (1.0, None), - u1: tuple[float | Callable[..., float], Perturbation] = (0.0, None), - u2: tuple[float | Callable[..., float], Perturbation] = (0.0, None), - u3: tuple[float | Callable[..., float], Perturbation] = (0.0, None), - vth1: tuple[float | Callable[..., float], Perturbation] = (1.0, None), - vth2: tuple[float | Callable[..., float], Perturbation] = (1.0, None), - vth3: tuple[float | Callable[..., float], Perturbation] = (1.0, None), + n: tuple[float | Callable, Perturbation] = (1.0, None), + u1: tuple[float | Callable, Perturbation] = (0.0, None), + u2: tuple[float | Callable, Perturbation] = (0.0, None), + u3: tuple[float | Callable, Perturbation] = (0.0, None), + vth1: tuple[float | Callable, Perturbation] = (1.0, None), + vth2: tuple[float | Callable, Perturbation] = (1.0, None), + vth3: tuple[float | Callable, Perturbation] = (1.0, None), ): self._maxw_params = {} self._maxw_params["n"] = n @@ -164,11 +163,11 @@ class GyroMaxwellian2D(Maxwellian): def __init__( self, - n: tuple[float | Callable[..., float], Perturbation] = (1.0, None), - u_para: tuple[float | Callable[..., float], Perturbation] = (0.0, None), - u_perp: tuple[float | Callable[..., float], Perturbation] = (0.0, None), - vth_para: tuple[float | Callable[..., float], Perturbation] = (1.0, None), - vth_perp: tuple[float | Callable[..., float], Perturbation] = (1.0, None), + n: tuple[float | Callable, Perturbation] = (1.0, None), + u_para: tuple[float | Callable, Perturbation] = (0.0, None), + u_perp: tuple[float | Callable, Perturbation] = (0.0, None), + vth_para: tuple[float | Callable, Perturbation] = (1.0, None), + vth_perp: tuple[float | Callable, Perturbation] = (1.0, None), equil: FluidEquilibriumWithB = None, volume_form: bool = True, ): @@ -302,8 +301,35 @@ def vth(self, eta1, eta2, eta3): return [ou * mom_fac for ou, mom_fac in zip(out, self.moment_factors["vth"])] -class CanonicalMaxwellian(CanonicalMaxwellian): - r"""A :class:`~struphy.kinetic_background.base.CanonicalMaxwellian`. +class CanonicalMaxwellian: + r"""canonical Maxwellian distribution function. + It is defined by three constants of motion in the axissymmetric toroidal system: + + - Shifted canonical toroidal momentum + + .. math:: + + \psi_c = \psi + \frac{m_s F}{q_s B}v_\parallel - \text{sign}(v_\parallel)\sqrt{2(\epsilon - \mu B)}\frac{m_sF}{q_sB} \mathcal{H}(\epsilon - \mu B), + + - Energy + + .. math:: + + \epsilon = \frac{1}{2}m_sv_\parallel² + \mu B, + + - Magnetic moment + + .. math:: + + \mu = \frac{m_s v_\perp²}{2B}, + + where :math:`\psi` is the poloidal magnetic flux function, :math:`F=F(\psi)` is the poloidal current function and :math:`\mathcal{H}` is the Heaviside function. + + With the three constants of motion, a canonical Maxwellian distribution function is defined as + + .. math:: + + F(\psi_c, \epsilon, \mu) = \frac{n(\psi_c)}{(2\pi)^{3/2}v_\text{th}³(\psi_c)} \text{exp}\left[ - \frac{\epsilon}{v_\text{th}²(\psi_c)}\right]. Parameters ---------- @@ -322,46 +348,22 @@ class CanonicalMaxwellian(CanonicalMaxwellian): of the polar coordinate transofrmation (default = False). """ - @classmethod - def default_maxw_params(cls): - """Default parameters dictionary defining constant moments of the Maxwellian.""" - return { - "n": 1.0, - "vth": 1.0, - "type": "Particles5D", - } - def __init__( self, - maxw_params: dict = None, - pert_params: dict = None, + n: tuple[float | Callable, Perturbation] = (1.0, None), + vth: tuple[float | Callable, Perturbation] = (1.0, None), equil: FluidEquilibriumWithB = None, volume_form: bool = True, ): - # Set background parameters - self._maxw_params = self.default_maxw_params() - - if maxw_params is not None: - assert isinstance(maxw_params, dict) - self._maxw_params = set_defaults( - maxw_params, - self.default_maxw_params(), - ) - - # Set parameters for perturbation - self._pert_params = pert_params - - if self.pert_params is not None: - assert isinstance(pert_params, dict) - assert "type" in self.pert_params, '"type" is mandatory in perturbation dictionary.' - ptype = self.pert_params["type"] - assert ptype in self.pert_params, f"{ptype} is mandatory in perturbation dictionary." - self._pert_type = ptype + self._maxw_params = {} + self._maxw_params["n"] = n + self._maxw_params["vth"] = vth - self._equil = equil + self.check_maxw_params() # volume form represenation self._volume_form = volume_form + self._equil = equil # factors multiplied onto the defined moments n and vth (can be set via setter) self._moment_factors = { @@ -380,17 +382,21 @@ def maxw_params(self): return self._maxw_params @property - def pert_params(self): - """Parameters dictionary defining the perturbations of the :meth:`~Maxwellian5D.maxw_params`.""" - return self._pert_params - - @property - def equil(self): + def equil(self) -> FluidEquilibriumWithB: """One of :mod:`~struphy.fields_background.equils` in case that moments are to be set in that way, None otherwise. """ return self._equil + def check_maxw_params(self): + for k, v in self.maxw_params.items(): + assert isinstance(k, str) + assert isinstance(v, tuple), f"Maxwallian parameter {k} must be tuple, but is {v}" + assert len(v) == 2 + + assert isinstance(v[0], (float, int, Callable)) + assert isinstance(v[1], Perturbation) or v[1] is None + def velocity_jacobian_det(self, eta1, eta2, eta3, energy): r"""TODO""" @@ -408,6 +414,91 @@ def velocity_jacobian_det(self, eta1, eta2, eta3, energy): return np.sqrt(energy) * 2.0 * np.sqrt(2.0) / absB0 + def gaussian(self, e, vth=1.0): + """3-dim. normal distribution, to which array-valued thermal velocities can be passed. + + Parameters + ---------- + e : float | array-like + Energy. + + vth : float | array-like + Thermal velocity evaluated at psic. + + Returns + ------- + An array of size(e). + """ + + if isinstance(vth, np.ndarray): + assert e.shape == vth.shape, f"{e.shape = } but {vth.shape = }" + + return 2.0 * np.sqrt(e / np.pi) / vth**3 * np.exp(-e / vth**2) + + def __call__(self, *args): + """Evaluates the canonical Maxwellian distribution function. + + There are two use-cases for this function in the code: + + 1. Evaluating for particles ("flat evaluation", inputs are all 1D of length N_p) + 2. Evaluating the function on a meshgrid (in phase space). + + Hence all arguments must always have + + 1. the same shape + 2. either ndim = 1 or ndim = 3. + + Parameters + ---------- + *args : array_like + Position-velocity arguments in the order energy, magnetic moment, canonical toroidal momentum. + + Returns + ------- + f : np.ndarray + The evaluated Maxwellian. + """ + + # Check that all args have the same shape + shape0 = np.shape(args[0]) + for i, arg in enumerate(args): + assert np.shape(arg) == shape0, f"Argument {i} has {np.shape(arg) = }, but must be {shape0 = }." + assert np.ndim(arg) == 1 or np.ndim(arg) == 3, ( + f"{np.ndim(arg) = } not allowed for canonical Maxwellian evaluation." + ) # flat or meshgrid evaluation + + # Get result evaluated with each particles' psic + res = self.n(args[2]) + vths = self.vth(args[2]) + + # take care of correct broadcasting, assuming args come from phase space meshgrid + if np.ndim(args[0]) == 3: + # move eta axes to the back + arg_t = np.moveaxis(args[0], 0, -1) + arg_t = np.moveaxis(arg_t, 0, -1) + arg_t = np.moveaxis(arg_t, 0, -1) + + # broadcast + res_broad = res + 0.0 * arg_t + + # move eta axes to the front + res = np.moveaxis(res_broad, -1, 0) + res = np.moveaxis(res, -1, 0) + res = np.moveaxis(res, -1, 0) + + # Multiply result with gaussian in energy + if np.ndim(args[0]) == 3: + vth_broad = vths + 0.0 * arg_t + vth = np.moveaxis(vth_broad, -1, 0) + vth = np.moveaxis(vth, -1, 0) + vth = np.moveaxis(vth, -1, 0) + else: + vth = vths + + res *= self.gaussian(args[0], vth=vth) + + return res + @property def volume_form(self): """Boolean. True if the background is represented as a volume form (thus including the velocity Jacobian |v_perp|).""" @@ -464,7 +555,7 @@ def rc(self, psic): return rc - def n(self, psic): + def n(self, psic, add_perturbation: bool = None): """Density as background + perturbation. Parameters @@ -476,7 +567,6 @@ def n(self, psic): ------- A float (background value) or a numpy.array of the evaluated density. """ - # collect arguments assert isinstance(psic, np.ndarray) @@ -485,15 +575,26 @@ def n(self, psic): psic = psic[0, 0, :] # set background density - if isinstance(self.maxw_params["n"], dict): - mom_funcs = self.maxw_params["n"] - for typ, params in mom_funcs.items(): - nfun = getattr(moment_functions, typ)(**params) - res = nfun(eta1=self.rc(psic)) + if isinstance(self.maxw_params["n"][0], (float, int)): + res = self.maxw_params["n"][0] + 0.0 * psic else: - res = self.maxw_params["n"] + 0.0 * psic + nfun = self.maxw_params["n"][1] + # for typ, params in mom_funcs.items(): + # nfun = getattr(moment_functions, typ)(**params) + res = nfun(eta1=self.rc(psic)) - # TODO: add perturbation + # add perturbation + if add_perturbation is None: + add_perturbation = self.add_perturbation + + perturbation = self.maxw_params["n"][1] + if perturbation is not None and add_perturbation: + assert isinstance(perturbation, Perturbation) + res = perturbation(eta1=self.rc(psic)) + # if eta1.ndim == 1: + # out += perturbation(eta1, eta2, eta3) + # else: + # out += perturbation(*etas) return res * self.moment_factors["n"] @@ -517,12 +618,23 @@ def vth(self, psic): if psic.ndim == 3: psic = psic[0, 0, :] - res = self.maxw_params["vth"] + 0.0 * psic + res = self.maxw_params["vth"][0] + 0.0 * psic # TODO: add perturbation return res * self.moment_factors["vth"] + @property + def add_perturbation(self) -> bool: + if not hasattr(self, "_add_perturbation"): + self._add_perturbation = True + return self._add_perturbation + + @add_perturbation.setter + def add_perturbation(self, new): + assert isinstance(new, bool) + self._add_perturbation = new + class ColdPlasma(Maxwellian): r"""Base class for a distribution as a Dirac-delta in velocity (vth = 0). @@ -606,3 +718,14 @@ def vth(self, eta1, eta2, eta3): def __call__(self, eta1, eta2, eta3): return self.n(eta1, eta2, eta3) + + @property + def add_perturbation(self) -> bool: + if not hasattr(self, "_add_perturbation"): + self._add_perturbation = True + return self._add_perturbation + + @add_perturbation.setter + def add_perturbation(self, new): + assert isinstance(new, bool) + self._add_perturbation = new diff --git a/src/struphy/kinetic_background/moment_functions.py b/src/struphy/kinetic_background/moment_functions.py deleted file mode 100644 index 629958bfe..000000000 --- a/src/struphy/kinetic_background/moment_functions.py +++ /dev/null @@ -1,54 +0,0 @@ -#!/usr/bin/env python3 -"Analytical moment functions." - -import numpy as np - - -class ITPA_density: - r"""ITPA radial density profile in `A. Könies et al. 2018 `_ - - .. math:: - - n(\eta_1) = n_0*c_3\exp\left[-\frac{c_2}{c_1}\tanh\left(\frac{\eta_1 - c_0}{c_2}\right)\right]\,. - - Note - ---- - In the parameter .yml, use the following template in the section ``kinetic/``:: - - ITPA_density : - given_in_basis : '0' - n0 : 0.00720655 - c : [0.491230, 0.298228, 0.198739, 0.521298] - """ - - def __init__(self, n0=0.00720655, c=(0.491230, 0.298228, 0.198739, 0.521298)): - """ - Parameters - ---------- - n0 : float - ITPA profile density - - c : tuple | list - 4 ITPA profile coefficients - """ - - assert len(c) == 4 - - self._n0 = n0 - self._c = c - - def __call__(self, eta1, eta2=None, eta3=None): - val = 0.0 - - if self._c[2] == 0.0: - val = self._c[3] - 0 * eta1 - else: - val = ( - self._n0 - * self._c[3] - * np.exp( - -self._c[2] / self._c[1] * np.tanh((eta1 - self._c[0]) / self._c[2]), - ) - ) - - return val diff --git a/src/struphy/kinetic_background/tests/test_base.py b/src/struphy/kinetic_background/tests/test_base.py index 425dad0ff..09172c866 100644 --- a/src/struphy/kinetic_background/tests/test_base.py +++ b/src/struphy/kinetic_background/tests/test_base.py @@ -15,8 +15,8 @@ def test_kinetic_background_magics(show_plot=False): m1_params = {"n": 0.5, "u1": 3.0} m2_params = {"n": 0.5, "u1": -3.0} - m1 = Maxwellian3D(maxw_params=m1_params) - m2 = Maxwellian3D(maxw_params=m2_params) + m1 = Maxwellian3D(n=(0.5, None), u1=(3.0, None)) + m2 = Maxwellian3D(n=(0.5, None), u1=(-3.0, None)) m_add = m1 + m2 m_rmul_int = 2 * m1 diff --git a/src/struphy/kinetic_background/tests/test_maxwellians.py b/src/struphy/kinetic_background/tests/test_maxwellians.py index 0c061ed58..387ba8f50 100644 --- a/src/struphy/kinetic_background/tests/test_maxwellians.py +++ b/src/struphy/kinetic_background/tests/test_maxwellians.py @@ -20,9 +20,7 @@ def test_maxwellian_3d_uniform(Nel, show_plot=False): # ========================================================== # ==== Test uniform non-shifted, isothermal Maxwellian ===== # ========================================================== - maxw_params = {"n": 2.0} - - maxwellian = Maxwellian3D(maxw_params=maxw_params) + maxwellian = Maxwellian3D(n=(2.0, None)) meshgrids = np.meshgrid(e1, e2, e3, [0.0], [0.0], [0.0]) @@ -56,9 +54,16 @@ def test_maxwellian_3d_uniform(Nel, show_plot=False): vth1 = 1.2 vth2 = 0.5 vth3 = 0.3 - maxw_params = {"n": n, "u1": u1, "u2": u2, "u3": u3, "vth1": vth1, "vth2": vth2, "vth3": vth3} - maxwellian = Maxwellian3D(maxw_params=maxw_params) + maxwellian = Maxwellian3D( + n=(2.0, None), + u1=(1.0, None), + u2=(-0.2, None), + u3=(0.1, None), + vth1=(1.2, None), + vth2=(0.5, None), + vth3=(0.3, None), + ) # test Maxwellian profile in v for i in range(3): @@ -91,6 +96,7 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): import matplotlib.pyplot as plt import numpy as np + from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import Maxwellian3D e1 = np.linspace(0.0, 1.0, Nel[0]) @@ -102,18 +108,9 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): amp = 0.1 mode = 1 - maxw_params = {"n": 2.0} - pert_params = { - "n": { - "ModesCos": { - "given_in_basis": "0", - "ls": [mode], - "amps": [amp], - } - } - } + pert = perturbations.ModesCos(ls=(mode,), amps=(amp,)) - maxwellian = Maxwellian3D(maxw_params=maxw_params, pert_params=pert_params) + maxwellian = Maxwellian3D(n=(2.0, pert)) meshgrids = np.meshgrid(e1, [0.0], [0.0], [0.0], [0.0], [0.0]) @@ -139,18 +136,9 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): n = 2.0 u1 = 1.2 - maxw_params = {"n": n, "u1": u1} - pert_params = { - "u1": { - "ModesCos": { - "given_in_basis": "0", - "ls": [mode], - "amps": [amp], - } - } - } + pert = perturbations.ModesCos(ls=(mode,), amps=(amp,)) - maxwellian = Maxwellian3D(maxw_params=maxw_params, pert_params=pert_params) + maxwellian = Maxwellian3D(n=(n, None), u1=(u1, pert)) meshgrids = np.meshgrid( e1, @@ -195,18 +183,9 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): n = 2.0 vth1 = 1.2 - maxw_params = {"n": n, "vth1": vth1} - pert_params = { - "vth1": { - "ModesCos": { - "given_in_basis": "0", - "ls": [mode], - "amps": [amp], - } - } - } + pert = perturbations.ModesCos(ls=(mode,), amps=(amp,)) - maxwellian = Maxwellian3D(maxw_params=maxw_params, pert_params=pert_params) + maxwellian = Maxwellian3D(n=(n, None), vth1=(vth1, pert)) meshgrids = np.meshgrid( e1, @@ -247,19 +226,11 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): # ===== Test ITPA perturbation in density ===== # ============================================= n0 = 0.00720655 - c = [0.491230, 0.298228, 0.198739, 0.521298] + c = (0.491230, 0.298228, 0.198739, 0.521298) - maxw_params = { - "n": { - "ITPA_density": { - "given_in_basis": "0", - "n0": n0, - "c": c, - } - } - } + pert = perturbations.ITPA_density(n0=n0, c=c) - maxwellian = Maxwellian3D(maxw_params=maxw_params) + maxwellian = Maxwellian3D(n=(0.0, pert)) meshgrids = np.meshgrid(e1, [0.0], [0.0], [0.0], [0.0], [0.0]) @@ -288,30 +259,12 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): import numpy as np from struphy.fields_background import equils + from struphy.fields_background.base import FluidEquilibrium from struphy.geometry import domains from struphy.initial import perturbations + from struphy.initial.base import Perturbation from struphy.kinetic_background.maxwellians import Maxwellian3D - maxw_params_mhd = { - "n": "fluid_background", - "u1": "fluid_background", - "u2": "fluid_background", - "u3": "fluid_background", - "vth1": "fluid_background", - "vth2": "fluid_background", - "vth3": "fluid_background", - } - - maxw_params_1 = { - "n": 1.0, - "u1": "fluid_background", - "u2": "fluid_background", - "u3": "fluid_background", - "vth1": "fluid_background", - "vth2": "fluid_background", - "vth3": "fluid_background", - } - e1 = np.linspace(0.0, 1.0, Nel[0]) e2 = np.linspace(0.0, 1.0, Nel[1]) e3 = np.linspace(0.0, 1.0, Nel[2]) @@ -344,6 +297,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): print(f"Attention: flat (marker) evaluation not tested for GVEC at the moment.") mhd_equil = val() + assert isinstance(mhd_equil, FluidEquilibrium) print(f"{mhd_equil.params = }") if "AdhocTorus" in key: mhd_equil.domain = domains.HollowTorus( @@ -375,9 +329,25 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): except: print(f"Not setting domain for {key}.") - maxwellian = Maxwellian3D(maxw_params=maxw_params_mhd, equil=mhd_equil) + maxwellian = Maxwellian3D( + n=(mhd_equil.n0, None), + u1=(mhd_equil.u_cart_1, None), + u2=(mhd_equil.u_cart_2, None), + u3=(mhd_equil.u_cart_3, None), + vth1=(mhd_equil.vth0, None), + vth2=(mhd_equil.vth0, None), + vth3=(mhd_equil.vth0, None), + ) - maxwellian_1 = Maxwellian3D(maxw_params=maxw_params_1, equil=mhd_equil) + maxwellian_1 = Maxwellian3D( + n=(1.0, None), + u1=(mhd_equil.u_cart_1, None), + u2=(mhd_equil.u_cart_2, None), + u3=(mhd_equil.u_cart_3, None), + vth1=(mhd_equil.vth0, None), + vth2=(mhd_equil.vth0, None), + vth3=(mhd_equil.vth0, None), + ) # test meshgrid evaluation n0 = mhd_equil.n0(*e_meshgrids) @@ -521,23 +491,22 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): maxw_params_zero = {"n": 0.0, "vth1": 0.0, "vth2": 0.0, "vth3": 0.0} for key_2, val_2 in inspect.getmembers(perturbations): - if inspect.isclass(val_2): - print(f"{key_2 = }") + if inspect.isclass(val_2) and val_2.__module__ == perturbations.__name__: pert = val_2() + assert isinstance(pert, Perturbation) print(f"{pert = }") - pert_params = { - "n": {key_2: {"given_in_basis": "0"}}, - "u1": {key_2: {"given_in_basis": "0"}}, - "u2": {key_2: {"given_in_basis": "0"}}, - "u3": {key_2: {"given_in_basis": "0"}}, - "vth1": {key_2: {"given_in_basis": "0"}}, - "vth2": {key_2: {"given_in_basis": "0"}}, - "vth3": {key_2: {"given_in_basis": "0"}}, - } + if isinstance(pert, perturbations.Noise): + continue # background + perturbation maxwellian_perturbed = Maxwellian3D( - maxw_params=maxw_params_mhd, pert_params=pert_params, equil=mhd_equil + n=(mhd_equil.n0, pert), + u1=(mhd_equil.u_cart_1, pert), + u2=(mhd_equil.u_cart_2, pert), + u3=(mhd_equil.u_cart_3, pert), + vth1=(mhd_equil.vth0, pert), + vth2=(mhd_equil.vth0, pert), + vth3=(mhd_equil.vth0, pert), ) # test meshgrid evaluation @@ -548,7 +517,13 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): # pure perturbation maxwellian_zero_bckgr = Maxwellian3D( - maxw_params=maxw_params_zero, pert_params=pert_params, equil=mhd_equil + n=(0.0, pert), + u1=(0.0, pert), + u2=(0.0, pert), + u3=(0.0, pert), + vth1=(0.0, pert), + vth2=(0.0, pert), + vth3=(0.0, pert), ) assert np.allclose(maxwellian_zero_bckgr.n(*e_meshgrids), pert(*e_meshgrids)) @@ -701,9 +676,7 @@ def test_maxwellian_2d_uniform(Nel, show_plot=False): # =========================================================== # ===== Test uniform non-shifted, isothermal Maxwellian ===== # =========================================================== - maxw_params = {"n": 2.0} - - maxwellian = GyroMaxwellian2D(maxw_params=maxw_params, volume_form=False) + maxwellian = GyroMaxwellian2D(n=(2.0, None), volume_form=False) meshgrids = np.meshgrid(e1, e2, e3, [0.01], [0.01]) @@ -738,9 +711,15 @@ def test_maxwellian_2d_uniform(Nel, show_plot=False): u_perp = 0.2 vth_para = 1.2 vth_perp = 0.5 - maxw_params = {"n": n, "u_para": u_para, "u_perp": u_perp, "vth_para": vth_para, "vth_perp": vth_perp} - maxwellian = GyroMaxwellian2D(maxw_params=maxw_params, volume_form=False) + maxwellian = GyroMaxwellian2D( + n=(n, None), + u_para=(u_para, None), + u_perp=(u_perp, None), + vth_para=(vth_para, None), + vth_perp=(vth_perp, None), + volume_form=False, + ) # test Maxwellian profile in v v_para = np.linspace(-5, 5, 64) @@ -781,6 +760,7 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): import matplotlib.pyplot as plt import numpy as np + from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import GyroMaxwellian2D e1 = np.linspace(0.0, 1.0, Nel[0]) @@ -792,10 +772,9 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): # =============================================== amp = 0.1 mode = 1 - maxw_params = {"n": 2.0} - pert_params = {"n": {"ModesCos": {"given_in_basis": "0", "ls": [mode], "amps": [amp]}}} + pert = perturbations.ModesCos(ls=(mode,), amps=(amp,)) - maxwellian = GyroMaxwellian2D(maxw_params=maxw_params, pert_params=pert_params, volume_form=False) + maxwellian = GyroMaxwellian2D(n=(2.0, pert), volume_form=False) v_perp = 0.1 meshgrids = np.meshgrid(e1, [0.0], [0.0], [0.0], v_perp) @@ -822,10 +801,13 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): mode = 1 n = 2.0 u_para = 1.2 - maxw_params = {"n": n, "u_para": u_para} - pert_params = {"u_para": {"ModesCos": {"given_in_basis": "0", "ls": [mode], "amps": [amp]}}} + pert = perturbations.ModesCos(ls=(mode,), amps=(amp,)) - maxwellian = GyroMaxwellian2D(maxw_params=maxw_params, pert_params=pert_params, volume_form=False) + maxwellian = GyroMaxwellian2D( + n=(2.0, None), + u_para=(u_para, pert), + volume_form=False, + ) v_perp = 0.1 meshgrids = np.meshgrid(e1, [0.0], [0.0], v1, v_perp) @@ -863,10 +845,13 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): mode = 1 n = 2.0 u_perp = 1.2 - maxw_params = {"n": n, "u_perp": u_perp} - pert_params = {"u_perp": {"ModesCos": {"given_in_basis": "0", "ls": [mode], "amps": [amp]}}} + pert = perturbations.ModesCos(ls=(mode,), amps=(amp,)) - maxwellian = GyroMaxwellian2D(maxw_params=maxw_params, pert_params=pert_params, volume_form=False) + maxwellian = GyroMaxwellian2D( + n=(2.0, None), + u_perp=(u_perp, pert), + volume_form=False, + ) meshgrids = np.meshgrid(e1, [0.0], [0.0], 0.0, v2) @@ -903,10 +888,13 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): mode = 1 n = 2.0 vth_para = 1.2 - maxw_params = {"n": n, "vth_para": vth_para} - pert_params = {"vth_para": {"ModesCos": {"given_in_basis": "0", "ls": [mode], "amps": [amp]}}} + pert = perturbations.ModesCos(ls=(mode,), amps=(amp,)) - maxwellian = GyroMaxwellian2D(maxw_params=maxw_params, pert_params=pert_params, volume_form=False) + maxwellian = GyroMaxwellian2D( + n=(2.0, None), + vth_para=(vth_para, pert), + volume_form=False, + ) v_perp = 0.1 meshgrids = np.meshgrid( @@ -951,10 +939,13 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): mode = 1 n = 2.0 vth_perp = 1.2 - maxw_params = {"n": n, "vth_perp": vth_perp} - pert_params = {"vth_perp": {"ModesCos": {"given_in_basis": "0", "ls": [mode], "amps": [amp]}}} + pert = perturbations.ModesCos(ls=(mode,), amps=(amp,)) - maxwellian = GyroMaxwellian2D(maxw_params=maxw_params, pert_params=pert_params, volume_form=False) + maxwellian = GyroMaxwellian2D( + n=(2.0, None), + vth_perp=(vth_perp, pert), + volume_form=False, + ) meshgrids = np.meshgrid( e1, @@ -995,17 +986,9 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): # ============================================= n0 = 0.00720655 c = [0.491230, 0.298228, 0.198739, 0.521298] - maxw_params = { - "n": { - "ITPA_density": { - "given_in_basis": "0", - "n0": n0, - "c": c, - } - } - } + pert = perturbations.ITPA_density(n0=n0, c=c) - maxwellian = GyroMaxwellian2D(maxw_params=maxw_params, volume_form=False) + maxwellian = GyroMaxwellian2D(n=(0.0, pert), volume_form=False) v_perp = 0.1 meshgrids = np.meshgrid(e1, [0.0], [0.0], [0.0], v_perp) @@ -1039,22 +1022,9 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): from struphy.fields_background.base import FluidEquilibriumWithB from struphy.geometry import domains from struphy.initial import perturbations + from struphy.initial.base import Perturbation from struphy.kinetic_background.maxwellians import GyroMaxwellian2D - maxw_params_mhd = { - "n": "fluid_background", - "u_para": "fluid_background", - "vth_para": "fluid_background", - "vth_perp": "fluid_background", - } - - maxw_params_1 = { - "n": 1.0, - "u_para": "fluid_background", - "vth_para": "fluid_background", - "vth_perp": "fluid_background", - } - e1 = np.linspace(0.0, 1.0, Nel[0]) e2 = np.linspace(0.0, 1.0, Nel[1]) e3 = np.linspace(0.0, 1.0, Nel[2]) @@ -1119,9 +1089,21 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): except: print(f"Not setting domain for {key}.") - maxwellian = GyroMaxwellian2D(maxw_params=maxw_params_mhd, equil=mhd_equil, volume_form=False) + maxwellian = GyroMaxwellian2D( + n=(mhd_equil.n0, None), + u_para=(mhd_equil.u_para0, None), + vth_para=(mhd_equil.vth0, None), + vth_perp=(mhd_equil.vth0, None), + volume_form=False, + ) - maxwellian_1 = GyroMaxwellian2D(maxw_params=maxw_params_1, equil=mhd_equil, volume_form=False) + maxwellian_1 = GyroMaxwellian2D( + n=(1.0, None), + u_para=(mhd_equil.u_para0, None), + vth_para=(mhd_equil.vth0, None), + vth_perp=(mhd_equil.vth0, None), + volume_form=False, + ) # test meshgrid evaluation n0 = mhd_equil.n0(*e_meshgrids) @@ -1261,24 +1243,22 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): # test perturbations if "EQDSKequilibrium" in key: - maxw_params_zero = {"n": 0.0, "vth_para": 0.0, "vth_perp": 0.0} - for key_2, val_2 in inspect.getmembers(perturbations): - if inspect.isclass(val_2): - print(f"{key_2 = }") + if inspect.isclass(val_2) and val_2.__module__ == perturbations.__name__: pert = val_2() print(f"{pert = }") - pert_params = { - "n": {key_2: {"given_in_basis": "0"}}, - "u_para": {key_2: {"given_in_basis": "0"}}, - "u_perp": {key_2: {"given_in_basis": "0"}}, - "vth_para": {key_2: {"given_in_basis": "0"}}, - "vth_perp": {key_2: {"given_in_basis": "0"}}, - } + assert isinstance(pert, Perturbation) + + if isinstance(pert, perturbations.Noise): + continue # background + perturbation maxwellian_perturbed = GyroMaxwellian2D( - maxw_params=maxw_params_mhd, pert_params=pert_params, equil=mhd_equil, volume_form=False + n=(mhd_equil.n0, pert), + u_para=(mhd_equil.u_para0, pert), + vth_para=(mhd_equil.vth0, pert), + vth_perp=(mhd_equil.vth0, pert), + volume_form=False, ) # test meshgrid evaluation @@ -1289,9 +1269,11 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): # pure perturbation maxwellian_zero_bckgr = GyroMaxwellian2D( - maxw_params=maxw_params_zero, - pert_params=pert_params, - equil=mhd_equil, + n=(0.0, pert), + u_para=(0.0, pert), + u_perp=(0.0, pert), + vth_para=(0.0, pert), + vth_perp=(0.0, pert), volume_form=False, ) @@ -1436,6 +1418,7 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): from struphy.fields_background import equils from struphy.geometry import domains + from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import CanonicalMaxwellian e1 = np.linspace(0.0, 1.0, Nel[0]) @@ -1496,7 +1479,7 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): # =========================================================== maxw_params = {"n": 2.0, "vth": 1.0} - maxwellian = CanonicalMaxwellian(maxw_params=maxw_params) + maxwellian = CanonicalMaxwellian(n=(2.0, None), vth=(1.0, None)) # Test constant value at v_para = v_perp = 0.01 res = maxwellian(energy, mu, psic).squeeze() @@ -1604,8 +1587,9 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): "n": {"ITPA_density": {"n0": n0, "c": c}}, "vth": 1.0, } + pert = perturbations.ITPA_density(n0=n0, c=c) - maxwellian = CanonicalMaxwellian(maxw_params=maxw_params, equil=mhd_equil) + maxwellian = CanonicalMaxwellian(n=(0.0, pert), equil=mhd_equil, volume_form=False) e1 = np.linspace(0.0, 1.0, Nel[0]) e2 = np.linspace(0.0, 1.0, Nel[1]) @@ -1654,10 +1638,10 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): if __name__ == "__main__": - # test_maxwellian_3d_uniform(Nel=[64, 1, 1], show_plot=False) - # test_maxwellian_3d_perturbed(Nel=[64, 1, 1], show_plot=False) + # test_maxwellian_3d_uniform(Nel=[64, 1, 1], show_plot=True) + # test_maxwellian_3d_perturbed(Nel=[64, 1, 1], show_plot=True) # test_maxwellian_3d_mhd(Nel=[8, 11, 12], with_desc=None, show_plot=False) # test_maxwellian_2d_uniform(Nel=[64, 1, 1], show_plot=True) # test_maxwellian_2d_perturbed(Nel=[64, 1, 1], show_plot=True) - test_maxwellian_2d_mhd(Nel=[8, 12, 12], with_desc=None, show_plot=False) - # test_canonical_maxwellian_uniform(Nel=[64, 1, 1], show_plot=True) + # test_maxwellian_2d_mhd(Nel=[8, 12, 12], with_desc=None, show_plot=False) + test_canonical_maxwellian_uniform(Nel=[64, 1, 1], show_plot=True) diff --git a/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py b/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py index e0b4143d7..74ee6293b 100644 --- a/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py +++ b/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py @@ -1,11 +1,12 @@ import pytest +@pytest.mark.skip @pytest.mark.mpi_skip @pytest.mark.parametrize("Nel", [[16, 1, 1], [32, 1, 1]]) @pytest.mark.parametrize("p", [[1, 1, 1], [2, 1, 1]]) @pytest.mark.parametrize("spl_kind", [[True, True, True]]) -@pytest.mark.parametrize("dirichlet_bc", [[[False, False], [False, False], [False, False]]]) +@pytest.mark.parametrize("dirichlet_bc", [((False, False), (False, False), (False, False))]) @pytest.mark.parametrize("mapping", [["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}]]) @pytest.mark.parametrize("epsilon", [0.000000001]) @pytest.mark.parametrize("dt", [0.001]) @@ -59,10 +60,10 @@ def test_propagator1D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): bas_ops = BasisProjectionOperators(derham, domain, eq_mhd=eq_mhd) # Manufactured solutions - uvec = FEECVariable("u", "Hdiv") - u_evec = FEECVariable("u_e", "Hdiv") - potentialvec = FEECVariable("potential", "L2") - uinitial = FEECVariable("u", "Hdiv") + uvec = FEECVariable(space="Hdiv") + u_evec = FEECVariable(space="Hdiv") + potentialvec = FEECVariable(space="L2") + uinitial = FEECVariable(space="Hdiv") pp_u = perturbations.ManufacturedSolutionVelocity() pp_ue = perturbations.ManufacturedSolutionVelocity(species="Electrons") @@ -211,11 +212,12 @@ def test_propagator1D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): import pytest +@pytest.mark.skip @pytest.mark.mpi_skip @pytest.mark.parametrize("Nel", [[16, 16, 1], [32, 32, 1]]) @pytest.mark.parametrize("p", [[1, 1, 1], [2, 2, 1]]) @pytest.mark.parametrize("spl_kind", [[True, True, True]]) -@pytest.mark.parametrize("dirichlet_bc", [[[False, False], [False, False], [False, False]]]) +@pytest.mark.parametrize("dirichlet_bc", [((False, False), (False, False), (False, False))]) @pytest.mark.parametrize("mapping", [["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}]]) @pytest.mark.parametrize("epsilon", [0.001]) @pytest.mark.parametrize("dt", [0.01]) @@ -230,6 +232,7 @@ def test_propagator2D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): from struphy.feec.utilities import compare_arrays from struphy.fields_background.equils import HomogenSlab from struphy.geometry import domains + from struphy.models.variables import FEECVariable from struphy.propagators.propagators_fields import TwoFluidQuasiNeutralFull mpi_comm = MPI.COMM_WORLD @@ -267,9 +270,9 @@ def test_propagator2D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): bas_ops = BasisProjectionOperators(derham, domain, eq_mhd=eq_mhd) # Manufactured solutions - uvec = FEECVariable("u", "Hdiv") - u_evec = FEECVariable("u_e", "Hdiv") - potentialvec = FEECVariable("potential", "L2") + uvec = FEECVariable(space="Hdiv") + u_evec = FEECVariable(space="Hdiv") + potentialvec = FEECVariable(space="L2") pp_u = { "ManufacturedSolutionVelocity": { diff --git a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py index c2ee8268d..31f54a491 100644 --- a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py +++ b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py @@ -6,7 +6,7 @@ @pytest.mark.parametrize("Nel", [[12, 8, 1]]) @pytest.mark.parametrize("p", [[3, 3, 1]]) @pytest.mark.parametrize("spl_kind", [[False, True, True]]) -@pytest.mark.parametrize("dirichlet_bc", [[[False, False], [False, False], [False, False]]]) +@pytest.mark.parametrize("dirichlet_bc", [((False, False), (False, False), (False, False))]) @pytest.mark.parametrize("mapping", [["Cuboid", {"l1": 0.0, "r1": 2.0, "l2": 0.0, "r2": 3.0, "l3": 0.0, "r3": 6.0}]]) def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): """Test saddle-point-solver with manufactured solutions.""" diff --git a/src/struphy/main.py b/src/struphy/main.py index bbca0c526..9cfa80c92 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -30,6 +30,7 @@ create_femfields, create_vtk, eval_femfields, + get_params_of_run, post_process_f, post_process_markers, post_process_n_sph, @@ -148,8 +149,8 @@ def run( pickle.dump(grid, f, pickle.HIGHEST_PROTOCOL) with open(os.path.join(path_out, "derham_opts.bin"), "wb") as f: pickle.dump(derham_opts, f, pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_out, "model.bin"), "wb") as f: - pickle.dump(model, f, pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_out, "model_class.bin"), "wb") as f: + pickle.dump(model.__class__, f, pickle.HIGHEST_PROTOCOL) # config clones if comm is None: @@ -193,16 +194,18 @@ def run( # default grid if grid is None: Nel = (16, 16, 16) - print(f"\nNo grid specified - using TensorProductGrid with {Nel = }.") + if rank == 0: + print(f"\nNo grid specified - using TensorProductGrid with {Nel = }.") grid = grids.TensorProductGrid(Nel=Nel) # allocate derham-related objects if derham_opts is None: p = (3, 3, 3) spl_kind = (False, False, False) - print( - f"\nNo Derham options specified - creating Derham with {p = } and {spl_kind = } for projecting equilibrium." - ) + if rank == 0: + print( + f"\nNo Derham options specified - creating Derham with {p = } and {spl_kind = } for projecting equilibrium." + ) derham_opts = DerhamOptions(p=p, spl_kind=spl_kind) model.allocate_feec(grid, derham_opts) @@ -446,20 +449,10 @@ def pproc( if MPI.COMM_WORLD.Get_rank() == 0: print(f"\n*** Start post-processing of {path}:") - # load light-weight model instance from simulation - try: - params_in = import_parameters_py(os.path.join(path, "parameters.py")) - model: StruphyModel = params_in.model - domain: Domain = params_in.domain - except FileNotFoundError: - with open(os.path.join(path, "model.bin"), "rb") as f: - model: StruphyModel = pickle.load(f) - with open(os.path.join(path, "domain.bin"), "rb") as f: - # domain: Domain = pickle.load(f) - # print(f"{domain = }") - # print(f"{domain.params = }") - domain_dct = pickle.load(f) - domain: Domain = getattr(domains, domain_dct["name"])(**domain_dct["params"]) + # import parameters + params_in = get_params_of_run(path) + model = params_in.model + domain = params_in.domain # create post-processing folder path_pproc = os.path.join(path, "post_processing") @@ -513,13 +506,13 @@ def pproc( # field post-processing if exist_fields: - fields, t_grid = create_femfields(path, step=step) + fields, t_grid = create_femfields(path, params_in=params_in, step=step) - point_data, grids_log, grids_phy = eval_femfields(path, fields, celldivide=[celldivide, celldivide, celldivide]) + point_data, grids_log, grids_phy = eval_femfields(params_in, fields, celldivide=[celldivide] * 3) if physical: point_data_phy, grids_log, grids_phy = eval_femfields( - path, fields, celldivide=[celldivide, celldivide, celldivide], physical=True + params_in, fields, celldivide=[celldivide] * 3, physical=True ) # directory for field data @@ -606,11 +599,24 @@ def pproc( else: compute_bckgr = False - post_process_f(path, path_kinetics_species, species, step, compute_bckgr=compute_bckgr) + post_process_f( + path, + params_in, + path_kinetics_species, + species, + step, + compute_bckgr=compute_bckgr, + ) # sph density if exist_kinetic["n_sph"]: - post_process_n_sph(path, path_kinetics_species, species, step, compute_bckgr=compute_bckgr) + post_process_n_sph( + path, + path_kinetics_species, + species, + step, + compute_bckgr=compute_bckgr, + ) class SimData: diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 289c5c4b8..b3da1e2f8 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1342,13 +1342,22 @@ def generate_default_parameter_file( exclude = f"# model.{sn}.{vn}.save_data = False\n" elif isinstance(var, PICVariable): has_pic = True - init_pert_pic = f"\nperturbation = perturbations.TorusModesCos()" + init_pert_pic = f"\n# if .add_initial_condition is not called, the background is the kinetic initial condition\n" + init_pert_pic += f"perturbation = perturbations.TorusModesCos()\n" if "6D" in var.space: - init_bckgr_pic = f"\nmaxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n" + init_bckgr_pic = f"maxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, None))\n" init_bckgr_pic += f"maxwellian_2 = maxwellians.Maxwellian3D(n=(0.1, None))\n" + init_pert_pic += f"maxwellian_1pt = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n" + init_pert_pic += f"init = maxwellian_1pt + maxwellian_2\n" + init_pert_pic += f"model.kinetic_ions.var.add_initial_condition(init)\n" elif "5D" in var.space: - init_bckgr_pic = f"\nmaxwellian_1 = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation))\n" - init_bckgr_pic += f"maxwellian_2 = maxwellians.GyroMaxwellian2D(n=(0.1, None))\n" + init_bckgr_pic = f"maxwellian_1 = maxwellians.GyroMaxwellian2D(n=(1.0, None), equil=equil)\n" + init_bckgr_pic += f"maxwellian_2 = maxwellians.GyroMaxwellian2D(n=(0.1, None), equil=equil)\n" + init_pert_pic += ( + f"maxwellian_1pt = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), equil=equil)\n" + ) + init_pert_pic += f"init = maxwellian_1pt + maxwellian_2\n" + init_pert_pic += f"model.kinetic_ions.var.add_initial_condition(init)\n" init_bckgr_pic += f"background = maxwellian_1 + maxwellian_2\n" init_bckgr_pic += f"model.{sn}.{vn}.add_background(background)\n" @@ -1412,13 +1421,13 @@ def generate_default_parameter_file( for prop in self.propagators.__dict__: file.write(f"model.propagators.{prop}.options = model.propagators.{prop}.Options()\n") - file.write("\n# initial conditions (background + perturbation)\n") + file.write("\n# background and initial conditions\n") if has_feec: file.write(init_bckgr_feec) file.write(init_pert_feec) if has_pic: - file.write(init_pert_pic) file.write(init_bckgr_pic) + file.write(init_pert_pic) file.write("\n# optional: exclude variables from saving\n") file.write(exclude) @@ -1442,8 +1451,8 @@ def generate_default_parameter_file( file.close() print( - f"\nDefault parameter file for '{self.__class__.__name__}' has been created in {path}.\n\ -You can now launch with 'struphy run {self.__class__.__name__}' or with 'struphy run -i params_{self.__class__.__name__}.py'" + f"\nDefault parameter file for '{self.__class__.__name__}' has been created in the cwd ({path}).\n\ +You can now launch a simlaltion with 'python params_{self.__class__.__name__}.py'" ) return path diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 814554465..125c3e5fd 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -187,7 +187,7 @@ def set_sorting_boxes( self, do_sort: bool = False, sorting_frequency: int = 0, - boxes_per_dim: tuple = (16, 1, 1), + boxes_per_dim: tuple = (12, 12, 1), box_bufsize: float = 2.0, dims_maks: tuple = (True, True, True), ): diff --git a/src/struphy/models/tests/test_fluid_models.py b/src/struphy/models/tests/test_fluid_models.py deleted file mode 100644 index 44473e163..000000000 --- a/src/struphy/models/tests/test_fluid_models.py +++ /dev/null @@ -1,28 +0,0 @@ -import inspect - -import pytest - -from struphy.models.tests.util import wrapper_for_testing - - -@pytest.mark.parametrize( - "map_and_equil", [("Cuboid", "HomogenSlab"), ("HollowTorus", "AdhocTorus"), ("Tokamak", "EQDSKequilibrium")] -) -def test_fluid( - map_and_equil: tuple | list, - fast: bool, - vrbose: bool, - verification: bool, - nclones: int, - show_plots: bool, -): - """Tests all models in models/fluid.py.""" - wrapper_for_testing( - mtype="fluid", - map_and_equil=map_and_equil, - fast=fast, - vrbose=vrbose, - verification=verification, - nclones=nclones, - show_plots=show_plots, - ) diff --git a/src/struphy/models/tests/test_hybrid_models.py b/src/struphy/models/tests/test_hybrid_models.py deleted file mode 100644 index fb056a86e..000000000 --- a/src/struphy/models/tests/test_hybrid_models.py +++ /dev/null @@ -1,28 +0,0 @@ -import inspect - -import pytest - -from struphy.models.tests.util import wrapper_for_testing - - -@pytest.mark.parametrize( - "map_and_equil", [("Cuboid", "HomogenSlab"), ("HollowTorus", "AdhocTorus"), ("Tokamak", "EQDSKequilibrium")] -) -def test_hybrid( - map_and_equil: tuple | list, - fast: bool, - vrbose: bool, - verification: bool, - nclones: int, - show_plots: bool, -): - """Tests all models in models/hybrid.py.""" - wrapper_for_testing( - mtype="hybrid", - map_and_equil=map_and_equil, - fast=fast, - vrbose=vrbose, - verification=verification, - nclones=nclones, - show_plots=show_plots, - ) diff --git a/src/struphy/models/tests/test_kinetic_models.py b/src/struphy/models/tests/test_kinetic_models.py deleted file mode 100644 index 33180b74a..000000000 --- a/src/struphy/models/tests/test_kinetic_models.py +++ /dev/null @@ -1,28 +0,0 @@ -import inspect - -import pytest - -from struphy.models.tests.util import wrapper_for_testing - - -@pytest.mark.parametrize( - "map_and_equil", [("Cuboid", "HomogenSlab"), ("HollowTorus", "AdhocTorus"), ("Tokamak", "EQDSKequilibrium")] -) -def test_kinetic( - map_and_equil: tuple | list, - fast: bool, - vrbose: bool, - verification: bool, - nclones: int, - show_plots: bool, -): - """Tests models in models/kinetic.py.""" - wrapper_for_testing( - mtype="kinetic", - map_and_equil=map_and_equil, - fast=fast, - vrbose=vrbose, - verification=verification, - nclones=nclones, - show_plots=show_plots, - ) diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index c0c455799..4bb145f3e 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -1,8 +1,8 @@ -import inspect import os from types import ModuleType import pytest +from mpi4py import MPI from struphy import main from struphy.io.options import EnvironmentOptions @@ -10,30 +10,44 @@ from struphy.models import fluid, hybrid, kinetic, toy from struphy.models.base import StruphyModel +rank = MPI.COMM_WORLD.Get_rank() + # available models -toy_models = [] -for name, obj in inspect.getmembers(toy): - if inspect.isclass(obj) and "models.toy" in obj.__module__: - toy_models += [name] -print(f"\n{toy_models = }") - -fluid_models = [] -for name, obj in inspect.getmembers(fluid): - if inspect.isclass(obj) and "models.fluid" in obj.__module__: - fluid_models += [name] -print(f"\n{fluid_models = }") - -kinetic_models = [] -for name, obj in inspect.getmembers(kinetic): - if inspect.isclass(obj) and "models.kinetic" in obj.__module__: - kinetic_models += [name] -print(f"\n{kinetic_models = }") +toy_models = [ + "Maxwell", + "Vlasov", + "GuidingCenter", +] +# for name, obj in inspect.getmembers(toy): +# if inspect.isclass(obj) and "models.toy" in obj.__module__: +# toy_models += [name] +if rank == 0: + print(f"\n{toy_models = }") + +fluid_models = [ + "LinearMHD", +] +# for name, obj in inspect.getmembers(fluid): +# if inspect.isclass(obj) and "models.fluid" in obj.__module__: +# fluid_models += [name] +if rank == 0: + print(f"\n{fluid_models = }") + +kinetic_models = [ + "VlasovAmpereOneSpecies", +] +# for name, obj in inspect.getmembers(kinetic): +# if inspect.isclass(obj) and "models.kinetic" in obj.__module__: +# kinetic_models += [name] +if rank == 0: + print(f"\n{kinetic_models = }") hybrid_models = [] -for name, obj in inspect.getmembers(hybrid): - if inspect.isclass(obj) and "models.hybrid" in obj.__module__: - hybrid_models += [name] -print(f"\n{hybrid_models = }") +# for name, obj in inspect.getmembers(hybrid): +# if inspect.isclass(obj) and "models.hybrid" in obj.__module__: +# hybrid_models += [name] +if rank == 0: + print(f"\n{hybrid_models = }") # folder for test simulations @@ -41,36 +55,49 @@ # generic function for calling model tests -def call_test(model_name: str, module: ModuleType, verbose=True): - print(f"\n*** Testing '{model_name}':") - model = getattr(module, model_name)() +def call_test(model_name: str, module: ModuleType = None, verbose=True): + if rank == 0: + print(f"\n*** Testing '{model_name}':") + + if module is None: + submods = [toy, fluid, kinetic, hybrid] + for submod in submods: + try: + model = getattr(submod, model_name)() + except: + continue + + else: + model = getattr(module, model_name)() + assert isinstance(model, StruphyModel) # generate paramater file for testing path = os.path.join(test_folder, f"params_{model_name}.py") - model.generate_default_parameter_file(path=path, prompt=False) - del model - print("\nDeleting light-weight instance ...") + if rank == 0: + model.generate_default_parameter_file(path=path, prompt=False) + del model + MPI.COMM_WORLD.Barrier() # set environment options - env = EnvironmentOptions(out_folders=test_folder, sim_folder=f"{model_name}_test") + env = EnvironmentOptions(out_folders=test_folder, sim_folder=f"{model_name}") # read parameters params_in = import_parameters_py(path) - units = params_in.units + base_units = params_in.base_units time_opts = params_in.time_opts domain = params_in.domain equil = params_in.equil grid = params_in.grid derham_opts = params_in.derham_opts + model = params_in.model # test - model = params_in.model main.run( model, params_path=path, env=env, - units=units, + base_units=base_units, time_opts=time_opts, domain=domain, equil=equil, @@ -79,26 +106,71 @@ def call_test(model_name: str, module: ModuleType, verbose=True): verbose=verbose, ) + MPI.COMM_WORLD.Barrier() + if rank == 0: + path_out = os.path.join(test_folder, model_name) + main.pproc(path=path_out) + main.load_data(path=path_out) + MPI.COMM_WORLD.Barrier() -# specific tests -@pytest.mark.parametrize("model_name", toy_models) -def test_toy(model_name: str, verbose=True): - call_test(model_name=model_name, module=toy, verbose=verbose) - -@pytest.mark.parametrize("model_name", fluid_models) -def test_fluid(model_name: str, verbose=True): - call_test(model_name=model_name, module=fluid, verbose=verbose) - - -@pytest.mark.parametrize("model_name", kinetic_models) -def test_kinetic(model_name: str, verbose=True): - call_test(model_name=model_name, module=kinetic, verbose=verbose) - - -@pytest.mark.parametrize("model_name", hybrid_models) -def test_hybrid(model_name: str, verbose=True): - call_test(model_name=model_name, module=hybrid, verbose=verbose) +# specific tests +@pytest.mark.models +@pytest.mark.toy +@pytest.mark.parametrize("model", toy_models) +def test_toy( + model: str, + vrbose: bool, + nclones: int, + show_plots: bool, +): + call_test(model_name=model, module=toy, verbose=vrbose) + + +@pytest.mark.models +@pytest.mark.fluid +@pytest.mark.parametrize("model", fluid_models) +def test_fluid( + model: str, + vrbose: bool, + nclones: int, + show_plots: bool, +): + call_test(model_name=model, module=fluid, verbose=vrbose) + + +@pytest.mark.models +@pytest.mark.kinetic +@pytest.mark.parametrize("model", kinetic_models) +def test_kinetic( + model: str, + vrbose: bool, + nclones: int, + show_plots: bool, +): + call_test(model_name=model, module=kinetic, verbose=vrbose) + + +@pytest.mark.models +@pytest.mark.hybrid +@pytest.mark.parametrize("model", hybrid_models) +def test_hybrid( + model: str, + vrbose: bool, + nclones: int, + show_plots: bool, +): + call_test(model_name=model, module=hybrid, verbose=vrbose) + + +@pytest.mark.single +def test_single_model( + model_name: str, + vrbose: bool, + nclones: int, + show_plots: bool, +): + call_test(model_name=model_name, module=None, verbose=vrbose) if __name__ == "__main__": diff --git a/src/struphy/models/tests/test_toy_models.py b/src/struphy/models/tests/test_toy_models.py deleted file mode 100644 index 8b1f03456..000000000 --- a/src/struphy/models/tests/test_toy_models.py +++ /dev/null @@ -1,28 +0,0 @@ -import inspect - -import pytest - -from struphy.models.tests.util import wrapper_for_testing - - -@pytest.mark.parametrize( - "map_and_equil", [("Cuboid", "HomogenSlab"), ("HollowTorus", "AdhocTorus"), ("Tokamak", "EQDSKequilibrium")] -) -def test_toy( - map_and_equil: tuple | list, - fast: bool, - vrbose: bool, - verification: bool, - nclones: int, - show_plots: bool, -): - """Tests models in models/toy.py.""" - wrapper_for_testing( - mtype="toy", - map_and_equil=map_and_equil, - fast=fast, - vrbose=vrbose, - verification=verification, - nclones=nclones, - show_plots=show_plots, - ) diff --git a/src/struphy/models/tests/test_LinearMHD.py b/src/struphy/models/tests/test_verif_LinearMHD.py similarity index 100% rename from src/struphy/models/tests/test_LinearMHD.py rename to src/struphy/models/tests/test_verif_LinearMHD.py diff --git a/src/struphy/models/tests/test_Maxwell.py b/src/struphy/models/tests/test_verif_Maxwell.py similarity index 100% rename from src/struphy/models/tests/test_Maxwell.py rename to src/struphy/models/tests/test_verif_Maxwell.py diff --git a/src/struphy/models/tests/test_xxpproc.py b/src/struphy/models/tests/test_xxpproc.py deleted file mode 100644 index bfd72b7ac..000000000 --- a/src/struphy/models/tests/test_xxpproc.py +++ /dev/null @@ -1,69 +0,0 @@ -def test_pproc_codes(model: str = None, group: str = None): - """Tests the post processing of runs in test_codes.py""" - - import inspect - import os - - from mpi4py import MPI - - import struphy - from struphy.models import fluid, hybrid, kinetic, toy - from struphy.post_processing import pproc_struphy - - comm = MPI.COMM_WORLD - - libpath = struphy.__path__[0] - - list_fluid = [] - for name, obj in inspect.getmembers(fluid): - if inspect.isclass(obj): - if name not in {"StruphyModel", "Propagator"}: - list_fluid += [name] - - list_kinetic = [] - for name, obj in inspect.getmembers(kinetic): - if inspect.isclass(obj): - if name not in {"StruphyModel", "KineticBackground", "Propagator"}: - list_kinetic += [name] - - list_hybrid = [] - for name, obj in inspect.getmembers(hybrid): - if inspect.isclass(obj): - if name not in {"StruphyModel", "Propagator"}: - list_hybrid += [name] - - list_toy = [] - for name, obj in inspect.getmembers(toy): - if inspect.isclass(obj): - if name not in {"StruphyModel", "Propagator"}: - list_toy += [name] - - if group is None: - list_models = list_fluid + list_kinetic + list_hybrid + list_toy - elif group == "fluid": - list_models = list_fluid - elif group == "kinetic": - list_models = list_kinetic - elif group == "hybrid": - list_models = list_hybrid - elif group == "toy": - list_models = list_toy - else: - raise ValueError(f"{group = } is not a valid group specification.") - - if comm.Get_rank() == 0: - if model is None: - for model in list_models: - if "Variational" in model or "Visco" in model: - print(f"Model {model} is currently excluded from tests.") - continue - - path_out = os.path.join(libpath, "io/out/test_" + model) - pproc_struphy.main(path_out) - else: - path_out = os.path.join(libpath, "io/out/test_" + model) - pproc_struphy.main(path_out) - - -if __name__ == "__main__": - test_pproc_codes() diff --git a/src/struphy/models/tests/util.py b/src/struphy/models/tests/util.py deleted file mode 100644 index 18dcd8086..000000000 --- a/src/struphy/models/tests/util.py +++ /dev/null @@ -1,426 +0,0 @@ -import copy -import inspect -import os -import sys - -import yaml -from mpi4py import MPI - -import struphy -from struphy.console.main import recursive_get_files -from struphy.io.setup import descend_options_dict -from struphy.models.base import StruphyModel -from struphy.struphy import run - -libpath = struphy.__path__[0] - - -def wrapper_for_testing( - mtype: str = "fluid", - map_and_equil: tuple | list = ("Cuboid", "HomogenSlab"), - fast: bool = True, - vrbose: bool = False, - verification: bool = False, - nclones: int = 1, - show_plots: bool = False, - model: str = None, - Tend: float = None, -): - """Wrapper for testing Struphy models. - - If model is not None, tests the specified model. - The argument "fast" is a pytest option that can be specified at the command line (see conftest.py). - """ - - if mtype == "fluid": - from struphy.models import fluid as modmod - elif mtype == "kinetic": - from struphy.models import kinetic as modmod - elif mtype == "hybrid": - from struphy.models import hybrid as modmod - elif mtype == "toy": - from struphy.models import toy as modmod - else: - raise ValueError(f'{mtype} must be either "fluid", "kinetic", "hybrid" or "toy".') - - comm = MPI.COMM_WORLD - - if model is None: - for key, val in inspect.getmembers(modmod): - if inspect.isclass(val) and val.__module__ == modmod.__name__: - # TODO: remove if-clauses - if "LinearExtendedMHD" in key and "HomogenSlab" not in map_and_equil[1]: - print(f"Model {key} is currently excluded from tests with mhd_equil other than HomogenSlab.") - continue - - if fast and "Cuboid" not in map_and_equil[0]: - print(f"Fast is enabled, mapping {map_and_equil[0]} skipped ...") - continue - - call_test( - key, - val, - map_and_equil, - Tend=Tend, - verbose=vrbose, - comm=comm, - verification=verification, - nclones=nclones, - show_plots=show_plots, - ) - else: - assert model in modmod.__dir__(), f"{model} not in {modmod.__name__}, please specify correct model type." - val = getattr(modmod, model) - - # TODO: remove if-clause - if "LinearExtendedMHD" in model and "HomogenSlab" not in map_and_equil[1]: - print(f"Model {model} is currently excluded from tests with mhd_equil other than HomogenSlab.") - sys.exit(0) - - call_test( - model, - val, - map_and_equil, - Tend=Tend, - verbose=vrbose, - comm=comm, - verification=verification, - nclones=nclones, - show_plots=show_plots, - ) - - -def call_test( - model_name: str, - model: StruphyModel, - map_and_equil: tuple, - *, - Tend: float = None, - verbose: bool = True, - comm=None, - verification: bool = False, - nclones: int = 1, - show_plots: bool = False, -): - """Does testing of one model, either all options or verification. - - Parameters - ---------- - model_name : str - Model name. - - model : StruphyModel - Instance of model base class. - - map_and_equil : tuple[str] - Name of mapping and MHD equilibirum. - - nclones : int - Number of domain clones. - - Tend : float - End time of simulation other than default. - - verbose : bool - Show info on screen. - - verification : bool - Do verifiaction runs. - - show_plots: bool - Show plots of verification tests. - """ - if "SPH" in model_name: - nclones = 1 - rank = comm.Get_rank() - - if verification: - ver_path = os.path.join(libpath, "io", "inp", "verification") - yml_files = recursive_get_files(ver_path, contains=(model_name,)) - if len(yml_files) == 0: - if MPI.COMM_WORLD.Get_rank() == 0: - print(f"\nVerification run not started: no .yml files for model {model_name} in {ver_path}.") - return - params_list = [] - paths_out = [] - py_scripts = [] - if MPI.COMM_WORLD.Get_rank() == 0: - print("\nThe following verification tests will be run:") - for n, file in enumerate(yml_files): - ref = file.split("_")[0] - if ref != model_name: - continue - if MPI.COMM_WORLD.Get_rank() == 0: - print(file) - with open(os.path.join(ver_path, file)) as tmp: - params_list += [yaml.load(tmp, Loader=yaml.FullLoader)] - paths_out += [os.path.join(libpath, "io", "out", "verification", model_name, f"{n + 1}")] - - # python scripts for data verification after the run below - from struphy.models.tests import verification as verif - - tname = file.split(".")[0] - try: - py_scripts += [getattr(verif, tname)] - except: - if MPI.COMM_WORLD.Get_rank() == 0: - print(f"A Python script for {model_name} is missing in models/tests/verification.py, exiting ...") - sys.exit(1) - else: - params = model.generate_default_parameter_file(save=False) - params["geometry"]["type"] = map_and_equil[0] - params["geometry"][map_and_equil[0]] = {} - params["fluid_background"][map_and_equil[1]] = {} - params_list = [params] - paths_out = [os.path.join(libpath, "io/out/test_" + model_name)] - py_scripts = [None] - - # run model - for parameters, path_out, py_script in zip(params_list, paths_out, py_scripts): - if Tend is not None: - parameters["time"]["Tend"] = Tend - if MPI.COMM_WORLD.Get_rank() == 0: - print_test_params(parameters) - run( - model_name, - parameters, - path_out, - save_step=int( - Tend / parameters["time"]["dt"], - ), - num_clones=nclones, - verbose=verbose, - ) - return - else: - # run with default - if MPI.COMM_WORLD.Get_rank() == 0: - print_test_params(parameters) - run( - model_name, - parameters, - path_out, - num_clones=nclones, - verbose=verbose, - ) - - # run the verification script on the output data - if verification: - py_script( - path_out, - rank, - show_plots=show_plots, - ) - - # run available options (if present) - if not verification: - d_opts, test_list = find_model_options(model, parameters) - params_default = copy.deepcopy(parameters) - - if len(d_opts["em_fields"]) > 0: - for opts_dict in d_opts["em_fields"]: - parameters = copy.deepcopy(params_default) - for opt in opts_dict: - parameters["em_fields"]["options"] = opt - - # test only if not aready tested - if any([opt == i for i in test_list]): - continue - else: - test_list += [opt] - if MPI.COMM_WORLD.Get_rank() == 0: - print_test_params(parameters) - run( - model_name, - parameters, - path_out, - num_clones=nclones, - verbose=verbose, - ) - - if len(d_opts["fluid"]) > 0: - for species, opts_dicts in d_opts["fluid"].items(): - for opts_dict in opts_dicts: - parameters = copy.deepcopy(params_default) - for opt in opts_dict: - parameters["fluid"][species]["options"] = opt - - # test only if not aready tested - if any([opt == i for i in test_list]): - continue - else: - test_list += [opt] - if MPI.COMM_WORLD.Get_rank() == 0: - print_test_params(parameters) - run( - model_name, - parameters, - path_out, - num_clones=nclones, - verbose=verbose, - ) - - if len(d_opts["kinetic"]) > 0: - for species, opts_dicts in d_opts["kinetic"].items(): - for opts_dict in opts_dicts: - parameters = copy.deepcopy(params_default) - for opt in opts_dict: - parameters["kinetic"][species]["options"] = opt - - # test only if not aready tested - if any([opt == i for i in test_list]): - continue - else: - test_list += [opt] - if MPI.COMM_WORLD.Get_rank() == 0: - print_test_params(parameters) - run( - model_name, - parameters, - path_out, - num_clones=nclones, - verbose=verbose, - ) - - -def print_test_params(parameters): - print("\nOptions of this test run:") - for k, v in parameters.items(): - if k == "em_fields": - if "options" in v: - print("\nem_fields:") - for kk, vv in v["options"].items(): - print(" " * 4, kk) - print(" " * 8, vv) - elif k in ("fluid", "kinetic"): - print(f"\n{k}:") - for kk, vv in v.items(): - if "options" in vv: - for kkk, vvv in vv["options"].items(): - print(" " * 4, kkk) - print(" " * 8, vvv) - - -def find_model_options( - model: StruphyModel, - parameters: dict, -): - """Find all options of a model and store them in d_opts. - The default options are also stored in test_list.""" - - d_opts = {"em_fields": [], "fluid": {}, "kinetic": {}} - # find out the em_fields options of the model - if "em_fields" in parameters: - if "options" in parameters["em_fields"]: - # create the default options parameters - d_default = parameters["em_fields"]["options"] - - # create a list of parameter dicts for the different options - descend_options_dict( - model.options()["em_fields"]["options"], - d_opts["em_fields"], - d_default=d_default, - ) - - for name in model.species()["fluid"]: - # find out the fluid options of the model - if "options" in parameters["fluid"][name]: - # create the default options parameters - d_default = parameters["fluid"][name]["options"] - - d_opts["fluid"][name] = [] - - # create a list of parameter dicts for the different options - descend_options_dict( - model.options()["fluid"][name]["options"], - d_opts["fluid"][name], - d_default=d_default, - ) - - for name in model.species()["kinetic"]: - # find out the kinetic options of the model - if "options" in parameters["kinetic"][name]: - # create the default options parameters - d_default = parameters["kinetic"][name]["options"] - - d_opts["kinetic"][name] = [] - - # create a list of parameter dicts for the different options - descend_options_dict( - model.options()["kinetic"][name]["options"], - d_opts["kinetic"][name], - d_default=d_default, - ) - - # store default options - test_list = [] - if "options" in model.options()["em_fields"]: - test_list += [parameters["em_fields"]["options"]] - if "fluid" in parameters: - for species in parameters["fluid"]: - if "options" in model.options()["fluid"][species]: - test_list += [parameters["fluid"][species]["options"]] - if "kinetic" in parameters: - for species in parameters["kinetic"]: - if "options" in model.options()["kinetic"][species]: - test_list += [parameters["kinetic"][species]["options"]] - - return d_opts, test_list - - -if __name__ == "__main__": - # This is called in struphy_test in case "group" is a model name - mtype = sys.argv[1] - group = sys.argv[2] - if sys.argv[3] == "None": - Tend = None - else: - Tend = float(sys.argv[3]) - fast = sys.argv[4] == "True" - vrbose = sys.argv[5] == "True" - verification = sys.argv[6] == "True" - if sys.argv[7] == "None": - nclones = 1 - else: - nclones = int(sys.argv[7]) - show_plots = sys.argv[8] == "True" - - map_and_equil = ("Cuboid", "HomogenSlab") - wrapper_for_testing( - mtype, - map_and_equil, - fast, - vrbose, - verification, - nclones, - show_plots, - model=group, - Tend=Tend, - ) - - if not fast and not verification: - map_and_equil = ("HollowTorus", "AdhocTorus") - wrapper_for_testing( - mtype, - map_and_equil, - fast, - vrbose, - verification, - nclones, - show_plots, - model=group, - Tend=Tend, - ) - - map_and_equil = ("Tokamak", "EQDSKequilibrium") - wrapper_for_testing( - mtype, - map_and_equil, - fast, - vrbose, - verification, - nclones, - show_plots, - model=group, - Tend=Tend, - ) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index ef4cdb6b2..3839f0c93 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -175,6 +175,21 @@ def add_background(self, background: KineticBackground, n_as_volume_form: bool = self._n_as_volume_form = n_as_volume_form super().add_background(background, verbose=verbose) + def add_initial_condition(self, init: KineticBackground, verbose=True): + self._initial_condition = init + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print( + f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added initial condition '{init.__class__.__name__}' with:" + ) + for k, v in init.__dict__.items(): + print(f" {k}: {v}") + + @property + def initial_condition(self) -> KineticBackground: + if not hasattr(self, "_initial_condition"): + self._initial_condition = self.backgrounds + return self._initial_condition + def allocate( self, clone_config: CloneConfig = None, @@ -218,6 +233,7 @@ def allocate( equil=equil, projected_equil=projected_equil, background=self.backgrounds, + initial_condition=self.initial_condition, n_as_volume_form=self.n_as_volume_form, # perturbations=self.perturbations, equation_params=self.species.equation_params, diff --git a/src/struphy/ode/tests/test_ode_feec.py b/src/struphy/ode/tests/test_ode_feec.py index 2df6aca95..d03af5a95 100644 --- a/src/struphy/ode/tests/test_ode_feec.py +++ b/src/struphy/ode/tests/test_ode_feec.py @@ -1,3 +1,5 @@ +from typing import get_args + import pytest from struphy.ode.utils import OptsButcher @@ -14,7 +16,7 @@ ("1", "0", "2"), ], ) -@pytest.mark.parametrize("algo", OptsButcher) +@pytest.mark.parametrize("algo", get_args(OptsButcher)) def test_exp_growth(spaces, algo, show_plots=False): """Solve dy/dt = omega*y for different feec variables y and with all available solvers from the ButcherTableau.""" @@ -27,6 +29,7 @@ def test_exp_growth(spaces, algo, show_plots=False): from struphy.feec.psydac_derham import Derham from struphy.ode.solvers import ODEsolverFEEC + from struphy.ode.utils import ButcherTableau comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -99,9 +102,10 @@ def f(t, y1, y2, y3, out=out): vector_field[var] = f print(f"{vector_field = }") - print(f"{algo = }") + butcher = ButcherTableau(algo=algo) + print(f"{butcher = }") - solver = ODEsolverFEEC(vector_field, algo=algo) + solver = ODEsolverFEEC(vector_field, butcher=butcher) hs = [0.1] n_hs = 6 diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 6bcacaad3..9ab323163 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -118,7 +118,10 @@ class Particles(metaclass=ABCMeta): Struphy fluid equilibrium projected into a discrete Derham complex. background : KineticBackground - Kinetic background parameters. + Kinetic background. + + initial_condition : KineticBackground + Kinetic initial condition. n_as_volume_form: bool Whether the number density n is given as a volume form or scalar function (=default). @@ -151,6 +154,8 @@ def __init__( equil: FluidEquilibrium = None, projected_equil: ProjectedFluidEquilibrium = None, background: KineticBackground = None, + initial_condition: KineticBackground = None, + perturbations: dict[str, Perturbation] = None, n_as_volume_form: bool = False, equation_params: dict = None, verbose: bool = False, @@ -289,11 +294,11 @@ def __init__( else: self._background = background - # background p-form description in [eta, v] (None means 0-form, "vol" means volume form -> divide by det) + # background p-form description in [eta, v] (False means 0-form, True means volume form -> divide by det) if isinstance(background, FluidEquilibrium): - self._pforms = (False, False) + self._is_volume_form = (n_as_volume_form, False) else: - self._pforms = ( + self._is_volume_form = ( n_as_volume_form, self.background.volume_form, ) @@ -302,8 +307,14 @@ def __init__( self._set_background_function() self._set_background_coordinates() - # perturbation parameters - # self._perturbations = perturbations + # perturbation parameters (needed for fluid background) + self._perturbations = perturbations + + # initial condition + if initial_condition is None: + self._initial_condition = self.background + else: + self._initial_condition = initial_condition # for loading # if self.loading_params["moments"] is None and self.type != "sph" and isinstance(self.bckgr_params, dict): @@ -494,10 +505,10 @@ def background(self) -> KineticBackground: """Kinetic background.""" return self._background - # @property - # def perturbations(self): - # """Kinetic perturbations.""" - # return self._perturbations + @property + def perturbations(self) -> dict[str, Perturbation]: + """Kinetic perturbations, keys are the names of moments of the distribution function ("n", "u1", etc.).""" + return self._perturbations @property def loading_params(self) -> LoadingParameters: @@ -537,6 +548,11 @@ def equation_params(self): """Parameters appearing in model equation due to Struphy normalization.""" return self._equation_params + @property + def initial_condition(self) -> KineticBackground: + """Kinetic initial condition""" + return self._initial_condition + @property def f_init(self): """Callable initial condition (background + perturbation). @@ -794,11 +810,9 @@ def marker_ids(self, new): self._markers[self.valid_mks, self.index["ids"]] = new @property - def pforms(self): - """Tuple of size 2; each entry must be either "vol" or None, defining the p-form - (space and velocity, respectively) of f_init. - """ - return self._pforms + def is_volume_form(self): + """Tuple of size 2 for (position, velocity), defining the p-form representation of f_init: True means volume-form, False means 0-form.""" + return self._is_volume_form @property def spatial(self): @@ -952,8 +966,14 @@ def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): return dom_arr, tuple(nprocs) def _set_background_function(self): - self._f0 = copy.deepcopy(self.background) - self.f0.add_perturbation = False + self._f0 = self.background + + # if isinstance(self.background, FluidEquilibrium): + # self._f0 = self.background + # else: + # self._f0 = copy.deepcopy(self.background) + # self.f0.add_perturbation = False + # self._f0 = None # if isinstance(self.bckgr_params, FluidEquilibrium): # self._f0 = self.bckgr_params @@ -1170,8 +1190,8 @@ def _generate_sampling_moments(self): # self.loading_params["moments"] = new_moments - def _set_initial_condition(self, bp_copy=None, pp_copy=None): - self._f_init = self.background + def _set_initial_condition(self): + self._f_init = self.initial_condition def _load_external( self, @@ -1656,7 +1676,7 @@ def initialize_weights( """ if self.loading == "tesselation": - if self.pforms[0] is None: + if not self.is_volume_form[0]: fvol = TransformedPformComponent([self.f_init], "0", "3", domain=self.domain) else: fvol = self.f_init @@ -1682,10 +1702,10 @@ def initialize_weights( f_init = self.f_init(*self.f_coords.T) # if f_init is vol-form, transform to 0-form - if self.pforms[0] == "vol": + if self.is_volume_form[0]: f_init /= self.domain.jacobian_det(self.positions) - if self.pforms[1] == "vol": + if self.is_volume_form[1]: f_init /= self.f_init.velocity_jacobian_det( *self.f_jacobian_coords.T, ) @@ -1728,10 +1748,10 @@ def update_weights(self): f0 = self.f0(*self.f_coords.T) # if f_init is vol-form, transform to 0-form - if self.pforms[0] == "vol": + if self.is_volume_form[0]: f0 /= self.domain.jacobian_det(self.positions) - if self.pforms[1] == "vol": + if self.is_volume_form[1]: f0 /= self.f0.velocity_jacobian_det(*self.f_jacobian_coords.T) self.weights = self.weights0 - f0 / self.sampling_density diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 1a8560375..ae12b5d5a 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -1,9 +1,13 @@ import copy -from struphy.fields_background.base import FluidEquilibriumWithB +from struphy.fields_background import equils +from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB from struphy.geometry.base import Domain +from struphy.geometry.utilities import TransformedPformComponent +from struphy.initial.base import Perturbation from struphy.kinetic_background import maxwellians +from struphy.kinetic_background.base import Maxwellian, SumKineticBackground from struphy.pic import utilities_kernels from struphy.pic.base import Particles @@ -31,10 +35,10 @@ def __init__( ): kwargs["type"] = "full_f" - # if "backgrounds" not in kwargs: - # kwargs["backgrounds"] = self.default_background() - # elif kwargs["backgrounds"] is None: - # kwargs["backgrounds"] = self.default_background() + if "background" not in kwargs: + kwargs["background"] = self.default_background() + elif kwargs["background"] is None: + kwargs["background"] = self.default_background() # default number of diagnostics and auxiliary columns self._n_cols_diagnostics = kwargs.pop("n_cols_diagn", 0) @@ -238,35 +242,45 @@ class DeltaFParticles6D(Particles6D): """ @classmethod - def default_bckgr_params(cls): - return {"Maxwellian3D": {}} + def default_background(cls): + return maxwellians.Maxwellian3D() def __init__( self, **kwargs, ): kwargs["type"] = "delta_f" - kwargs["control_variate"] = False + if "weights_params" in kwargs: + kwargs["weights_params"].control_variate = False super().__init__(**kwargs) def _set_initial_condition(self): - bp_copy = copy.deepcopy(self.bckgr_params) - pp_copy = copy.deepcopy(self.pert_params) - - # Prepare delta-f perturbation parameters - if pp_copy is not None: - for fi in bp_copy: - # Set background to zero (if "use_background_n" in perturbation params is set to false or not in keys) - if fi in pp_copy: - if "use_background_n" in pp_copy[fi]: - if not pp_copy[fi]["use_background_n"]: - bp_copy[fi]["n"] = 0.0 - else: - bp_copy[fi]["n"] = 0.0 - else: - bp_copy[fi]["n"] = 0.0 - - super()._set_initial_condition(bp_copy=bp_copy, pp_copy=pp_copy) + # bp_copy = copy.deepcopy(self.bckgr_params) + # pp_copy = copy.deepcopy(self.pert_params) + + # # Prepare delta-f perturbation parameters + # if pp_copy is not None: + # for fi in bp_copy: + # # Set background to zero (if "use_background_n" in perturbation params is set to false or not in keys) + # if fi in pp_copy: + # if "use_background_n" in pp_copy[fi]: + # if not pp_copy[fi]["use_background_n"]: + # bp_copy[fi]["n"] = 0.0 + # else: + # bp_copy[fi]["n"] = 0.0 + # else: + # bp_copy[fi]["n"] = 0.0 + self.set_n_to_zero(self.initial_condition) + + super()._set_initial_condition() + + def set_n_to_zero(self, background: Maxwellian | SumKineticBackground): + if isinstance(background, Maxwellian): + background.maxw_params["n"] = (0.0, background.maxw_params["n"][1]) + else: + assert isinstance(background, SumKineticBackground) + self.set_n_to_zero(background._f1) + self.set_n_to_zero(background._f2) class Particles5D(Particles): @@ -754,8 +768,8 @@ class ParticlesSPH(Particles): """ @classmethod - def default_bckgr_params(cls): - return {"ConstantVelocity": {}} + def default_background(cls): + return equils.ConstantVelocity() def __init__( self, @@ -763,8 +777,14 @@ def __init__( ): kwargs["type"] = "sph" - if "bckgr_params" not in kwargs: - kwargs["bckgr_params"] = self.default_bckgr_params() + if "background" not in kwargs: + bckgr = self.default_background() + bckgr.domain = kwargs["domain"] + kwargs["background"] = bckgr + elif kwargs["background"] is None: + bckgr = self.default_background() + bckgr.domain = kwargs["domain"] + kwargs["background"] = bckgr if "boxes_per_dim" not in kwargs: boxes_per_dim = (1, 1, 1) @@ -869,10 +889,6 @@ def s0(self, eta1, eta2, eta3, *v, flat_eval=False, remove_holes=True): def _set_initial_condition(self): """Set a callable initial condition f_init as a 0-form (scalar), and u_init in Cartesian coordinates.""" - from struphy.feec.psydac_derham import transform_perturbation - from struphy.fields_background.base import FluidEquilibrium - - pp_copy = copy.deepcopy(self.pert_params) # Get the initialization function and pass the correct arguments self._f_init = None @@ -880,31 +896,43 @@ def _set_initial_condition(self): self._f_init = self.f0.n0 self._u_init = self.f0.u_cart - if pp_copy is not None: - if "n" in pp_copy: - for _type, _params in pp_copy["n"].items(): # only one perturbation is taken into account at the moment - _fun = transform_perturbation(_type, _params, "0", self.domain) - - def _f_init(*etas): - if len(etas) == 1: - return self.f0.n0(etas[0]) + _fun(*etas[0].T) - else: - assert len(etas) == 3 - E1, E2, E3, is_sparse_meshgrid = Domain.prepare_eval_pts( - etas[0], - etas[1], - etas[2], - flat_eval=False, - ) - return self.f0.n0(E1, E2, E3) + _fun(E1, E2, E3) - - self._f_init = _f_init - - if "u1" in pp_copy: - for _type, _params in pp_copy[ - "u1" - ].items(): # only one perturbation is taken into account at the moment - _fun = transform_perturbation(_type, _params, "v", self.domain) + if self.perturbations is not None: + for moment, pert in self.perturbations.items(): # only one perturbation is taken into account at the moment + assert isinstance(moment, str) + assert isinstance(pert, Perturbation) + + if moment == "n": + _fun = TransformedPformComponent( + pert, + pert.given_in_basis, + "0", + comp=pert.comp, + domain=self.domain, + ) + + def _f_init(*etas): + if len(etas) == 1: + return self.f0.n0(etas[0]) + _fun(*etas[0].T) + else: + assert len(etas) == 3 + E1, E2, E3, is_sparse_meshgrid = Domain.prepare_eval_pts( + etas[0], + etas[1], + etas[2], + flat_eval=False, + ) + return self.f0.n0(E1, E2, E3) + _fun(E1, E2, E3) + + self._f_init = _f_init + + elif moment == "u1": + _fun = TransformedPformComponent( + pert, + pert.given_in_basis, + "v", + comp=pert.comp, + domain=self.domain, + ) _fun_cart = lambda e1, e2, e3: self.domain.push(_fun, e1, e2, e3, kind="v") - self._u_init = lambda e1, e2, e3: self.f0.u_cart(e1, e2, e3)[0] + _fun_cart(e1, e2, e3) - # TODO: add other velocity components + self._u_init = lambda e1, e2, e3: self.f0.u_cart(e1, e2, e3)[0] + _fun_cart(e1, e2, e3) + # TODO: add other velocity components diff --git a/src/struphy/pic/tests/test_accum_vec_H1.py b/src/struphy/pic/tests/test_accum_vec_H1.py index 6d1e90208..3f2ce4923 100644 --- a/src/struphy/pic/tests/test_accum_vec_H1.py +++ b/src/struphy/pic/tests/test_accum_vec_H1.py @@ -55,6 +55,7 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): from struphy.pic.accumulation import accum_kernels from struphy.pic.accumulation.particles_to_grid import AccumulatorVector from struphy.pic.particles import Particles6D + from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters from struphy.utils.clone_config import CloneConfig mpi_comm = MPI.COMM_WORLD @@ -89,17 +90,16 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): print("Domain decomposition according to", derham.domain_array) # load distributed markers first and use Send/Receive to make global marker copies for the legacy routines - loading_params = { - "seed": 1607, - "moments": [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], - "spatial": "uniform", - } + loading_params = LoadingParameters( + Np=Np, + seed=1607, + moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), + spatial="uniform", + ) particles = Particles6D( comm_world=mpi_comm, clone_config=clone_config, - Np=Np, - bc=["periodic"] * 3, loading_params=loading_params, domain=domain, domain_decomp=domain_decomp, diff --git a/src/struphy/pic/tests/test_accumulation.py b/src/struphy/pic/tests/test_accumulation.py index d7c5ac32b..7fcd8f039 100644 --- a/src/struphy/pic/tests/test_accumulation.py +++ b/src/struphy/pic/tests/test_accumulation.py @@ -59,6 +59,7 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): from struphy.pic.accumulation.particles_to_grid import Accumulator from struphy.pic.particles import Particles6D from struphy.pic.tests.test_pic_legacy_files.accumulation_kernels_3d import kernel_step_ph_full + from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters mpi_comm = MPI.COMM_WORLD # assert mpi_comm.size >= 2 @@ -84,12 +85,10 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): print(derham.domain_array) # load distributed markers first and use Send/Receive to make global marker copies for the legacy routines - loading_params = {"seed": 1607, "moments": [0.0, 0.0, 0.0, 1.0, 2.0, 3.0], "spatial": "uniform"} + loading_params = LoadingParameters(Np=Np, seed=1607, moments=(0.0, 0.0, 0.0, 1.0, 2.0, 3.0), spatial="uniform") particles = Particles6D( comm_world=mpi_comm, - Np=Np, - bc=["periodic"] * 3, loading_params=loading_params, domain=domain, domain_decomp=domain_decomp, diff --git a/src/struphy/pic/tests/test_binning.py b/src/struphy/pic/tests/test_binning.py index 3d4875cf2..2e7500e01 100644 --- a/src/struphy/pic/tests/test_binning.py +++ b/src/struphy/pic/tests/test_binning.py @@ -40,8 +40,14 @@ def test_binning_6D_full_f(mapping, show_plot=False): from mpi4py import MPI from struphy.geometry import domains + from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import Maxwellian3D from struphy.pic.particles import Particles6D + from struphy.pic.utilities import ( + BoundaryParameters, + LoadingParameters, + WeightsParameters, + ) # Set seed seed = 1234 @@ -54,19 +60,17 @@ def test_binning_6D_full_f(mapping, show_plot=False): domain = domain_class(**mapping[1]) # create particles - loading_params = { - "seed": seed, - "spatial": "uniform", - } - bc_params = ["periodic", "periodic", "periodic"] + bc_params = ("periodic", "periodic", "periodic") # =========================================== # ===== Test Maxwellian in v1 direction ===== # =========================================== + loading_params = LoadingParameters(Np=Np, seed=seed, spatial="uniform") + boundary_params = BoundaryParameters(bc=bc_params) + particles = Particles6D( - Np=Np, - bc=bc_params, loading_params=loading_params, + boundary_params=boundary_params, domain=domain, ) @@ -106,22 +110,14 @@ def test_binning_6D_full_f(mapping, show_plot=False): # test weights amp_n = 0.1 l_n = 2 - pert_params = { - "n": { - "ModesCos": { - "given_in_basis": "0", - "ls": [l_n], - "amps": [amp_n], - } - } - } + pert = perturbations.ModesCos(ls=(l_n,), amps=(amp_n,)) + maxwellian = Maxwellian3D(n=(1.0, pert)) particles = Particles6D( - Np=Np, - bc=bc_params, loading_params=loading_params, + boundary_params=boundary_params, domain=domain, - pert_params=pert_params, + background=maxwellian, ) particles.draw_markers() particles.initialize_weights() @@ -154,55 +150,34 @@ def test_binning_6D_full_f(mapping, show_plot=False): # ============================================================== # ===== Test cosines for two backgrounds in eta1 direction ===== # ============================================================== - loading_params = { - "seed": seed, - "spatial": "uniform", - } n1 = 0.8 n2 = 0.2 - bckgr_params = { - "Maxwellian3D_1": { - "n": n1, - }, - "Maxwellian3D_2": { - "n": n2, - "vth1": 0.5, - "u1": 4.5, - }, - } + # test weights amp_n1 = 0.1 amp_n2 = 0.1 l_n1 = 2 l_n2 = 4 - pert_params = { - "Maxwellian3D_1": { - "n": { - "ModesCos": { - "given_in_basis": "0", - "ls": [l_n], - "amps": [amp_n], - } - } - }, - "Maxwellian3D_2": { - "n": { - "ModesCos": { - "given_in_basis": "0", - "ls": [l_n2], - "amps": [amp_n2], - } - } - }, - } - particles = Particles6D( + pert_1 = perturbations.ModesCos(ls=(l_n,), amps=(amp_n,)) + pert_2 = perturbations.ModesCos(ls=(l_n2,), amps=(amp_n2,)) + maxw_1 = Maxwellian3D(n=(n1, pert_1)) + maxw_2 = Maxwellian3D(n=(n2, pert_2), u1=(4.5, None), vth1=(0.5, None)) + background = maxw_1 + maxw_2 + + # adapt s0 for importance sampling + loading_params = LoadingParameters( Np=Np, - bc=bc_params, + seed=seed, + spatial="uniform", + moments=(2.5, 0, 0, 3, 1, 1), + ) + + particles = Particles6D( loading_params=loading_params, + boundary_params=boundary_params, domain=domain, - bckgr_params=bckgr_params, - pert_params=pert_params, + background=background, ) particles.draw_markers() particles.initialize_weights() @@ -221,16 +196,15 @@ def test_binning_6D_full_f(mapping, show_plot=False): # Compare s0 and the sum of two Maxwellians if show_plot: - s0_dict = { - "n": 1.0, - "u1": particles.loading_params["moments"][0], - "u2": particles.loading_params["moments"][1], - "u3": particles.loading_params["moments"][2], - "vth1": particles.loading_params["moments"][3], - "vth2": particles.loading_params["moments"][4], - "vth3": particles.loading_params["moments"][5], - } - s0 = Maxwellian3D(maxw_params=s0_dict) + s0 = Maxwellian3D( + n=(1.0, None), + u1=(particles.loading_params.moments[0], None), + u2=(particles.loading_params.moments[1], None), + u3=(particles.loading_params.moments[2], None), + vth1=(particles.loading_params.moments[3], None), + vth2=(particles.loading_params.moments[4], None), + vth3=(particles.loading_params.moments[5], None), + ) v1 = np.linspace(-10.0, 10.0, 400) phase_space = np.meshgrid( @@ -299,8 +273,14 @@ def test_binning_6D_delta_f(mapping, show_plot=False): from mpi4py import MPI from struphy.geometry import domains + from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import Maxwellian3D from struphy.pic.particles import DeltaFParticles6D + from struphy.pic.utilities import ( + BoundaryParameters, + LoadingParameters, + WeightsParameters, + ) # Set seed seed = 1234 @@ -313,34 +293,25 @@ def test_binning_6D_delta_f(mapping, show_plot=False): domain = domain_class(**mapping[1]) # create particles - loading_params = { - "seed": seed, - "spatial": "uniform", - } - bc_params = ["periodic", "periodic", "periodic"] + bc_params = ("periodic", "periodic", "periodic") # ========================================= # ===== Test cosine in eta1 direction ===== # ========================================= + loading_params = LoadingParameters(Np=Np, seed=seed, spatial="uniform") + boundary_params = BoundaryParameters(bc=bc_params) + # test weights amp_n = 0.1 l_n = 2 - pert_params = { - "n": { - "ModesCos": { - "given_in_basis": "0", - "ls": [l_n], - "amps": [amp_n], - }, - } - } + pert = perturbations.ModesCos(ls=(l_n,), amps=(amp_n,)) + background = Maxwellian3D(n=(1.0, pert)) particles = DeltaFParticles6D( - Np=Np, - bc=bc_params, loading_params=loading_params, + boundary_params=boundary_params, domain=domain, - pert_params=pert_params, + background=background, ) particles.draw_markers() particles.initialize_weights() @@ -373,57 +344,34 @@ def test_binning_6D_delta_f(mapping, show_plot=False): # ============================================================== # ===== Test cosines for two backgrounds in eta1 direction ===== # ============================================================== - loading_params = { - "seed": seed, - "spatial": "uniform", - } n1 = 0.8 n2 = 0.2 - bckgr_params = { - "Maxwellian3D_1": { - "n": n1, - }, - "Maxwellian3D_2": { - "n": n2, - "vth1": 0.5, - "u1": 4.5, - }, - } + # test weights amp_n1 = 0.1 amp_n2 = 0.1 l_n1 = 2 l_n2 = 4 - pert_params = { - "Maxwellian3D_1": { - "use_background_n": False, - "n": { - "ModesCos": { - "given_in_basis": "0", - "ls": [l_n1], - "amps": [amp_n1], - } - }, - }, - "Maxwellian3D_2": { - "use_background_n": True, - "n": { - "ModesCos": { - "given_in_basis": "0", - "ls": [l_n2], - "amps": [amp_n2], - } - }, - }, - } - particles = DeltaFParticles6D( + pert_1 = perturbations.ModesCos(ls=(l_n,), amps=(amp_n,)) + pert_2 = perturbations.ModesCos(ls=(l_n2,), amps=(amp_n2,)) + maxw_1 = Maxwellian3D(n=(n1, pert_1)) + maxw_2 = Maxwellian3D(n=(n2, pert_2), u1=(4.5, None), vth1=(0.5, None)) + background = maxw_1 + maxw_2 + + # adapt s0 for importance sampling + loading_params = LoadingParameters( Np=Np, - bc=bc_params, + seed=seed, + spatial="uniform", + moments=(2.5, 0, 0, 2, 1, 1), + ) + + particles = DeltaFParticles6D( loading_params=loading_params, + boundary_params=boundary_params, domain=domain, - bckgr_params=bckgr_params, - pert_params=pert_params, + background=background, ) particles.draw_markers() particles.initialize_weights() @@ -438,20 +386,19 @@ def test_binning_6D_delta_f(mapping, show_plot=False): e1_plot = e1_bins[:-1] + de / 2 - ana_res = amp_n1 * np.cos(2 * np.pi * l_n1 * e1_plot) + n2 + amp_n2 * np.cos(2 * np.pi * l_n2 * e1_plot) + ana_res = amp_n1 * np.cos(2 * np.pi * l_n1 * e1_plot) + amp_n2 * np.cos(2 * np.pi * l_n2 * e1_plot) # Compare s0 and the sum of two Maxwellians if show_plot: - s0_dict = { - "n": 1.0, - "u1": particles.loading_params["moments"][0], - "u2": particles.loading_params["moments"][1], - "u3": particles.loading_params["moments"][2], - "vth1": particles.loading_params["moments"][3], - "vth2": particles.loading_params["moments"][4], - "vth3": particles.loading_params["moments"][5], - } - s0 = Maxwellian3D(maxw_params=s0_dict) + s0 = Maxwellian3D( + n=(1.0, None), + u1=(particles.loading_params.moments[0], None), + u2=(particles.loading_params.moments[1], None), + u3=(particles.loading_params.moments[2], None), + vth1=(particles.loading_params.moments[3], None), + vth2=(particles.loading_params.moments[4], None), + vth3=(particles.loading_params.moments[5], None), + ) v1 = np.linspace(-10.0, 10.0, 400) phase_space = np.meshgrid( @@ -523,8 +470,14 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): from mpi4py import MPI from struphy.geometry import domains + from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import Maxwellian3D from struphy.pic.particles import Particles6D + from struphy.pic.utilities import ( + BoundaryParameters, + LoadingParameters, + WeightsParameters, + ) # Set seed seed = 1234 @@ -543,19 +496,17 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): assert size > 1 # create particles - loading_params = { - "seed": seed, - "spatial": "uniform", - } - bc_params = ["periodic", "periodic", "periodic"] + bc_params = ("periodic", "periodic", "periodic") # =========================================== # ===== Test Maxwellian in v1 direction ===== # =========================================== + loading_params = LoadingParameters(Np=Np, seed=seed, spatial="uniform") + boundary_params = BoundaryParameters(bc=bc_params) + particles = Particles6D( - Np=Np, - bc=bc_params, loading_params=loading_params, + boundary_params=boundary_params, comm_world=comm, domain=domain, ) @@ -600,23 +551,15 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): # test weights amp_n = 0.1 l_n = 2 - pert_params = { - "n": { - "ModesCos": { - "given_in_basis": "0", - "ls": [l_n], - "amps": [amp_n], - } - } - } + pert = perturbations.ModesCos(ls=(l_n,), amps=(amp_n,)) + maxwellian = Maxwellian3D(n=(1.0, pert)) particles = Particles6D( - Np=Np, - bc=bc_params, loading_params=loading_params, + boundary_params=boundary_params, comm_world=comm, domain=domain, - pert_params=pert_params, + background=maxwellian, ) particles.draw_markers() particles.initialize_weights() @@ -654,10 +597,6 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): # ============================================================== # ===== Test cosines for two backgrounds in eta1 direction ===== # ============================================================== - loading_params = { - "seed": seed, - "spatial": "uniform", - } n1 = 0.8 n2 = 0.2 bckgr_params = { @@ -695,15 +634,26 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): } }, } + pert_1 = perturbations.ModesCos(ls=(l_n1,), amps=(amp_n1,)) + pert_2 = perturbations.ModesCos(ls=(l_n2,), amps=(amp_n2,)) + maxw_1 = Maxwellian3D(n=(n1, pert_1)) + maxw_2 = Maxwellian3D(n=(n2, pert_2), u1=(4.5, None), vth1=(0.5, None)) + background = maxw_1 + maxw_2 + + # adapt s0 for importance sampling + loading_params = LoadingParameters( + Np=Np, + seed=seed, + spatial="uniform", + moments=(2.5, 0, 0, 2, 1, 1), + ) particles = Particles6D( - Np=Np, - bc=bc_params, loading_params=loading_params, + boundary_params=boundary_params, comm_world=comm, domain=domain, - bckgr_params=bckgr_params, - pert_params=pert_params, + background=background, ) particles.draw_markers() particles.initialize_weights() @@ -727,16 +677,15 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): # Compare s0 and the sum of two Maxwellians if show_plot and rank == 0: - s0_dict = { - "n": 1.0, - "u1": particles.loading_params["moments"][0], - "u2": particles.loading_params["moments"][1], - "u3": particles.loading_params["moments"][2], - "vth1": particles.loading_params["moments"][3], - "vth2": particles.loading_params["moments"][4], - "vth3": particles.loading_params["moments"][5], - } - s0 = Maxwellian3D(maxw_params=s0_dict) + s0 = Maxwellian3D( + n=(1.0, None), + u1=(particles.loading_params.moments[0], None), + u2=(particles.loading_params.moments[1], None), + u3=(particles.loading_params.moments[2], None), + vth1=(particles.loading_params.moments[3], None), + vth2=(particles.loading_params.moments[4], None), + vth3=(particles.loading_params.moments[5], None), + ) v1 = np.linspace(-10.0, 10.0, 400) phase_space = np.meshgrid( @@ -805,8 +754,14 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): from mpi4py import MPI from struphy.geometry import domains + from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import Maxwellian3D from struphy.pic.particles import DeltaFParticles6D + from struphy.pic.utilities import ( + BoundaryParameters, + LoadingParameters, + WeightsParameters, + ) # Set seed seed = 1234 @@ -825,15 +780,14 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): assert size > 1 # create particles - loading_params = { - "seed": seed, - "spatial": "uniform", - } - bc_params = ["periodic", "periodic", "periodic"] + bc_params = ("periodic", "periodic", "periodic") # ========================================= # ===== Test cosine in eta1 direction ===== # ========================================= + loading_params = LoadingParameters(Np=Np, seed=seed, spatial="uniform") + boundary_params = BoundaryParameters(bc=bc_params) + # test weights amp_n = 0.1 l_n = 2 @@ -846,14 +800,15 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): } } } + pert = perturbations.ModesCos(ls=(l_n,), amps=(amp_n,)) + background = Maxwellian3D(n=(1.0, pert)) particles = DeltaFParticles6D( - Np=Np, - bc=bc_params, loading_params=loading_params, + boundary_params=boundary_params, comm_world=comm, domain=domain, - pert_params=pert_params, + background=background, ) particles.draw_markers() particles.initialize_weights() @@ -891,10 +846,6 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): # ============================================================== # ===== Test cosines for two backgrounds in eta1 direction ===== # ============================================================== - loading_params = { - "seed": seed, - "spatial": "uniform", - } n1 = 0.8 n2 = 0.2 bckgr_params = { @@ -934,15 +885,26 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): }, }, } + pert_1 = perturbations.ModesCos(ls=(l_n1,), amps=(amp_n1,)) + pert_2 = perturbations.ModesCos(ls=(l_n2,), amps=(amp_n2,)) + maxw_1 = Maxwellian3D(n=(n1, pert_1)) + maxw_2 = Maxwellian3D(n=(n2, pert_2), u1=(4.5, None), vth1=(0.5, None)) + background = maxw_1 + maxw_2 + + # adapt s0 for importance sampling + loading_params = LoadingParameters( + Np=Np, + seed=seed, + spatial="uniform", + moments=(2.5, 0, 0, 2, 1, 1), + ) particles = DeltaFParticles6D( - Np=Np, - bc=bc_params, loading_params=loading_params, + boundary_params=boundary_params, comm_world=comm, domain=domain, - bckgr_params=bckgr_params, - pert_params=pert_params, + background=background, ) particles.draw_markers() particles.initialize_weights() @@ -962,20 +924,19 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): e1_plot = e1_bins[:-1] + de / 2 - ana_res = amp_n1 * np.cos(2 * np.pi * l_n1 * e1_plot) + n2 + amp_n2 * np.cos(2 * np.pi * l_n2 * e1_plot) + ana_res = amp_n1 * np.cos(2 * np.pi * l_n1 * e1_plot) + amp_n2 * np.cos(2 * np.pi * l_n2 * e1_plot) # Compare s0 and the sum of two Maxwellians if show_plot and rank == 0: - s0_dict = { - "n": 1.0, - "u1": particles.loading_params["moments"][0], - "u2": particles.loading_params["moments"][1], - "u3": particles.loading_params["moments"][2], - "vth1": particles.loading_params["moments"][3], - "vth2": particles.loading_params["moments"][4], - "vth3": particles.loading_params["moments"][5], - } - s0 = Maxwellian3D(maxw_params=s0_dict) + s0 = Maxwellian3D( + n=(1.0, None), + u1=(particles.loading_params.moments[0], None), + u2=(particles.loading_params.moments[1], None), + u3=(particles.loading_params.moments[2], None), + vth1=(particles.loading_params.moments[3], None), + vth2=(particles.loading_params.moments[4], None), + vth3=(particles.loading_params.moments[5], None), + ) v1 = np.linspace(-10.0, 10.0, 400) phase_space = np.meshgrid( diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index 7efad5f2b..c4520eeab 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -42,6 +42,7 @@ def test_draw(Nel, p, spl_kind, mapping, ppc=10): from struphy.feec.psydac_derham import Derham from struphy.geometry import domains from struphy.pic.particles import Particles6D + from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters comm = MPI.COMM_WORLD assert comm.size >= 2 @@ -66,17 +67,16 @@ def test_draw(Nel, p, spl_kind, mapping, ppc=10): print(derham.domain_array) # create particles - loading_params = { - "seed": seed, - "moments": [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], - "spatial": "uniform", - } + loading_params = LoadingParameters( + ppc=ppc, + seed=seed, + moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), + spatial="uniform", + ) particles = Particles6D( comm_world=comm, - ppc=ppc, domain_decomp=domain_decomp, - bc=["periodic", "periodic", "periodic"], loading_params=loading_params, domain=domain, ) diff --git a/src/struphy/pic/tests/test_pushers.py b/src/struphy/pic/tests/test_pushers.py index 218d65f29..acb89f2e9 100644 --- a/src/struphy/pic/tests/test_pushers.py +++ b/src/struphy/pic/tests/test_pushers.py @@ -33,6 +33,7 @@ def test_push_vxb_analytic(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str + from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -58,14 +59,12 @@ def test_push_vxb_analytic(Nel, p, spl_kind, mapping, show_plots=False): # particle loading and sorting seed = 1234 - loader_params = {"seed": seed, "moments": [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], "spatial": "uniform"} + loading_params = LoadingParameters(ppc=2, seed=seed, moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), spatial="uniform") particles = Particles6D( comm_world=comm, - ppc=2, domain_decomp=domain_decomp, - bc=["periodic", "periodic", "periodic"], - loading_params=loader_params, + loading_params=loading_params, ) particles.draw_markers() @@ -170,6 +169,7 @@ def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str + from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -195,14 +195,12 @@ def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): # particle loading and sorting seed = 1234 - loader_params = {"seed": seed, "moments": [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], "spatial": "uniform"} + loading_params = LoadingParameters(ppc=2, seed=seed, moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), spatial="uniform") particles = Particles6D( comm_world=comm, - ppc=2, domain_decomp=domain_decomp, - bc=["periodic", "periodic", "periodic"], - loading_params=loader_params, + loading_params=loading_params, ) particles.draw_markers() @@ -318,6 +316,7 @@ def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str + from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -343,14 +342,12 @@ def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): # particle loading and sorting seed = 1234 - loader_params = {"seed": seed, "moments": [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], "spatial": "uniform"} + loading_params = LoadingParameters(ppc=2, seed=seed, moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), spatial="uniform") particles = Particles6D( comm_world=comm, - ppc=2, domain_decomp=domain_decomp, - bc=["periodic", "periodic", "periodic"], - loading_params=loader_params, + loading_params=loading_params, ) particles.draw_markers() @@ -466,6 +463,7 @@ def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str + from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -491,14 +489,12 @@ def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): # particle loading and sorting seed = 1234 - loader_params = {"seed": seed, "moments": [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], "spatial": "uniform"} + loading_params = LoadingParameters(ppc=2, seed=seed, moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), spatial="uniform") particles = Particles6D( comm_world=comm, - ppc=2, domain_decomp=domain_decomp, - bc=["periodic", "periodic", "periodic"], - loading_params=loader_params, + loading_params=loading_params, ) particles.draw_markers() @@ -614,6 +610,7 @@ def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str + from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -639,14 +636,12 @@ def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): # particle loading and sorting seed = 1234 - loader_params = {"seed": seed, "moments": [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], "spatial": "uniform"} + loading_params = LoadingParameters(ppc=2, seed=seed, moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), spatial="uniform") particles = Particles6D( comm_world=comm, - ppc=2, domain_decomp=domain_decomp, - bc=["periodic", "periodic", "periodic"], - loading_params=loader_params, + loading_params=loading_params, ) particles.draw_markers() @@ -765,6 +760,7 @@ def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing import pusher_kernels from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str + from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -791,14 +787,12 @@ def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): # particle loading and sorting seed = 1234 - loader_params = {"seed": seed, "moments": [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], "spatial": "uniform"} + loading_params = LoadingParameters(ppc=2, seed=seed, moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), spatial="uniform") particles = Particles6D( comm_world=comm, - ppc=2, domain_decomp=domain_decomp, - bc=["periodic", "periodic", "periodic"], - loading_params=loader_params, + loading_params=loading_params, ) particles.draw_markers() diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index 151bf852f..d4c050554 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -7,6 +7,7 @@ from struphy.feec.psydac_derham import Derham from struphy.geometry import domains from struphy.pic.particles import Particles6D +from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters @pytest.mark.mpi(min_size=2) @@ -50,13 +51,11 @@ def test_sorting(Nel, p, spl_kind, mapping, Np, verbose=False): nprocs = derham.domain_decomposition.nprocs domain_decomp = (domain_array, nprocs) - loading_params = {"seed": 1607, "moments": [0.0, 0.0, 0.0, 1.0, 2.0, 3.0], "spatial": "uniform"} + loading_params = LoadingParameters(Np=Np, seed=1607, moments=(0.0, 0.0, 0.0, 1.0, 2.0, 3.0), spatial="uniform") boxes_per_dim = (3, 3, 6) particles = Particles6D( comm_world=mpi_comm, - Np=Np, - bc=["periodic", "periodic", "periodic"], loading_params=loading_params, domain_decomp=domain_decomp, boxes_per_dim=boxes_per_dim, diff --git a/src/struphy/pic/tests/test_sph.py b/src/struphy/pic/tests/test_sph.py index 0ab7caa16..73f8b4ec8 100644 --- a/src/struphy/pic/tests/test_sph.py +++ b/src/struphy/pic/tests/test_sph.py @@ -5,8 +5,11 @@ from mpi4py import MPI from struphy.feec.psydac_derham import Derham +from struphy.fields_background.equils import ConstantVelocity from struphy.geometry import domains +from struphy.initial import perturbations from struphy.pic.particles import ParticlesSPH +from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters @pytest.mark.mpi(min_size=2) @@ -22,27 +25,26 @@ def test_evaluation_mc(Np, bc_x, show_plot=False): domain = domain_class(**dom_params) boxes_per_dim = (16, 1, 1) - loading_params = {"seed": 1607} + loading_params = LoadingParameters(Np=Np, seed=1607) + boundary_params = BoundaryParameters(bc=(bc_x, "periodic", "periodic")) - cst_vel = {"density_profile": "constant", "n": 1.0} - bckgr_params = {"ConstantVelocity": cst_vel, "pforms": ["vol", None]} + background = ConstantVelocity(n=1.0, density_profile="constant") + background.domain = domain - mode_params = {"given_in_basis": "0", "ls": [1], "amps": [1e-0]} - modes = {"ModesSin": mode_params} - pert_params = {"n": modes} + pert = {"n": perturbations.ModesSin(ls=(1,), amps=(1e-0,))} fun_exact = lambda e1, e2, e3: 1.0 + np.sin(2 * np.pi * e1) particles = ParticlesSPH( comm_world=comm, - Np=Np, + loading_params=loading_params, + boundary_params=boundary_params, boxes_per_dim=boxes_per_dim, - bc=[bc_x, "periodic", "periodic"], bufsize=1.0, - loading_params=loading_params, domain=domain, - bckgr_params=bckgr_params, - pert_params=pert_params, + background=background, + perturbations=pert, + n_as_volume_form=True, ) particles.draw_markers(sort=False) @@ -66,6 +68,7 @@ def test_evaluation_mc(Np, bc_x, show_plot=False): plt.figure(figsize=(12, 8)) plt.plot(ee1.squeeze(), fun_exact(ee1, ee2, ee3).squeeze(), label="exact") plt.plot(ee1.squeeze(), all_eval.squeeze(), "--.", label="eval_sph") + plt.legend() plt.show() @@ -88,29 +91,26 @@ def test_evaluation_tesselation(boxes_per_dim, ppb, bc_x, show_plot=False): domain_class = getattr(domains, dom_type) domain = domain_class(**dom_params) - loading = "tesselation" - loading_params = {"n_quad": 1} + loading_params = LoadingParameters(ppb=ppb, loading="tesselation") + boundary_params = BoundaryParameters(bc=(bc_x, "periodic", "periodic")) - cst_vel = {"density_profile": "constant", "n": 1.0} - bckgr_params = {"ConstantVelocity": cst_vel, "pforms": ["vol", None]} + background = ConstantVelocity(n=1.0, density_profile="constant") + background.domain = domain - mode_params = {"given_in_basis": "0", "ls": [1], "amps": [1e-0]} - modes = {"ModesSin": mode_params} - pert_params = {"n": modes} + pert = {"n": perturbations.ModesSin(ls=(1,), amps=(1e-0,))} fun_exact = lambda e1, e2, e3: 1.0 + np.sin(2 * np.pi * e1) particles = ParticlesSPH( comm_world=comm, - ppb=ppb, + loading_params=loading_params, + boundary_params=boundary_params, boxes_per_dim=boxes_per_dim, - bc=[bc_x, "periodic", "periodic"], bufsize=1.0, - loading=loading, - loading_params=loading_params, domain=domain, - bckgr_params=bckgr_params, - pert_params=pert_params, + background=background, + perturbations=pert, + n_as_volume_form=True, verbose=True, ) @@ -149,10 +149,5 @@ def test_evaluation_tesselation(boxes_per_dim, ppb, bc_x, show_plot=False): if __name__ == "__main__": - test_evaluation_mc(40000, "periodic", show_plot=True) - # test_evaluation_tesselation( - # (8, 1, 1), - # 4, - # "periodic", - # show_plot=True - # ) + # test_evaluation_mc(40000, "periodic", show_plot=True) + test_evaluation_tesselation((8, 1, 1), 4, "periodic", show_plot=True) diff --git a/src/struphy/pic/tests/test_tesselation.py b/src/struphy/pic/tests/test_tesselation.py index 145931d15..adc67447b 100644 --- a/src/struphy/pic/tests/test_tesselation.py +++ b/src/struphy/pic/tests/test_tesselation.py @@ -6,8 +6,11 @@ from mpi4py import MPI from struphy.feec.psydac_derham import Derham +from struphy.fields_background.equils import ConstantVelocity from struphy.geometry import domains +from struphy.initial import perturbations from struphy.pic.particles import ParticlesSPH +from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters @pytest.mark.mpi(min_size=2) @@ -25,17 +28,14 @@ def test_draw(ppb, nx, ny, nz): domain = domain_class(**dom_params) boxes_per_dim = (nx, ny, nz) - bc = ["periodic"] * 3 - loading = "tesselation" bufsize = 0.5 + loading_params = LoadingParameters(ppb=ppb, loading="tesselation") # instantiate Particle object particles = ParticlesSPH( comm_world=comm, - ppb=ppb, + loading_params=loading_params, boxes_per_dim=boxes_per_dim, - bc=bc, - loading=loading, domain=domain, verbose=False, bufsize=bufsize, @@ -89,31 +89,24 @@ def test_cell_average(ppb, nx, ny, nz, n_quad, show_plot=False): domain = domain_class(**dom_params) boxes_per_dim = (nx, ny, nz) - bc = ["periodic"] * 3 - loading = "tesselation" - loading_params = {"n_quad": n_quad} + loading_params = LoadingParameters(ppb=ppb, loading="tesselation", n_quad=n_quad) bufsize = 0.5 - cst_vel = {"ux": 0.0, "uy": 0.0, "uz": 0.0, "density_profile": "constant"} - bckgr_params = {"ConstantVelocity": cst_vel} + background = ConstantVelocity(n=1.0, ux=0.0, uy=0.0, uz=0.0, density_profile="constant") + background.domain = domain - mode_params = {"given_in_basis": "0", "ls": [1], "amps": [1e-0]} - modes = {"ModesSin": mode_params} - pert_params = {"n": modes} + pert = {"n": perturbations.ModesSin(ls=(1,), amps=(1e-0,))} # instantiate Particle object particles = ParticlesSPH( comm_world=comm, - ppb=ppb, boxes_per_dim=boxes_per_dim, - bc=bc, - loading=loading, loading_params=loading_params, domain=domain, verbose=False, bufsize=bufsize, - bckgr_params=bckgr_params, - pert_params=pert_params, + background=background, + perturbations=pert, ) particles.draw_markers(sort=False) @@ -190,5 +183,5 @@ def test_cell_average(ppb, nx, ny, nz, n_quad, show_plot=False): if __name__ == "__main__": - # test_draw(8, 16, 1, 1) + test_draw(8, 16, 1, 1) test_cell_average(8, 6, 16, 14, n_quad=2, show_plot=True) diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 14c9bb8dc..39ac730ac 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -93,8 +93,9 @@ def get_params_of_run(path: str) -> ParamsIn: grid = pickle.load(f) with open(os.path.join(path, "derham_opts.bin"), "rb") as f: derham_opts = pickle.load(f) - with open(os.path.join(path, "model.bin"), "rb") as f: - model = pickle.load(f) + with open(os.path.join(path, "model_class.bin"), "rb") as f: + model_class: StruphyModel = pickle.load(f) + model = model_class() else: raise FileNotFoundError(f"Neither of the paths {params_path} or {bin_path} exists.") @@ -115,6 +116,7 @@ def get_params_of_run(path: str) -> ParamsIn: def create_femfields( path: str, + params_in: ParamsIn, *, step: int = 1, ): @@ -125,6 +127,9 @@ def create_femfields( path : str Absolute path of simulation output folder. + params_in : ParamsIn + Simulation parameters. + step : int Whether to create FEM fields at every time step (step=1, default), every second time step (step=2), etc. @@ -141,9 +146,6 @@ def create_femfields( meta = yaml.load(f, Loader=yaml.FullLoader) nproc = meta["MPI processes"] - # import parameters - params_in = get_params_of_run(path) - derham = setup_derham( params_in.grid, params_in.derham_opts, @@ -249,7 +251,7 @@ def create_femfields( def eval_femfields( - path: str, + params_in: ParamsIn, fields: dict, *, celldivide: list = [1, 1, 1], @@ -259,8 +261,8 @@ def eval_femfields( Parameters ---------- - path : str - Absolute path of simulation output folder. + params_in : ParamsIn + Simulation parameters. fields : dict Obtained from struphy.diagnostics.post_processing.create_femfields. @@ -287,9 +289,6 @@ def eval_femfields( Mapped (physical) grids obtained by domain(*grids_log). """ - # import parameters - params_in = get_params_of_run(path) - # get domain domain = params_in.domain @@ -618,7 +617,14 @@ def post_process_markers( file.close() -def post_process_f(path_in, path_out, species, step=1, compute_bckgr=False): +def post_process_f( + path_in, + params_in: ParamsIn, + path_out, + species, + step=1, + compute_bckgr=False, +): """Computes and saves distribution functions of saved binning data during a simulation. Parameters @@ -626,6 +632,9 @@ def post_process_f(path_in, path_out, species, step=1, compute_bckgr=False): path_in : str Absolute path of simulation output folder. + params_in : ParamsIn + Simulation parameters. + path_out : str Absolute path of where to store the .txt files. Will be in path_out/orbits. @@ -657,9 +666,6 @@ def post_process_f(path_in, path_out, species, step=1, compute_bckgr=False): for i in range(int(nproc)) ] - # import parameters - params = get_params_of_run(path_in) - # directory for .npy files path_distr = os.path.join(path_out, "distribution_function") @@ -729,7 +735,7 @@ def post_process_f(path_in, path_out, species, step=1, compute_bckgr=False): # maxw_params=maxw_params, # ) - spec: KineticSpecies = getattr(params.model, species) + spec: KineticSpecies = getattr(params_in.model, species) var: PICVariable = spec.var f_bckgr: KineticBackground = var.backgrounds diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 6c12fb5a3..71a36358e 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -2738,7 +2738,8 @@ def allocate(self): # always stabilize if np.abs(self.options.sigma_1) < 1e-14: self.options.sigma_1 = 1e-14 - print(f"Stabilizing Poisson solve with {self.options.sigma_1 = }") + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"Stabilizing Poisson solve with {self.options.sigma_1 = }") # model parameters self._sigma_1 = self.options.sigma_1 diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index cce0e45c4..d5cfc2554 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -7,6 +7,8 @@ from struphy.feec.projectors import L2Projector from struphy.feec.psydac_derham import Derham from struphy.geometry import domains +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable from struphy.propagators.base import Propagator from struphy.propagators.propagators_fields import ImplicitDiffusion @@ -31,15 +33,6 @@ def test_poisson_M1perp_1d(direction, bc_type, mapping, show_plot=False): in 1D by means of manufactured solutions. """ - solver_params = { - "type": ("pcg", "MassMatrixPreconditioner"), - "tol": 1.0e-13, - "maxiter": 3000, - "info": True, - "verbose": False, - "recycle": False, - } - # create domain object dom_type = mapping[0] dom_params = mapping[1] @@ -91,7 +84,8 @@ def rho1_xyz(x, y, z): else: if bc_type == "dirichlet": spl_kind = [False, True, True] - dirichlet_bc = [[not kd] * 2 for kd in spl_kind] + dirichlet_bc = [(not kd,) * 2 for kd in spl_kind] + dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): return np.sin(2 * np.pi / Lx * x) @@ -115,7 +109,8 @@ def rho1_xyz(x, y, z): else: if bc_type == "dirichlet": spl_kind = [True, False, True] - dirichlet_bc = [[not kd] * 2 for kd in spl_kind] + dirichlet_bc = [(not kd,) * 2 for kd in spl_kind] + dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): return np.sin(2 * np.pi / Ly * y) @@ -126,6 +121,7 @@ def rho1_xyz(x, y, z): print("Direction should be either 0 or 1") # create derham object + print(f"{dirichlet_bc = }") derham = Derham(Nel, p, spl_kind, dirichlet_bc=dirichlet_bc, comm=comm) # mass matrices @@ -142,24 +138,40 @@ def rho1(e1, e2, e3): rho_vec = L2Projector("H1", mass_ops).get_dofs(rho1, apply_bc=True) # create Poisson solver - _phi = derham.create_spline_function("phi", "H1") - poisson_solver = ImplicitDiffusion( - _phi.vector, + solver_params = SolverParameters( + tol=1.0e-13, + maxiter=3000, + info=True, + verbose=False, + recycle=False, + ) + + _phi = FEECVariable(space="H1") + _phi.allocate(derham=derham, domain=domain) + + poisson_solver = ImplicitDiffusion() + poisson_solver.variables.phi = _phi + + poisson_solver.options = poisson_solver.Options( sigma_1=1e-12, sigma_2=0.0, sigma_3=1.0, divide_by_dt=True, diffusion_mat="M1perp", rho=rho_vec, - solver=solver_params, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, ) + poisson_solver.allocate() + # Solve Poisson (call propagator with dt=1.) dt = 1.0 poisson_solver(dt) # push numerical solution and compare - sol_val1 = domain.push(_phi, e1, e2, e3, kind="0") + sol_val1 = domain.push(_phi.spline, e1, e2, e3, kind="0") x, y, z = domain(e1, e2, e3) analytic_value1 = sol1_xyz(x, y, z) @@ -227,14 +239,6 @@ def test_poisson_M1perp_2d(Nel, p, bc_type, mapping, show_plot=False): Test the Poisson solver with M1perp diffusion matrix by means of manufactured solutions in 2D . """ - solver_params = { - "type": ("pcg", "MassMatrixPreconditioner"), - "tol": 1.0e-13, - "maxiter": 3000, - "info": True, - "verbose": False, - "recycle": False, - } # create domain object dom_type = mapping[0] @@ -274,7 +278,8 @@ def rho2_xyz(x, y, z): elif bc_type == "dirichlet": spl_kind = [False, True, True] - dirichlet_bc = [[not kd] * 2 for kd in spl_kind] + dirichlet_bc = [(not kd,) * 2 for kd in spl_kind] + dirichlet_bc = tuple(dirichlet_bc) print(f"{dirichlet_bc = }") # manufactured solution in 2D @@ -333,24 +338,62 @@ def rho2(e1, e2, e3): rho_vec2 = l2_proj.get_dofs(rho2, apply_bc=True) # Create Poisson solvers - _phi1 = derham.create_spline_function("test1", "H1") - poisson_solver1 = ImplicitDiffusion( - _phi1.vector, sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, diffusion_mat="M1perp", rho=rho_vec1, solver=solver_params + solver_params = SolverParameters( + tol=1.0e-13, + maxiter=3000, + info=True, + verbose=False, + recycle=False, ) - _phi2 = derham.create_spline_function("test2", "H1") - poisson_solver2 = ImplicitDiffusion( - _phi2.vector, sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, diffusion_mat="M1perp", rho=rho_vec2, solver=solver_params + _phi1 = FEECVariable(space="H1") + _phi1.allocate(derham=derham, domain=domain) + + poisson_solver1 = ImplicitDiffusion() + poisson_solver1.variables.phi = _phi1 + + poisson_solver1.options = poisson_solver1.Options( + sigma_1=1e-8, + sigma_2=0.0, + sigma_3=1.0, + divide_by_dt=True, + diffusion_mat="M1perp", + rho=rho_vec1, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, ) + poisson_solver1.allocate() + + _phi2 = FEECVariable(space="H1") + _phi2.allocate(derham=derham, domain=domain) + + poisson_solver2 = ImplicitDiffusion() + poisson_solver2.variables.phi = _phi2 + + poisson_solver2.options = poisson_solver2.Options( + sigma_1=1e-8, + sigma_2=0.0, + sigma_3=1.0, + divide_by_dt=True, + diffusion_mat="M1perp", + rho=rho_vec2, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, + ) + + poisson_solver2.allocate() + # Solve Poisson equation (call propagator with dt=1.) dt = 1.0 poisson_solver1(dt) poisson_solver2(dt) # push numerical solutions - sol_val1 = domain.push(_phi1, e1, e2, e3, kind="0") - sol_val2 = domain.push(_phi2, e1, e2, e3, kind="0") + sol_val1 = domain.push(_phi1.spline, e1, e2, e3, kind="0") + sol_val2 = domain.push(_phi2.spline, e1, e2, e3, kind="0") x, y, z = domain(e1, e2, e3) analytic_value1 = sol1_xyz(x, y, z) @@ -411,15 +454,6 @@ def test_poisson_M1perp_3d_compare_2p5d(Nel, p, mapping, show_plot=False): from time import time - solver_params = { - "type": ("pcg", "MassMatrixPreconditioner"), - "tol": 1.0e-13, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": False, - } - # create domain object dom_type = mapping[0] dom_params = mapping[1] @@ -429,7 +463,7 @@ def test_poisson_M1perp_3d_compare_2p5d(Nel, p, mapping, show_plot=False): # boundary conditions spl_kind = [False, True, True] - dirichlet_bc = [[True, True], [False, False], [False, False]] + dirichlet_bc = ((True, True), (False, False), (False, False)) # evaluation grid e1 = np.linspace(0.0, 1.0, 50) @@ -458,20 +492,44 @@ def rho(e1, e2, e3): print(f"{rho_vec[:].shape = }") # Create 3d Poisson solver - _phi = derham.create_spline_function("test2", "H1") - _phi_2p5d = derham.create_spline_function("sol_2p5d", "H1") - poisson_solver_3d = ImplicitDiffusion( - _phi.vector, sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, diffusion_mat="M1perp", rho=rho_vec, solver=solver_params + solver_params = SolverParameters( + tol=1.0e-13, + maxiter=3000, + info=True, + verbose=False, + recycle=False, + ) + + _phi = FEECVariable(space="H1") + _phi.allocate(derham=derham, domain=domain) + + _phi_2p5d = FEECVariable(space="H1") + _phi_2p5d.allocate(derham=derham, domain=domain) + + poisson_solver_3d = ImplicitDiffusion() + poisson_solver_3d.variables.phi = _phi + + poisson_solver_3d.options = poisson_solver_3d.Options( + sigma_1=1e-8, + sigma_2=0.0, + sigma_3=1.0, + divide_by_dt=True, + diffusion_mat="M1perp", + rho=rho_vec, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, ) - s = _phi.starts - e = _phi.ends + poisson_solver_3d.allocate() + + s = _phi.spline.starts + e = _phi.spline.ends # create 2.5d deRham object Nel_new = [Nel[0], Nel[1], 1] p[2] = 1 spl_kind[2] = True - dirichlet_bc[2] = [False, False] derham = Derham(Nel_new, p, spl_kind, dirichlet_bc=dirichlet_bc, comm=comm) mass_ops = WeightedMassOperators(derham, domain) @@ -479,18 +537,27 @@ def rho(e1, e2, e3): Propagator.derham = derham Propagator.mass_ops = mass_ops - _phi_small = derham.create_spline_function("test_small", "H1") + _phi_small = FEECVariable(space="H1") + _phi_small.allocate(derham=derham, domain=domain) rhs = derham.create_spline_function("rhs", "H1") - poisson_solver_2p5d = ImplicitDiffusion( - _phi_small.vector, + + poisson_solver_2p5d = ImplicitDiffusion() + poisson_solver_2p5d.variables.phi = _phi_small + + poisson_solver_2p5d.options = poisson_solver_2p5d.Options( sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, + divide_by_dt=True, diffusion_mat="M1perp", rho=rhs.vector, - solver=solver_params, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, ) + poisson_solver_2p5d.allocate() + # Solve Poisson equation (call propagator with dt=1.) dt = 1.0 t0 = time() @@ -508,16 +575,16 @@ def rho(e1, e2, e3): poisson_solver_2p5d(dt) t1i = time() t_inner += t1i - t0i - _tmp = _phi_small.vector.copy() - _phi_2p5d.vector[s[0] : e[0] + 1, s[1] : e[1] + 1, n] = _tmp[s[0] : e[0] + 1, s[1] : e[1] + 1, 0] + _tmp = _phi_small.spline.vector.copy() + _phi_2p5d.spline.vector[s[0] : e[0] + 1, s[1] : e[1] + 1, n] = _tmp[s[0] : e[0] + 1, s[1] : e[1] + 1, 0] t1 = time() print(f"rank {rank}, 2.5d pure solve time (without copy) = {t_inner}") print(f"rank {rank}, 2.5d solve time = {t1 - t0}") # push numerical solutions - sol_val = domain.push(_phi, e1, e2, e3, kind="0") - sol_val_2p5d = domain.push(_phi_2p5d, e1, e2, e3, kind="0") + sol_val = domain.push(_phi.spline, e1, e2, e3, kind="0") + sol_val_2p5d = domain.push(_phi_2p5d.spline, e1, e2, e3, kind="0") x, y, z = domain(e1, e2, e3) print("max diff:", np.max(np.abs(sol_val - sol_val_2p5d))) @@ -558,9 +625,9 @@ def rho(e1, e2, e3): if __name__ == "__main__": direction = 0 - bc_type = "periodic" + bc_type = "dirichlet" mapping = ["Cuboid", {"l1": 0.0, "r1": 4.0, "l2": 0.0, "r2": 2.0, "l3": 0.0, "r3": 3.0}] - # mapping = ['Orthogonal', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 3.}] + mapping = ["Orthogonal", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 3.0}] test_poisson_M1perp_1d(direction, bc_type, mapping, show_plot=True) # Nel = [64, 64, 1] @@ -570,8 +637,7 @@ def rho(e1, e2, e3): # mapping = ['Orthogonal', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 1.}] # test_poisson_M1perp_2d(Nel, p, bc_type, mapping, show_plot=True) - # Nel = [64, 64, 16] - # p = [2, 2, 1] - # mapping = ['Cuboid', {'l1': 0., 'r1': 1., - # 'l2': 0., 'r2': 1., 'l3': 0., 'r3': 1.}] + Nel = [64, 64, 16] + p = [2, 2, 1] + mapping = ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}] # test_poisson_M1perp_3d_compare_2p5d(Nel, p, mapping, show_plot=True) diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 8281f15da..2460968b6 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -7,6 +7,8 @@ from struphy.feec.projectors import L2Projector from struphy.feec.psydac_derham import Derham from struphy.geometry import domains +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable from struphy.propagators.base import Propagator from struphy.propagators.propagators_fields import ImplicitDiffusion @@ -30,15 +32,6 @@ def test_poisson_1d(direction, bc_type, mapping, show_plot=False): Test the convergence of Poisson solver in 1D by means of manufactured solutions. """ - solver_params = { - "type": ("pcg", "MassMatrixPreconditioner"), - "tol": 1.0e-13, - "maxiter": 3000, - "info": True, - "verbose": False, - "recycle": False, - } - # create domain object dom_type = mapping[0] dom_params = mapping[1] @@ -90,7 +83,8 @@ def rho1_xyz(x, y, z): else: if bc_type == "dirichlet": spl_kind = [False, True, True] - dirichlet_bc = [[not kd] * 2 for kd in spl_kind] + dirichlet_bc = [(not kd,) * 2 for kd in spl_kind] + dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): return np.sin(2 * np.pi / Lx * x) @@ -114,7 +108,8 @@ def rho1_xyz(x, y, z): else: if bc_type == "dirichlet": spl_kind = [True, False, True] - dirichlet_bc = [[not kd] * 2 for kd in spl_kind] + dirichlet_bc = [(not kd,) * 2 for kd in spl_kind] + dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): return np.sin(2 * np.pi / Ly * y) @@ -138,7 +133,8 @@ def rho1_xyz(x, y, z): else: if bc_type == "dirichlet": spl_kind = [True, True, False] - dirichlet_bc = [[not kd] * 2 for kd in spl_kind] + dirichlet_bc = [(not kd,) * 2 for kd in spl_kind] + dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): return np.sin(2 * np.pi / Lz * z) @@ -165,17 +161,38 @@ def rho1(e1, e2, e3): rho_vec = L2Projector("H1", mass_ops).get_dofs(rho1, apply_bc=True) # create Poisson solver - _phi = derham.create_spline_function("phi", "H1") - poisson_solver = ImplicitDiffusion( - _phi.vector, sigma_1=1e-12, sigma_2=0.0, sigma_3=1.0, rho=rho_vec, solver=solver_params + solver_params = SolverParameters( + tol=1.0e-13, + maxiter=3000, + info=True, + verbose=False, + recycle=False, + ) + + _phi = FEECVariable(space="H1") + _phi.allocate(derham=derham, domain=domain) + + poisson_solver = ImplicitDiffusion() + poisson_solver.variables.phi = _phi + + poisson_solver.options = poisson_solver.Options( + sigma_1=1e-12, + sigma_2=0.0, + sigma_3=1.0, + rho=rho_vec, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, ) + poisson_solver.allocate() + # Solve Poisson (call propagator with dt=1.) dt = 1.0 poisson_solver(dt) # push numerical solution and compare - sol_val1 = domain.push(_phi, e1, e2, e3, kind="0") + sol_val1 = domain.push(_phi.spline, e1, e2, e3, kind="0") x, y, z = domain(e1, e2, e3) analytic_value1 = sol1_xyz(x, y, z) @@ -246,14 +263,6 @@ def test_poisson_2d(Nel, p, bc_type, mapping, show_plot=False): """ Test the Poisson solver by means of manufactured solutions in 2D . """ - solver_params = { - "type": ("pcg", "MassMatrixPreconditioner"), - "tol": 1.0e-13, - "maxiter": 3000, - "info": True, - "verbose": False, - "recycle": False, - } # create domain object dom_type = mapping[0] @@ -293,7 +302,8 @@ def rho2_xyz(x, y, z): elif bc_type == "dirichlet": spl_kind = [False, True, True] - dirichlet_bc = [[not kd] * 2 for kd in spl_kind] + dirichlet_bc = [(not kd,) * 2 for kd in spl_kind] + dirichlet_bc = tuple(dirichlet_bc) print(f"{dirichlet_bc = }") # manufactured solution in 2D @@ -352,24 +362,68 @@ def rho2(e1, e2, e3): rho_vec2 = l2_proj.get_dofs(rho2, apply_bc=True) # Create Poisson solvers - _phi1 = derham.create_spline_function("test1", "H1") - poisson_solver1 = ImplicitDiffusion( - _phi1.vector, sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, rho=rho_vec1, solver=solver_params + solver_params = SolverParameters( + tol=1.0e-13, + maxiter=3000, + info=True, + verbose=False, + recycle=False, + ) + + _phi1 = FEECVariable(space="H1") + _phi1.allocate(derham=derham, domain=domain) + + poisson_solver1 = ImplicitDiffusion() + poisson_solver1.variables.phi = _phi1 + + poisson_solver1.options = poisson_solver1.Options( + sigma_1=1e-8, + sigma_2=0.0, + sigma_3=1.0, + rho=rho_vec1, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, ) - _phi2 = derham.create_spline_function("test2", "H1") - poisson_solver2 = ImplicitDiffusion( - _phi2.vector, sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, rho=rho_vec2, solver=solver_params + poisson_solver1.allocate() + + # _phi1 = derham.create_spline_function("test1", "H1") + # poisson_solver1 = ImplicitDiffusion( + # _phi1.vector, sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, rho=rho_vec1, solver=solver_params + # ) + + _phi2 = FEECVariable(space="H1") + _phi2.allocate(derham=derham, domain=domain) + + poisson_solver2 = ImplicitDiffusion() + poisson_solver2.variables.phi = _phi2 + + poisson_solver2.options = poisson_solver2.Options( + sigma_1=1e-8, + sigma_2=0.0, + sigma_3=1.0, + rho=rho_vec2, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, ) + poisson_solver2.allocate() + + # _phi2 = derham.create_spline_function("test2", "H1") + # poisson_solver2 = ImplicitDiffusion( + # _phi2.vector, sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, rho=rho_vec2, solver=solver_params + # ) + # Solve Poisson equation (call propagator with dt=1.) dt = 1.0 poisson_solver1(dt) poisson_solver2(dt) # push numerical solutions - sol_val1 = domain.push(_phi1, e1, e2, e3, kind="0") - sol_val2 = domain.push(_phi2, e1, e2, e3, kind="0") + sol_val1 = domain.push(_phi1.spline, e1, e2, e3, kind="0") + sol_val2 = domain.push(_phi2.spline, e1, e2, e3, kind="0") x, y, z = domain(e1, e2, e3) analytic_value1 = sol1_xyz(x, y, z) @@ -414,7 +468,7 @@ def rho2(e1, e2, e3): if __name__ == "__main__": - direction = 0 + direction = 2 bc_type = "dirichlet" mapping = ["Cuboid", {"l1": 0.0, "r1": 4.0, "l2": 0.0, "r2": 2.0, "l3": 0.0, "r3": 3.0}] # mapping = ['Orthogonal', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 3.}] diff --git a/src/struphy/tutorials/tests/test_tutorials.py b/src/struphy/tutorials/tests/test_tutorials.py deleted file mode 100644 index 17cf2069d..000000000 --- a/src/struphy/tutorials/tests/test_tutorials.py +++ /dev/null @@ -1,175 +0,0 @@ -import os - -import pytest -import yaml -from mpi4py import MPI - -import struphy -from struphy.post_processing import pproc_struphy -from struphy.struphy import run - -comm = MPI.COMM_WORLD -rank = comm.Get_rank() - -libpath = struphy.__path__[0] -i_path = os.path.join(libpath, "io", "inp") -o_path = os.path.join(libpath, "io", "out") - - -@pytest.mark.mpi(min_size=2) -def test_tutorial_02(): - run( - "LinearMHDVlasovCC", - os.path.join(i_path, "tutorials", "params_02.yml"), - os.path.join(o_path, "tutorial_02"), - supress_out=True, - ) - - -@pytest.mark.mpi(min_size=2) -def test_tutorial_03(): - run( - "LinearMHD", - os.path.join(i_path, "tutorials", "params_03.yml"), - os.path.join(o_path, "tutorial_03"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_03"), physical=True) - - -@pytest.mark.mpi(min_size=2) -def test_tutorial_04(fast): - run( - "Maxwell", - os.path.join(i_path, "tutorials", "params_04a.yml"), - os.path.join(o_path, "tutorial_04a"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_04a")) - - run( - "LinearMHD", - os.path.join(i_path, "tutorials", "params_04b.yml"), - os.path.join(o_path, "tutorial_04b"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_04b")) - - if not fast: - run( - "VariationalMHD", - os.path.join(i_path, "tutorials", "params_04c.yml"), - os.path.join(o_path, "tutorial_04c"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_04c")) - - -def test_tutorial_05(): - run( - "Vlasov", - os.path.join(i_path, "tutorials", "params_05a.yml"), - os.path.join(o_path, "tutorial_05a"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05a")) - - run( - "Vlasov", - os.path.join(i_path, "tutorials", "params_05b.yml"), - os.path.join(o_path, "tutorial_05b"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05b")) - - run( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_05c.yml"), - os.path.join(o_path, "tutorial_05c"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05c")) - - run( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_05d.yml"), - os.path.join(o_path, "tutorial_05d"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05d")) - - run( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_05e.yml"), - os.path.join(o_path, "tutorial_05e"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05e")) - - run( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_05f.yml"), - os.path.join(o_path, "tutorial_05f"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05f")) - - -def test_tutorial_12(): - run( - "Vlasov", - os.path.join(i_path, "tutorials", "params_12a.yml"), - os.path.join(o_path, "tutorial_12a"), - save_step=100, - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_12a")) - - run( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_12b.yml"), - os.path.join(o_path, "tutorial_12b"), - save_step=10, - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_12b")) - - -if __name__ == "__main__": - test_tutorial_04(True) diff --git a/src/struphy/utils/utils.py b/src/struphy/utils/utils.py index f52fb3f85..aac94fd10 100644 --- a/src/struphy/utils/utils.py +++ b/src/struphy/utils/utils.py @@ -111,15 +111,15 @@ def refresh_models(): list_fluid = [] fluid_string = "" for name, obj in inspect.getmembers(fluid): - if inspect.isclass(obj): - if name not in {"StruphyModel", "Propagator"}: - list_fluid += [name] - fluid_string += '"' + name + '"\n' + if inspect.isclass(obj) and obj.__module__ == fluid.__name__: + # if name not in {"StruphyModel", "Propagator"}: + list_fluid += [name] + fluid_string += '"' + name + '"\n' list_kinetic = [] kinetic_string = "" for name, obj in inspect.getmembers(kinetic): - if inspect.isclass(obj): + if inspect.isclass(obj) and obj.__module__ == kinetic.__name__: if name not in {"StruphyModel", "Propagator"}: list_kinetic += [name] kinetic_string += '"' + name + '"\n' @@ -127,7 +127,7 @@ def refresh_models(): list_hybrid = [] hybrid_string = "" for name, obj in inspect.getmembers(hybrid): - if inspect.isclass(obj): + if inspect.isclass(obj) and obj.__module__ == hybrid.__name__: if name not in {"StruphyModel", "Propagator"}: list_hybrid += [name] hybrid_string += '"' + name + '"\n' @@ -135,7 +135,7 @@ def refresh_models(): list_toy = [] toy_string = "" for name, obj in inspect.getmembers(toy): - if inspect.isclass(obj): + if inspect.isclass(obj) and obj.__module__ == toy.__name__: if name not in {"StruphyModel", "Propagator"}: list_toy += [name] toy_string += '"' + name + '"\n' diff --git a/tutorials/tutorial_02_test_particles.ipynb b/tutorials/tutorial_02_test_particles.ipynb index d2711ffbd..3c4b05b56 100644 --- a/tutorials/tutorial_02_test_particles.ipynb +++ b/tutorials/tutorial_02_test_particles.ipynb @@ -1129,7 +1129,7 @@ "\n", "# initial conditions (background + perturbation)\n", "perturbation = None\n", - "background = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation))\n", + "background = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), equil=equil)\n", "\n", "model.kinetic_ions.var.add_background(background)" ] diff --git a/tutorials/tutorial_05_vlasov_maxwell.ipynb b/tutorials/tutorial_05_vlasov_maxwell.ipynb index 149d60427..fd8bcac5b 100644 --- a/tutorials/tutorial_05_vlasov_maxwell.ipynb +++ b/tutorials/tutorial_05_vlasov_maxwell.ipynb @@ -167,9 +167,12 @@ "metadata": {}, "outputs": [], "source": [ + "background = maxwellians.Maxwellian3D(n=(1.0, None))\n", + "model.kinetic_ions.var.add_background(background)\n", + "\n", "perturbation = perturbations.ModesCos(ls=[1], amps=[0.001])\n", - "background = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", - "model.kinetic_ions.var.add_background(background)" + "init = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n", + "model.kinetic_ions.var.add_initial_condition(init)" ] }, { From af3668b4e39b17ffa291632e7b98632e4ee283f9 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 25 Sep 2025 15:44:34 +0000 Subject: [PATCH 123/292] Resolve "Port toy model PressurelessSPH to 318" --- .gitlab-ci.yml | 39 +- Testcopy.ipynb | 765 ----------- pyproject.toml | 6 +- src/struphy/console/main.py | 13 +- src/struphy/console/params.py | 26 +- src/struphy/console/pproc.py | 53 +- .../examples/restelli2018/callables.py | 157 ++- src/struphy/feec/linear_operators.py | 4 + src/struphy/feec/tests/test_field_init.py | 7 +- src/struphy/fields_background/equils.py | 18 +- src/struphy/initial/perturbations.py | 800 ++++++++++- .../initial/tests/test_init_perturbations.py | 2 +- .../tests/test_maxwellians.py | 8 +- src/struphy/linear_algebra/saddle_point.py | 42 +- .../tests/test_saddlepoint_massmatrices.py | 28 +- src/struphy/main.py | 2 +- src/struphy/models/base.py | 206 +-- src/struphy/models/fluid.py | 17 +- src/struphy/models/kinetic.py | 4 +- src/struphy/models/species.py | 7 +- src/struphy/models/tests/test_models.py | 1 + src/struphy/models/toy.py | 94 +- src/struphy/models/variables.py | 244 +++- .../pic/accumulation/particles_to_grid.py | 31 +- src/struphy/pic/base.py | 1224 ++++++++++++----- src/struphy/pic/particles.py | 56 +- src/struphy/pic/pushing/pusher.py | 16 +- src/struphy/pic/sorting_kernels.py | 135 +- src/struphy/pic/sph_eval_kernels.py | 14 +- src/struphy/pic/tests/test_sorting.py | 55 +- src/struphy/pic/tests/test_sph.py | 862 +++++++++++- src/struphy/pic/utilities.py | 20 + .../likwid/plot_time_traces.py | 431 ++++-- .../post_processing/post_processing_tools.py | 35 +- src/struphy/post_processing/pproc_struphy.py | 12 - src/struphy/profiling/profiling.py | 31 +- src/struphy/propagators/propagators_fields.py | 664 +++++++-- .../propagators/propagators_markers.py | 8 +- 38 files changed, 4200 insertions(+), 1937 deletions(-) delete mode 100644 Testcopy.ipynb diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 580ef9eee..d26ced2dc 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -510,6 +510,23 @@ inspect_repo: ### PUSH PIPELINE ### ##################### +# check_struphy_simulation_params: +# stage: test +# extends: +# - .rules_mr_to_devel +# - .image_gitlab_mpcdf_struphy +# - .before_script_load_modules +# - .variables_push +# script: +# - !reference [.scripts, install_on_push] +# - git clone https://gitlab.mpcdf.mpg.de/struphy/struphy-simulations.git +# - | +# for file in struphy-simulations/**/*.yml; do +# echo "Checking $file" +# # TODO: How to deduce the model? +# struphy params LinearVlasovAmpereOneSpecies --check-file $file +# done + install_tests: stage: test extends: @@ -588,14 +605,19 @@ pages_tests: - !reference [.scripts, inspect_struphy] - module load pandoc - module list - - pip install -e .[doc] + - pip install .[doc] - | struphy compile -y || ( echo "Initial compile failed. Removing compiled kernels and trying again..." && struphy compile -d -y && struphy compile -y ) - - cd doc ; make html; cd .. + - cd doc ; make html + - mv _build/html/ $CI_PROJECT_DIR/documentation/ + artifacts: + expose_as: 'Documentation' + paths: + - documentation/ ubuntu_latest: stage: test @@ -667,6 +689,7 @@ cleanup_macos: ### SCHEDULED PIPELINE ### ########################## + install_scheduled: stage: install needs: [] @@ -941,12 +964,11 @@ lint_full_repo_report: - pip install -e .[dev] # We have to install struphy in editable mode for struphy lint to work - !reference [.scripts, inspect_struphy] - struphy lint all --output-format report - - mkdir -p lint_reports - - mv *.html lint_reports/ 2>/dev/null || echo "No HTML files found." allow_failure: true artifacts: + expose_as: 'Branch lint report' paths: - - lint_reports/ + - code_analysis_report.html expire_in: 1 month # Lint struphy with `struphy lint` @@ -991,12 +1013,11 @@ lint_branch_report: - pip install -e .[dev] # We have to install struphy in editable mode for struphy lint to work - !reference [.scripts, inspect_struphy] - struphy lint branch --output-format report - - mkdir -p lint_reports - - mv *.html lint_reports/ 2>/dev/null || echo "No HTML files found." allow_failure: true artifacts: + expose_as: 'Branch lint report' paths: - - lint_reports/ + - code_analysis_report.html expire_in: 1 month ################### @@ -1018,7 +1039,7 @@ pages: - !reference [.scripts, inspect_struphy] - !reference [.scripts, compile] - cd doc ; make html - - mv _build/html/ ../public/ + - mv _build/html/ $CI_PROJECT_DIR/public/ artifacts: name: 'pages' paths: diff --git a/Testcopy.ipynb b/Testcopy.ipynb deleted file mode 100644 index eefc4d406..000000000 --- a/Testcopy.ipynb +++ /dev/null @@ -1,765 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from mpi4py import MPI" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from struphy.geometry.domains import Cuboid\n", - "\n", - "l1 = -.5\n", - "r1 = .5\n", - "l2 = -.5\n", - "r2 = .5\n", - "l3 = 0.\n", - "r3 = 1.\n", - "domain = Cuboid(l1=l1, r1=r1, l2=l2, r2=r2, l3=l3, r3=r3)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from struphy.pic.particles import ParticlesSPH\n", - "\n", - "# mandatory parameters\n", - "name = 'test'\n", - "Np = 1000\n", - "bc = ['reflect', 'reflect', 'periodic']\n", - "loading = 'pseudo_random'\n", - "\n", - "# optional\n", - "loading_params = {'seed': None}\n", - "params_sorting = {\"nx\": 1, \"ny\": 1, \"nz\": 1, \"eps\": 0.25}\n", - "params_loading = {\"seed\": 1234, \"moments\": \"degenerate\", \"spatial\": \"uniform\"}\n", - "\n", - "bckgr_params = {\n", - " \"type\": \"BeltramiFlow\",\n", - " \"BeltramiFlow\": {},\n", - " \"pforms\": [None, None],\n", - " }\n", - "\n", - "mpi_comm = MPI.COMM_WORLD\n", - "\n", - "# instantiate Particle object\n", - "particles = ParticlesSPH(\n", - " name=name,\n", - " Np=Np,\n", - " bc=bc,\n", - " loading=loading,\n", - " bufsize=10.0, # Lots a buffering needed since only 3*3*3 box\n", - " comm= mpi_comm,\n", - " loading_params=params_loading,\n", - " domain=domain,\n", - " bckgr_params=bckgr_params,\n", - " sorting_params=params_sorting,\n", - " )\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "particles.draw_markers(sort=False)\n", - "particles.mpi_sort_markers()\n", - "particles.initialize_weights()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "particles.positions\n", - "# positions on the physical domain Omega\n", - "pushed_pos = domain(particles.positions).T\n", - "pushed_pos" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from struphy.propagators.propagators_markers import PushEta\n", - "\n", - "# default parameters of Propagator\n", - "opts_eta = PushEta.options(default=False)\n", - "print(opts_eta)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# pass simulation parameters to Propagator class\n", - "PushEta.domain = domain" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# instantiate Propagator object\n", - "prop_eta = PushEta(particles, algo = \"forward_euler\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from struphy.feec.psydac_derham import Derham\n", - "\n", - "Nel = [64, 64, 1] # Number of grid cells\n", - "p = [3, 3, 1] # spline degrees\n", - "spl_kind = [False, False, True] # spline types (clamped vs. periodic)\n", - "\n", - "derham = Derham(Nel, p, spl_kind)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from struphy.propagators.propagators_markers import PushVinEfield\n", - "\n", - "# instantiate Propagator object\n", - "PushVinEfield.domain = domain\n", - "PushVinEfield.derham = derham" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from struphy.fields_background.fluid_equils.equils import BeltramiFlow\n", - "bel_flow = BeltramiFlow()\n", - "p_xyz = bel_flow.p_xyz\n", - "p = lambda eta1, eta2, eta3: domain.pull(p_xyz, eta1, eta2, eta3)\n", - "p_coeffs = derham.P[\"0\"](p)\n", - "p = lambda eta1, eta2, eta3: domain.pull(p_xyz, eta1, eta2, eta3, squeeze_out=True)\n", - "p_coeffs" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "p_h = derham.create_field('pressure', 'H1')\n", - "p_h.vector = p_coeffs" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "\n", - "plt.figure(figsize=(12, 12))\n", - "x = np.linspace(-.5, .5, 100)\n", - "y = np.linspace(-.5, .5, 90)\n", - "xx, yy = np.meshgrid(x, y)\n", - "eta1 = np.linspace(0, 1, 100)\n", - "eta2 = np.linspace(0, 1, 90)\n", - "\n", - "plt.subplot(2, 2, 1)\n", - "plt.pcolor(xx, yy, p_xyz(xx, yy, 0))\n", - "plt.axis('square')\n", - "plt.title('p_xyz')\n", - "plt.colorbar()\n", - "\n", - "plt.subplot(2, 2, 2)\n", - "p_vals = p(eta1, eta2, 0).T\n", - "plt.pcolor(eta1, eta2, p_vals)\n", - "plt.axis('square')\n", - "plt.title('p logical')\n", - "plt.colorbar()\n", - "\n", - "plt.subplot(2, 2, 3)\n", - "p_h_vals = p_h(eta1, eta2, 0, squeeze_out=True).T\n", - "plt.pcolor(eta1, eta2, p_h_vals)\n", - "plt.axis('square')\n", - "plt.title('p_h (logical)')\n", - "plt.colorbar()\n", - "\n", - "plt.subplot(2, 2, 4)\n", - "plt.pcolor(eta1, eta2, np.abs(p_vals - p_h_vals))\n", - "plt.axis('square')\n", - "plt.title('difference')\n", - "plt.colorbar()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "grad_p_ana = lambda x, y: -1.*np.array([np.pi * np.sin(np.pi*x)*np.cos(np.pi*x),\n", - " np.pi * np.sin(np.pi*y)*np.cos(np.pi*y)])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "grad_p = derham.grad.dot(p_coeffs)\n", - "grad_p.update_ghost_regions() # very important, we will move it inside grad\n", - "grad_p *= -1.\n", - "prop_v = PushVinEfield(particles, e_field=grad_p)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "grad_p_h = derham.create_field('pressure gradient', 'Hcurl')\n", - "grad_p_h.vector = grad_p" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "\n", - "plt.figure(figsize=(12, 18))\n", - "\n", - "plt.subplot(3, 2, 1)\n", - "plt.pcolor(xx, yy, grad_p_ana(xx, yy)[0])\n", - "plt.axis('square')\n", - "plt.title('grad_p x')\n", - "plt.colorbar()\n", - "\n", - "plt.subplot(3, 2, 2)\n", - "plt.pcolor(xx, yy, grad_p_ana(xx, yy)[1])\n", - "plt.axis('square')\n", - "plt.title('grad_p y')\n", - "plt.colorbar()\n", - "\n", - "plt.subplot(3, 2, 3)\n", - "grad_p_vals = grad_p_h(eta1, eta2, 0, squeeze_out=True)\n", - "plt.pcolor(eta1, eta2, grad_p_vals[0].T)\n", - "plt.axis('square')\n", - "plt.title('grad_p_h x')\n", - "plt.colorbar()\n", - "\n", - "plt.subplot(3, 2, 4)\n", - "plt.pcolor(eta1, eta2, grad_p_vals[1].T)\n", - "plt.axis('square')\n", - "plt.title('grad_p_h x')\n", - "plt.colorbar()\n", - "\n", - "plt.subplot(3, 2, 5)\n", - "plt.pcolor(eta1, eta2, np.abs(grad_p_vals[0].T - grad_p_ana(xx, yy)[0]))\n", - "plt.axis('square')\n", - "plt.title('diff x')\n", - "plt.colorbar()\n", - "\n", - "plt.subplot(3, 2, 6)\n", - "plt.pcolor(eta1, eta2, np.abs(grad_p_vals[1].T - grad_p_ana(xx, yy)[1]))\n", - "plt.axis('square')\n", - "plt.title('diff y')\n", - "plt.colorbar()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "import math\n", - "import tqdm\n", - "\n", - "# time stepping\n", - "dt = 0.02\n", - "Nt = 200\n", - "\n", - "pos = np.zeros((Nt + 1, Np, 3), dtype=float)\n", - "velo = np.zeros((Nt + 1, Np, 3), dtype=float)\n", - "energy = np.zeros((Nt + 1, Np), dtype=float)\n", - "\n", - "particles.draw_markers(sort=False)\n", - "particles.mpi_sort_markers()\n", - "particles.initialize_weights()\n", - "\n", - "pos[0] = domain(particles.positions).T\n", - "velo[0] = particles.velocities\n", - "energy[0] = .5*(velo[0, : , 0]**2 + velo[0, : , 1]**2) + p_h(particles.positions)\n", - "\n", - "time = 0.\n", - "time_vec = np.zeros(Nt + 1, dtype=float)\n", - "n = 0\n", - "while n < Nt:\n", - " time += dt\n", - " n += 1\n", - " time_vec[n] = time\n", - " \n", - " # advance in time\n", - " prop_eta(dt/2)\n", - " prop_v(dt)\n", - " prop_eta(dt/2)\n", - " \n", - " # positions on the physical domain Omega\n", - " pos[n] = domain(particles.positions).T\n", - " velo[n] = particles.velocities\n", - " \n", - " energy[n] = .5*(velo[n, : , 0]**2 + velo[n, : , 1]**2) + p_h(particles.positions)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# energy plots\n", - "fig = plt.figure(figsize = (13, 6))\n", - "\n", - "plt.subplot(2, 2, 1)\n", - "plt.plot(time_vec, energy[:, 0])\n", - "plt.title('particle 1')\n", - "plt.xlabel('time')\n", - "plt.ylabel('energy')\n", - "\n", - "plt.subplot(2, 2, 2)\n", - "plt.plot(time_vec, energy[:, 1])\n", - "plt.title('particle 2')\n", - "plt.xlabel('time')\n", - "plt.ylabel('energy')\n", - "\n", - "plt.subplot(2, 2, 3)\n", - "plt.plot(time_vec, energy[:, 2])\n", - "plt.title('particle 3')\n", - "plt.xlabel('time')\n", - "plt.ylabel('energy')\n", - "\n", - "plt.subplot(2, 2, 4)\n", - "plt.plot(time_vec, energy[:, 3])\n", - "plt.title('particle 4')\n", - "plt.xlabel('time')\n", - "plt.ylabel('energy')" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "fig = plt.figure()\n", - "ax = fig.add_subplot(projection=\"3d\")\n", - "ax.scatter(pos[-1,:,0],pos[-1,:,1],pos[-1,:,2])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import matplotlib.animation as animation\n", - "n_frame = 200\n", - "fig, ax = plt.subplots()\n", - "\n", - "coloring = np.select([pos[0,:,0]<=-0.2, np.abs(pos[0,:,0]) < +0.2, pos[0,:,0] >= 0.2],\n", - " [-1.0, 0.0, +1.0])\n", - "scat = ax.scatter(pos[0,:,0], pos[0,:,1], c=coloring)\n", - "ax.set_xlim([-0.5,0.5])\n", - "ax.set_ylim([-0.5,0.5])\n", - "ax.set_aspect('equal')\n", - "\n", - "f = lambda x, y: np.cos(np.pi*x)*np.cos(np.pi*y)\n", - "ax.contour(xx, yy, f(xx, yy))\n", - "ax.set_title(f'time = {time_vec[0]:4.2f}')\n", - "\n", - "def update_frame(frame):\n", - " scat.set_offsets(pos[frame,:,:2])\n", - " ax.set_title(f'time = {time_vec[frame]:4.2f}')\n", - " return scat\n", - "\n", - "ani = animation.FuncAnimation(fig=fig, func=update_frame, frames = n_frame)\n", - "ani.save(\"Beltramiaktuell.gif\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import matplotlib.animation as animation\n", - "n_frame = 10\n", - "fig, ax = plt.subplots()\n", - "\n", - "coloring = np.select([pos[0]<=-0.2, np.abs(pos[0]) < +0.2, pos[0] >= 0.2],\n", - " [-1.0, 0.0, +1.0])\n", - "scat = ax.scatter(pos[0], pos[1])\n", - "def update_frame(frame):\n", - " scat.set_offsets(pos[frame])\n", - " return scat\n", - "\n", - "ani = animation.FuncAnimation(fig=fig, func=update_frame, frames = n_frame)\n", - "ani.save(\"Beltramiaktuell2.gif\")\n", - "np.shape(pos[0])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "fig = plt.figure()\n", - "ax = fig.add_subplot(111, aspect='equal')\n", - "\n", - "coloring = np.select([pos[0]<=-0.2, np.abs(pos[0]) < +0.2, pos[0] >= 0.2],\n", - " [-1.0, 0.0, +1.0])\n", - "scatter = ax.scatter(pos[0], pos[1])\n", - "ax.set_xlim([-0.5, +0.5])\n", - "ax.set_ylim([-0.5, +0.5])\n", - "def animate(i):\n", - " pos[i]\n", - " ##scatter.set_offsets(Xi.T)\n", - " \n", - " anim = animation.FuncAnimation(fig, animate, frames=Nt, interval=1,\n", - " repeat_delay=5000) \n", - "plt.show() \n", - "\n", - "##anim.save(\"test.gif\", writer='pillow', fps=60)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np \n", - "import matplotlib.pyplot as plt \n", - "from matplotlib.animation import FuncAnimation\n", - "\n", - "t = np.linspace(0,10,100)\n", - "y = np.sin(t)\n", - "\n", - "fig, axis = plt.subplots()\n", - "axis.set_xlim([min(t), max(t)])\n", - "axis.set_ylim([-2,2])\n", - "animated_plot, = axis.plot([],[])\n", - "\n", - "def update_data(frame):\n", - " animated_plot.set_data(t,y)\n", - " return animated_plot\n", - "\n", - "animation = FuncAnimation(fig=fig, func=update_data, frames = len(t), interval = 0.01)\n", - "\n", - "plt.show()\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "import matplotlib.pyplot as plt\n", - "from matplotlib.animation import FuncAnimation\n", - "\n", - "# Create time and y data\n", - "t = np.linspace(0, 10, 100)\n", - "y = np.sin(t)\n", - "\n", - "# Create a figure and an axis\n", - "fig, axis = plt.subplots()\n", - "axis.set_xlim([min(t), max(t)])\n", - "axis.set_ylim([-2, 2])\n", - "\n", - "# Initialize the plot with empty data\n", - "animated_plot, = axis.plot([], [], lw=2)\n", - "\n", - "# Update function to animate the plot\n", - "def update_data(frame):\n", - " # Set the data for the current frame\n", - " animated_plot.set_data(t[:frame], y[:frame]) # Display data up to the current frame\n", - " return animated_plot, # Return the plot object (blit requires an iterable)\n", - "\n", - "# Create the animation\n", - "animation = FuncAnimation(fig, update_data, frames=len(t), interval=25, blit=True)\n", - "\n", - "# Show the plot\n", - "plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "import matplotlib.pyplot as plt\n", - "import matplotlib.animation as animation\n", - "\n", - "\n", - "#--- Parameters\n", - "N = 1 # number of particles per direction\n", - "dt = 1./50. # time step\n", - "n_steps = 10 ###10000 # max. number of steps\n", - "x_offset = 0.05 # off-set at the boundary\n", - "\n", - "\n", - "#--- Velocity field\n", - "u = lambda x: np.array([-np.cos(np.pi*x[0])*np.sin(np.pi*x[1]),\n", - " np.sin(np.pi*x[0])*np.cos(np.pi*x[1])])\n", - "\n", - "#--- Pressure field\n", - "p = lambda x: 0.5*(np.sin(np.pi*x[0])**2 + np.sin(np.pi*x[1])**2)\n", - "\n", - "#--- Analytical pressure gradient\n", - "grad_p = lambda x: np.array([np.pi * np.sin(np.pi*x[0])*np.cos(np.pi*x[0]),\n", - " np.pi * np.sin(np.pi*x[1])*np.cos(np.pi*x[1])])\n", - "\n", - "#--- Perturbed analytical pressure gradient\n", - "# grad_p = lambda x: np.array([np.pi * np.sin(np.pi*x[0])*np.cos(np.pi*x[0]) + 0.1*np.random.random(),\n", - "# np.pi * np.sin(np.pi*x[1])*np.cos(np.pi*x[1]) + 0.1*np.random.random()]) \n", - "\n", - "#--- Invariant and Hamiltonian\n", - "f = lambda x: np.cos(np.pi*x[0])*np.cos(np.pi*x[1])\n", - "k = lambda v: 0.5*(v[0]**2 + v[1]**2)\n", - "h = lambda x, v: 0.5*(v[0]**2 + v[1]**2) + p(x)\n", - "\n", - "#--- Update function for the Lagrangian trajectory\n", - "# Hamiltonian system where p acts as a potential energy\n", - "# Symplectic Euler scheme\n", - "def update(x, v, dt):\n", - " \"\"\"\n", - " Symplectic Euler for separable Hamiltonians.\n", - " \"\"\"\n", - "\n", - " # Classical symplectic Euler\n", - " # (equivalent to Hamiltonian Lie splitting)\n", - " # x += dt * v\n", - " # v -= dt * grad_p(x)\n", - "\n", - " # Hamiltonian Strang splitting\n", - " # x += 0.5 * dt * v\n", - " print(x + .5)\n", - " print(dt, -grad_p(x))\n", - " v -= dt * grad_p(x)\n", - " # x += 0.5 * dt * v \n", - " \n", - " return (x,v)\n", - "\n", - "#--- Initialization of the particles on a uniform grid\n", - "# dx = x_offset\n", - "# s = np.linspace(-0.5+dx, +0.5-dx, N)\n", - "# XX, YY = np.meshgrid(s, s)\n", - "# X0 = XX.flatten()\n", - "# X1 = YY.flatten()\n", - "# X = np.row_stack((X0,X1))\n", - "\n", - "#--- Random initialization of particles\n", - "np.random.seed(1234)\n", - "X = -0.5 + np.random.random((2,N*N))\n", - "\n", - "print(X)\n", - "\n", - "#--- Initialization of velocities\n", - "V = u(X)\n", - "\n", - "print(V)\n", - "\n", - "Lagr_trajectories = np.empty(X.shape+(n_steps,))\n", - "Lagr_velocities = np.empty(V.shape+(n_steps,))\n", - "Lagr_trajectories[...,0] = X\n", - "Lagr_velocities[...,0] = V\n", - "for i in range(1,n_steps):\n", - " X, V = update(X, V, dt)\n", - " Lagr_trajectories[...,i] = X\n", - " Lagr_velocities[...,i] = V\n", - "\n", - "print('Computation completed. Plotting...')\n", - " \n", - "#--- Finer grid for the contours of the invariant\n", - "Sg = np.linspace(-0.5, +0.5, 200)\n", - "Xg, Yg = np.meshgrid(Sg, Sg)\n", - "\n", - "#--- Animation\n", - "fig = plt.figure()\n", - "ax = fig.add_subplot(111, aspect='equal')\n", - "ax.contour(Xg, Yg, f([Xg,Yg]))\n", - "Xi = Lagr_trajectories[...,0]\n", - "coloring = np.select([Xi[0]<=-0.2, np.abs(Xi[0]) < +0.2, Xi[0] >= 0.2],\n", - " [-1.0, 0.0, +1.0])\n", - "scatter = ax.scatter(Xi[0], Xi[1], c=coloring, marker='.', cmap='rainbow')\n", - "ax.set_xlim([-0.5, +0.5])\n", - "ax.set_ylim([-0.5, +0.5])\n", - "def animate(i):\n", - " Xi = Lagr_trajectories[...,i]\n", - " scatter.set_offsets(Xi.T)\n", - " \n", - "anim = animation.FuncAnimation(fig, animate, frames=n_steps, interval=1,\n", - " repeat_delay=5000) \n", - "\n", - "#--- Hamiltonian, invariant, and calculation of the total energies\n", - "figh = plt.figure()\n", - "axh1 = figh.add_subplot(121)\n", - "axh2 = figh.add_subplot(122)\n", - "tot_Ham = np.zeros_like(Lagr_trajectories[0,0,:])\n", - "tot_Kin = np.zeros_like(Lagr_trajectories[0,0,:])\n", - "tot_Ein = np.zeros_like(Lagr_trajectories[0,0,:])\n", - "for ip in range(Lagr_trajectories.shape[1]):\n", - " particle_h = h(Lagr_trajectories[:,ip,:], Lagr_velocities[:,ip,:])\n", - " particle_k = k(Lagr_velocities[:,ip,:])\n", - " particle_p = p(Lagr_trajectories[:,ip,:])\n", - " particle_f = f(Lagr_trajectories[:,ip,:])\n", - " tot_Ham += particle_h\n", - " tot_Kin += particle_k\n", - " tot_Ein += particle_p\n", - " if ip < 100: # only for the first 100 particles\n", - " axh1.plot(particle_h[:] - particle_h[0])\n", - " axh2.plot(particle_f[:] - particle_f[0])\n", - "\n", - "#--- Normalization to the number of particles \n", - "tot_Ham /= N*N\n", - "tot_Kin /= N*N\n", - "tot_Ein /= N*N\n", - "\n", - "print('... h and total h computed and plotted ...')\n", - "print('Initial kinetic energy = {}'.format(tot_Kin[0]))\n", - "print('Initial potential energy = {}'.format(tot_Ein[0]))\n", - "print('Initial Hamiltonian = {}'.format(tot_Ham[0]))\n", - "\n", - "#--- Total Hamiltonian conservation error\n", - "figth = plt.figure()\n", - "axth1 = figth.add_subplot(121)\n", - "axth1.plot(tot_Ham[:] - tot_Ham[0])\n", - "axth2 = figth.add_subplot(122)\n", - "axth2.plot(tot_Kin[:] - tot_Kin[0])\n", - "axth2.plot(tot_Ein[:] - tot_Ein[0])\n", - "\n", - "\n", - "plt.show() \n", - "\n", - "###anim.save(\"test.gif\", writer='pillow', fps=60)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "X = -0.5 + np.random.random((2,N*N))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "N= 2" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "X = np.random.random((2,3))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "X\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "X.shape" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "Lagr_trajectories = np.empty(X.shape+(n_steps,))" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "env", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.12" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/pyproject.toml b/pyproject.toml index e3c087296..e06b57a0b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,9 @@ dev = [ "nbstripout", "tabulate", "tomli-w", + "kaleido", + "pandas", + "plotly", ] doc = [ "struphy[phys]", @@ -75,10 +78,7 @@ doc = [ ] likwid = [ "struphy[dev]", - "kaleido", "pylikwid", - "pandas", - "plotly", ] all = [ "struphy[phys]", diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index b12fc9ad8..ee3f46c41 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -706,6 +706,13 @@ def add_parser_params(subparsers, list_models, model_message): help="Absolute path to the parameter file (default is getcwd()/params_MODEL.py)", ) + parser_params.add_argument( + "--check-file", + type=str, + metavar="FILE", + help="check if the parameters in the .yml file are valid", + ) + parser_params.add_argument( "-y", "--yes", @@ -898,8 +905,10 @@ def add_parser_pproc(subparsers, out_folders): parser_pproc.add_argument( "--time-trace", - help="whether to plot the time traces", - action="store_true", + help="List of regions to include in time trace plot\n(options: propagators, kernels, any profiling region name).", + type=str, + nargs="+", + default=[], ) diff --git a/src/struphy/console/params.py b/src/struphy/console/params.py index 1648c7383..bbe33ace3 100644 --- a/src/struphy/console/params.py +++ b/src/struphy/console/params.py @@ -1,8 +1,13 @@ +import sys + +import yaml +from mpi4py import MPI + from struphy.models import fluid, hybrid, kinetic, toy from struphy.models.base import StruphyModel -def struphy_params(model_name: str, params_path: str, yes: bool = False): +def struphy_params(model_name: str, params_path: str, yes: bool = False, check_file: bool = False): """Create a model's default parameter file and save in current input path. Parameters @@ -24,5 +29,20 @@ def struphy_params(model_name: str, params_path: str, yes: bool = False): except AttributeError: pass - prompt = not yes - model.generate_default_parameter_file(path=params_path, prompt=prompt) + # print units + if check_file: + print(f"Checking {check_file} with model {model_class}") + with open(check_file) as file: + params = yaml.load(file, Loader=yaml.FullLoader) + # TODO: Enable running struphy without any communicators + comm = MPI.COMM_WORLD + try: + model = model_class(params=params, comm=MPI.COMM_WORLD) + print("Model initialized successfully.") + except Exception as e: + print(f"Failed to initialize model: {e}") + sys.exit(1) + + else: + prompt = not yes + model.generate_default_parameter_file(path=params_path, prompt=prompt) diff --git a/src/struphy/console/pproc.py b/src/struphy/console/pproc.py index 4d80082a8..5ae99a733 100644 --- a/src/struphy/console/pproc.py +++ b/src/struphy/console/pproc.py @@ -1,3 +1,14 @@ +import inspect + +import struphy.post_processing.pproc_struphy as pproc_struphy +import struphy.propagators.propagators_coupling as propagators_coupling +import struphy.propagators.propagators_fields as propagators_fields +import struphy.propagators.propagators_markers as propagators_markers +from struphy.post_processing.likwid.plot_time_traces import ( + plot_gantt_chart, + plot_gantt_chart_plotly, + plot_time_vs_duration, +) from struphy.utils.utils import subp_run @@ -10,7 +21,7 @@ def struphy_pproc( guiding_center=False, classify=False, no_vtk=False, - time_trace=False, + time_trace=[], ): """Post process data from finished Struphy runs. @@ -53,13 +64,15 @@ def struphy_pproc( dirs = [dir_abs] use_state_o_path = False + absolute_paths = [] for dir in dirs: # create absolute path if use_state_o_path: - path_to_simulation = os.path.join(o_path, dir) + absolute_paths.append(os.path.join(o_path, dir)) else: - path_to_simulation = dir + absolute_paths.append(dir) + for path_to_simulation in absolute_paths: print(f"Post processing data in {path_to_simulation}") command = [ @@ -85,7 +98,35 @@ def struphy_pproc( if no_vtk: command += ["--no-vtk"] - if time_trace: - command += ["--time-trace"] - subp_run(command) + + if len(time_trace) > 0: + print(f"Plotting time trace for the following regions: {', '.join(time_trace)}") + for path in absolute_paths: + path_time_trace = os.path.join(path, "profiling_time_trace.pkl") + if not os.path.isfile(path_time_trace): + raise FileNotFoundError(f"No profiling time trace found at {path_time_trace}") + + # plot_time_vs_duration(path_time_trace, output_path=path_pproc) + # plot_gantt_chart(path_time_trace, output_path=path_pproc) + + propagators = [] + for module in [propagators_coupling, propagators_markers, propagators_fields]: + propagators += [ + name + for name, obj in inspect.getmembers(module, inspect.isclass) + if obj.__module__ == module.__name__ + ] + groups_include = time_trace + + if "kernels" in groups_include: + groups_include += ["kernel:*"] + if "propagators" in groups_include: + groups_include += propagators + + plot_gantt_chart_plotly( + path_time_trace, + output_path=path, + groups_include=groups_include, + show=False, + ) diff --git a/src/struphy/examples/restelli2018/callables.py b/src/struphy/examples/restelli2018/callables.py index adf76d09f..89ce4bfda 100644 --- a/src/struphy/examples/restelli2018/callables.py +++ b/src/struphy/examples/restelli2018/callables.py @@ -43,7 +43,7 @@ class RestelliForcingTerm: in plasma physics, Journal of Computational Physics 2018. """ - def __init__(self, nu=1.0, R0=2.0, a=1.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0): + def __init__(self, nu=1.0, R0=2.0, a=1.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0, eps=1.0): r""" Parameters ---------- @@ -60,6 +60,8 @@ def __init__(self, nu=1.0, R0=2.0, a=1.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0) alpha: 0.1 beta: 1. + + eps: 1. # Normalization """ self._nu = nu @@ -69,11 +71,12 @@ def __init__(self, nu=1.0, R0=2.0, a=1.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0) self._Bp = Bp self._alpha = alpha self._beta = beta + self._eps_norm = eps def __call__(self, x, y, z): R = np.sqrt(x**2 + y**2) R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(y, x) + phi = np.arctan2(-y, x) force_Z = self._nu * ( self._alpha * (self._R0 - 4 * R) / (self._a * self._R0 * R) - self._beta * self._Bp * self._R0**2 / (self._B0 * self._a * R**3) @@ -95,22 +98,53 @@ class ManufacturedSolutionForceterm: .. math:: - f = \left[\begin{array}{c} 2 \pi cos(2 \pi x) + \nu 4 \pi^2 sin(2\pi x) + \frac{sin(2 \pi x) + 1.0}{dt} \\ B_0 (sin(2 \pi x) + 1.9) \\ 0 \end{array} \right] \,, + f = \left[\begin{array}{c} 2 \pi cos(2 \pi x) + \nu 4 \pi^2 sin(2\pi x) + \frac{sin(2 \pi x) + 1.0}{dt} \\ \frac{B_0}{\epsilon} (sin(2 \pi x) + 1.9) \\ 0 \end{array} \right] \,, \\[2mm] - f_e = \left[\begin{array}{c} -2 \pi cos(2 \pi x) + \nu_e 4 \pi^2 sin(2\pi x) - \epsilon sin(2 \pi x) \\ -B_0 sin(2 \pi x) \\ 0 \end{array} \right] \,. + f_e = \left[\begin{array}{c} -2 \pi cos(2 \pi x) + \nu_e 4 \pi^2 sin(2\pi x) - \sigma sin(2 \pi x) \\ -\frac{B_0}{\epsilon} sin(2 \pi x) \\ 0 \end{array} \right] \,. In 2D it is defined as follows: .. math:: - f = \left[\begin{array}{c} -2\pi sin(2\pi x) + B_0 cos(2\pi x)cos(2\pi y) - \nu 8 \pi^2 sin(2\pi x)sin(2\pi y) \\ 2\pi cos(2\pi y) - B_0 sin(2\pi x)sin(2\pi y) - \nu 8 \pi^2 cos(2\pi x)cos(2\pi y) \\ 0 \end{array} \right] \,, + f = \left[\begin{array}{c} -2\pi sin(2\pi x) + \frac{B_0}{\epsilon} cos(2\pi x)cos(2\pi y) - \nu 8 \pi^2 sin(2\pi x)sin(2\pi y) \\ 2\pi cos(2\pi y) - \frac{B_0}{\epsilon} sin(2\pi x)sin(2\pi y) - \nu 8 \pi^2 cos(2\pi x)cos(2\pi y) \\ 0 \end{array} \right] \,, + \\[2mm] + f_e = \left[\begin{array}{c} 2\pi sin(2\pi x) -\frac{B_0}{\epsilon} cos(4\pi x)cos(4\pi y) - \nu_e 32 \pi^2 sin(4\pi x)sin(4\pi y) + \sigma sin(4\pi x) sin(4 \pi y) \\ -2\pi cos(2\pi y) +\frac{B_0}{\epsilon} sin(4\pi x)sin(4\pi y) - \nu_e 32 \pi^2 cos(4\pi x)cos(4\pi y) + \sigma cos(4\pi x) cos(4 \pi y) \\ 0 \end{array} \right] \,. + + In Tokamak geometry it is defined as follows: + + .. math:: + + f = \left[\begin{array}{c} \alpha \frac{B_0}{a}(R-R_0) - \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu \alpha \frac{1}{a R_0} \frac{R_0}{R^2} \\ + \alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} + \alpha \frac{B_0Z}{a} \\ + \alpha \frac{1}{a R_0} \frac{R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right) \end{array} \right] \,, + \\[2mm] + f = \left[\begin{array}{c} -\alpha \frac{B_0}{a}(R-R_0) + \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu_e \alpha \frac{1}{a R_0} \frac{R_0}{R^2} \\ + -\alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} - \alpha \frac{B_0 Z}{a} \\ + -\alpha \frac{1}{a R_0} \frac{ R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right) \end{array} \right] \,, \\[2mm] - f_e = \left[\begin{array}{c} 2\pi sin(2\pi x) -B_0 cos(4\pi x)cos(4\pi y) - \nu_e 32 \pi^2 sin(4\pi x)sin(4\pi y) + \epsilon sin(4\pi x) sin(4 \pi y) \\ -2\pi cos(2\pi y) +B_0 sin(4\pi x)sin(4\pi y) - \nu_e 32 \pi^2 cos(4\pi x)cos(4\pi y) + \epsilon cos(4\pi x) cos(4 \pi y) \\ 0 \end{array} \right] \,. + R = \sqrt{x^2 + y^2} \,. + Can only be defined in Cartesian coordinates. """ - def __init__(self, species, comp, dimension, epsilon, dt, b0=1.0, nu=1.0, nu_e=0.01): + def __init__( + self, + species, + comp, + dimension, + stab_sigma, + eps, + dt, + b0=1.0, + nu=1.0, + nu_e=0.01, + R0=2.0, + a=1.0, + Bp=12.5, + alpha=0.1, + beta=1.0, + ): """ Parameters ---------- @@ -120,57 +154,89 @@ def __init__(self, species, comp, dimension, epsilon, dt, b0=1.0, nu=1.0, nu_e=0 Which component of the solution ('0', '1' or '2'). dimension: string Defines the manufactured solution to be selected ('1D' or '2D'). - epsilon : float + stab_sigma : float Stabilization parameter. + eps : float + Normalization parameter. dt : float Time step. b0 : float Magnetic field (default: 1.0). nu : float - Viscosity of ions (default: 1.0) + Viscosity of ions (default: 1.0). nu_e : float - Viscosity of electrons (default: 0.01) + Viscosity of electrons (default: 0.01). + R0 : float + Major radius of torus (default: 2.). + a : float + Minor radius of torus (default: 1.). + Bp : float + Poloidal magnetic field (default: 12.5). + alpha : float + (default: 0.1) + beta : float + (default: 1.0) """ - self._b = b0 + self._B0 = b0 self._nu = nu self._nu_e = nu_e self._comp = comp self._species = species self._dimension = dimension - self._epsilon = epsilon + self._eps_norm = eps + self._stab_sigma = stab_sigma self._dt = dt + self._R0 = R0 + self._Bp = Bp + self._alpha = alpha + self._beta = beta + self._a = a # equilibrium ion velocity def __call__(self, x, y, z): + A = self._alpha / (self._a * self._R0) + C = self._beta * self._Bp * self._R0 / (self._B0 * self._a) + R = np.sqrt(x**2 + y**2) + R = np.where(R == 0.0, 1e-9, R) + phi = np.arctan2(-y, x) if self._species == "Ions": """Forceterm for ions on the right hand side.""" - """x component""" if self._dimension == "2D": fx = ( -2.0 * np.pi * np.sin(2 * np.pi * x) - + np.cos(2 * np.pi * x) * np.cos(2 * np.pi * y) * self._b + + np.cos(2 * np.pi * x) * np.cos(2 * np.pi * y) * self._B0 / self._eps_norm - self._nu * 8.0 * np.pi**2 * np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y) ) + fy = ( + 2.0 * np.pi * np.cos(2 * np.pi * y) + - np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y) * self._B0 / self._eps_norm + - self._nu * 8.0 * np.pi**2 * np.cos(2 * np.pi * x) * np.cos(2 * np.pi * y) + ) + fz = 0.0 * x + elif self._dimension == "1D": fx = ( 2.0 * np.pi * np.cos(2 * np.pi * x) + self._nu * 4.0 * np.pi**2 * np.sin(2 * np.pi * x) + (np.sin(2 * np.pi * x) + 1.0) / self._dt ) - - """y component""" - if self._dimension == "2D": - fy = ( - 2.0 * np.pi * np.cos(2 * np.pi * y) - - np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y) * self._b - - self._nu * 8.0 * np.pi**2 * np.cos(2 * np.pi * x) * np.cos(2 * np.pi * y) + fy = (np.sin(2 * np.pi * x) + 1.0) * self._B0 / self._eps_norm + fz = 0.0 * x + + elif self._dimension == "Tokamak": + # Covariant basis with transfo DF u withoud phi dependency + fR = ( + self._alpha * self._B0 / self._a * (R - self._R0) + - A * self._R0 / R * (z * self._B0) + + self._nu * A / R**2 * self._R0 ) - elif self._dimension == "1D": - fy = (np.sin(2 * np.pi * x) + 1.0) * self._b + fZ = self._alpha * self._B0 * z / self._a + A * self._R0 / R * ((R - self._R0) * self._B0) + fphi = A * self._R0 * self._Bp / (self._a * R**2) * ((R - self._R0) ** 2 + z**2) - """z component""" - fz = 0.0 * x + fx = np.cos(phi) * fR - R * np.sin(phi) * fphi + fy = -np.sin(phi) * fR - R * np.cos(phi) * fphi + fz = fZ if self._comp == "0": return fx @@ -183,34 +249,43 @@ def __call__(self, x, y, z): elif self._species == "Electrons": """Forceterm for electrons on the right hand side.""" - """x component""" if self._dimension == "2D": fx = ( 2.0 * np.pi * np.sin(2 * np.pi * x) - - np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y) * self._b + - np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y) * self._B0 / self._eps_norm - self._nu_e * 32.0 * np.pi**2 * np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y) - - self._epsilon * (-np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y)) + - self._stab_sigma * (-np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y)) ) + fy = ( + -2.0 * np.pi * np.cos(2 * np.pi * y) + + np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y) * self._B0 / self._eps_norm + - self._nu_e * 32.0 * np.pi**2 * np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y) + - self._stab_sigma * (-np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y)) + ) + fz = 0.0 * x + elif self._dimension == "1D": fx = ( -2.0 * np.pi * np.cos(2 * np.pi * x) + self._nu_e * 4.0 * np.pi**2 * np.sin(2 * np.pi * x) - - self._epsilon * np.sin(2 * np.pi * x) + - self._stab_sigma * np.sin(2 * np.pi * x) ) - - """y component""" - if self._dimension == "2D": - fy = ( - -2.0 * np.pi * np.cos(2 * np.pi * y) - + np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y) * self._b - - self._nu_e * 32.0 * np.pi**2 * np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y) - - self._epsilon * (-np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y)) + fy = -np.sin(2 * np.pi * x) * self._B0 / self._eps_norm + fz = 0.0 * x + + elif self._dimension == "Tokamak": + # Covariant basis with transfo DF u Solution without phi dependency + fR = ( + -self._alpha * self._B0 / self._a * (R - self._R0) + + A * self._R0 / R * (z * self._B0) + + self._nu_e * A * self._R0 / R**2 ) - elif self._dimension == "1D": - fy = -np.sin(2 * np.pi * x) * self._b + fZ = -self._alpha * self._B0 * z / self._a - A * self._R0 / R * ((R - self._R0) * self._B0) + fphi = -A * self._R0 * self._Bp / (self._a * R**2) * ((R - self._R0) ** 2 + z**2) - """z component""" - fz = 0.0 * x + fx = np.cos(phi) * fR - R * np.sin(phi) * fphi + fy = -np.sin(phi) * fR - R * np.cos(phi) * fphi + fz = fZ if self._comp == "0": return fx diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 64191d38d..0fbcee763 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -52,8 +52,12 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): if isinstance(self.domain, BlockVectorSpace): comm = self.domain.spaces[0].cart.comm + if comm is None: + comm = MPI.COMM_SELF elif isinstance(self.domain, StencilVectorSpace): comm = self.domain.cart.comm + if comm is None: + comm = MPI.COMM_SELF rank = comm.Get_rank() size = comm.Get_size() diff --git a/src/struphy/feec/tests/test_field_init.py b/src/struphy/feec/tests/test_field_init.py index 0e8b38e48..a6d5ca815 100644 --- a/src/struphy/feec/tests/test_field_init.py +++ b/src/struphy/feec/tests/test_field_init.py @@ -123,7 +123,12 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show elif "EQDSKequilibrium" in key: mhd_equil.domain = domains.Tokamak(equilibrium=mhd_equil) elif "CircularTokamak" in key: - mhd_equil.domain = domains.Tokamak(equilibrium=mhd_equil) + mhd_equil.domain = domains.HollowTorus( + a1=1e-3, + a2=mhd_equil.params["a"], + R0=mhd_equil.params["R0"], + tor_period=1, + ) elif "HomogenSlab" in key: mhd_equil.domain = domains.Cuboid() elif "ShearedSlab" in key: diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index af11467fe..96febc74b 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -2900,12 +2900,12 @@ class ConstantVelocity(CartesianFluidEquilibrium): def __init__( self, - ux: float = 1.0, - uy: float = 1.0, - uz: float = 1.0, + ux: float = 0.0, + uy: float = 0.0, + uz: float = 0.0, n: float = 1.0, n1: float = 0.0, - density_profile: str = "affine", + density_profile: str = "constant", p0: float = 1.0, ): # use params setter @@ -2938,7 +2938,11 @@ def n_xyz(self, x, y, z): return self.params["n"] * np.exp(-(x**2 + y**2) / self.params["p0"]) elif self.params["density_profile"] == "step_function_x": out = 1e-8 + 0 * x - out[x < 0.0] = self.params["n"] + # mask_x = np.logical_and(x < .6, x > .4) + # mask_y = np.logical_and(y < .6, y > .4) + # mask = np.logical_and(mask_x, mask_y) + mask = x < -2.0 + out[mask] = self.params["n"] return out @@ -3154,7 +3158,7 @@ def psi(self, R, Z, dR=0, dZ=0): "Only combinations (dR=0, dZ=0), (dR=1, dZ=0), (dR=0, dZ=1), (dR=2, dZ=0), (dR=0, dZ=2) and (dR=1, dZ=1) possible!", ) - return out + return -out def g_tor(self, R, Z, dR=0, dZ=0): """Toroidal field function g = g(R, Z).""" @@ -3170,7 +3174,7 @@ def g_tor(self, R, Z, dR=0, dZ=0): "Only combinations (dR=0, dZ=0), (dR=1, dZ=0) and (dR=0, dZ=1) possible!", ) - return out + return -out def p_xyz(self, x, y, z): """Pressure p = p(x, y, z).""" diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 1ff365c8b..f37e38584 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -423,6 +423,431 @@ def __call__(self, eta1, eta2, eta3): return val +class ModesCosCos(Perturbation): + r""" + + .. math:: + + u(x, y, z) = \sum_s A_s \, \chi_s(z) + \cos \!\left(l_s \tfrac{2\pi}{L_x} x + \theta_{x,s}\right) + \cos \!\left(m_s \tfrac{2\pi}{L_y} y + \theta_{y,s}\right) + + where :math:`\chi_s(z)` can be either 1 or localized in z. + + """ + + def __init__( + self, + ls=None, + ms=None, + amps=(1e-4,), + theta_x=None, + theta_y=None, + pfuns=("Id",), + pfuns_params=(0.0,), + Lx=1.0, + Ly=1.0, + Lz=1.0, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): + if ls is not None: + n_modes = len(ls) + elif ms is not None: + n_modes = len(ms) + ls = [0] * n_modes + else: + n_modes = 1 + ls = [0] + ms = [0] + + if ms is None: + ms = [0] * n_modes + else: + assert len(ms) == n_modes + + if len(amps) == 1: + amps = [amps[0]] * n_modes + else: + assert len(amps) == n_modes + + if theta_x is None: + theta_x = [0] * n_modes + elif len(theta_x) == 1: + theta_x = [theta_x[0]] * n_modes + else: + assert len(theta_x) == n_modes + + if theta_y is None: + theta_y = [0] * n_modes + elif len(theta_y) == 1: + theta_y = [theta_y[0]] * n_modes + else: + assert len(theta_y) == n_modes + + if len(pfuns) == 1: + pfuns = [pfuns[0]] * n_modes + else: + assert len(pfuns) == n_modes + + if len(pfuns_params) == 1: + pfuns_params = [pfuns_params[0]] * n_modes + else: + assert len(pfuns_params) == n_modes + + self._ls = ls + self._ms = ms + self._amps = amps + self._theta_x = theta_x + self._theta_y = theta_y + self._Lx = Lx + self._Ly = Ly + self._Lz = Lz + + self._pfuns = [] + for pfun, params in zip(pfuns, pfuns_params): + if pfun == "Id": + self._pfuns += [lambda z: 1.0] + elif pfun == "localize": + self._pfuns += [lambda z, p=params: np.tanh((z - 0.5) / p) / np.cosh((z - 0.5) / p)] + else: + raise ValueError(f"Profile function {pfun} is not defined..") + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp + + def __call__(self, x, y, z): + val = 0.0 + for amp, l, m, thx, thy, pfun in zip(self._amps, self._ls, self._ms, self._theta_x, self._theta_y, self._pfuns): + val += ( + amp + * pfun(z) + * np.cos(l * 2.0 * np.pi / self._Lx * x + thx) + * np.cos(m * 2.0 * np.pi / self._Ly * y + thy) + ) + return val + + +class ModesSinSin(Perturbation): + r""" + + .. math:: + + u(x, y, z) = \sum_s A_s \, \chi_s(z) + \sin \!\left(l_s \tfrac{2\pi}{L_x} x + \theta_{x,s}\right) + \sin \!\left(m_s \tfrac{2\pi}{L_y} y + \theta_{y,s}\right) + + where :math:`\chi_s(z)` can be either 1 or localized in z. + """ + + def __init__( + self, + ls=None, + ms=None, + amps=(1e-4,), + theta_x=None, + theta_y=None, + pfuns=("Id",), + pfuns_params=(0.0,), + Lx=1.0, + Ly=1.0, + Lz=1.0, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): + if ls is not None: + n_modes = len(ls) + elif ms is not None: + n_modes = len(ms) + ls = [0] * n_modes + else: + n_modes = 1 + ls = [0] + ms = [0] + + if ms is None: + ms = [0] * n_modes + else: + assert len(ms) == n_modes + + if len(amps) == 1: + amps = [amps[0]] * n_modes + else: + assert len(amps) == n_modes + + if theta_x is None: + theta_x = [0] * n_modes + elif len(theta_x) == 1: + theta_x = [theta_x[0]] * n_modes + else: + assert len(theta_x) == n_modes + + if theta_y is None: + theta_y = [0] * n_modes + elif len(theta_y) == 1: + theta_y = [theta_y[0]] * n_modes + else: + assert len(theta_y) == n_modes + + if len(pfuns) == 1: + pfuns = [pfuns[0]] * n_modes + else: + assert len(pfuns) == n_modes + + if len(pfuns_params) == 1: + pfuns_params = [pfuns_params[0]] * n_modes + else: + assert len(pfuns_params) == n_modes + + self._ls = ls + self._ms = ms + self._amps = amps + self._theta_x = theta_x + self._theta_y = theta_y + self._Lx = Lx + self._Ly = Ly + self._Lz = Lz + + self._pfuns = [] + for pfun, params in zip(pfuns, pfuns_params): + if pfun == "Id": + self._pfuns += [lambda z: 1.0] + elif pfun == "localize": + self._pfuns += [lambda z, p=params: np.tanh((z - 0.5) / p) / np.cosh((z - 0.5) / p)] + else: + raise ValueError(f"Profile function {pfun} is not defined..") + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp + + def __call__(self, x, y, z): + val = 0.0 + for amp, l, m, thx, thy, pfun in zip(self._amps, self._ls, self._ms, self._theta_x, self._theta_y, self._pfuns): + val += ( + amp + * pfun(z) + * np.sin(l * 2.0 * np.pi / self._Lx * x + thx) + * np.sin(m * 2.0 * np.pi / self._Ly * y + thy) + ) + return val + + +class ModesSinCos(Perturbation): + r""" + + .. math:: + + u(x, y, z) = \sum_s A_s \, \chi_s(z) + \sin \!\left(l_s \tfrac{2\pi}{L_x} x + \theta_{x,s}\right) + \cos \!\left(m_s \tfrac{2\pi}{L_y} y + \theta_{y,s}\right) + + where :math:`\chi_s(z)` can be either 1 or localized in z. + """ + + def __init__( + self, + ls=None, + ms=None, + amps=(1e-4,), + theta_x=None, + theta_y=None, + pfuns=("Id",), + pfuns_params=(0.0,), + Lx=1.0, + Ly=1.0, + Lz=1.0, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): + # number of modes + if ls is not None: + n_modes = len(ls) + elif ms is not None: + n_modes = len(ms) + ls = [0] * n_modes + else: + n_modes = 1 + ls = [0] + ms = [0] + + if ms is None: + ms = [0] * n_modes + else: + assert len(ms) == n_modes + + if len(amps) == 1: + amps = [amps[0]] * n_modes + else: + assert len(amps) == n_modes + + if theta_x is None: + theta_x = [0] * n_modes + elif len(theta_x) == 1: + theta_x = [theta_x[0]] * n_modes + else: + assert len(theta_x) == n_modes + + if theta_y is None: + theta_y = [0] * n_modes + elif len(theta_y) == 1: + theta_y = [theta_y[0]] * n_modes + else: + assert len(theta_y) == n_modes + + if len(pfuns) == 1: + pfuns = [pfuns[0]] * n_modes + else: + assert len(pfuns) == n_modes + + if len(pfuns_params) == 1: + pfuns_params = [pfuns_params[0]] * n_modes + else: + assert len(pfuns_params) == n_modes + + # store + self._ls = ls + self._ms = ms + self._amps = amps + self._theta_x = theta_x + self._theta_y = theta_y + self._Lx = Lx + self._Ly = Ly + self._Lz = Lz + + self._pfuns = [] + for pfun, params in zip(pfuns, pfuns_params): + if pfun == "Id": + self._pfuns += [lambda z: 1.0] + elif pfun == "localize": + self._pfuns += [lambda z, p=params: np.tanh((z - 0.5) / p) / np.cosh((z - 0.5) / p)] + else: + raise ValueError(f"Profile function {pfun} is not defined..") + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp + + def __call__(self, x, y, z): + val = 0.0 + for amp, l, m, thx, thy, pfun in zip(self._amps, self._ls, self._ms, self._theta_x, self._theta_y, self._pfuns): + val += ( + amp + * pfun(z) + * np.sin(l * 2.0 * np.pi / self._Lx * x + thx) + * np.cos(m * 2.0 * np.pi / self._Ly * y + thy) + ) + return val + + +class ModesCosSin(Perturbation): + r""" + + .. math:: + + u(x, y, z) = \sum_s A_s \, \chi_s(z) + \cos \!\left(l_s \tfrac{2\pi}{L_x} x + \theta_{x,s}\right) + \sin \!\left(m_s \tfrac{2\pi}{L_y} y + \theta_{y,s}\right) + + where :math:`\chi_s(z)` can be either 1 or localized in z. + """ + + def __init__( + self, + ls=None, + ms=None, + amps=(1e-4,), + theta_x=None, + theta_y=None, + pfuns=("Id",), + pfuns_params=(0.0,), + Lx=1.0, + Ly=1.0, + Lz=1.0, + given_in_basis: GivenInBasis = "0", + comp: int = 0, + ): + # number of modes + if ls is not None: + n_modes = len(ls) + elif ms is not None: + n_modes = len(ms) + ls = [0] * n_modes + else: + n_modes = 1 + ls = [0] + ms = [0] + + if ms is None: + ms = [0] * n_modes + else: + assert len(ms) == n_modes + + if len(amps) == 1: + amps = [amps[0]] * n_modes + else: + assert len(amps) == n_modes + + if theta_x is None: + theta_x = [0] * n_modes + elif len(theta_x) == 1: + theta_x = [theta_x[0]] * n_modes + else: + assert len(theta_x) == n_modes + + if theta_y is None: + theta_y = [0] * n_modes + elif len(theta_y) == 1: + theta_y = [theta_y[0]] * n_modes + else: + assert len(theta_y) == n_modes + + if len(pfuns) == 1: + pfuns = [pfuns[0]] * n_modes + else: + assert len(pfuns) == n_modes + + if len(pfuns_params) == 1: + pfuns_params = [pfuns_params[0]] * n_modes + else: + assert len(pfuns_params) == n_modes + + # store + self._ls = ls + self._ms = ms + self._amps = amps + self._theta_x = theta_x + self._theta_y = theta_y + self._Lx = Lx + self._Ly = Ly + self._Lz = Lz + + self._pfuns = [] + for pfun, params in zip(pfuns, pfuns_params): + if pfun == "Id": + self._pfuns += [lambda z: 1.0] + elif pfun == "localize": + self._pfuns += [lambda z, p=params: np.tanh((z - 0.5) / p) / np.cosh((z - 0.5) / p)] + else: + raise ValueError(f"Profile function {pfun} is not defined..") + + # use the setters + self.given_in_basis = given_in_basis + self.comp = comp + + def __call__(self, x, y, z): + val = 0.0 + for amp, l, m, thx, thy, pfun in zip(self._amps, self._ls, self._ms, self._theta_x, self._theta_y, self._pfuns): + val += ( + amp + * pfun(z) + * np.cos(l * 2.0 * np.pi / self._Lx * x + thx) + * np.sin(m * 2.0 * np.pi / self._Ly * y + thy) + ) + return val + + class TorusModesSin(Perturbation): r"""Sinusoidal function in the periodic coordinates of a Torus. @@ -953,23 +1378,28 @@ def __call__(self, x, y, z): """Velocity of ions and electrons.""" R = np.sqrt(x**2 + y**2) R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(y, x) - uR = ( + phi = np.arctan2(-y, x) + ustarR = ( self._alpha * R / (self._a * self._R0) * (-z) + self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * z ) - uZ = self._alpha * R / (self._a * self._R0) * (R - self._R0) + self._beta * self._Bp * self._R0 / ( + ustarZ = self._alpha * R / (self._a * self._R0) * (R - self._R0) + self._beta * self._Bp * self._R0 / ( self._B0 * self._a * R ) * (-(R - self._R0)) - uphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp + ustarphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp + + # form normalized to cylindrical coordinates: + uR = ustarR + uphi = ustarphi / R + uZ = ustarZ + + # from cylindrical to cartesian: if self.comp == 0: - # ux = np.cos(phi) * uR - R * np.sin(phi) * uphi - ux = np.cos(phi) * uR - np.sin(phi) * uphi + ux = np.cos(phi) * uR - R * np.sin(phi) * uphi return ux elif self.comp == 1: - # uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi - uy = np.sin(phi) * uR + np.cos(phi) * uphi + uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi return uy elif self.comp == 2: uz = uZ @@ -1054,23 +1484,28 @@ def __call__(self, x, y, z): """Velocity of ions and electrons.""" R = np.sqrt(x**2 + y**2) R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(y, x) - uR = ( + phi = np.arctan2(-y, x) + ustarR = ( self._alpha * R / (self._a * self._R0) * (-z) + self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * z ) - uZ = self._alpha * R / (self._a * self._R0) * (R - self._R0) + self._beta * self._Bp * self._R0 / ( + ustarZ = self._alpha * R / (self._a * self._R0) * (R - self._R0) + self._beta * self._Bp * self._R0 / ( self._B0 * self._a * R ) * (-(R - self._R0)) - uphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp + ustarphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp + + # form normalized to cylindrical coordinates: + uR = ustarR + uphi = ustarphi / R + uZ = ustarZ + + # from cylindrical to cartesian: if self.comp == 0: - # ux = np.cos(phi) * uR - R * np.sin(phi) * uphi - ux = np.cos(phi) * uR - np.sin(phi) * uphi + ux = np.cos(phi) * uR - R * np.sin(phi) * uphi return ux elif self.comp == 1: - # uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi - uy = np.sin(phi) * uR + np.cos(phi) * uphi + uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi return uy elif self.comp == 2: uz = uZ @@ -1155,23 +1590,28 @@ def __call__(self, x, y, z): """Velocity of ions and electrons.""" R = np.sqrt(x**2 + y**2) R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(y, x) - uR = ( + phi = np.arctan2(-y, x) + ustarR = ( self._alpha * R / (self._a * self._R0) * (-z) + self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * z ) - uZ = self._alpha * R / (self._a * self._R0) * (R - self._R0) + self._beta * self._Bp * self._R0 / ( + ustarZ = self._alpha * R / (self._a * self._R0) * (R - self._R0) + self._beta * self._Bp * self._R0 / ( self._B0 * self._a * R ) * (-(R - self._R0)) - uphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp + ustarphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp + + # form normalized to cylindrical coordinates: + uR = ustarR + uphi = ustarphi / R + uZ = ustarZ + + # from cylindrical to cartesian: if self.comp == 0: - # ux = np.cos(phi) * uR - R * np.sin(phi) * uphi - ux = np.cos(phi) * uR - np.sin(phi) * uphi + ux = np.cos(phi) * uR - R * np.sin(phi) * uphi return ux elif self.comp == 1: - # uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi - uy = np.sin(phi) * uR + np.cos(phi) * uphi + uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi return uy elif self.comp == 2: uz = uZ @@ -1567,3 +2007,315 @@ def __call__(self, eta1, eta2=None, eta3=None): ) return val + + +class TokamakManufacturedSolutionVelocity(Perturbation): + r"""Analytic solution :math:`u=u_e` of the system: + + .. math:: + + \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\ + 0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\ + \nabla \cdot (u-u_e) = 0 \,. + + where :math:`f` is defined as follows: + + .. math:: + + f = \left[\begin{array}{c} \alpha \frac{B_0}{a}(R-R_0) - \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu \alpha \frac{1}{a R_0} \frac{R_0}{R^2} \\ + \alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} + \alpha \frac{B_0Z}{a} \\ + \alpha \frac{1}{a R_0} \frac{R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right) \end{array} \right] \,, + \\[2mm] + f = \left[\begin{array}{c} -\alpha \frac{B_0}{a}(R-R_0) + \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu_e \alpha \frac{1}{a R_0} \frac{R_0}{R^2} \\ + -\alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} - \alpha \frac{B_0 Z}{a} \\ + -\alpha \frac{1}{a R_0} \frac{ R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right) \end{array} \right] \,, + \\[2mm] + R = \sqrt{x^2 + y^2} \,. + + Can only be defined in Cartesian coordinates. + The solution is given by: + + .. math:: + \mathbf{u} = \alpha \frac{1}{a R_0} \left[\begin{array}{c} R-R_0 \\ z \\ 0 \end{array} \right] \,, + \\[2mm] + R = \sqrt{x^2 + y^2} \,. + + Parameters + ---------- + comp : string + Which component of the solution ('0', '1' or '2'). + a : float + Minor radius of torus (default: 1.). + R0 : float + Major radius of torus (default: 2.). + B0 : float + On-axis (r=0) toroidal magnetic field (default: 10.). + Bp : float + Poloidal magnetic field (default: 12.5). + alpha : float + (default: 0.1) + beta : float + (default: 1.0) + + References + ---------- + [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising + in plasma physics, Journal of Computational Physics 2018. + """ + + def __init__( + self, + comp=0, + a=1.0, + R0=2.0, + B0=10.0, + Bp=12.5, + alpha=0.1, + beta=1.0, + ): + self._comp = comp + self._a = a + self._R0 = R0 + self._B0 = B0 + self._Bp = Bp + self._alpha = alpha + self._beta = beta + + # use the setters + self.given_in_basis = "physical" + self.comp = comp + + # equilibrium ion velocity + def __call__(self, x, y, z): + """Velocity of ions and electrons.""" + R = np.sqrt(x**2 + y**2) + R = np.where(R == 0.0, 1e-9, R) + phi = np.arctan2(-y, x) + A = self._alpha / (self._a * self._R0) + C = self._beta * self._Bp * self._R0 / (self._B0 * self._a) + + uR = A * (R - self._R0) + uZ = A * z + uphi = 0 + + # from cylindrical to cartesian: + + if self.comp == 0: + ux = np.cos(phi) * uR - R * np.sin(phi) * uphi + return ux + elif self.comp == 1: + uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi + return uy + elif self.comp == 2: + uz = uZ + return uz + else: + raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") + + +class TokamakManufacturedSolutionVelocity_1(Perturbation): + r"""Analytic solution :math:`u=u_e` of the system: + + .. math:: + + \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\ + 0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\ + \nabla \cdot (u-u_e) = 0 \,. + + where :math:`f` is defined as follows: + + .. math:: + + f = \left[\begin{array}{c} \alpha \frac{B_0}{a}(R-R_0) - \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu \alpha \frac{1}{a R_0} \frac{R_0}{R^2} \\ + \alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} + \alpha \frac{B_0Z}{a} \\ + \alpha \frac{1}{a R_0} \frac{R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right) \end{array} \right] \,, + \\[2mm] + f = \left[\begin{array}{c} -\alpha \frac{B_0}{a}(R-R_0) + \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu_e \alpha \frac{1}{a R_0} \frac{R_0}{R^2} \\ + -\alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} - \alpha \frac{B_0 Z}{a} \\ + -\alpha \frac{1}{a R_0} \frac{ R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right) \end{array} \right] \,, + \\[2mm] + R = \sqrt{x^2 + y^2} \,. + + Can only be defined in Cartesian coordinates. + The solution is given by: + + .. math:: + \mathbf{u} = \alpha \frac{1}{a R_0} \left[\begin{array}{c} R-R_0 \\ z \\ 0 \end{array} \right] \,, + \\[2mm] + R = \sqrt{x^2 + y^2} \,. + + Parameters + ---------- + comp : string + Which component of the solution ('0', '1' or '2'). + a : float + Minor radius of torus (default: 1.). + R0 : float + Major radius of torus (default: 2.). + B0 : float + On-axis (r=0) toroidal magnetic field (default: 10.). + Bp : float + Poloidal magnetic field (default: 12.5). + alpha : float + (default: 0.1) + beta : float + (default: 1.0) + + References + ---------- + [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising + in plasma physics, Journal of Computational Physics 2018. + """ + + def __init__( + self, + comp=0, + a=1.0, + R0=2.0, + B0=10.0, + Bp=12.5, + alpha=0.1, + beta=1.0, + ): + self._comp = comp + self._a = a + self._R0 = R0 + self._B0 = B0 + self._Bp = Bp + self._alpha = alpha + self._beta = beta + + # use the setters + self.given_in_basis = "physical" + self.comp = comp + + # equilibrium ion velocity + def __call__(self, x, y, z): + """Velocity of ions and electrons.""" + R = np.sqrt(x**2 + y**2) + R = np.where(R == 0.0, 1e-9, R) + phi = np.arctan2(-y, x) + A = self._alpha / (self._a * self._R0) + C = self._beta * self._Bp * self._R0 / (self._B0 * self._a) + + uR = A * (R - self._R0) + uZ = A * z + uphi = 0 + + # from cylindrical to cartesian: + + if self.comp == 0: + ux = np.cos(phi) * uR - R * np.sin(phi) * uphi + return ux + elif self.comp == 1: + uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi + return uy + elif self.comp == 2: + uz = uZ + return uz + else: + raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") + + +class TokamakManufacturedSolutionVelocity_2(Perturbation): + r"""Analytic solution :math:`u=u_e` of the system: + + .. math:: + + \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\ + 0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\ + \nabla \cdot (u-u_e) = 0 \,. + + where :math:`f` is defined as follows: + + .. math:: + + f = \left[\begin{array}{c} \alpha \frac{B_0}{a}(R-R_0) - \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu \alpha \frac{1}{a R_0} \frac{R_0}{R^2} \\ + \alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} + \alpha \frac{B_0Z}{a} \\ + \alpha \frac{1}{a R_0} \frac{R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right) \end{array} \right] \,, + \\[2mm] + f = \left[\begin{array}{c} -\alpha \frac{B_0}{a}(R-R_0) + \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu_e \alpha \frac{1}{a R_0} \frac{R_0}{R^2} \\ + -\alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} - \alpha \frac{B_0 Z}{a} \\ + -\alpha \frac{1}{a R_0} \frac{ R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right) \end{array} \right] \,, + \\[2mm] + R = \sqrt{x^2 + y^2} \,. + + Can only be defined in Cartesian coordinates. + The solution is given by: + + .. math:: + \mathbf{u} = \alpha \frac{1}{a R_0} \left[\begin{array}{c} R-R_0 \\ z \\ 0 \end{array} \right] \,, + \\[2mm] + R = \sqrt{x^2 + y^2} \,. + + Parameters + ---------- + comp : string + Which component of the solution ('0', '1' or '2'). + a : float + Minor radius of torus (default: 1.). + R0 : float + Major radius of torus (default: 2.). + B0 : float + On-axis (r=0) toroidal magnetic field (default: 10.). + Bp : float + Poloidal magnetic field (default: 12.5). + alpha : float + (default: 0.1) + beta : float + (default: 1.0) + + References + ---------- + [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising + in plasma physics, Journal of Computational Physics 2018. + """ + + def __init__( + self, + comp=0, + a=1.0, + R0=2.0, + B0=10.0, + Bp=12.5, + alpha=0.1, + beta=1.0, + ): + self._comp = comp + self._a = a + self._R0 = R0 + self._B0 = B0 + self._Bp = Bp + self._alpha = alpha + self._beta = beta + + # use the setters + self.given_in_basis = "physical" + self.comp = comp + + # equilibrium ion velocity + def __call__(self, x, y, z): + """Velocity of ions and electrons.""" + R = np.sqrt(x**2 + y**2) + R = np.where(R == 0.0, 1e-9, R) + phi = np.arctan2(-y, x) + A = self._alpha / (self._a * self._R0) + C = self._beta * self._Bp * self._R0 / (self._B0 * self._a) + + uR = A * (R - self._R0) + uZ = A * z + uphi = 0 + + # from cylindrical to cartesian: + + if self.comp == 0: + ux = np.cos(phi) * uR - R * np.sin(phi) * uphi + return ux + elif self.comp == 1: + uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi + return uy + elif self.comp == 2: + uz = uZ + return uz + else: + raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.") diff --git a/src/struphy/initial/tests/test_init_perturbations.py b/src/struphy/initial/tests/test_init_perturbations.py index 105043c01..4a012f08b 100644 --- a/src/struphy/initial/tests/test_init_perturbations.py +++ b/src/struphy/initial/tests/test_init_perturbations.py @@ -83,7 +83,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False if inspect.isclass(val) and val.__module__ == perturbations.__name__: print(key, val) - if "Modes" not in key: + if key not in ("ModesCos", "ModesSin", "TorusModesCos", "TorusModesSin"): continue # skip impossible combinations diff --git a/src/struphy/kinetic_background/tests/test_maxwellians.py b/src/struphy/kinetic_background/tests/test_maxwellians.py index 387ba8f50..b64a5bbe3 100644 --- a/src/struphy/kinetic_background/tests/test_maxwellians.py +++ b/src/struphy/kinetic_background/tests/test_maxwellians.py @@ -306,7 +306,9 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): elif "EQDSKequilibrium" in key: mhd_equil.domain = domains.Tokamak(equilibrium=mhd_equil) elif "CircularTokamak" in key: - mhd_equil.domain = domains.Tokamak(equilibrium=mhd_equil) + mhd_equil.domain = domains.HollowTorus( + a1=1e-3, a2=mhd_equil.params["a"], R0=mhd_equil.params["R0"], tor_period=1 + ) elif "HomogenSlab" in key: mhd_equil.domain = domains.Cuboid() elif "ShearedSlab" in key: @@ -1066,7 +1068,9 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): elif "EQDSKequilibrium" in key: mhd_equil.domain = domains.Tokamak(equilibrium=mhd_equil) elif "CircularTokamak" in key: - mhd_equil.domain = domains.Tokamak(equilibrium=mhd_equil) + mhd_equil.domain = domains.HollowTorus( + a1=1e-3, a2=mhd_equil.params["a"], R0=mhd_equil.params["R0"], tor_period=1 + ) elif "HomogenSlab" in key: mhd_equil.domain = domains.Cuboid() elif "ShearedSlab" in key: diff --git a/src/struphy/linear_algebra/saddle_point.py b/src/struphy/linear_algebra/saddle_point.py index 1a9685a0b..15ca4aac4 100644 --- a/src/struphy/linear_algebra/saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point.py @@ -7,6 +7,8 @@ from psydac.linalg.direct_solvers import SparseSolver from psydac.linalg.solvers import inverse +from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms + class SaddlePointSolver: r"""Solves for :math:`(x, y)` in the saddle point problem @@ -135,6 +137,7 @@ def __init__( self._max_iter = max_iter self._spectralanalysis = spectralanalysis self._dimension = dimension + self._verbose = solver_params["verbose"] if self._variant == "Inverse_Solver": self._BT = B.transpose() @@ -284,6 +287,7 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): self._blockU = BlockVector(self._A.domain, blocks=[self._U1, self._U2]) self._solblocks = [self._blockU, self._P1] + # comment out the next two lines if working with lifting and GMRES x0 = BlockVector(self._block_domainM, blocks=self._solblocks) self._solverMinv._options["x0"] = x0 @@ -316,6 +320,13 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): self._Unp = U_init.toarray() if U_init is not None else self._Unp self._Uenp = Ue_init.toarray() if U_init is not None else self._Uenp + if self._verbose: + print("Uzawa solver:") + print("+---------+---------------------+") + print("+ Iter. # | L2-norm of residual |") + print("+---------+---------------------+") + template = "| {:7d} | {:19.2e} |" + for iteration in range(self._max_iter): # Step 1: Compute velocity U by solving A U = -Bᵀ P + F -A Un self._rhs0np *= 0 @@ -347,8 +358,13 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): self._residual_norms.append(residual_normR1) # Store residual norm # Check for convergence based on residual norm if residual_norm < self._tol: + if self._verbose: + print(template.format(iteration + 1, residual_norm)) + print("+---------+---------------------+") info["success"] = True info["niter"] = iteration + 1 + if self._verbose: + _plot_residual_norms(self._residual_norms) return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult # Steepest gradient @@ -357,9 +373,17 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): # alpha = ((self._Precnp.dot(R)).dot(R)) / ((self._Precnp.dot(R)).dot(self._Precnp.dot(R))) self._Pnp += alpha.real * R.real + if self._verbose: + print(template.format(iteration + 1, residual_norm)) + + if self._verbose: + print("+---------+---------------------+") + # Return with info if maximum iterations reached info["success"] = False info["niter"] = iteration + 1 + if self._verbose == True: + _plot_residual_norms(self._residual_norms) return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult def _setup_inverses(self): @@ -368,18 +392,18 @@ def _setup_inverses(self): # === Preconditioner inverses, if used if self._preconditioner: - A11 = self._Apre[0] - A22 = self._Apre[1] + A11_pre = self._Apre[0] + A22_pre = self._Apre[1] - if hasattr(self, "_A11npinv") and self._is_inverse_still_valid(self._A11npinv, A11, "A11 pre"): + if hasattr(self, "_A11npinv") and self._is_inverse_still_valid(self._A11npinv, A11_pre, "A11 pre"): pass else: - self._A11npinv = self._compute_inverse(A11, which="A11 pre") + self._A11npinv = self._compute_inverse(A11_pre, which="A11 pre") - if hasattr(self, "_A22npinv") and self._is_inverse_still_valid(self._A22npinv, A22, "A22 pre"): + if hasattr(self, "_A22npinv") and self._is_inverse_still_valid(self._A22npinv, A22_pre, "A22 pre"): pass else: - self._A22npinv = self._compute_inverse(A22, which="A22 pre") + self._A22npinv = self._compute_inverse(A22_pre, which="A22 pre") # === Inverse for A[0] if preconditioned if hasattr(self, "_Anpinv") and self._is_inverse_still_valid(self._Anpinv, A0, "A[0]", pre=self._A11npinv): @@ -458,8 +482,10 @@ def _spectral_analysis(self): # A11 before if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): eigvalsA11_before, eigvecs_before = np.linalg.eig(self._A[0]) + condA11_before = np.linalg.cond(self._A[0]) elif self._method_to_solve in ("SparseSolver", "ScipySparse"): eigvalsA11_before, eigvecs_before = np.linalg.eig(self._A[0].toarray()) + condA11_before = np.linalg.cond(self._A[0].toarray()) maxbeforeA11 = max(eigvalsA11_before) maxbeforeA11_abs = np.max(np.abs(eigvalsA11_before)) minbeforeA11_abs = np.min(np.abs(eigvalsA11_before)) @@ -533,7 +559,7 @@ def _spectral_analysis(self): # print(f'{specA22_aft_prec = }') print(f"{specA22_aft_abs_prec = }") - return condA22_before, specA22_bef_abs, condA22_after, specA22_aft_abs_prec + return condA22_before, specA22_bef_abs, condA11_before, condA22_after, specA22_aft_abs_prec else: - return condA22_before, specA22_bef_abs + return condA22_before, specA22_bef_abs, condA11_before diff --git a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py index 31f54a491..67de947e0 100644 --- a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py +++ b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py @@ -71,11 +71,10 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m nue = 0.01 * 100 nu = 1.0 dt = 0.001 - eps = 1e-6 - eps2 = eps # 1e-5#1. #Preconditioner Ae + stab_sigma = 1e-4 method_to_solve = "DirectNPInverse" # 'ScipySparse', 'DirectNPInverse', 'InexactNPInverse', , 'SparseSolver' preconditioner = True - spectralanalysis = True + spectralanalysis = False # Create the solver rho = 0.0005 # Example descent parameter @@ -92,16 +91,13 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m A11 = M2 / dt + nu * (D.T @ M3 @ D + S21.T @ C.T @ M2 @ C @ S21) - M2R A12 = None A21 = A12 - A22 = eps * IdentityOperator(A11.domain) + nue * (D.T @ M3 @ D + S21.T @ C.T @ M2 @ C @ S21) + M2R + A22 = stab_sigma * IdentityOperator(A11.domain) + nue * (D.T @ M3 @ D + S21.T @ C.T @ M2 @ C @ S21) + M2R B1 = -M3 @ D B1T = B1.T B2 = M3 @ D B2T = B2.T F1 = A11.dot(x1) + B1T.dot(y1_rdm) F2 = A22.dot(x2) + B2T.dot(y1_rdm) - # Preconditioner - _A11 = M2 / dt + nu * (D.T @ M3 @ D) # + S21.T @ C.T @ M2 @ C @ S21 - _A22 = nue * (D.T @ M3 @ D) + M2 # +eps2*IdentityOperator(A22.domain) # elif method_for_solving in ("SaddlePointSolverUzawaNumpy"): # Change to numpy if method_to_solve in ("DirectNPInverse", "InexactNPInverse"): @@ -136,20 +132,25 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m A11np = M2np / dt + nu * (Dnp.T @ M3np @ Dnp + S21np.T @ Cnp.T @ M2np @ Cnp @ S21np) - M2Bnp if method_to_solve in ("DirectNPInverse", "InexactNPInverse"): A22np = ( - eps * np.identity(A11np.shape[0]) + stab_sigma * np.identity(A11np.shape[0]) + nue * (Dnp.T @ M3np @ Dnp + S21np.T @ Cnp.T @ M2np @ Cnp @ S21np) + M2Bnp ) - _A22np_pre = eps * np.identity(A22np.shape[0]) # + nue*(Dnp.T @ M3np @ Dnp) + # Preconditioner + _A22np_pre = stab_sigma * np.identity(A22np.shape[0]) # + nue*(Dnp.T @ M3np @ Dnp) + _A11np_pre = M2np / dt # + nu * (Dnp.T @ M3np @ Dnp) elif method_to_solve in ("SparseSolver", "ScipySparse"): A22np = ( - eps * sc.sparse.identity(A11np.shape[0], format="csr") + stab_sigma * sc.sparse.identity(A11np.shape[0], format="csr") + nue * (Dnp.T @ M3np @ Dnp + S21np.T @ Cnp.T @ M2np @ Cnp @ S21np) + M2Bnp ) - +nue * (Dnp.T @ M3np @ Dnp) + eps * sc.sparse.identity(A22np.shape[0], format="csr") # - _A22np_pre = eps * sc.sparse.identity(A22np.shape[0], format="csr") # + nue*(Dnp.T @ M3np @ Dnp) + +nue * (Dnp.T @ M3np @ Dnp) + stab_sigma * sc.sparse.identity(A22np.shape[0], format="csr") # + # Preconditioner + _A22np_pre = stab_sigma * sc.sparse.identity(A22np.shape[0], format="csr") # + nue*(Dnp.T @ M3np @ Dnp) _A22np_pre = _A22np_pre.tocsr() + _A11np_pre = M2np / dt # + nu * (Dnp.T @ M3np @ Dnp) + _A11np_pre = _A11np_pre.tocsr() B1np = -M3np @ Dnp B2np = M3np @ Dnp ynp = y1_rdm.toarray() @@ -159,7 +160,7 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m Anp = [A11np, A22np] Bnp = [B1np, B2np] Fnp = [F1np, F2np] - _A11np_pre = M2np / dt + nu * (Dnp.T @ M3np @ Dnp) + # Preconditioner not inverted Anppre = [_A11np_pre, _A22np_pre] if method_for_solving in ("SaddlePointSolverGMRES", "SaddlePointSolverGMRESwithPC"): @@ -259,6 +260,7 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m spectralanalysis=spectralanalysis, tol=tol, max_iter=max_iter, + verbose=verbose, ) solver.A = Anp solver.B = Bnp diff --git a/src/struphy/main.py b/src/struphy/main.py index 9cfa80c92..74bdd33a8 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -165,7 +165,7 @@ def run( # between the clones : inter_comm clone_config = CloneConfig(comm=comm, params=None, num_clones=num_clones) clone_config.print_clone_config() - if model.kinetic_species: + if model.particle_species: clone_config.print_particle_config() model.clone_config = clone_config diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index b3da1e2f8..479c5b63a 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -28,7 +28,7 @@ from struphy.io.output_handling import DataContainer from struphy.io.setup import descend_options_dict, setup_derham from struphy.kinetic_background import maxwellians -from struphy.models.species import DiagnosticSpecies, FieldSpecies, FluidSpecies, KineticSpecies, Species +from struphy.models.species import DiagnosticSpecies, FieldSpecies, FluidSpecies, ParticleSpecies, Species from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.pic import particles from struphy.pic.base import Particles @@ -86,8 +86,8 @@ def setup_equation_params(self, units: Units, verbose=False): assert isinstance(species, FluidSpecies) species.setup_equation_params(units=units, verbose=verbose) - for _, species in self.kinetic_species.items(): - assert isinstance(species, KineticSpecies) + for _, species in self.particle_species.items(): + assert isinstance(species, ParticleSpecies) species.setup_equation_params(units=units, verbose=verbose) def setup_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): @@ -139,13 +139,13 @@ def fluid_species(self) -> dict: return self._fluid_species @property - def kinetic_species(self) -> dict: - if not hasattr(self, "_kinetic_species"): - self._kinetic_species = {} + def particle_species(self) -> dict: + if not hasattr(self, "_particle_species"): + self._particle_species = {} for k, v in self.__dict__.items(): - if isinstance(v, KineticSpecies): - self._kinetic_species[k] = v - return self._kinetic_species + if isinstance(v, ParticleSpecies): + self._particle_species[k] = v + return self._particle_species @property def diagnostic_species(self) -> dict: @@ -159,7 +159,7 @@ def diagnostic_species(self) -> dict: @property def species(self): if not hasattr(self, "_species"): - self._species = self.field_species | self.fluid_species | self.kinetic_species + self._species = self.field_species | self.fluid_species | self.particle_species return self._species ## allocate methods @@ -601,19 +601,27 @@ def allocate_variables(self, verbose: bool = False): ) # allocate memory for marker arrays of kinetic variables - if self.kinetic_species: - for species, spec in self.kinetic_species.items(): - assert isinstance(spec, KineticSpecies) + if self.particle_species: + for species, spec in self.particle_species.items(): + assert isinstance(spec, ParticleSpecies) for k, v in spec.variables.items(): - assert isinstance(v, (PICVariable, SPHVariable)) - v.allocate( - clone_config=self.clone_config, - derham=self.derham, - domain=self.domain, - equil=self.equil, - projected_equil=self.projected_equil, - verbose=verbose, - ) + if isinstance(v, PICVariable): + v.allocate( + clone_config=self.clone_config, + derham=self.derham, + domain=self.domain, + equil=self.equil, + projected_equil=self.projected_equil, + verbose=verbose, + ) + if isinstance(v, SPHVariable): + v.allocate( + derham=self.derham, + domain=self.domain, + equil=self.equil, + projected_equil=self.projected_equil, + verbose=verbose, + ) # TODO: allocate memory for FE coeffs of diagnostics # if self.params.diagnostic_fields is not None: @@ -682,8 +690,8 @@ def update_markers_to_be_saved(self): Writes markers with IDs that are supposed to be saved into corresponding array. """ - for name, species in self.kinetic_species.items(): - assert isinstance(species, KineticSpecies) + for name, species in self.particle_species.items(): + assert isinstance(species, ParticleSpecies) assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." for _, var in species.variables.items(): assert isinstance(var, PICVariable | SPHVariable) @@ -703,10 +711,10 @@ def update_markers_to_be_saved(self): self._n_markers_saved = n_markers assert self._n_markers_saved <= obj.Np, ( - f"The number of markers for which data should be stored (={n_markers}) murst be <= than the total number of markers (={obj.Np})" + f"The number of markers for which data should be stored (={self._n_markers_saved}) murst be <= than the total number of markers (={obj.Np})" ) if self._n_markers_saved > 0: - var.kinetic_data["markers"] = np.zeros( + var.particle_data["markers"] = np.zeros( (self._n_markers_saved, obj.markers.shape[1]), dtype=float, ) @@ -717,8 +725,8 @@ def update_markers_to_be_saved(self): obj.markers[:, -1] < self._n_markers_saved, ) n_markers_on_proc = np.count_nonzero(markers_on_proc) - var.kinetic_data["markers"][:] = -1.0 - var.kinetic_data["markers"][:n_markers_on_proc] = obj.markers[markers_on_proc] + var.particle_data["markers"][:] = -1.0 + var.particle_data["markers"][:n_markers_on_proc] = obj.markers[markers_on_proc] @profile def update_distr_functions(self): @@ -728,53 +736,53 @@ def update_distr_functions(self): dim_to_int = {"e1": 0, "e2": 1, "e3": 2, "v1": 3, "v2": 4, "v3": 5} - for name, species in self.kinetic_species.items(): - assert isinstance(species, KineticSpecies) + for name, species in self.particle_species.items(): + assert isinstance(species, ParticleSpecies) assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." for _, var in species.variables.items(): assert isinstance(var, PICVariable | SPHVariable) obj = var.particles assert isinstance(obj, Particles) - if obj.n_cols_diagnostics > 0: - for i in range(obj.n_cols_diagnostics): - str_dn = f"d{i + 1}" - dim_to_int[str_dn] = 3 + obj.vdim + 3 + i + if obj.n_cols_diagnostics > 0: + for i in range(obj.n_cols_diagnostics): + str_dn = f"d{i + 1}" + dim_to_int[str_dn] = 3 + obj.vdim + 3 + i - if species.binning_plots: - for slice_i, edges in var.kinetic_data["bin_edges"].items(): - comps = slice_i.split("_") - components = [False] * (3 + obj.vdim + 3 + obj.n_cols_diagnostics) + if species.binning_plots: + for slice_i, edges in var.particle_data["bin_edges"].items(): + comps = slice_i.split("_") + components = [False] * (3 + obj.vdim + 3 + obj.n_cols_diagnostics) - for comp in comps: - components[dim_to_int[comp]] = True + for comp in comps: + components[dim_to_int[comp]] = True - f_slice, df_slice = obj.binning(components, edges) + f_slice, df_slice = obj.binning(components, edges) - var.kinetic_data["f"][slice_i][:] = f_slice - var.kinetic_data["df"][slice_i][:] = df_slice + var.particle_data["f"][slice_i][:] = f_slice + var.particle_data["df"][slice_i][:] = df_slice - if species.n_sph is not None: - h1 = 1 / obj.boxes_per_dim[0] - h2 = 1 / obj.boxes_per_dim[1] - h3 = 1 / obj.boxes_per_dim[2] + if species.kernel_density_plots: + h1 = 1 / obj.boxes_per_dim[0] + h2 = 1 / obj.boxes_per_dim[1] + h3 = 1 / obj.boxes_per_dim[2] - ndim = np.count_nonzero([d > 1 for d in obj.boxes_per_dim]) - if ndim == 0: - kernel_type = "gaussian_3d" - else: - kernel_type = "gaussian_" + str(ndim) + "d" - - for i, pts in enumerate(val["plot_pts"]): - n_sph = obj.eval_density( - *pts, - h1=h1, - h2=h2, - h3=h3, - kernel_type=kernel_type, - fast=True, - ) - val["kinetic_data"]["n_sph"][i][:] = n_sph + ndim = np.count_nonzero([d > 1 for d in obj.boxes_per_dim]) + if ndim == 0: + kernel_type = "gaussian_3d" + else: + kernel_type = "gaussian_" + str(ndim) + "d" + + for i, pts in enumerate(var.particle_data["plot_pts"]): + n_sph = obj.eval_density( + *pts, + h1=h1, + h2=h2, + h3=h3, + kernel_type=kernel_type, + fast=True, + ) + var.particle_data["n_sph"][i][:] = n_sph def print_scalar_quantities(self): """ @@ -1086,8 +1094,8 @@ def initialize_data_output(self, data: DataContainer, size): ) # save kinetic data in group 'kinetic/' - for name, species in self.kinetic_species.items(): - assert isinstance(species, KineticSpecies) + for name, species in self.particle_species.items(): + assert isinstance(species, ParticleSpecies) assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." for varname, var in species.variables.items(): assert isinstance(var, PICVariable | SPHVariable) @@ -1099,8 +1107,8 @@ def initialize_data_output(self, data: DataContainer, size): data.add_data({key_spec_restart: obj._markers}) - # TODO: kinetic_data should be a KineticData object, not a dict - for key1, val1 in var.kinetic_data.items(): + # TODO: particle_data should be a KineticData object, not a dict + for key1, val1 in var.particle_data.items(): key_dat = os.path.join(key_spec, key1) if key1 == "bin_edges": @@ -1114,10 +1122,10 @@ def initialize_data_output(self, data: DataContainer, size): dims = (len(key2) - 2) // 3 + 1 for dim in range(dims): data.file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = ( - var.kinetic_data["bin_edges"][key2][dim][:-1] + var.particle_data["bin_edges"][key2][dim][:-1] + ( - var.kinetic_data["bin_edges"][key2][dim][1] - - var.kinetic_data["bin_edges"][key2][dim][0] + var.particle_data["bin_edges"][key2][dim][1] + - var.particle_data["bin_edges"][key2][dim][0] ) / 2 ) @@ -1127,9 +1135,9 @@ def initialize_data_output(self, data: DataContainer, size): key_n = os.path.join(key_dat, "view_", str(i)) data.add_data({key_n: v1}) # save 1d point values, not meshgrids, because attrs size is limited - eta1 = var.kinetic_data["plot_pts"][i][0][:, 0, 0] - eta2 = var.kinetic_data["plot_pts"][i][1][0, :, 0] - eta3 = var.kinetic_data["plot_pts"][i][2][0, 0, :] + eta1 = var.particle_data["plot_pts"][i][0][:, 0, 0] + eta2 = var.particle_data["plot_pts"][i][1][0, :, 0] + eta3 = var.particle_data["plot_pts"][i][2][0, 0, :] data.file[key_n].attrs["eta1"] = eta1 data.file[key_n].attrs["eta2"] = eta2 data.file[key_n].attrs["eta3"] = eta3 @@ -1306,7 +1314,7 @@ def generate_default_parameter_file( file.write("from struphy.fields_background import equils\n") species_params = "\n# species parameters\n" - kinetic_params = "" + particle_params = "" has_plasma = False has_feec = False has_pic = False @@ -1314,16 +1322,16 @@ def generate_default_parameter_file( for sn, species in self.species.items(): assert isinstance(species, Species) - if isinstance(species, (FluidSpecies, KineticSpecies)): + if isinstance(species, (FluidSpecies, ParticleSpecies)): has_plasma = True species_params += f"model.{sn}.set_phys_params()\n" - if isinstance(species, KineticSpecies): - kinetic_params += f"\nloading_params = LoadingParameters()\n" - kinetic_params += f"weights_params = WeightsParameters()\n" - kinetic_params += f"boundary_params = BoundaryParameters()\n" - kinetic_params += f"model.{sn}.set_markers(loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params)\n" - kinetic_params += f"model.{sn}.set_sorting_boxes()\n" - kinetic_params += f"model.{sn}.set_save_data()\n" + if isinstance(species, ParticleSpecies): + particle_params += f"\nloading_params = LoadingParameters()\n" + particle_params += f"weights_params = WeightsParameters()\n" + particle_params += f"boundary_params = BoundaryParameters()\n" + particle_params += f"model.{sn}.set_markers(loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params)\n" + particle_params += f"model.{sn}.set_sorting_boxes()\n" + particle_params += f"model.{sn}.set_save_data()\n" for vn, var in species.variables.items(): if isinstance(var, FEECVariable): @@ -1339,7 +1347,6 @@ def generate_default_parameter_file( model.{sn}.{vn}.add_perturbation(perturbations.TorusModesCos(given_in_basis='v', comp=2))\n" ) - exclude = f"# model.{sn}.{vn}.save_data = False\n" elif isinstance(var, PICVariable): has_pic = True init_pert_pic = f"\n# if .add_initial_condition is not called, the background is the kinetic initial condition\n" @@ -1362,8 +1369,14 @@ def generate_default_parameter_file( init_bckgr_pic += f"model.{sn}.{vn}.add_background(background)\n" exclude = f"# model.....save_data = False\n" + elif isinstance(var, SPHVariable): has_sph = True + init_bckgr_sph = f"background = equils.ConstantVelocity()\n" + init_bckgr_sph += f"model.{sn}.{vn}.add_background(background)\n" + init_pert_sph = f"perturbation = perturbations.TorusModesCos()\n" + init_pert_sph += f"model.{sn}.{vn}.add_perturbation(del_n=perturbation)\n" + exclude = f"# model.{sn}.{vn}.save_data = False\n" file.write("from struphy.topology import grids\n") file.write("from struphy.io.options import DerhamOptions\n") @@ -1414,8 +1427,8 @@ def generate_default_parameter_file( if has_plasma: file.write(species_params) - if has_pic: - file.write(kinetic_params) + if has_pic or has_sph: + file.write(particle_params) file.write("\n# propagator options\n") for prop in self.propagators.__dict__: @@ -1428,6 +1441,9 @@ def generate_default_parameter_file( if has_pic: file.write(init_bckgr_pic) file.write(init_pert_pic) + if has_sph: + file.write(init_bckgr_sph) + file.write(init_pert_sph) file.write("\n# optional: exclude variables from saving\n") file.write(exclude) @@ -1436,16 +1452,16 @@ def generate_default_parameter_file( file.write(" # start run\n") file.write( " main.run(model, \n\ - params_path=__file__, \n\ - env=env, \n\ - base_units=base_units, \n\ - time_opts=time_opts, \n\ - domain=domain, \n\ - equil=equil, \n\ - grid=grid, \n\ - derham_opts=derham_opts, \n\ - verbose=verbose, \n\ - )" + params_path=__file__, \n\ + env=env, \n\ + base_units=base_units, \n\ + time_opts=time_opts, \n\ + domain=domain, \n\ + equil=equil, \n\ + grid=grid, \n\ + derham_opts=derham_opts, \n\ + verbose=verbose, \n\ + )" ) file.close() diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 5b88a1fab..6b1047541 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -3,7 +3,7 @@ from psydac.linalg.block import BlockVector from struphy.models.base import StruphyModel -from struphy.models.species import FieldSpecies, FluidSpecies, KineticSpecies +from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.polar.basic import PolarVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers @@ -2207,7 +2207,8 @@ class IsothermalEulerSPH(StruphyModel): :ref:`propagators` (called in sequence): 1. :class:`~struphy.propagators.propagators_markers.PushEta` - 2. :class:`~struphy.propagators.propagators_markers.PushVinSPHpressure` + 2. :class:`~struphy.propagators.propagators_markers.PushVxB` + 3. :class:`~struphy.propagators.propagators_markers.PushVinSPHpressure` :ref:`Model info `: """ @@ -2237,6 +2238,7 @@ def velocity_scale(): def propagators_dct(): return { propagators_markers.PushEta: ["euler_fluid"], + # propagators_markers.PushVxB: ["euler_fluid"], propagators_markers.PushVinSPHpressure: ["euler_fluid"], } @@ -2253,17 +2255,28 @@ def __init__(self, params, comm, clone_config=None): # prelim _p = params["kinetic"]["euler_fluid"] algo_eta = _p["options"]["PushEta"]["algo"] + # algo_vxb = _p["options"]["PushVxB"]["algo"] kernel_type = _p["options"]["PushVinSPHpressure"]["kernel_type"] algo_sph = _p["options"]["PushVinSPHpressure"]["algo"] gravity = _p["options"]["PushVinSPHpressure"]["gravity"] thermodynamics = _p["options"]["PushVinSPHpressure"]["thermodynamics"] + # magnetic field + # self._b_eq = self.projected_equil.b2 + # set keyword arguments for propagators self._kwargs[propagators_markers.PushEta] = { "algo": algo_eta, # "density_field": self.pointer["projected_density"], } + # self._kwargs[propagators_markers.PushVxB] = { + # "algo": algo_vxb, + # "kappa": 1.0, + # "b2": self._b_eq, + # "b2_add": None, + # } + self._kwargs[propagators_markers.PushVinSPHpressure] = { "kernel_type": kernel_type, "algo": algo_sph, diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 0c19aadef..4f5dd89bd 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -4,7 +4,7 @@ from struphy.feec.projectors import L2Projector from struphy.kinetic_background.base import KineticBackground from struphy.models.base import StruphyModel -from struphy.models.species import FieldSpecies, FluidSpecies, KineticSpecies +from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.pic.accumulation.particles_to_grid import AccumulatorVector @@ -90,7 +90,7 @@ def __init__(self): self.phi = FEECVariable(space="H1") self.init_variables() - class KineticIons(KineticSpecies): + class KineticIons(ParticleSpecies): def __init__(self): self.var = PICVariable(space="Particles6D") self.init_variables() diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 125c3e5fd..0c1a2311e 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -15,6 +15,7 @@ from struphy.pic.utilities import ( BinningPlot, BoundaryParameters, + KernelDensityPlot, LoadingParameters, WeightsParameters, ) @@ -144,7 +145,7 @@ class FluidSpecies(Species): pass -class KineticSpecies(Species): +class ParticleSpecies(Species): """Single kinetic species in 3d + vdim phase space.""" def set_markers( @@ -202,12 +203,12 @@ def set_save_data( self, n_markers: int | float = 3, binning_plots: tuple[BinningPlot] = (), - n_sph: dict = None, + kernel_density_plots: tuple[KernelDensityPlot] = (), ): """Saving marker orits, binned data and kernel density reconstructions.""" self.n_markers = n_markers self.binning_plots = binning_plots - self.n_sph = n_sph + self.kernel_density_plots = kernel_density_plots class DiagnosticSpecies(Species): diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index 4bb145f3e..af5c28410 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -17,6 +17,7 @@ "Maxwell", "Vlasov", "GuidingCenter", + "PressureLessSPH", ] # for name, obj in inspect.getmembers(toy): # if inspect.isclass(obj) and "models.toy" in obj.__module__: diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index ce1bf9035..1a8d3562e 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -4,7 +4,7 @@ from mpi4py import MPI from struphy.models.base import StruphyModel -from struphy.models.species import FieldSpecies, FluidSpecies, KineticSpecies +from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers from struphy.propagators.base import Propagator @@ -116,7 +116,7 @@ class Vlasov(StruphyModel): ## species - class KineticIons(KineticSpecies): + class KineticIons(ParticleSpecies): def __init__(self): self.var = PICVariable(space="Particles6D") self.init_variables() @@ -204,7 +204,7 @@ class GuidingCenter(StruphyModel): ## species - class KineticIons(KineticSpecies): + class KineticIons(ParticleSpecies): def __init__(self): self.var = PICVariable(space="Particles5D") self.init_variables() @@ -1059,65 +1059,61 @@ class PressureLessSPH(StruphyModel): This is discretized by particles going in straight lines. """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species - dct["kinetic"]["p_fluid"] = "ParticlesSPH" - return dct + class ColdFluid(ParticleSpecies): + def __init__(self): + self.var = SPHVariable() + self.init_variables() - @staticmethod - def bulk_species(): - return "p_fluid" + ## propagators - @staticmethod - def velocity_scale(): - return None + class Propagators: + def __init__(self): + self.push_eta = propagators_markers.PushEta() - @staticmethod - def diagnostics_dct(): - dct = {} - dct["projected_density"] = "L2" - return dct + ## abstract methods - @staticmethod - def propagators_dct(): - return {propagators_markers.PushEta: ["p_fluid"]} + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + # 1. instantiate all species + self.cold_fluid = self.ColdFluid() - def __init__(self, params, comm, clone_config=None): - super().__init__(params, comm=comm, clone_config=clone_config) + # 2. instantiate all propagators + self.propagators = self.Propagators() - from mpi4py.MPI import IN_PLACE, SUM + # 3. assign variables to propagators + self.propagators.push_eta.variables.var = self.cold_fluid.var - # prelim - p_fluid_params = self.kinetic["p_fluid"]["params"] - algo_eta = params["kinetic"]["p_fluid"]["options"]["PushEta"]["algo"] + # define scalars for update_scalar_quantities + self.add_scalar("en_kin", compute="from_particles", variable=self.cold_fluid.var) - # set keyword arguments for propagators - self._kwargs[propagators_markers.PushEta] = { - "algo": algo_eta, - "density_field": self.pointer["projected_density"], - } + @property + def bulk_species(self): + return self.cold_fluid - # Initialize propagators used in splitting substeps - self.init_propagators() + @property + def velocity_scale(self): + return None - # Scalar variables to be saved during simulation - self.add_scalar("en_kin", compute="from_particles", species="p_fluid") + # @staticmethod + # def diagnostics_dct(): + # dct = {} + # dct["projected_density"] = "L2" + # return dct + + def allocate_helpers(self): + pass def update_scalar_quantities(self): - en_kin = self.pointer["p_fluid"].markers_wo_holes_and_ghost[:, 6].dot( - self.pointer["p_fluid"].markers_wo_holes_and_ghost[:, 3] ** 2 - + self.pointer["p_fluid"].markers_wo_holes_and_ghost[:, 4] ** 2 - + self.pointer["p_fluid"].markers_wo_holes_and_ghost[:, 5] ** 2 - ) / (2.0 * self.pointer["p_fluid"].Np) + particles = self.cold_fluid.var.particles + en_kin = particles.markers_wo_holes_and_ghost[:, 6].dot( + particles.markers_wo_holes_and_ghost[:, 3] ** 2 + + particles.markers_wo_holes_and_ghost[:, 4] ** 2 + + particles.markers_wo_holes_and_ghost[:, 5] ** 2 + ) / (2.0 * particles.Np) self.update_scalar("en_kin", en_kin) @@ -1235,6 +1231,7 @@ def __init__(self, params, comm, clone_config=None): stokes_spectralanalysis = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"][ "spectralanalysis" ] + stokes_lifting = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["lifting"] stokes_dimension = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["dimension"] stokes_1D_dt = params["time"]["dt"] @@ -1264,6 +1261,7 @@ def __init__(self, params, comm, clone_config=None): "spectralanalysis": stokes_spectralanalysis, "dimension": stokes_dimension, "D1_dt": stokes_1D_dt, + "lifting": stokes_lifting, } # Initialize propagators used in splitting substeps diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 3839f0c93..01482dc55 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -21,10 +21,11 @@ from struphy.kinetic_background.base import KineticBackground from struphy.pic import particles from struphy.pic.base import Particles +from struphy.pic.particles import ParticlesSPH from struphy.utils.clone_config import CloneConfig if TYPE_CHECKING: - from struphy.models.species import FieldSpecies, FluidSpecies, KineticSpecies, Species + from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies, Species class Variable(metaclass=ABCMeta): @@ -86,21 +87,6 @@ def add_background(self, background, verbose=True): for k, v in background.__dict__.items(): print(f" {k}: {v}") - def add_perturbation(self, perturbation: Perturbation, verbose=True): - if not hasattr(self, "_perturbations") or self.perturbations is None: - self._perturbations = perturbation - else: - if not isinstance(self.perturbations, list): - self._perturbations = [self.perturbations] - self._perturbations += [perturbation] - - if verbose and MPI.COMM_WORLD.Get_rank() == 0: - print( - f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added perturbation '{perturbation.__class__.__name__}' with:" - ) - for k, v in perturbation.__dict__.items(): - print(f" {k}: {v}") - class FEECVariable(Variable): def __init__(self, space: OptsFEECSpace = "H1"): @@ -124,6 +110,21 @@ def species(self) -> FieldSpecies | FluidSpecies: def add_background(self, background: FieldsBackground, verbose=True): super().add_background(background, verbose=verbose) + def add_perturbation(self, perturbation: Perturbation, verbose=True): + if not hasattr(self, "_perturbations") or self.perturbations is None: + self._perturbations = perturbation + else: + if not isinstance(self.perturbations, list): + self._perturbations = [self.perturbations] + self._perturbations += [perturbation] + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print( + f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added perturbation '{perturbation.__class__.__name__}' with:" + ) + for k, v in perturbation.__dict__.items(): + print(f" {k}: {v}") + def allocate( self, derham: Derham, @@ -144,7 +145,7 @@ class PICVariable(Variable): def __init__(self, space: OptsPICSpace = "Particles6D"): check_option(space, OptsPICSpace) self._space = space - self._kinetic_data = {} + self._particle_data = {} @property def space(self): @@ -155,11 +156,11 @@ def particles(self) -> Particles: return self._particles @property - def kinetic_data(self): - return self._kinetic_data + def particle_data(self) -> dict: + return self._particle_data @property - def species(self) -> KineticSpecies: + def species(self) -> ParticleSpecies: if not hasattr(self, "_species"): self._species = None return self._species @@ -248,9 +249,9 @@ def allocate( self.particles.initialize_weights() # for storing the binned distribution function - self.kinetic_data["bin_edges"] = {} - self.kinetic_data["f"] = {} - self.kinetic_data["df"] = {} + self.particle_data["bin_edges"] = {} + self.particle_data["f"] = {} + self.particle_data["df"] = {} for bin_plot in self.species.binning_plots: sli = bin_plot.slice @@ -261,42 +262,193 @@ def allocate( assert len(sli.split("_")) == len(ranges) == len(n_bins), ( f"Number of slices names ({len(sli.split('_'))}), number of bins ({len(n_bins)}), and number of ranges ({len(ranges)}) are inconsistent with each other!\n\n" ) - self.kinetic_data["bin_edges"][sli] = [] + self.particle_data["bin_edges"][sli] = [] dims = (len(sli) - 2) // 3 + 1 for j in range(dims): - self.kinetic_data["bin_edges"][sli] += [ + self.particle_data["bin_edges"][sli] += [ np.linspace( ranges[j][0], ranges[j][1], n_bins[j] + 1, ), ] - self.kinetic_data["f"][sli] = np.zeros(n_bins, dtype=float) - self.kinetic_data["df"][sli] = np.zeros(n_bins, dtype=float) + self.particle_data["f"][sli] = np.zeros(n_bins, dtype=float) + self.particle_data["df"][sli] = np.zeros(n_bins, dtype=float) # for storing an sph evaluation of the density n - if self.species.n_sph is not None: - plot_pts = self.species.n_sph["plot_pts"] - - self.kinetic_data["n_sph"] = [] - self.kinetic_data["plot_pts"] = [] - for i, pts in enumerate(plot_pts): - assert len(pts) == 3 - eta1 = np.linspace(0.0, 1.0, pts[0]) - eta2 = np.linspace(0.0, 1.0, pts[1]) - eta3 = np.linspace(0.0, 1.0, pts[2]) - ee1, ee2, ee3 = np.meshgrid( - eta1, - eta2, - eta3, - indexing="ij", - ) - self.kinetic_data["plot_pts"] += [(ee1, ee2, ee3)] - self.kinetic_data["n_sph"] += [np.zeros(ee1.shape, dtype=float)] + self.particle_data["n_sph"] = [] + self.particle_data["plot_pts"] = [] + + for kd_plot in self.species.kernel_density_plots: + eta1 = np.linspace(0.0, 1.0, kd_plot.pts_e1) + eta2 = np.linspace(0.0, 1.0, kd_plot.pts_e2) + eta3 = np.linspace(0.0, 1.0, kd_plot.pts_e3) + ee1, ee2, ee3 = np.meshgrid( + eta1, + eta2, + eta3, + indexing="ij", + ) + self.particle_data["plot_pts"] += [(ee1, ee2, ee3)] + self.particle_data["n_sph"] += [np.zeros(ee1.shape, dtype=float)] # other data (wave-particle power exchange, etc.) # TODO class SPHVariable(Variable): - pass + def __init__(self): + self._space = "ParticlesSPH" + self._n_as_volume_form = True + self._particle_data = {} + + @property + def space(self): + return self._space + + @property + def particles(self) -> ParticlesSPH: + return self._particles + + @property + def particle_data(self): + return self._particle_data + + @property + def species(self) -> ParticleSpecies: + if not hasattr(self, "_species"): + self._species = None + return self._species + + @property + def n_as_volume_form(self) -> bool: + """Whether the number density n is given as a volume form or scalar function (=default).""" + return self._n_as_volume_form + + def add_background(self, background: FluidEquilibrium, verbose=True): + super().add_background(background, verbose=verbose) + + def add_perturbation( + self, + del_n: Perturbation = None, + del_u1: Perturbation = None, + del_u2: Perturbation = None, + del_u3: Perturbation = None, + verbose=True, + ): + self._perturbations = {} + self._perturbations["n"] = del_n + self._perturbations["u1"] = del_u1 + self._perturbations["u2"] = del_u2 + self._perturbations["u3"] = del_u3 + + if verbose and MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added perturbation:") + for k, v in self._perturbations.items(): + print(f" {k}: {v}") + + @property + def perturbations(self) -> dict[str, Perturbation]: + if not hasattr(self, "_perturbations"): + self._perturbations = None + return self._perturbations + + def allocate( + self, + derham: Derham = None, + domain: Domain = None, + equil: FluidEquilibrium = None, + projected_equil: ProjectedFluidEquilibrium = None, + verbose: bool = False, + ): + assert isinstance(self.backgrounds, FluidEquilibrium), ( + f"List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." + ) + + self.backgrounds.domain = domain + + if derham is None: + domain_decomp = None + else: + domain_array = derham.domain_array + nprocs = derham.domain_decomposition.nprocs + domain_decomp = (domain_array, nprocs) + + comm_world = MPI.COMM_WORLD + if comm_world.Get_size() == 1: + comm_world = None + + self._particles = ParticlesSPH( + comm_world=comm_world, + domain_decomp=domain_decomp, + mpi_dims_mask=self.species.dims_mask, + boxes_per_dim=self.species.boxes_per_dim, + box_bufsize=self.species.box_bufsize, + name=self.species.__class__.__name__, + loading_params=self.species.loading_params, + weights_params=self.species.weights_params, + boundary_params=self.species.boundary_params, + bufsize=self.species.bufsize, + domain=domain, + equil=equil, + projected_equil=projected_equil, + background=self.backgrounds, + n_as_volume_form=self.n_as_volume_form, + perturbations=self.perturbations, + equation_params=self.species.equation_params, + verbose=verbose, + ) + + if self.species.do_sort: + sort = True + else: + sort = False + self.particles.draw_markers(sort=sort, verbose=verbose) + self.particles.initialize_weights() + + # for storing the binned distribution function + self.particle_data["bin_edges"] = {} + self.particle_data["f"] = {} + self.particle_data["df"] = {} + + for bin_plot in self.species.binning_plots: + sli = bin_plot.slice + n_bins = bin_plot.n_bins + ranges = bin_plot.ranges + + assert ((len(sli) - 2) / 3).is_integer(), f"Binning coordinates must be separated by '_', but reads {sli}." + assert len(sli.split("_")) == len(ranges) == len(n_bins), ( + f"Number of slices names ({len(sli.split('_'))}), number of bins ({len(n_bins)}), and number of ranges ({len(ranges)}) are inconsistent with each other!\n\n" + ) + self.particle_data["bin_edges"][sli] = [] + dims = (len(sli) - 2) // 3 + 1 + for j in range(dims): + self.particle_data["bin_edges"][sli] += [ + np.linspace( + ranges[j][0], + ranges[j][1], + n_bins[j] + 1, + ), + ] + self.particle_data["f"][sli] = np.zeros(n_bins, dtype=float) + self.particle_data["df"][sli] = np.zeros(n_bins, dtype=float) + + # for storing an sph evaluation of the density n + self.particle_data["n_sph"] = [] + self.particle_data["plot_pts"] = [] + + for kd_plot in self.species.kernel_density_plots: + eta1 = np.linspace(0.0, 1.0, kd_plot.pts_e1) + eta2 = np.linspace(0.0, 1.0, kd_plot.pts_e2) + eta3 = np.linspace(0.0, 1.0, kd_plot.pts_e3) + ee1, ee2, ee3 = np.meshgrid( + eta1, + eta2, + eta3, + indexing="ij", + ) + self.particle_data["plot_pts"] += [(ee1, ee2, ee3)] + self.particle_data["n_sph"] += [np.zeros(ee1.shape, dtype=float)] + + # other data (wave-particle power exchange, etc.) + # TODO diff --git a/src/struphy/pic/accumulation/particles_to_grid.py b/src/struphy/pic/accumulation/particles_to_grid.py index 2092db2a4..23345df23 100644 --- a/src/struphy/pic/accumulation/particles_to_grid.py +++ b/src/struphy/pic/accumulation/particles_to_grid.py @@ -12,6 +12,7 @@ from struphy.feec.psydac_derham import Derham from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments from struphy.pic.base import Particles +from struphy.profiling.profiling import ProfileManager class Accumulator: @@ -203,13 +204,14 @@ def __call__(self, *optional_args, **args_control): dat[:] = 0.0 # accumulate into matrix (and vector) with markers - self.kernel( - self.particles.args_markers, - self.derham.args_derham, - self.args_domain, - *self._args_data, - *optional_args, - ) + with ProfileManager.profile_region("kernel: " + self.kernel.__name__): + self.kernel( + self.particles.args_markers, + self.derham.args_derham, + self.args_domain, + *self._args_data, + *optional_args, + ) # apply filter if self.filter_params["use_filter"] is not None: @@ -607,13 +609,14 @@ def __call__(self, *optional_args, **args_control): dat[:] = 0.0 # accumulate into matrix (and vector) with markers - self.kernel( - self.particles.args_markers, - self.derham._args_derham, - self.args_domain, - *self._args_data, - *optional_args, - ) + with ProfileManager.profile_region("kernel: " + self.kernel.__name__): + self.kernel( + self.particles.args_markers, + self.derham._args_derham, + self.args_domain, + *self._args_data, + *optional_args, + ) if self.particles.clone_config is None: num_clones = 1 diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 9ab323163..4408f2593 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -26,10 +26,10 @@ from struphy.pic import sampling_kernels, sobol_seq from struphy.pic.pushing.pusher_utilities_kernels import reflect from struphy.pic.sorting_kernels import ( + assign_box_to_each_particle, + assign_particles_to_boxes, flatten_index, initialize_neighbours, - put_particles_in_boxes_kernel, - reassign_boxes, sort_boxed_particles, ) from struphy.pic.sph_eval_kernels import ( @@ -153,7 +153,7 @@ def __init__( domain: Domain = None, equil: FluidEquilibrium = None, projected_equil: ProjectedFluidEquilibrium = None, - background: KineticBackground = None, + background: KineticBackground | FluidEquilibrium = None, initial_condition: KineticBackground = None, perturbations: dict[str, Perturbation] = None, n_as_volume_form: bool = False, @@ -170,16 +170,6 @@ def __init__( self._num_clones = self.clone_config.num_clones self._clone_id = self.clone_config.clone_id - # other parameters - self._name = name - self._loading_params = loading_params - self._weights_params = weights_params - self._boundary_params = boundary_params - self._domain = domain - self._equil = equil - self._projected_equil = projected_equil - self._equation_params = equation_params - # defaults if loading_params is None: loading_params = LoadingParameters() @@ -190,6 +180,16 @@ def __init__( if boundary_params is None: boundary_params = BoundaryParameters() + # other parameters + self._name = name + self._loading_params = loading_params + self._weights_params = weights_params + self._boundary_params = boundary_params + self._domain = domain + self._equil = equil + self._projected_equil = projected_equil + self._equation_params = equation_params + # check for mpi communicator (i.e. sub_comm of clone) if self.mpi_comm is None: self._mpi_size = 1 @@ -210,8 +210,8 @@ def __init__( # total number of cells (equal to mpi_size if no grid) n_cells = np.sum(np.prod(self.domain_array[:, 2::3], axis=1, dtype=int)) * self.num_clones - if verbose: - print(f"\n{self.mpi_rank = }, {n_cells = }") + # if verbose: + # print(f"\n{self.mpi_rank = }, {n_cells = }") # total number of boxes if self.boxes_per_dim is None: @@ -225,8 +225,8 @@ def __init__( ) n_boxes = np.prod(self.boxes_per_dim, dtype=int) * self.num_clones - if verbose: - print(f"\n{self.mpi_rank = }, {n_boxes = }") + # if verbose: + # print(f"\n{self.mpi_rank = }, {n_boxes = }") # total number of markers (Np) and particles per cell (ppc) Np = self.loading_params.Np @@ -265,12 +265,21 @@ def __init__( if bc_refill is not None: for bc_refilli in bc_refill: assert bc_refilli in ("outer", "inner") + self._bc = bc self._periodic_axes = [axis for axis, b_c in enumerate(bc) if b_c == "periodic"] self._reflect_axes = [axis for axis, b_c in enumerate(bc) if b_c == "reflect"] self._remove_axes = [axis for axis, b_c in enumerate(bc) if b_c == "remove"] self._bc_refill = bc_refill + bc_sph = boundary_params.bc_sph + if bc_sph is None: + bc_sph = [bci if bci == "periodic" else "mirror" for bci in self.bc] + + for bci in bc_sph: + assert bci in ("periodic", "mirror", "fixed") + self._bc_sph = bc_sph + # particle type assert type in ("full_f", "delta_f", "sph") self._type = type @@ -444,6 +453,11 @@ def bc_refill(self): """How to re-enter particles if bc is 'refill'.""" return self._bc_refill + @property + def bc_sph(self): + """List of boundary conditions for sph evaluation in each direction.""" + return self._bc_sph + @property def Np(self): """Total number of markers/particles, from user input.""" @@ -1105,25 +1119,52 @@ def _allocate_marker_array(self): ) def _initialize_sorting_boxes(self): + """Initializes the sorting boxes. + + Each MPI process has exactly the same box structure and numbering. + For instance, if boxes_per_dim = (16, 1, 1) and there are 2 MPI processes, + each process would get 8 boxes in the first direction. + Hence boxes_per_dim has to be divisible by the number of ranks in each direction. + """ + self._initialized_sorting = False if self.boxes_per_dim is not None: + # split boxes across MPI processes nboxes = [nboxes // nproc for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)] + # check whether this process touches the domain boundary + is_domain_boundary = {} + x_l = self.domain_array[self.mpi_rank, 0] + x_r = self.domain_array[self.mpi_rank, 1] + y_l = self.domain_array[self.mpi_rank, 3] + y_r = self.domain_array[self.mpi_rank, 4] + z_l = self.domain_array[self.mpi_rank, 6] + z_r = self.domain_array[self.mpi_rank, 7] + is_domain_boundary["x_m"] = x_l == 0.0 + is_domain_boundary["x_p"] = x_r == 1.0 + is_domain_boundary["y_m"] = y_l == 0.0 + is_domain_boundary["y_p"] = y_r == 1.0 + is_domain_boundary["z_m"] = z_l == 0.0 + is_domain_boundary["z_p"] = z_r == 1.0 + self._sorting_boxes = self.SortingBoxes( self.markers.shape, self.type == "sph", nx=nboxes[0], ny=nboxes[1], nz=nboxes[2], + bc_sph=self.bc_sph, + is_domain_boundary=is_domain_boundary, comm=self.mpi_comm, - verbose=self.verbose, + verbose=False, box_bufsize=self._box_bufsize, ) + if self.sorting_boxes.communicate: self._get_neighbouring_proc() self._initialized_sorting = True - self._argsort_array = np.zeros(self.markers.shape[0], dtype=int) + else: self._sorting_boxes = None @@ -1191,7 +1232,75 @@ def _generate_sampling_moments(self): # self.loading_params["moments"] = new_moments def _set_initial_condition(self): - self._f_init = self.initial_condition + if self.type != "sph": + self._f_init = self.initial_condition + else: + # Get the initialization function and pass the correct arguments + assert isinstance(self.f0, FluidEquilibrium) + self._u_init = self.f0.u_cart + + if self.perturbations is not None: + for ( + moment, + pert, + ) in self.perturbations.items(): # only one perturbation is taken into account at the moment + assert isinstance(moment, str) + if pert is None: + continue + assert isinstance(pert, Perturbation) + + if moment == "n": + _fun = TransformedPformComponent( + pert, + pert.given_in_basis, + "0", + comp=pert.comp, + domain=self.domain, + ) + elif moment == "u1": + _fun = TransformedPformComponent( + pert, + pert.given_in_basis, + "v", + comp=pert.comp, + domain=self.domain, + ) + _fun_cart = lambda e1, e2, e3: self.domain.push(_fun, e1, e2, e3, kind="v") + self._u_init = lambda e1, e2, e3: self.f0.u_cart(e1, e2, e3)[0] + _fun_cart(e1, e2, e3) + # TODO: add other velocity components + else: + _fun = None + + def _f_init(*etas, flat_eval=False): + if len(etas) == 1: + if _fun is None: + out = self.f0.n0(etas[0]) + else: + out = self.f0.n0(etas[0]) + _fun(*etas[0].T) + else: + assert len(etas) == 3 + E1, E2, E3, is_sparse_meshgrid = Domain.prepare_eval_pts( + etas[0], + etas[1], + etas[2], + flat_eval=flat_eval, + ) + + out0 = self.f0.n0(E1, E2, E3) + + if _fun is None: + out = out0 + else: + out1 = _fun(E1, E2, E3) + assert out0.shape == out1.shape + out = out0 + out1 + + if flat_eval: + out = np.squeeze(out) + + return out + + self._f_init = _f_init def _load_external( self, @@ -1960,8 +2069,9 @@ def apply_kinetic_bc(self, newton=False): for axis in self._reflect_axes: outside_inds = self._find_outside_particles(axis) - self.markers[self._is_outside_left, axis] = 1e-4 - self.markers[self._is_outside_right, axis] = 1 - 1e-4 + self.markers[self._is_outside_left, axis] *= -1.0 + self.markers[self._is_outside_right, axis] *= -1.0 + self.markers[self._is_outside_right, axis] += 2.0 self.markers[self._is_outside, self.first_pusher_idx] = -1.0 @@ -1970,7 +2080,7 @@ def apply_kinetic_bc(self, newton=False): for axis in self._reflect_axes: if len(outside_inds_per_axis[axis]) == 0: continue - + # flip velocity reflect( self.markers, self.domain.args_domain, @@ -2119,9 +2229,9 @@ def gyro_transfer(self, outside_inds): class SortingBoxes: """Boxes used for the sorting of the particles. - Represented as a 2D array of integers, - each line of the array corespond to one box, - and all the non (-1) entries of line i are the particles in the i-th box + Boxes are represented as a 2D array of integers, where + each line coresponds to one box, and all entries of line i that are not -1 + correspond to a particles in the i-th box. Parameters ---------- @@ -2140,6 +2250,13 @@ class SortingBoxes: nz : int number of boxes in the z direction. + bc_sph : list + Boundary condition for sph density evaluation. + Either 'periodic', 'mirror' or 'fixed' in each direction. + + is_domain_boundary: dict + Has two booleans for each direction; True when the boundary of the MPI process is a domain boundary. + comm : Intracomm MPI communicator or None. @@ -2158,6 +2275,8 @@ def __init__( nx: int = 1, ny: int = 1, nz: int = 1, + bc_sph: list = None, + is_domain_boundary: dict = None, comm: Intracomm = None, box_index: "int" = -2, box_bufsize: "float" = 2.0, @@ -2172,6 +2291,21 @@ def __init__( self._box_bufsize = box_bufsize self._verbose = verbose + if bc_sph is None: + bc_sph = ["periodic"] * 3 + self._bc_sph = bc_sph + + if is_domain_boundary is None: + is_domain_boundary = {} + is_domain_boundary["x_m"] = True + is_domain_boundary["x_p"] = True + is_domain_boundary["y_m"] = True + is_domain_boundary["y_p"] = True + is_domain_boundary["z_m"] = True + is_domain_boundary["z_p"] = True + + self._is_domain_boundary = is_domain_boundary + if comm is None: self._rank = 0 else: @@ -2220,6 +2354,51 @@ def neighbours(self): def communicate(self): return self._communicate + @property + def is_domain_boundary(self): + """Dict with two booleans for each direction (e.g. 'x_m' and 'x_p'); True when the boundary of the MPI process is a domain boundary (0.0 or 1.0).""" + return self._is_domain_boundary + + @property + def bc_sph(self): + """List of boundary conditions for sph evaluation in each direction.""" + return self._bc_sph + + @property + def bc_sph_index_shifts(self): + """Dictionary holding the index shifts of box number for ghost particles in each direction.""" + if not hasattr(self, "_bc_sph_index_shifts"): + self._compute_sph_index_shifts() + return self._bc_sph_index_shifts + + def _compute_sph_index_shifts(self): + """The index shifts are applied to ghost particles to indicate their new box after sending.""" + self._bc_sph_index_shifts = {} + self._bc_sph_index_shifts["x_m"] = flatten_index(self.nx, 0, 0, self.nx, self.ny, self.nz) + self._bc_sph_index_shifts["x_p"] = flatten_index(self.nx, 0, 0, self.nx, self.ny, self.nz) + self._bc_sph_index_shifts["y_m"] = flatten_index(0, self.ny, 0, self.nx, self.ny, self.nz) + self._bc_sph_index_shifts["y_p"] = flatten_index(0, self.ny, 0, self.nx, self.ny, self.nz) + self._bc_sph_index_shifts["z_m"] = flatten_index(0, 0, self.nz, self.nx, self.ny, self.nz) + self._bc_sph_index_shifts["z_p"] = flatten_index(0, 0, self.nz, self.nx, self.ny, self.nz) + + if self.bc_sph[0] in ("mirror", "fixed"): + if self.is_domain_boundary["x_m"]: + self._bc_sph_index_shifts["x_m"] = flatten_index(-1, 0, 0, self.nx, self.ny, self.nz) + if self.is_domain_boundary["x_p"]: + self._bc_sph_index_shifts["x_p"] = flatten_index(-1, 0, 0, self.nx, self.ny, self.nz) + + if self.bc_sph[1] in ("mirror", "fixed"): + if self.is_domain_boundary["y_m"]: + self._bc_sph_index_shifts["y_m"] = flatten_index(0, -1, 0, self.nx, self.ny, self.nz) + if self.is_domain_boundary["y_p"]: + self._bc_sph_index_shifts["y_p"] = flatten_index(0, -1, 0, self.nx, self.ny, self.nz) + + if self.bc_sph[2] in ("mirror", "fixed"): + if self.is_domain_boundary["z_m"]: + self._bc_sph_index_shifts["z_m"] = flatten_index(0, 0, -1, self.nx, self.ny, self.nz) + if self.is_domain_boundary["z_p"]: + self._bc_sph_index_shifts["z_p"] = flatten_index(0, 0, -1, self.nx, self.ny, self.nz) + def _set_boxes(self): """ "(Re)set the box structure.""" self._n_boxes = (self._nx + 2) * (self._ny + 2) * (self._nz + 2) @@ -2403,10 +2582,14 @@ def _set_boundary_boxes(self): ) ) - def sort_boxed_particles_numpy(self): - """Sort the particles by box using numpy.sort.""" + def _sort_boxed_particles_numpy(self): + """Sort the particles by box using numpy.argsort.""" sorting_axis = self._sorting_boxes.box_index + + if not hasattr(self, "_argsort_array"): + self._argsort_array = np.zeros(self.markers.shape[0], dtype=int) self._argsort_array[:] = self._markers[:, sorting_axis].argsort() + self._markers[:, :] = self._markers[self._argsort_array] @profile @@ -2416,29 +2599,53 @@ def put_particles_in_boxes(self): neighbouring boxes of neighbours processors or also communicated""" self.remove_ghost_particles() - put_particles_in_boxes_kernel( - self._markers, + assign_box_to_each_particle( + self.markers, self.holes, self._sorting_boxes.nx, self._sorting_boxes.ny, self._sorting_boxes.nz, - self._sorting_boxes._boxes, - self._sorting_boxes._next_index, self.domain_array[self.mpi_rank], ) + self.check_and_assign_particles_to_boxes() + if self.sorting_boxes.communicate: - self.communicate_boxes() - reassign_boxes( - self._markers, - self.holes, - self._sorting_boxes._boxes, - self._sorting_boxes._next_index, - ) + self.communicate_boxes(verbose=self.verbose) + self.check_and_assign_particles_to_boxes() self.update_ghost_particles() + # if self.verbose: + # valid_box_ids = np.nonzero(self._sorting_boxes._boxes[:, 0] != -1)[0] + # print(f"Boxes holding at least one particle: {valid_box_ids}") + # for i in valid_box_ids: + # n_mks_box = np.count_nonzero(self._sorting_boxes._boxes[i] != -1) + # print(f"Number of markers in box {i} is {n_mks_box}") + + def check_and_assign_particles_to_boxes(self): + """Check whether the box array has enough columns (detect load imbalance wrt to sorting boxes), + and then assigne the particles to boxes.""" + + bcount = np.bincount(np.int64(self.markers_wo_holes[:, -2])) + max_in_box = np.max(bcount) + if max_in_box > self._sorting_boxes.boxes.shape[1]: + warnings.warn( + f'Strong load imbalance detected in sorting boxes: \ +max number of markers in a box ({max_in_box}) on rank {self.mpi_rank} \ +exceeds the column-size of the box array ({self._sorting_boxes.boxes.shape[1]}). \ +Increasing the value of "box_bufsize" in the markers parameters for the next run.' + ) + self.mpi_comm.Abort() + + assign_particles_to_boxes( + self.markers, + self.holes, + self._sorting_boxes._boxes, + self._sorting_boxes._next_index, + ) + @profile - def do_sort(self): + def do_sort(self, use_numpy_argsort=False): """Assign the particles to boxes and then sort them.""" nx = self._sorting_boxes.nx ny = self._sorting_boxes.ny @@ -2447,18 +2654,17 @@ def do_sort(self): self.put_particles_in_boxes() - # We could either use numpy routine or kernel to sort - # Kernel seems to be 3x faster - # self.sort_boxed_particles_numpy() - - sort_boxed_particles( - self._markers, - self._sorting_boxes._swap_line_1, - self._sorting_boxes._swap_line_2, - nboxes + 1, - self._sorting_boxes._next_index, - self._sorting_boxes._cumul_next_index, - ) + if use_numpy_argsort: + self._sort_boxed_particles_numpy() + else: + sort_boxed_particles( + self._markers, + self._sorting_boxes._swap_line_1, + self._sorting_boxes._swap_line_2, + nboxes + 1, + self._sorting_boxes._next_index, + self._sorting_boxes._cumul_next_index, + ) if self.sorting_boxes.communicate: self.update_ghost_particles() @@ -2469,27 +2675,28 @@ def remove_ghost_particles(self): self._markers[new_holes] = -1.0 self.update_holes() - def determine_send_markers_box(self): - """Determine which markers belong to boxes that are at the boundary and put them in a new array""" - # Faces - self._markers_x_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_m) - self._markers_x_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_p) - self._markers_y_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_y_m) - self._markers_y_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_y_p) - self._markers_z_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_z_m) - self._markers_z_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_z_p) + def prepare_ghost_particles(self): + """Markers for boundary conditions and MPI communication. - # Adjust box number - self._markers_x_m[:, self._sorting_boxes.box_index] += self._sorting_boxes.nx - self._markers_x_p[:, self._sorting_boxes.box_index] -= self._sorting_boxes.nx - self._markers_y_m[:, self._sorting_boxes.box_index] += (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - self._markers_y_p[:, self._sorting_boxes.box_index] -= (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - self._markers_z_m[:, self._sorting_boxes.box_index] += ( - (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_z_p[:, self._sorting_boxes.box_index] -= ( - (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) + Does the following: + 1. determine which markers belong to boxes that are at the boundary and put these markers in a new array (e.g. markers_x_m) + 2. set their last index to -2 to indicate that they will be "ghost particles" after sending + 3. set their new box number (boundary conditions enter here) + 4. optional: mirror position for boundary conditions + """ + shifts = self.sorting_boxes.bc_sph_index_shifts + # if self.verbose: + # print(f"{self.sorting_boxes.bc_sph_index_shifts = }") + + ## Faces + + # ghost marker arrays + self._markers_x_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m) + self._markers_x_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p) + self._markers_y_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_m) + self._markers_y_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_p) + self._markers_z_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_z_m) + self._markers_z_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_z_p) # Put last index to -2 to indicate that they are ghosts on the new process self._markers_x_m[:, -1] = -2.0 @@ -2499,147 +2706,133 @@ def determine_send_markers_box(self): self._markers_z_m[:, -1] = -2.0 self._markers_z_p[:, -1] = -2.0 - # Edges x-y - self._markers_x_m_y_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_m_y_m) - self._markers_x_m_y_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_m_y_p) - self._markers_x_p_y_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_p_y_m) - self._markers_x_p_y_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p) - # Adjust box number - self._markers_x_m_y_m[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx + (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - ) - self._markers_x_m_y_p[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx - (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - ) - self._markers_x_p_y_m[:, self._sorting_boxes.box_index] += ( - -self._sorting_boxes.nx + (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - ) - self._markers_x_p_y_p[:, self._sorting_boxes.box_index] += ( - -self._sorting_boxes.nx - (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - ) + self._markers_x_m[:, self._sorting_boxes.box_index] += shifts["x_m"] + self._markers_x_p[:, self._sorting_boxes.box_index] -= shifts["x_p"] + self._markers_y_m[:, self._sorting_boxes.box_index] += shifts["y_m"] + self._markers_y_p[:, self._sorting_boxes.box_index] -= shifts["y_p"] + self._markers_z_m[:, self._sorting_boxes.box_index] += shifts["z_m"] + self._markers_z_p[:, self._sorting_boxes.box_index] -= shifts["z_p"] + + # Mirror position for boundary condition + if self.bc_sph[0] in ("mirror", "fixed"): + self._mirror_particles( + "_markers_x_m", "_markers_x_p", is_domain_boundary=self.sorting_boxes.is_domain_boundary + ) + + if self.bc_sph[1] in ("mirror", "fixed"): + self._mirror_particles( + "_markers_y_m", "_markers_y_p", is_domain_boundary=self.sorting_boxes.is_domain_boundary + ) + + if self.bc_sph[2] in ("mirror", "fixed"): + self._mirror_particles( + "_markers_z_m", "_markers_z_p", is_domain_boundary=self.sorting_boxes.is_domain_boundary + ) - # Put first last index to -2 to indicate that they are ghosts on the new process + ## Edges x-y + + # ghost marker arrays + self._markers_x_m_y_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_m) + self._markers_x_m_y_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_p) + self._markers_x_p_y_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_m) + self._markers_x_p_y_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p) + + # Put last index to -2 to indicate that they are ghosts on the new process self._markers_x_m_y_m[:, -1] = -2.0 self._markers_x_m_y_p[:, -1] = -2.0 self._markers_x_p_y_m[:, -1] = -2.0 self._markers_x_p_y_p[:, -1] = -2.0 - # Edges x-z - self._markers_x_m_z_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_m_z_m) - self._markers_x_m_z_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_m_z_p) - self._markers_x_p_z_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_p_z_m) - self._markers_x_p_z_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_p_z_p) - # Adjust box number - self._markers_x_m_z_m[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx - + (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_m_z_p[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx - - (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_p_z_m[:, self._sorting_boxes.box_index] += ( - -self._sorting_boxes.nx - + (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_p_z_p[:, self._sorting_boxes.box_index] += ( - -self._sorting_boxes.nx - - (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) + self._markers_x_m_y_m[:, self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] + self._markers_x_m_y_p[:, self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] + self._markers_x_p_y_m[:, self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] + self._markers_x_p_y_p[:, self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] + + # Mirror position for boundary condition + if self.bc_sph[0] in ("mirror", "fixed") or self.bc_sph[1] in ("mirror", "fixed"): + self._mirror_particles( + "_markers_x_m_y_m", + "_markers_x_m_y_p", + "_markers_x_p_y_m", + "_markers_x_p_y_p", + is_domain_boundary=self.sorting_boxes.is_domain_boundary, + ) + + ## Edges x-z - # Put first last index to -2 to indicate that they are ghosts on the new process + # ghost marker arrays + self._markers_x_m_z_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_z_m) + self._markers_x_m_z_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_z_p) + self._markers_x_p_z_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_z_m) + self._markers_x_p_z_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_z_p) + + # Put last index to -2 to indicate that they are ghosts on the new process self._markers_x_m_z_m[:, -1] = -2.0 self._markers_x_m_z_p[:, -1] = -2.0 self._markers_x_p_z_m[:, -1] = -2.0 self._markers_x_p_z_p[:, -1] = -2.0 - # Edges y-z - self._markers_y_m_z_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_y_m_z_m) - self._markers_y_m_z_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_y_m_z_p) - self._markers_y_p_z_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_y_p_z_m) - self._markers_y_p_z_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_y_p_z_p) - # Adjust box number - self._markers_y_m_z_m[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx + 2 - ) * self._sorting_boxes.ny + (self._sorting_boxes.nx + 2) * ( - self._sorting_boxes.ny + 2 - ) * self._sorting_boxes.nz - self._markers_y_m_z_p[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx + 2 - ) * self._sorting_boxes.ny - (self._sorting_boxes.nx + 2) * ( - self._sorting_boxes.ny + 2 - ) * self._sorting_boxes.nz - self._markers_y_p_z_m[:, self._sorting_boxes.box_index] += ( - -(self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - + (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_y_p_z_p[:, self._sorting_boxes.box_index] += ( - -(self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - - (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) + self._markers_x_m_z_m[:, self._sorting_boxes.box_index] += shifts["x_m"] + shifts["z_m"] + self._markers_x_m_z_p[:, self._sorting_boxes.box_index] += shifts["x_m"] - shifts["z_p"] + self._markers_x_p_z_m[:, self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["z_m"] + self._markers_x_p_z_p[:, self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["z_p"] + + # Mirror position for boundary condition + if self.bc_sph[0] in ("mirror", "fixed") or self.bc_sph[2] in ("mirror", "fixed"): + self._mirror_particles( + "_markers_x_m_z_m", + "_markers_x_m_z_p", + "_markers_x_p_z_m", + "_markers_x_p_z_p", + is_domain_boundary=self.sorting_boxes.is_domain_boundary, + ) + + ## Edges y-z - # Put first last index to -2 to indicate that they are ghosts on the new process + # ghost marker arrays + self._markers_y_m_z_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_m_z_m) + self._markers_y_m_z_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_m_z_p) + self._markers_y_p_z_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_p_z_m) + self._markers_y_p_z_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_p_z_p) + + # Put last index to -2 to indicate that they are ghosts on the new process self._markers_y_m_z_m[:, -1] = -2.0 self._markers_y_m_z_p[:, -1] = -2.0 self._markers_y_p_z_m[:, -1] = -2.0 self._markers_y_p_z_p[:, -1] = -2.0 - # Corners - self._markers_x_m_y_m_z_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_m_y_m_z_m) - self._markers_x_m_y_m_z_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_m_y_m_z_p) - self._markers_x_m_y_p_z_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_m_y_p_z_m) - self._markers_x_m_y_p_z_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_m_y_p_z_p) - self._markers_x_p_y_m_z_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_p_y_m_z_m) - self._markers_x_p_y_m_z_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_p_y_m_z_p) - self._markers_x_p_y_p_z_m = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p_z_m) - self._markers_x_p_y_p_z_p = self.determine_marker_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p_z_p) - # Adjust box number - self._markers_x_m_y_m_z_m[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx - + (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - + (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_m_y_m_z_p[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx - + (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - - (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_m_y_p_z_m[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx - - (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - + (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_m_y_p_z_p[:, self._sorting_boxes.box_index] += ( - self._sorting_boxes.nx - - (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - - (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_p_y_m_z_m[:, self._sorting_boxes.box_index] += ( - -self._sorting_boxes.nx - + (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - + (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_p_y_m_z_p[:, self._sorting_boxes.box_index] += ( - -self._sorting_boxes.nx - + (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - - (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_p_y_p_z_m[:, self._sorting_boxes.box_index] += ( - -self._sorting_boxes.nx - - (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - + (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) - self._markers_x_p_y_p_z_p[:, self._sorting_boxes.box_index] += ( - -self._sorting_boxes.nx - - (self._sorting_boxes.nx + 2) * self._sorting_boxes.ny - - (self._sorting_boxes.nx + 2) * (self._sorting_boxes.ny + 2) * self._sorting_boxes.nz - ) + self._markers_y_m_z_m[:, self._sorting_boxes.box_index] += shifts["y_m"] + shifts["z_m"] + self._markers_y_m_z_p[:, self._sorting_boxes.box_index] += shifts["y_m"] - shifts["z_p"] + self._markers_y_p_z_m[:, self._sorting_boxes.box_index] += -shifts["y_p"] + shifts["z_m"] + self._markers_y_p_z_p[:, self._sorting_boxes.box_index] += -shifts["y_p"] - shifts["z_p"] + + # Mirror position for boundary condition + if self.bc_sph[1] in ("mirror", "fixed") or self.bc_sph[2] in ("mirror", "fixed"): + self._mirror_particles( + "_markers_y_m_z_m", + "_markers_y_m_z_p", + "_markers_y_p_z_m", + "_markers_y_p_z_p", + is_domain_boundary=self.sorting_boxes.is_domain_boundary, + ) + + ## Corners + + # ghost marker arrays + self._markers_x_m_y_m_z_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_m_z_m) + self._markers_x_m_y_m_z_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_m_z_p) + self._markers_x_m_y_p_z_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_p_z_m) + self._markers_x_m_y_p_z_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_p_z_p) + self._markers_x_p_y_m_z_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_m_z_m) + self._markers_x_p_y_m_z_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_m_z_p) + self._markers_x_p_y_p_z_m = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p_z_m) + self._markers_x_p_y_p_z_p = self.determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p_z_p) - # Put first last index to -2 to indicate that they are ghosts on the new process + # Put last index to -2 to indicate that they are ghosts on the new process self._markers_x_m_y_m_z_m[:, -1] = -2.0 self._markers_x_m_y_m_z_p[:, -1] = -2.0 self._markers_x_m_y_p_z_m[:, -1] = -2.0 @@ -2649,8 +2842,123 @@ def determine_send_markers_box(self): self._markers_x_p_y_p_z_m[:, -1] = -2.0 self._markers_x_p_y_p_z_p[:, -1] = -2.0 - def determine_marker_in_box(self, list_boxes): - """Determine the markers that belong to a certain box and put them in an array""" + # Adjust box number + self._markers_x_m_y_m_z_m[:, self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] + shifts["z_m"] + self._markers_x_m_y_m_z_p[:, self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] - shifts["z_p"] + self._markers_x_m_y_p_z_m[:, self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] + shifts["z_m"] + self._markers_x_m_y_p_z_p[:, self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] - shifts["z_p"] + self._markers_x_p_y_m_z_m[:, self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] + shifts["z_m"] + self._markers_x_p_y_m_z_p[:, self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] - shifts["z_p"] + self._markers_x_p_y_p_z_m[:, self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] + shifts["z_m"] + self._markers_x_p_y_p_z_p[:, self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] - shifts["z_p"] + + # Mirror position for boundary condition + if any([bci in ("mirror", "fixed") for bci in self.bc_sph]): + self._mirror_particles( + "_markers_x_m_y_m_z_m", + "_markers_x_m_y_m_z_p", + "_markers_x_m_y_p_z_m", + "_markers_x_m_y_p_z_p", + "_markers_x_p_y_m_z_m", + "_markers_x_p_y_m_z_p", + "_markers_x_p_y_p_z_m", + "_markers_x_p_y_p_z_p", + is_domain_boundary=self.sorting_boxes.is_domain_boundary, + ) + + def _mirror_particles(self, *marker_array_names, is_domain_boundary=None): + self._fixed_markers_set = {} + + for arr_name in marker_array_names: + assert isinstance(arr_name, str) + arr = getattr(self, arr_name) + + if arr.size == 0: + continue + + # x-direction + if self.bc_sph[0] in ("mirror", "fixed"): + if "x_m" in arr_name and is_domain_boundary["x_m"]: + arr[:, 0] *= -1.0 + if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set: + boundary_values = self.f_init( + *arr[:, :3].T, flat_eval=True + ) # evaluation outside of the unit cube - maybe not working for all f_init! + arr[:, self.index["weights"]] = -boundary_values / self.s0( + *arr[:, :3].T, + flat_eval=True, + remove_holes=False, + ) + self._fixed_markers_set[arr_name] = True + elif "x_p" in arr_name and is_domain_boundary["x_p"]: + arr[:, 0] = 2.0 - arr[:, 0] + if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set: + boundary_values = self.f_init( + *arr[:, :3].T, flat_eval=True + ) # evaluation outside of the unit cube - maybe not working for all f_init! + arr[:, self.index["weights"]] = -boundary_values / self.s0( + *arr[:, :3].T, + flat_eval=True, + remove_holes=False, + ) + self._fixed_markers_set[arr_name] = True + + # y-direction + if self.bc_sph[1] in ("mirror", "fixed"): + if "y_m" in arr_name and is_domain_boundary["y_m"]: + arr[:, 1] *= -1.0 + if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set: + boundary_values = self.f_init( + *arr[:, :3].T, flat_eval=True + ) # evaluation outside of the unit cube - maybe not working for all f_init! + arr[:, self.index["weights"]] = -boundary_values / self.s0( + *arr[:, :3].T, + flat_eval=True, + remove_holes=False, + ) + self._fixed_markers_set[arr_name] = True + elif "y_p" in arr_name and is_domain_boundary["y_p"]: + arr[:, 1] = 2.0 - arr[:, 1] + if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set: + boundary_values = self.f_init( + *arr[:, :3].T, flat_eval=True + ) # evaluation outside of the unit cube - maybe not working for all f_init! + arr[:, self.index["weights"]] = -boundary_values / self.s0( + *arr[:, :3].T, + flat_eval=True, + remove_holes=False, + ) + self._fixed_markers_set[arr_name] = True + + # z-direction + if self.bc_sph[2] in ("mirror", "fixed"): + if "z_m" in arr_name and is_domain_boundary["z_m"]: + arr[:, 2] *= -1.0 + if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set: + boundary_values = self.f_init( + *arr[:, :3].T, flat_eval=True + ) # evaluation outside of the unit cube - maybe not working for all f_init! + arr[:, self.index["weights"]] = -boundary_values / self.s0( + *arr[:, :3].T, + flat_eval=True, + remove_holes=False, + ) + self._fixed_markers_set[arr_name] = True + elif "z_p" in arr_name and is_domain_boundary["z_p"]: + arr[:, 2] = 2.0 - arr[:, 2] + if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set: + boundary_values = self.f_init( + *arr[:, :3].T, flat_eval=True + ) # evaluation outside of the unit cube - maybe not working for all f_init! + arr[:, self.index["weights"]] = -boundary_values / self.s0( + *arr[:, :3].T, + flat_eval=True, + remove_holes=False, + ) + self._fixed_markers_set[arr_name] = True + + def determine_markers_in_box(self, list_boxes): + """Determine the markers that belong to a certain box (list of boxes) and put them in an array""" indices = [] for i in list_boxes: indices += list(self._sorting_boxes._boxes[i][self._sorting_boxes._boxes[i] != -1]) @@ -2665,123 +2973,149 @@ def get_destinations_box(self): self._send_list_box = [np.zeros((0, self.n_cols))] * self.mpi_size # Faces + # if self._x_m_proc is not None: self._send_info_box[self._x_m_proc] += len(self._markers_x_m) self._send_list_box[self._x_m_proc] = np.concatenate((self._send_list_box[self._x_m_proc], self._markers_x_m)) + # if self._x_p_proc is not None: self._send_info_box[self._x_p_proc] += len(self._markers_x_p) self._send_list_box[self._x_p_proc] = np.concatenate((self._send_list_box[self._x_p_proc], self._markers_x_p)) + # if self._y_m_proc is not None: self._send_info_box[self._y_m_proc] += len(self._markers_y_m) self._send_list_box[self._y_m_proc] = np.concatenate((self._send_list_box[self._y_m_proc], self._markers_y_m)) + # if self._y_p_proc is not None: self._send_info_box[self._y_p_proc] += len(self._markers_y_p) self._send_list_box[self._y_p_proc] = np.concatenate((self._send_list_box[self._y_p_proc], self._markers_y_p)) + # if self._z_m_proc is not None: self._send_info_box[self._z_m_proc] += len(self._markers_z_m) self._send_list_box[self._z_m_proc] = np.concatenate((self._send_list_box[self._z_m_proc], self._markers_z_m)) + # if self._z_p_proc is not None: self._send_info_box[self._z_p_proc] += len(self._markers_z_p) self._send_list_box[self._z_p_proc] = np.concatenate((self._send_list_box[self._z_p_proc], self._markers_z_p)) # x-y edges + # if self._x_m_y_m_proc is not None: self._send_info_box[self._x_m_y_m_proc] += len(self._markers_x_m_y_m) self._send_list_box[self._x_m_y_m_proc] = np.concatenate( (self._send_list_box[self._x_m_y_m_proc], self._markers_x_m_y_m) ) + # if self._x_m_y_p_proc is not None: self._send_info_box[self._x_m_y_p_proc] += len(self._markers_x_m_y_p) self._send_list_box[self._x_m_y_p_proc] = np.concatenate( (self._send_list_box[self._x_m_y_p_proc], self._markers_x_m_y_p) ) + # if self._x_p_y_m_proc is not None: self._send_info_box[self._x_p_y_m_proc] += len(self._markers_x_p_y_m) self._send_list_box[self._x_p_y_m_proc] = np.concatenate( (self._send_list_box[self._x_p_y_m_proc], self._markers_x_p_y_m) ) + # if self._x_p_y_p_proc is not None: self._send_info_box[self._x_p_y_p_proc] += len(self._markers_x_p_y_p) self._send_list_box[self._x_p_y_p_proc] = np.concatenate( (self._send_list_box[self._x_p_y_p_proc], self._markers_x_p_y_p) ) # x-z edges + # if self._x_m_z_m_proc is not None: self._send_info_box[self._x_m_z_m_proc] += len(self._markers_x_m_z_m) self._send_list_box[self._x_m_z_m_proc] = np.concatenate( (self._send_list_box[self._x_m_z_m_proc], self._markers_x_m_z_m) ) + # if self._x_m_z_p_proc is not None: self._send_info_box[self._x_m_z_p_proc] += len(self._markers_x_m_z_p) self._send_list_box[self._x_m_z_p_proc] = np.concatenate( (self._send_list_box[self._x_m_z_p_proc], self._markers_x_m_z_p) ) + # if self._x_p_z_m_proc is not None: self._send_info_box[self._x_p_z_m_proc] += len(self._markers_x_p_z_m) self._send_list_box[self._x_p_z_m_proc] = np.concatenate( (self._send_list_box[self._x_p_z_m_proc], self._markers_x_p_z_m) ) + # if self._x_p_z_p_proc is not None: self._send_info_box[self._x_p_z_p_proc] += len(self._markers_x_p_z_p) self._send_list_box[self._x_p_z_p_proc] = np.concatenate( (self._send_list_box[self._x_p_z_p_proc], self._markers_x_p_z_p) ) # y-z edges + # if self._y_m_z_m_proc is not None: self._send_info_box[self._y_m_z_m_proc] += len(self._markers_y_m_z_m) self._send_list_box[self._y_m_z_m_proc] = np.concatenate( (self._send_list_box[self._y_m_z_m_proc], self._markers_y_m_z_m) ) + # if self._y_m_z_p_proc is not None: self._send_info_box[self._y_m_z_p_proc] += len(self._markers_y_m_z_p) self._send_list_box[self._y_m_z_p_proc] = np.concatenate( (self._send_list_box[self._y_m_z_p_proc], self._markers_y_m_z_p) ) + # if self._y_p_z_m_proc is not None: self._send_info_box[self._y_p_z_m_proc] += len(self._markers_y_p_z_m) self._send_list_box[self._y_p_z_m_proc] = np.concatenate( (self._send_list_box[self._y_p_z_m_proc], self._markers_y_p_z_m) ) + # if self._y_p_z_p_proc is not None: self._send_info_box[self._y_p_z_p_proc] += len(self._markers_y_p_z_p) self._send_list_box[self._y_p_z_p_proc] = np.concatenate( (self._send_list_box[self._y_p_z_p_proc], self._markers_y_p_z_p) ) # corners + # if self._x_m_y_m_z_m_proc is not None: self._send_info_box[self._x_m_y_m_z_m_proc] += len(self._markers_x_m_y_m_z_m) self._send_list_box[self._x_m_y_m_z_m_proc] = np.concatenate( (self._send_list_box[self._x_m_y_m_z_m_proc], self._markers_x_m_y_m_z_m) ) + # if self._x_m_y_m_z_p_proc is not None: self._send_info_box[self._x_m_y_m_z_p_proc] += len(self._markers_x_m_y_m_z_p) self._send_list_box[self._x_m_y_m_z_p_proc] = np.concatenate( (self._send_list_box[self._x_m_y_m_z_p_proc], self._markers_x_m_y_m_z_p) ) + # if self._x_m_y_p_z_m_proc is not None: self._send_info_box[self._x_m_y_p_z_m_proc] += len(self._markers_x_m_y_p_z_m) self._send_list_box[self._x_m_y_p_z_m_proc] = np.concatenate( (self._send_list_box[self._x_m_y_p_z_m_proc], self._markers_x_m_y_p_z_m) ) + # if self._x_m_y_p_z_p_proc is not None: self._send_info_box[self._x_m_y_p_z_p_proc] += len(self._markers_x_m_y_p_z_p) self._send_list_box[self._x_m_y_p_z_p_proc] = np.concatenate( (self._send_list_box[self._x_m_y_p_z_p_proc], self._markers_x_m_y_p_z_p) ) + # if self._x_p_y_m_z_m_proc is not None: self._send_info_box[self._x_p_y_m_z_m_proc] += len(self._markers_x_p_y_m_z_m) self._send_list_box[self._x_p_y_m_z_m_proc] = np.concatenate( (self._send_list_box[self._x_p_y_m_z_m_proc], self._markers_x_p_y_m_z_m) ) + # if self._x_p_y_m_z_p_proc is not None: self._send_info_box[self._x_p_y_m_z_p_proc] += len(self._markers_x_p_y_m_z_p) self._send_list_box[self._x_p_y_m_z_p_proc] = np.concatenate( (self._send_list_box[self._x_p_y_m_z_p_proc], self._markers_x_p_y_m_z_p) ) + # if self._x_p_y_p_z_m_proc is not None: self._send_info_box[self._x_p_y_p_z_m_proc] += len(self._markers_x_p_y_p_z_m) self._send_list_box[self._x_p_y_p_z_m_proc] = np.concatenate( (self._send_list_box[self._x_p_y_p_z_m_proc], self._markers_x_p_y_p_z_m) ) + # if self._x_p_y_p_z_p_proc is not None: self._send_info_box[self._x_p_y_p_z_p_proc] += len(self._markers_x_p_y_p_z_p) self._send_list_box[self._x_p_y_p_z_p_proc] = np.concatenate( (self._send_list_box[self._x_p_y_p_z_p_proc], self._markers_x_p_y_p_z_p) @@ -2821,13 +3155,13 @@ def self_communication_boxes(self): @profile def communicate_boxes(self, verbose=False): - if verbose: - n_valid = np.count_nonzero(self.valid_mks) - n_holes = np.count_nonzero(self.holes) - n_ghosts = np.count_nonzero(self.ghost_particles) - print(f"before communicate_boxes: {self.mpi_rank = }, {n_valid = } {n_holes = }, {n_ghosts = }") + # if verbose: + # n_valid = np.count_nonzero(self.valid_mks) + # n_holes = np.count_nonzero(self.holes) + # n_ghosts = np.count_nonzero(self.ghost_particles) + # print(f"before communicate_boxes: {self.mpi_rank = }, {n_valid = } {n_holes = }, {n_ghosts = }") - self.determine_send_markers_box() + self.prepare_ghost_particles() self.get_destinations_box() self.self_communication_boxes() self.update_holes() @@ -2838,11 +3172,11 @@ def communicate_boxes(self, verbose=False): self.update_holes() self.update_ghost_particles() - if verbose: - n_valid = np.count_nonzero(self.valid_mks) - n_holes = np.count_nonzero(self.holes) - n_ghosts = np.count_nonzero(self.ghost_particles) - print(f"after communicate_boxes: {self.mpi_rank = }, {n_valid = }, {n_holes = }, {n_ghosts = }") + # if verbose: + # n_valid = np.count_nonzero(self.valid_mks) + # n_holes = np.count_nonzero(self.holes) + # n_ghosts = np.count_nonzero(self.ghost_particles) + # print(f"after communicate_boxes: {self.mpi_rank = }, {n_valid = }, {n_holes = }, {n_ghosts = }") def sendrecv_all_to_all_boxes(self): """ @@ -2904,18 +3238,56 @@ def sendrecv_markers_boxes(self): self.mpi_comm.Barrier() def _get_neighbouring_proc(self): - """Find the neighbouring processes for the sending of boxes""" - dd = self.domain_array - periodic1, periodic2, periodic3 = ( - [True] * 3 - ) # for the moment we always assume periodicity for the evaluation near the boundary, TODO: fill ghost boxes with suitable markers for other bcs? + """Find the neighbouring processes for the sending of boxes. + + The left (right) neighbour in direction 1 is called x_m_proc (x_p_proc), etc. + By default every process is its own neighbour. + """ + # Faces + self._x_m_proc = None + self._x_p_proc = None + self._y_m_proc = None + self._y_p_proc = None + self._z_m_proc = None + self._z_p_proc = None + # Edges + self._x_m_y_m_proc = None + self._x_m_y_p_proc = None + self._x_p_y_m_proc = None + self._x_p_y_p_proc = None + self._x_m_z_m_proc = None + self._x_m_z_p_proc = None + self._x_p_z_m_proc = None + self._x_p_z_p_proc = None + self._y_m_z_m_proc = None + self._y_m_z_p_proc = None + self._y_p_z_m_proc = None + self._y_p_z_p_proc = None + # Corners + self._x_m_y_m_z_m_proc = None + self._x_m_y_m_z_p_proc = None + self._x_m_y_p_z_m_proc = None + self._x_p_y_m_z_m_proc = None + self._x_m_y_p_z_p_proc = None + self._x_p_y_m_z_p_proc = None + self._x_p_y_p_z_m_proc = None + self._x_p_y_p_z_p_proc = None + + # periodicitiy for distance computation + periodic1 = self.bc_sph[0] == "periodic" + periodic2 = self.bc_sph[1] == "periodic" + periodic3 = self.bc_sph[2] == "periodic" + # Determine which proc are on which side - x_l = dd[self.mpi_rank][0] - x_r = dd[self.mpi_rank][1] - y_l = dd[self.mpi_rank][3] - y_r = dd[self.mpi_rank][4] - z_l = dd[self.mpi_rank][6] - z_r = dd[self.mpi_rank][7] + dd = self.domain_array + rank = self.mpi_rank + + x_l = dd[rank][0] + x_r = dd[rank][1] + y_l = dd[rank][3] + y_r = dd[rank][4] + z_l = dd[rank][6] + z_r = dd[rank][7] for i in range(self.mpi_size): xl_i = dd[i][0] xr_i = dd[i][1] @@ -2924,244 +3296,360 @@ def _get_neighbouring_proc(self): zl_i = dd[i][6] zr_i = dd[i][7] + is_same_x_l = abs(distance(xl_i, x_l, periodic1)) < 1e-5 + is_same_x_r = abs(distance(xr_i, x_r, periodic1)) < 1e-5 + is_same_y_l = abs(distance(yl_i, y_l, periodic2)) < 1e-5 + is_same_y_r = abs(distance(yr_i, y_r, periodic2)) < 1e-5 + is_same_z_l = abs(distance(zl_i, z_l, periodic3)) < 1e-5 + is_same_z_r = abs(distance(zr_i, z_r, periodic3)) < 1e-5 + + is_neigh_x_l = abs(distance(xr_i, x_l, periodic1)) < 1e-5 + is_neigh_x_r = abs(distance(xl_i, x_r, periodic1)) < 1e-5 + is_neigh_y_l = abs(distance(yr_i, y_l, periodic2)) < 1e-5 + is_neigh_y_r = abs(distance(yl_i, y_r, periodic2)) < 1e-5 + is_neigh_z_l = abs(distance(zr_i, z_l, periodic3)) < 1e-5 + is_neigh_z_r = abs(distance(zl_i, z_r, periodic3)) < 1e-5 + # Faces # Process on the left (minus axis) in the x direction - if ( - abs(distance(yl_i, y_l, periodic2)) < 1e-5 - and abs(distance(yr_i, y_r, periodic2)) < 1e-5 - and abs(distance(zl_i, z_l, periodic3)) < 1e-5 - and abs(distance(zr_i, z_r, periodic3)) < 1e-5 - and abs(distance(xr_i, x_l, periodic1)) < 1e-5 - ): + if is_same_y_l and is_same_y_r and is_same_z_l and is_same_z_r and is_neigh_x_l: self._x_m_proc = i # Process on the right (plus axis) in the x direction - if ( - abs(distance(yl_i, y_l, periodic2)) < 1e-5 - and abs(distance(yr_i, y_r, periodic2)) < 1e-5 - and abs(distance(zl_i, z_l, periodic3)) < 1e-5 - and abs(distance(zr_i, z_r, periodic3)) < 1e-5 - and abs(distance(xl_i, x_r, periodic1)) < 1e-5 - ): + if is_same_y_l and is_same_y_r and is_same_z_l and is_same_z_r and is_neigh_x_r: self._x_p_proc = i # Process on the left (minus axis) in the y direction - if ( - abs(distance(xl_i, x_l, periodic1)) < 1e-5 - and abs(distance(xr_i, x_r, periodic1)) < 1e-5 - and abs(distance(zl_i, z_l, periodic3)) < 1e-5 - and abs(distance(zr_i, z_r, periodic3)) < 1e-5 - and abs(distance(yr_i, y_l, periodic2)) < 1e-5 - ): + if is_same_x_l and is_same_x_r and is_same_z_l and is_same_z_r and is_neigh_y_l: self._y_m_proc = i # Process on the right (plus axis) in the y direction - if ( - abs(distance(xl_i, x_l, periodic1)) < 1e-5 - and abs(distance(xr_i, x_r, periodic1)) < 1e-5 - and abs(distance(zl_i, z_l, periodic3)) < 1e-5 - and abs(distance(zr_i, z_r, periodic3)) < 1e-5 - and abs(distance(yl_i, y_r, periodic2)) < 1e-5 - ): + if is_same_x_l and is_same_x_r and is_same_z_l and is_same_z_r and is_neigh_y_r: self._y_p_proc = i # Process on the left (minus axis) in the z direction - if ( - abs(distance(xl_i, x_l, periodic1)) < 1e-5 - and abs(distance(xr_i, x_r, periodic1)) < 1e-5 - and abs(distance(yl_i, y_l, periodic2)) < 1e-5 - and abs(distance(yr_i, y_r, periodic2)) < 1e-5 - and abs(distance(zr_i, z_l, periodic3)) < 1e-5 - ): + if is_same_x_l and is_same_x_r and is_same_y_l and is_same_y_r and is_neigh_z_l: self._z_m_proc = i # Process on the right (plus axis) in the z direction - if ( - abs(distance(xl_i, x_l, periodic1)) < 1e-5 - and abs(distance(xr_i, x_r, periodic1)) < 1e-5 - and abs(distance(yl_i, y_l, periodic2)) < 1e-5 - and abs(distance(yr_i, y_r, periodic2)) < 1e-5 - and abs(distance(zl_i, z_r, periodic3)) < 1e-5 - ): + if is_same_x_l and is_same_x_r and is_same_y_l and is_same_y_r and is_neigh_z_r: self._z_p_proc = i # Edges # Process on the left in x and left in y axis - if ( - abs(distance(zl_i, z_l, periodic3)) < 1e-5 - and abs(distance(zr_i, z_r, periodic3)) < 1e-5 - and abs(distance(xr_i, x_l, periodic1)) < 1e-5 - and abs(distance(yr_i, y_l, periodic2)) < 1e-5 - ): + if is_same_z_l and is_same_z_r and is_neigh_x_l and is_neigh_y_l: self._x_m_y_m_proc = i # Process on the left in x and right in y axis - if ( - abs(distance(zl_i, z_l, periodic3)) < 1e-5 - and abs(distance(zr_i, z_r, periodic3)) < 1e-5 - and abs(distance(xr_i, x_l, periodic1)) < 1e-5 - and abs(distance(yl_i, y_r, periodic2)) < 1e-5 - ): + if is_same_z_l and is_same_z_r and is_neigh_x_l and is_neigh_y_r: self._x_m_y_p_proc = i # Process on the right in x and left in y axis - if ( - abs(distance(zl_i, z_l, periodic3)) < 1e-5 - and abs(distance(zr_i, z_r, periodic3)) < 1e-5 - and abs(distance(xl_i, x_r, periodic1)) < 1e-5 - and abs(distance(yr_i, y_l, periodic2)) < 1e-5 - ): + if is_same_z_l and is_same_z_r and is_neigh_x_r and is_neigh_y_l: self._x_p_y_m_proc = i # Process on the right in x and right in y axis - if ( - abs(distance(zl_i, z_l, periodic3)) < 1e-5 - and abs(distance(zr_i, z_r, periodic3)) < 1e-5 - and abs(distance(xl_i, x_r, periodic1)) < 1e-5 - and abs(distance(yl_i, y_r, periodic2)) < 1e-5 - ): + if is_same_z_l and is_same_z_r and is_neigh_x_r and is_neigh_y_r: self._x_p_y_p_proc = i # Process on the left in x and left in z axis - if ( - abs(distance(yl_i, y_l, periodic2)) < 1e-5 - and abs(distance(yr_i, y_r, periodic2)) < 1e-5 - and abs(distance(xr_i, x_l, periodic1)) < 1e-5 - and abs(distance(zr_i, z_l, periodic3)) < 1e-5 - ): + if is_same_y_l and is_same_y_r and is_neigh_x_l and is_neigh_z_l: self._x_m_z_m_proc = i # Process on the left in x and right in z axis - if ( - abs(distance(yl_i, y_l, periodic2)) < 1e-5 - and abs(distance(yr_i, y_r, periodic2)) < 1e-5 - and abs(distance(xr_i, x_l, periodic1)) < 1e-5 - and abs(distance(zl_i, z_r, periodic3)) < 1e-5 - ): + if is_same_y_l and is_same_y_r and is_neigh_x_l and is_neigh_z_r: self._x_m_z_p_proc = i # Process on the right in x and left in z axis - if ( - abs(distance(yl_i, y_l, periodic2)) < 1e-5 - and abs(distance(yr_i, y_r, periodic2)) < 1e-5 - and abs(distance(xl_i, x_r, periodic1)) < 1e-5 - and abs(distance(zr_i, z_l, periodic3)) < 1e-5 - ): + if is_same_y_l and is_same_y_r and is_neigh_x_r and is_neigh_z_l: self._x_p_z_m_proc = i # Process on the right in x and right in z axis - if ( - abs(distance(yl_i, y_l, periodic2)) < 1e-5 - and abs(distance(yr_i, y_r, periodic2)) < 1e-5 - and abs(distance(xl_i, x_r, periodic1)) < 1e-5 - and abs(distance(zl_i, z_r, periodic3)) < 1e-5 - ): + if is_same_y_l and is_same_y_r and is_neigh_x_r and is_neigh_z_r: self._x_p_z_p_proc = i # Process on the left in y and left in z axis - if ( - abs(distance(xl_i, x_l, periodic1)) < 1e-5 - and abs(distance(xr_i, x_r, periodic1)) < 1e-5 - and abs(distance(yr_i, y_l, periodic2)) < 1e-5 - and abs(distance(zr_i, z_l, periodic3)) < 1e-5 - ): + if is_same_x_l and is_same_x_r and is_neigh_y_l and is_neigh_z_l: self._y_m_z_m_proc = i # Process on the left in y and right in z axis - if ( - abs(distance(xl_i, x_l, periodic1)) < 1e-5 - and abs(distance(xr_i, x_r, periodic1)) < 1e-5 - and abs(distance(yr_i, y_l, periodic2)) < 1e-5 - and abs(distance(zl_i, z_r, periodic3)) < 1e-5 - ): + if is_same_x_l and is_same_x_r and is_neigh_y_l and is_neigh_z_r: self._y_m_z_p_proc = i # Process on the right in y and left in z axis - if ( - abs(distance(xl_i, x_l, periodic1)) < 1e-5 - and abs(distance(xr_i, x_r, periodic1)) < 1e-5 - and abs(distance(yl_i, y_r, periodic2)) < 1e-5 - and abs(distance(zr_i, z_l, periodic3)) < 1e-5 - ): + if is_same_x_l and is_same_x_r and is_neigh_y_r and is_neigh_z_l: self._y_p_z_m_proc = i # Process on the right in y and right in z axis - if ( - abs(distance(xl_i, x_l, periodic1)) < 1e-5 - and abs(distance(xr_i, x_r, periodic1)) < 1e-5 - and abs(distance(yl_i, y_r, periodic2)) < 1e-5 - and abs(distance(zl_i, z_r, periodic3)) < 1e-5 - ): + if is_same_x_l and is_same_x_r and is_neigh_y_r and is_neigh_z_r: self._y_p_z_p_proc = i # Corners # Process on the left in x, left in y and left in z axis - if ( - abs(distance(xr_i, x_l, periodic1)) < 1e-5 - and abs(distance(yr_i, y_l, periodic2)) < 1e-5 - and abs(distance(zr_i, z_l, periodic3)) < 1e-5 - ): + if is_neigh_x_l and is_neigh_y_l and is_neigh_z_l: self._x_m_y_m_z_m_proc = i # Process on the left in x, left in y and right in z axis - if ( - abs(distance(xr_i, x_l, periodic1)) < 1e-5 - and abs(distance(yr_i, y_l, periodic2)) < 1e-5 - and abs(distance(zl_i, z_r, periodic3)) < 1e-5 - ): + if is_neigh_x_l and is_neigh_y_l and is_neigh_z_r: self._x_m_y_m_z_p_proc = i # Process on the left in x, right in y and left in z axis - if ( - abs(distance(xr_i, x_l, periodic1)) < 1e-5 - and abs(distance(yl_i, y_r, periodic2)) < 1e-5 - and abs(distance(zr_i, z_l, periodic3)) < 1e-5 - ): + if is_neigh_x_l and is_neigh_y_r and is_neigh_z_l: self._x_m_y_p_z_m_proc = i # Process on the left in x, right in y and right in z axis - if ( - abs(distance(xr_i, x_l, periodic1)) < 1e-5 - and abs(distance(yl_i, y_r, periodic2)) < 1e-5 - and abs(distance(zl_i, z_r, periodic3)) < 1e-5 - ): + if is_neigh_x_l and is_neigh_y_r and is_neigh_z_r: self._x_m_y_p_z_p_proc = i # Process on the right in x, left in y and left in z axis - if ( - abs(distance(xl_i, x_r, periodic1)) < 1e-5 - and abs(distance(yr_i, y_l, periodic2)) < 1e-5 - and abs(distance(zr_i, z_l, periodic3)) < 1e-5 - ): + if is_neigh_x_r and is_neigh_y_l and is_neigh_z_l: self._x_p_y_m_z_m_proc = i # Process on the right in x, left in y and right in z axis - if ( - abs(distance(xl_i, x_r, periodic1)) < 1e-5 - and abs(distance(yr_i, y_l, periodic2)) < 1e-5 - and abs(distance(zl_i, z_r, periodic3)) < 1e-5 - ): + if is_neigh_x_r and is_neigh_y_l and is_neigh_z_r: self._x_p_y_m_z_p_proc = i # Process on the right in x, right in y and left in z axis - if ( - abs(distance(xl_i, x_r, periodic1)) < 1e-5 - and abs(distance(yl_i, y_r, periodic2)) < 1e-5 - and abs(distance(zr_i, z_l, periodic3)) < 1e-5 - ): + if is_neigh_x_r and is_neigh_y_r and is_neigh_z_l: self._x_p_y_p_z_m_proc = i # Process on the right in x, right in y and right in z axis - if ( - abs(distance(xl_i, x_r, periodic1)) < 1e-5 - and abs(distance(yl_i, y_r, periodic2)) < 1e-5 - and abs(distance(zl_i, z_r, periodic3)) < 1e-5 - ): + if is_neigh_x_r and is_neigh_y_r and is_neigh_z_r: self._x_p_y_p_z_p_proc = i + # set empty faces in x + if self._x_m_proc is None: + self._x_m_proc = rank + if self._x_p_proc is None: + self._x_p_proc = rank + + # set empty faces in y + if self._y_m_proc is None: + self._y_m_proc = rank + if self._y_p_proc is None: + self._y_p_proc = rank + + # set empty faces in z + if self._z_m_proc is None: + self._z_m_proc = rank + if self._z_p_proc is None: + self._z_p_proc = rank + + # set empty edges in xy + if self._x_m_y_m_proc is None: + if self._x_m_proc == rank: + self._x_m_y_m_proc = self._y_m_proc + elif self._y_m_proc == rank: + self._x_m_y_m_proc = self._x_m_proc + + if self._x_m_y_p_proc is None: + if self._x_m_proc == rank: + self._x_m_y_p_proc = self._y_p_proc + elif self._y_p_proc == rank: + self._x_m_y_p_proc = self._x_m_proc + + if self._x_p_y_m_proc is None: + if self._x_p_proc == rank: + self._x_p_y_m_proc = self._y_m_proc + elif self._y_m_proc == rank: + self._x_p_y_m_proc = self._x_p_proc + + if self._x_p_y_p_proc is None: + if self._x_p_proc == rank: + self._x_p_y_p_proc = self._y_p_proc + elif self._y_p_proc == rank: + self._x_p_y_p_proc = self._x_p_proc + + # set empty edges in xz + if self._x_m_z_m_proc is None: + if self._x_m_proc == rank: + self._x_m_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_m_z_m_proc = self._x_m_proc + + if self._x_m_z_p_proc is None: + if self._x_m_proc == rank: + self._x_m_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_m_z_p_proc = self._x_m_proc + + if self._x_p_z_m_proc is None: + if self._x_p_proc == rank: + self._x_p_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_p_z_m_proc = self._x_p_proc + + if self._x_p_z_p_proc is None: + if self._x_p_proc == rank: + self._x_p_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_p_z_p_proc = self._x_p_proc + + # set empty edges in yz + if self._y_m_z_m_proc is None: + if self._y_m_proc == rank: + self._y_m_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._y_m_z_m_proc = self._y_m_proc + + if self._y_m_z_p_proc is None: + if self._y_m_proc == rank: + self._y_m_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._y_m_z_p_proc = self._y_m_proc + + if self._y_p_z_m_proc is None: + if self._y_p_proc == rank: + self._y_p_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._y_p_z_m_proc = self._y_p_proc + + if self._y_p_z_p_proc is None: + if self._y_p_proc == rank: + self._y_p_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._y_p_z_p_proc = self._y_p_proc + + # set empty corners + if self._x_m_y_m_z_m_proc is None: + if self._x_m_proc == rank: + if self._y_m_proc == rank: + self._x_m_y_m_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_m_y_m_z_m_proc = self._y_m_proc + elif self._y_m_proc == rank: + if self._x_m_proc == rank: + self._x_m_y_m_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_m_y_m_z_m_proc = self._x_m_proc + elif self._z_m_proc == rank: + if self._x_m_proc == rank: + self._x_m_y_m_z_m_proc = self._y_m_proc + elif self._y_m_proc == rank: + self._x_m_y_m_z_m_proc = self._x_m_proc + + if self._x_m_y_m_z_p_proc is None: + if self._x_m_proc == rank: + if self._y_m_proc == rank: + self._x_m_y_m_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_m_y_m_z_p_proc = self._y_m_proc + elif self._y_m_proc == rank: + if self._x_m_proc == rank: + self._x_m_y_m_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_m_y_m_z_p_proc = self._x_m_proc + elif self._z_p_proc == rank: + if self._x_m_proc == rank: + self._x_m_y_m_z_p_proc = self._y_m_proc + elif self._y_m_proc == rank: + self._x_m_y_m_z_p_proc = self._x_m_proc + + if self._x_m_y_p_z_m_proc is None: + if self._x_m_proc == rank: + if self._y_p_proc == rank: + self._x_m_y_p_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_m_y_p_z_m_proc = self._y_p_proc + elif self._y_p_proc == rank: + if self._x_m_proc == rank: + self._x_m_y_p_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_m_y_p_z_m_proc = self._x_m_proc + elif self._z_m_proc == rank: + if self._x_m_proc == rank: + self._x_m_y_p_z_m_proc = self._y_p_proc + elif self._y_p_proc == rank: + self._x_m_y_p_z_m_proc = self._x_m_proc + + if self._x_m_y_p_z_p_proc is None: + if self._x_m_proc == rank: + if self._y_p_proc == rank: + self._x_m_y_p_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_m_y_p_z_p_proc = self._y_p_proc + elif self._y_p_proc == rank: + if self._x_m_proc == rank: + self._x_m_y_p_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_m_y_p_z_p_proc = self._x_m_proc + elif self._z_p_proc == rank: + if self._x_m_proc == rank: + self._x_m_y_p_z_p_proc = self._y_p_proc + elif self._y_p_proc == rank: + self._x_m_y_p_z_p_proc = self._x_m_proc + + if self._x_p_y_m_z_m_proc is None: + if self._x_p_proc == rank: + if self._y_m_proc == rank: + self._x_p_y_m_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_p_y_m_z_m_proc = self._y_m_proc + elif self._y_m_proc == rank: + if self._x_p_proc == rank: + self._x_p_y_m_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_p_y_m_z_m_proc = self._x_p_proc + elif self._z_m_proc == rank: + if self._x_p_proc == rank: + self._x_p_y_m_z_m_proc = self._y_m_proc + elif self._y_m_proc == rank: + self._x_p_y_m_z_m_proc = self._x_p_proc + + if self._x_p_y_m_z_p_proc is None: + if self._x_p_proc == rank: + if self._y_m_proc == rank: + self._x_p_y_m_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_p_y_m_z_p_proc = self._y_m_proc + elif self._y_m_proc == rank: + if self._x_p_proc == rank: + self._x_p_y_m_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_p_y_m_z_p_proc = self._x_p_proc + elif self._z_p_proc == rank: + if self._x_p_proc == rank: + self._x_p_y_m_z_p_proc = self._y_m_proc + elif self._y_m_proc == rank: + self._x_p_y_m_z_p_proc = self._x_p_proc + + if self._x_p_y_p_z_m_proc is None: + if self._x_p_proc == rank: + if self._y_p_proc == rank: + self._x_p_y_p_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_p_y_p_z_m_proc = self._y_p_proc + elif self._y_p_proc == rank: + if self._x_p_proc == rank: + self._x_p_y_p_z_m_proc = self._z_m_proc + elif self._z_m_proc == rank: + self._x_p_y_p_z_m_proc = self._x_p_proc + elif self._z_m_proc == rank: + if self._x_p_proc == rank: + self._x_p_y_p_z_m_proc = self._y_p_proc + elif self._y_p_proc == rank: + self._x_p_y_p_z_m_proc = self._x_p_proc + + if self._x_p_y_p_z_p_proc is None: + if self._x_p_proc == rank: + if self._y_p_proc == rank: + self._x_p_y_p_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_p_y_p_z_p_proc = self._y_p_proc + elif self._y_p_proc == rank: + if self._x_p_proc == rank: + self._x_p_y_p_z_p_proc = self._z_p_proc + elif self._z_p_proc == rank: + self._x_p_y_p_z_p_proc = self._x_p_proc + elif self._z_p_proc == rank: + if self._x_p_proc == rank: + self._x_p_y_p_z_p_proc = self._y_p_proc + elif self._y_p_proc == rank: + self._x_p_y_p_z_p_proc = self._x_p_proc + def eval_density( self, eta1, diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index ae12b5d5a..1252a530e 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -1,5 +1,7 @@ import copy +import numpy as np + from struphy.fields_background import equils from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB @@ -787,10 +789,10 @@ def __init__( kwargs["background"] = bckgr if "boxes_per_dim" not in kwargs: - boxes_per_dim = (1, 1, 1) + kwargs["boxes_per_dim"] = (1, 1, 1) else: if kwargs["boxes_per_dim"] is None: - boxes_per_dim = (1, 1, 1) + kwargs["boxes_per_dim"] = (1, 1, 1) # TODO: maybe this needs a fix # else: @@ -886,53 +888,3 @@ def s0(self, eta1, eta2, eta3, *v, flat_eval=False, remove_holes=True): kind="3_to_0", remove_outside=remove_holes, ) - - def _set_initial_condition(self): - """Set a callable initial condition f_init as a 0-form (scalar), and u_init in Cartesian coordinates.""" - - # Get the initialization function and pass the correct arguments - self._f_init = None - assert isinstance(self.f0, FluidEquilibrium) - self._f_init = self.f0.n0 - self._u_init = self.f0.u_cart - - if self.perturbations is not None: - for moment, pert in self.perturbations.items(): # only one perturbation is taken into account at the moment - assert isinstance(moment, str) - assert isinstance(pert, Perturbation) - - if moment == "n": - _fun = TransformedPformComponent( - pert, - pert.given_in_basis, - "0", - comp=pert.comp, - domain=self.domain, - ) - - def _f_init(*etas): - if len(etas) == 1: - return self.f0.n0(etas[0]) + _fun(*etas[0].T) - else: - assert len(etas) == 3 - E1, E2, E3, is_sparse_meshgrid = Domain.prepare_eval_pts( - etas[0], - etas[1], - etas[2], - flat_eval=False, - ) - return self.f0.n0(E1, E2, E3) + _fun(E1, E2, E3) - - self._f_init = _f_init - - elif moment == "u1": - _fun = TransformedPformComponent( - pert, - pert.given_in_basis, - "v", - comp=pert.comp, - domain=self.domain, - ) - _fun_cart = lambda e1, e2, e3: self.domain.push(_fun, e1, e2, e3, kind="v") - self._u_init = lambda e1, e2, e3: self.f0.u_cart(e1, e2, e3)[0] + _fun_cart(e1, e2, e3) - # TODO: add other velocity components diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index bb48607df..acaae39c2 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -6,6 +6,7 @@ from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments from struphy.pic.base import Particles +from struphy.profiling.profiling import ProfileManager class Pusher: @@ -279,13 +280,14 @@ def __call__(self, dt: float): ) # push markers - self.kernel( - dt, - stage, - self.particles.args_markers, - self._args_domain, - *self._args_kernel, - ) + with ProfileManager.profile_region("kernel: " + self.kernel.__name__): + self.kernel( + dt, + stage, + self.particles.args_markers, + self._args_domain, + *self._args_kernel, + ) self.particles.apply_kinetic_bc(newton=self._newton) self.particles.update_holes() diff --git a/src/struphy/pic/sorting_kernels.py b/src/struphy/pic/sorting_kernels.py index 1c1f2c394..46d84007f 100644 --- a/src/struphy/pic/sorting_kernels.py +++ b/src/struphy/pic/sorting_kernels.py @@ -8,12 +8,61 @@ def flatten_index( nx: "int", ny: "int", nz: "int", + algo_in: "str" = None, ): """Find the global index of a box based on its index in all 3 direction. At the moment this is simply sorted according to x then y then z but in the future more evolved indexing could be implemented. """ - return n1 + n2 * (nx + 2) + n3 * (nx + 2) * (ny + 2) + if algo_in is None: + algo = "fortran_ordering" + else: + algo = algo_in + + if algo == "fortran_ordering": + n_glob = n1 + n2 * (nx + 2) + n3 * (nx + 2) * (ny + 2) + elif algo == "c_ordering": + n_glob = n3 + n2 * (nz + 2) + n1 * (nz + 2) * (ny + 2) + else: + n_glob = -99 + print(algo, "is not implemented, n_glob set to -99 !!!") + + return n_glob + + +def unflatten_index( + n_glob: "int", + nx: "int", + ny: "int", + nz: "int", + algo_in: "str" = None, +): + """Find the multi-index (i, j, k) of a box based on its global flattened index. + At the moment this is simply sorted according to x then y then z but in the future + more evolved indexing could be implemented. + """ + if algo_in is None: + algo = "fortran_ordering" + else: + algo = algo_in + + if algo == "fortran_ordering": + n3 = n_glob // ((nx + 2) * (ny + 2)) + rest = n_glob - n3 * (nx + 2) * (ny + 2) + n2 = rest // (nx + 2) + n1 = rest - n2 * (nx + 2) + elif algo == "c_ordering": + n1 = n_glob // ((nz + 2) * (ny + 2)) + rest = n_glob - n1 * (nz + 2) * (ny + 2) + n2 = rest // (nz + 2) + n3 = rest - n2 * (nz + 2) + else: + n1 = -99 + n2 = -99 + n3 = -99 + print(algo, "is not implemented, n1, n2 and n3 set to -99 !!!") + + return n1, n2, n3 def initialize_neighbours( @@ -63,6 +112,14 @@ def find_box( domain_array : array Information of the domain on the current mpi process. """ + # offset if point is on left boundary + if eta1 == domain_array[0]: + eta1 += 1e-8 + if eta2 == domain_array[3]: + eta2 += 1e-8 + if eta3 == domain_array[6]: + eta3 += 1e-8 + # offset if point is on right boundary if eta1 == domain_array[1]: eta1 -= 1e-8 @@ -84,45 +141,8 @@ def find_box( n1 = int(floor((eta1 - x_l) / (x_r - x_l) * (nx + 2))) n2 = int(floor((eta2 - y_l) / (y_r - y_l) * (ny + 2))) n3 = int(floor((eta3 - z_l) / (z_r - z_l) * (nz + 2))) - return flatten_index(n1, n2, n3, nx, ny, nz) - -def put_particles_in_boxes_kernel( - markers: "float[:,:]", - holes: "bool[:]", - nx: "int", - ny: "int", - nz: "int", - boxes: "int[:,:]", - next_index: "int[:]", - domain_array: "float[:]", - box_index: "int" = -2, -): - """Assign the right box to all particles.""" - boxes[:, :] = -1 - next_index[:] = 0 - l = markers.shape[1] - for p in range(markers.shape[0]): - if holes[p]: - n_box = (nx + 2) * (ny + 2) * (nz + 2) - else: - a = find_box( - markers[p, 0], - markers[p, 1], - markers[p, 2], - nx, - ny, - nz, - domain_array, - ) - if a >= (nx + 2) * (ny + 2) * (nz + 2) or a < 0: - n_box = (nx + 2) * (ny + 2) * (nz + 2) - else: - n_box = a - boxes[n_box, next_index[n_box]] = p - next_index[n_box] += 1 - b = float(n_box) - markers[p, l + box_index] = b + return flatten_index(n1, n2, n3, nx, ny, nz) def get_next_index(box_nb: "float", next_index: "int[:]", cumul_next_index: "int[:]"): @@ -177,15 +197,46 @@ def sort_boxed_particles( loc_i += 1 -def reassign_boxes( +def assign_box_to_each_particle( + markers: "float[:,:]", + holes: "bool[:]", + nx: "int", + ny: "int", + nz: "int", + domain_array: "float[:]", + box_index: "int" = -2, +): + """Assign the right box to each particle, written into column -2 or marker array.""" + l = markers.shape[1] + for p in range(markers.shape[0]): + if holes[p]: + n_box = (nx + 2) * (ny + 2) * (nz + 2) + else: + a = find_box( + markers[p, 0], + markers[p, 1], + markers[p, 2], + nx, + ny, + nz, + domain_array, + ) + if a >= (nx + 2) * (ny + 2) * (nz + 2) or a < 0: + n_box = (nx + 2) * (ny + 2) * (nz + 2) + else: + n_box = a + b = float(n_box) + markers[p, l + box_index] = b + + +def assign_particles_to_boxes( markers: "float[:,:]", holes: "bool[:]", boxes: "int[:,:]", next_index: "int[:]", box_index: "int" = -2, ): - """Reloop over the particles after communication to update the neighbouring boxes - with the right particles and the next_index for later sorting.""" + """Write particle indices into box array.""" boxes[:, :] = -1 next_index[:] = 0 l = markers.shape[1] diff --git a/src/struphy/pic/sph_eval_kernels.py b/src/struphy/pic/sph_eval_kernels.py index e928f092d..37414f447 100644 --- a/src/struphy/pic/sph_eval_kernels.py +++ b/src/struphy/pic/sph_eval_kernels.py @@ -1,5 +1,3 @@ -from numpy import sqrt - import struphy.pic.sorting_kernels as sorting_kernels import struphy.pic.sph_smoothing_kernels as sph_smoothing_kernels from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments @@ -492,16 +490,22 @@ def box_based_evaluation_meshgrid( out[:] = 0.0 for i in range(n_eval_1): e1 = eta1[i, 0, 0] - if e1 < domain_array[0] or e1 > domain_array[1]: + + if e1 < domain_array[0] or e1 >= domain_array[1] and e1 != 1.0: continue + for j in range(n_eval_2): e2 = eta2[0, j, 0] - if e2 < domain_array[3] or e2 > domain_array[4]: + + if e2 < domain_array[3] or e2 >= domain_array[4] and e2 != 1.0: continue + for k in range(n_eval_3): e3 = eta3[0, 0, k] - if e3 < domain_array[6] or e3 > domain_array[7]: + + if e3 < domain_array[6] or e3 >= domain_array[7] and e3 != 1.0: continue + loc_box = sorting_kernels.find_box( e1, e2, diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index d4c050554..ce404d3e2 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -10,6 +10,26 @@ from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters +@pytest.mark.parametrize("nx", [8, 70]) +@pytest.mark.parametrize("ny", [16, 80]) +@pytest.mark.parametrize("nz", [32, 90]) +@pytest.mark.parametrize("algo", ["fortran_ordering", "c_ordering"]) +def test_flattening(nx, ny, nz, algo): + from struphy.pic.sorting_kernels import flatten_index, unflatten_index + + n1s = np.array(np.random.rand(10) * (nx + 1), dtype=int) + n2s = np.array(np.random.rand(10) * (ny + 1), dtype=int) + n3s = np.array(np.random.rand(10) * (nz + 1), dtype=int) + for n1 in n1s: + for n2 in n2s: + for n3 in n3s: + n_glob = flatten_index(int(n1), int(n2), int(n3), nx, ny, nz, algo) + n1n, n2n, n3n = unflatten_index(n_glob, nx, ny, nz, algo) + assert n1n == n1 + assert n2n == n2 + assert n3n == n3 + + @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @pytest.mark.parametrize("p", [[2, 3, 4]]) @@ -76,20 +96,21 @@ def test_sorting(Nel, p, spl_kind, mapping, Np, verbose=False): if __name__ == "__main__": - test_sorting( - [8, 9, 10], - [2, 3, 4], - [False, True, False], - [ - "Cuboid", - { - "l1": 1.0, - "r1": 2.0, - "l2": 10.0, - "r2": 20.0, - "l3": 100.0, - "r3": 200.0, - }, - ], - 1000000, - ) + test_flattening(8, 8, 8, "c_orderwding") + # test_sorting( + # [8, 9, 10], + # [2, 3, 4], + # [False, True, False], + # [ + # "Cuboid", + # { + # "l1": 1.0, + # "r1": 2.0, + # "l2": 10.0, + # "r2": 20.0, + # "l3": 100.0, + # "r3": 200.0, + # }, + # ], + # 1000000, + # ) diff --git a/src/struphy/pic/tests/test_sph.py b/src/struphy/pic/tests/test_sph.py index 73f8b4ec8..abb38aff6 100644 --- a/src/struphy/pic/tests/test_sph.py +++ b/src/struphy/pic/tests/test_sph.py @@ -1,10 +1,8 @@ -from time import time - import numpy as np import pytest +from matplotlib import pyplot as plt from mpi4py import MPI -from struphy.feec.psydac_derham import Derham from struphy.fields_background.equils import ConstantVelocity from struphy.geometry import domains from struphy.initial import perturbations @@ -12,10 +10,21 @@ from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters -@pytest.mark.mpi(min_size=2) -@pytest.mark.parametrize("Np", [40000, 46200]) -@pytest.mark.parametrize("bc_x", ["periodic", "reflect", "remove"]) -def test_evaluation_mc(Np, bc_x, show_plot=False): +@pytest.mark.parametrize("boxes_per_dim", [(24, 1, 1)]) +@pytest.mark.parametrize("kernel", ["trigonometric_1d", "gaussian_1d", "linear_1d"]) +@pytest.mark.parametrize("derivative", [0, 1]) +@pytest.mark.parametrize("bc_x", ["periodic", "mirror", "fixed"]) +@pytest.mark.parametrize("eval_pts", [11, 16]) +@pytest.mark.parametrize("tesselation", [False, True]) +def test_sph_evaluation_1d( + boxes_per_dim, + kernel, + derivative, + bc_x, + eval_pts, + tesselation, + show_plot=False, +): comm = MPI.COMM_WORLD # DOMAIN object @@ -24,16 +33,31 @@ def test_evaluation_mc(Np, bc_x, show_plot=False): domain_class = getattr(domains, dom_type) domain = domain_class(**dom_params) - boxes_per_dim = (16, 1, 1) - loading_params = LoadingParameters(Np=Np, seed=1607) - boundary_params = BoundaryParameters(bc=(bc_x, "periodic", "periodic")) - - background = ConstantVelocity(n=1.0, density_profile="constant") + if tesselation: + if kernel == "trigonometric_1d" and derivative == 1: + ppb = 100 + else: + ppb = 4 + loading_params = LoadingParameters(ppb=ppb, seed=1607, loading="tesselation") + else: + if derivative == 0: + ppb = 1000 + else: + ppb = 20000 + loading_params = LoadingParameters(ppb=ppb, seed=223) + + # background + background = ConstantVelocity(n=1.5, density_profile="constant") background.domain = domain - pert = {"n": perturbations.ModesSin(ls=(1,), amps=(1e-0,))} + pert = {"n": perturbations.ModesCos(ls=(1,), amps=(1e-0,))} + + if derivative == 0: + fun_exact = lambda e1, e2, e3: 1.5 + np.cos(2 * np.pi * e1) + else: + fun_exact = lambda e1, e2, e3: -2 * np.pi * np.sin(2 * np.pi * e1) - fun_exact = lambda e1, e2, e3: 1.0 + np.sin(2 * np.pi * e1) + boundary_params = BoundaryParameters(bc_sph=(bc_x, "periodic", "periodic")) particles = ParticlesSPH( comm_world=comm, @@ -47,107 +71,825 @@ def test_evaluation_mc(Np, bc_x, show_plot=False): n_as_volume_form=True, ) - particles.draw_markers(sort=False) + # eval points + eta1 = np.linspace(0, 1.0, eval_pts) + eta2 = np.array([0.0]) + eta3 = np.array([0.0]) + + particles.draw_markers(sort=False, verbose=False) particles.mpi_sort_markers() particles.initialize_weights() h1 = 1 / boxes_per_dim[0] h2 = 1 / boxes_per_dim[1] h3 = 1 / boxes_per_dim[2] - eta1 = np.linspace(0, 1.0, 100) # add offset for non-periodic boundary conditions, TODO: implement Neumann - eta2 = np.array([0.0]) - eta3 = np.array([0.0]) ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") - test_eval = particles.eval_density(ee1, ee2, ee3, h1=h1, h2=h2, h3=h3) + test_eval = particles.eval_density( + ee1, + ee2, + ee3, + h1=h1, + h2=h2, + h3=h3, + kernel_type=kernel, + derivative=derivative, + ) all_eval = np.zeros_like(test_eval) comm.Allreduce(test_eval, all_eval, op=MPI.SUM) - if show_plot and comm.Get_rank() == 0: - from matplotlib import pyplot as plt + exact_eval = fun_exact(ee1, ee2, ee3) + err_max_norm = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + + if comm.Get_rank() == 0: + print(f"\n{boxes_per_dim = }") + print(f"{kernel = }, {derivative =}") + print(f"{bc_x = }, {eval_pts = }, {tesselation = }, {err_max_norm = }") + if show_plot: + plt.figure(figsize=(12, 8)) + plt.plot(ee1.squeeze(), fun_exact(ee1, ee2, ee3).squeeze(), label="exact") + plt.plot(ee1.squeeze(), all_eval.squeeze(), "--.", label="eval_sph") + plt.xlabel("e1") + plt.legend() + plt.show() + + if tesselation: + if derivative == 0: + assert err_max_norm < 0.0081 + else: + assert err_max_norm < 0.027 + else: + if derivative == 0: + assert err_max_norm < 0.05 + else: + assert err_max_norm < 0.37 + + +@pytest.mark.parametrize("boxes_per_dim", [(12, 12, 1)]) +@pytest.mark.parametrize("kernel", ["trigonometric_2d", "gaussian_2d", "linear_2d"]) +@pytest.mark.parametrize("derivative", [0, 1, 2]) +@pytest.mark.parametrize("bc_x", ["periodic", "mirror", "fixed"]) +@pytest.mark.parametrize("bc_y", ["periodic", "mirror", "fixed"]) +@pytest.mark.parametrize("eval_pts", [11, 16]) +def test_sph_evaluation_2d( + boxes_per_dim, + kernel, + derivative, + bc_x, + bc_y, + eval_pts, + show_plot=False, +): + comm = MPI.COMM_WORLD - plt.figure(figsize=(12, 8)) - plt.plot(ee1.squeeze(), fun_exact(ee1, ee2, ee3).squeeze(), label="exact") - plt.plot(ee1.squeeze(), all_eval.squeeze(), "--.", label="eval_sph") - plt.legend() + tesselation = True - plt.show() + # DOMAIN object + dom_type = "Cuboid" + dom_params = {"l1": 1.0, "r1": 2.0, "l2": 0.0, "r2": 2.0, "l3": 100.0, "r3": 200.0} + domain_class = getattr(domains, dom_type) + domain = domain_class(**dom_params) + + if kernel == "trigonometric_2d" and derivative != 0: + ppb = 100 + else: + ppb = 16 + + loading_params = LoadingParameters(ppb=ppb, loading="tesselation") + + # background + background = ConstantVelocity(n=1.5, density_profile="constant") + background.domain = domain + + pert = {"n": perturbations.ModesCosCos(ls=(1,), ms=(1,), amps=(1e-0,))} + + if derivative == 0: + fun_exact = lambda e1, e2, e3: 1.5 + np.cos(2 * np.pi * e1) * np.cos(2 * np.pi * e2) + elif derivative == 1: + fun_exact = lambda e1, e2, e3: -2 * np.pi * np.sin(2 * np.pi * e1) * np.cos(2 * np.pi * e2) + else: + fun_exact = lambda e1, e2, e3: -2 * np.pi * np.cos(2 * np.pi * e1) * np.sin(2 * np.pi * e2) + + # boundary conditions + boundary_params = BoundaryParameters(bc_sph=(bc_x, bc_y, "periodic")) + + # eval points + eta1 = np.linspace(0, 1.0, eval_pts) + eta2 = np.linspace(0, 1.0, eval_pts) + eta3 = np.array([0.0]) + + # particles object + particles = ParticlesSPH( + comm_world=comm, + loading_params=loading_params, + boundary_params=boundary_params, + boxes_per_dim=boxes_per_dim, + bufsize=1.0, + domain=domain, + background=background, + perturbations=pert, + n_as_volume_form=True, + verbose=False, + ) - # print(f'{fun_exact(ee1, ee2, ee3) = }') - # print(f'{comm.Get_rank() = }, {all_eval = }') - # print(f'{np.max(np.abs(all_eval - fun_exact(ee1, ee2, ee3))) = }') - assert np.all(np.abs(all_eval - fun_exact(ee1, ee2, ee3)) < 0.065) + particles.draw_markers(sort=False, verbose=False) + particles.mpi_sort_markers() + particles.initialize_weights() + h1 = 1 / boxes_per_dim[0] + h2 = 1 / boxes_per_dim[1] + h3 = 1 / boxes_per_dim[2] + ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + test_eval = particles.eval_density( + ee1, + ee2, + ee3, + h1=h1, + h2=h2, + h3=h3, + kernel_type=kernel, + derivative=derivative, + ) + all_eval = np.zeros_like(test_eval) + comm.Allreduce(test_eval, all_eval, op=MPI.SUM) -@pytest.mark.mpi(min_size=2) -@pytest.mark.parametrize("boxes_per_dim", [(8, 1, 1), (10, 1, 1)]) -@pytest.mark.parametrize("ppb", [4, 9]) -@pytest.mark.parametrize("bc_x", ["periodic", "reflect", "remove"]) -def test_evaluation_tesselation(boxes_per_dim, ppb, bc_x, show_plot=False): + exact_eval = fun_exact(ee1, ee2, ee3) + err_max_norm = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + + if comm.Get_rank() == 0: + print(f"\n{boxes_per_dim = }") + print(f"{kernel = }, {derivative =}") + print(f"{bc_x = }, {bc_y = }, {eval_pts = }, {tesselation = }, {err_max_norm = }") + if show_plot: + plt.figure(figsize=(12, 24)) + plt.subplot(2, 1, 1) + plt.pcolor(ee1.squeeze(), ee2.squeeze(), fun_exact(ee1, ee2, ee3).squeeze()) + plt.title("exact") + plt.subplot(2, 1, 2) + plt.pcolor(ee1.squeeze(), ee2.squeeze(), all_eval.squeeze()) + plt.title("sph eval") + plt.xlabel("e1") + plt.xlabel("e2") + plt.show() + + if derivative == 0: + assert err_max_norm < 0.031 + else: + assert err_max_norm < 0.069 + + +@pytest.mark.parametrize("boxes_per_dim", [(12, 8, 8)]) +@pytest.mark.parametrize("kernel", ["trigonometric_3d", "gaussian_3d", "linear_3d", "linear_isotropic_3d"]) +@pytest.mark.parametrize("derivative", [0, 3]) +@pytest.mark.parametrize("bc_x", ["periodic"]) +@pytest.mark.parametrize("bc_y", ["periodic"]) +@pytest.mark.parametrize("bc_z", ["periodic", "mirror", "fixed"]) +@pytest.mark.parametrize("eval_pts", [11]) +def test_sph_evaluation_3d( + boxes_per_dim, + kernel, + derivative, + bc_x, + bc_y, + bc_z, + eval_pts, + show_plot=False, +): comm = MPI.COMM_WORLD + tesselation = True + # DOMAIN object dom_type = "Cuboid" - dom_params = {"l1": 1.0, "r1": 2.0, "l2": 10.0, "r2": 20.0, "l3": 100.0, "r3": 200.0} + dom_params = {"l1": 1.0, "r1": 2.0, "l2": 0.0, "r2": 2.0, "l3": -1.0, "r3": 2.0} domain_class = getattr(domains, dom_type) domain = domain_class(**dom_params) + if kernel in ("trigonometric_3d", "linear_isotropic_3d") and derivative != 0: + ppb = 100 + else: + ppb = 64 + loading_params = LoadingParameters(ppb=ppb, loading="tesselation") - boundary_params = BoundaryParameters(bc=(bc_x, "periodic", "periodic")) - background = ConstantVelocity(n=1.0, density_profile="constant") + # background + background = ConstantVelocity(n=1.5, density_profile="constant") background.domain = domain - pert = {"n": perturbations.ModesSin(ls=(1,), amps=(1e-0,))} + if derivative == 0: + fun_exact = lambda e1, e2, e3: 1.5 + 0.0 * e1 + else: + fun_exact = lambda e1, e2, e3: 0.0 * e1 + + # boundary conditions + boundary_params = BoundaryParameters(bc_sph=(bc_x, bc_y, bc_z)) - fun_exact = lambda e1, e2, e3: 1.0 + np.sin(2 * np.pi * e1) + # eval points + eta1 = np.linspace(0, 1.0, eval_pts) + eta2 = np.linspace(0, 1.0, eval_pts) + eta3 = np.linspace(0, 1.0, eval_pts) + # particles object particles = ParticlesSPH( comm_world=comm, loading_params=loading_params, boundary_params=boundary_params, boxes_per_dim=boxes_per_dim, - bufsize=1.0, + bufsize=2.0, domain=domain, background=background, - perturbations=pert, n_as_volume_form=True, - verbose=True, + verbose=False, ) - particles.draw_markers(sort=False) + particles.draw_markers(sort=False, verbose=False) particles.mpi_sort_markers() particles.initialize_weights() h1 = 1 / boxes_per_dim[0] h2 = 1 / boxes_per_dim[1] h3 = 1 / boxes_per_dim[2] - eta1 = np.linspace(0, 1.0, 10) # add offset for non-periodic boundary conditions, TODO: implement Neumann - eta2 = np.array([0.0]) - eta3 = np.array([0.0]) ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") - test_eval = particles.eval_density(ee1, ee2, ee3, h1=h1, h2=h2, h3=h3) + test_eval = particles.eval_density( + ee1, + ee2, + ee3, + h1=h1, + h2=h2, + h3=h3, + kernel_type=kernel, + derivative=derivative, + ) all_eval = np.zeros_like(test_eval) - # rank = comm.Get_rank() - # print(f'{rank = }, {test_eval.squeeze() = }') - comm.Allreduce(test_eval, all_eval, op=MPI.SUM) + exact_eval = fun_exact(ee1, ee2, ee3) + err_max_norm = np.max(np.abs(all_eval - exact_eval)) + + if comm.Get_rank() == 0: + print(f"\n{boxes_per_dim = }") + print(f"{kernel = }, {derivative =}") + print(f"{bc_x = }, {bc_y = }, {bc_z = }, {eval_pts = }, {tesselation = }, {err_max_norm = }") + if show_plot: + print(f"\n{fun_exact(ee1, ee2, ee3)[5, 5, 5] = }") + print(f"{ee1[5, 5, 5] = }, {ee2[5, 5, 5] = }, {ee3[5, 5, 5] = }") + print(f"{all_eval[5, 5, 5] = }") + + print(f"\n{ee1[4, 4, 4] = }, {ee2[4, 4, 4] = }, {ee3[4, 4, 4] = }") + print(f"{all_eval[4, 4, 4] = }") + + print(f"\n{ee1[3, 3, 3] = }, {ee2[3, 3, 3] = }, {ee3[3, 3, 3] = }") + print(f"{all_eval[3, 3, 3] = }") + + print(f"\n{ee1[2, 2, 2] = }, {ee2[2, 2, 2] = }, {ee3[2, 2, 2] = }") + print(f"{all_eval[2, 2, 2] = }") + + print(f"\n{ee1[1, 1, 1] = }, {ee2[1, 1, 1] = }, {ee3[1, 1, 1] = }") + print(f"{all_eval[1, 1, 1] = }") + + print(f"\n{ee1[0, 0, 0] = }, {ee2[0, 0, 0] = }, {ee3[0, 0, 0] = }") + print(f"{all_eval[0, 0, 0] = }") + # plt.figure(figsize=(12, 24)) + # plt.subplot(2, 1, 1) + # plt.pcolor(ee1[0, :, :], ee2[0, :, :], fun_exact(ee1, ee2, ee3)[0, :, :]) + # plt.title("exact") + # plt.subplot(2, 1, 2) + # plt.pcolor(ee1[0, :, :], ee2[0, :, :], all_eval[0, :, :]) + # plt.title("sph eval") + # plt.xlabel("e1") + # plt.xlabel("e2") + # plt.show() + + assert err_max_norm < 0.03 + + +@pytest.mark.parametrize("boxes_per_dim", [(12, 1, 1)]) +@pytest.mark.parametrize("bc_x", ["periodic", "mirror", "fixed"]) +@pytest.mark.parametrize("eval_pts", [11, 16]) +@pytest.mark.parametrize("tesselation", [False, True]) +def test_evaluation_SPH_Np_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselation, show_plot=False): + comm = MPI.COMM_WORLD + + # DOMAIN object + dom_type = "Cuboid" + dom_params = {"l1": 0.0, "r1": 3.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0} + domain_class = getattr(domains, dom_type) + domain = domain_class(**dom_params) + + if tesselation: + ppbs = [4, 8, 16, 32, 64] + Nps = [None] * len(ppbs) + else: + Nps = [(2**k) * 10**3 for k in range(-2, 9)] + ppbs = [None] * len(Nps) + + # background + background = ConstantVelocity(n=1.5, density_profile="constant") + background.domain = domain + + # perturbation]} + if bc_x in ("periodic", "fixed"): + fun_exact = lambda e1, e2, e3: 1.5 - np.sin(2 * np.pi * e1) + pert = {"n": perturbations.ModesSin(ls=(1,), amps=(-1e-0,))} + elif bc_x == "mirror": + fun_exact = lambda e1, e2, e3: 1.5 - np.cos(2 * np.pi * e1) + pert = {"n": perturbations.ModesCos(ls=(1,), amps=(-1e-0,))} + + # exact solution + eta1 = np.linspace(0, 1.0, eval_pts) # add offset for non-periodic boundary conditions, TODO: implement Neumann + eta2 = np.array([0.0]) + eta3 = np.array([0.0]) + ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + exact_eval = fun_exact(ee1, ee2, ee3) + + # boundary conditions + boundary_params = BoundaryParameters(bc_sph=(bc_x, "periodic", "periodic")) + + # loop + err_vec = [] + for Np, ppb in zip(Nps, ppbs): + if tesselation: + loading_params = LoadingParameters(ppb=ppb, loading="tesselation") + else: + loading_params = LoadingParameters(Np=Np, seed=1607) + + particles = ParticlesSPH( + comm_world=comm, + loading_params=loading_params, + boundary_params=boundary_params, + boxes_per_dim=boxes_per_dim, + bufsize=1.0, + domain=domain, + background=background, + perturbations=pert, + n_as_volume_form=True, + verbose=False, + ) + + particles.draw_markers(sort=False, verbose=False) + particles.mpi_sort_markers() + particles.initialize_weights() + h1 = 1 / boxes_per_dim[0] + h2 = 1 / boxes_per_dim[1] + h3 = 1 / boxes_per_dim[2] + + test_eval = particles.eval_density(ee1, ee2, ee3, h1=h1, h2=h2, h3=h3) + all_eval = np.zeros_like(test_eval) + + comm.Allreduce(test_eval, all_eval, op=MPI.SUM) + + if show_plot and comm.Get_rank() == 0: + plt.figure() + plt.plot(ee1.squeeze(), exact_eval.squeeze(), label="exact") + plt.plot(ee1.squeeze(), all_eval.squeeze(), "--.", label="eval_sph") + plt.title(f"{Np = }, {ppb = }") + # plt.savefig(f"fun_{Np}_{ppb}.png") + + diff = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + err_vec += [diff] + print(f"{Np = }, {ppb = }, {diff = }") + + if tesselation: + fit = np.polyfit(np.log(ppbs), np.log(err_vec), 1) + xvec = ppbs + else: + fit = np.polyfit(np.log(Nps), np.log(err_vec), 1) + xvec = Nps + if show_plot and comm.Get_rank() == 0: - from matplotlib import pyplot as plt + plt.figure(figsize=(12, 8)) + plt.loglog(xvec, err_vec, label="Convergence") + plt.loglog(xvec, np.exp(fit[1]) * np.array(xvec) ** (fit[0]), "--", label=f"fit with slope {fit[0]}") + plt.legend() + plt.show() + # plt.savefig(f"Convergence_SPH_{tesselation=}") + + if comm.Get_rank() == 0: + print(f"\n{bc_x = }, {eval_pts = }, {tesselation = }, {fit[0] = }") + + if tesselation: + assert fit[0] < 2e-3 + else: + assert np.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate + +@pytest.mark.parametrize("boxes_per_dim", [(12, 1, 1)]) +@pytest.mark.parametrize("bc_x", ["periodic", "fixed", "mirror"]) +@pytest.mark.parametrize("eval_pts", [11, 16]) +@pytest.mark.parametrize("tesselation", [False, True]) +def test_evaluation_SPH_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselation, show_plot=False): + comm = MPI.COMM_WORLD + + # DOMAIN object + dom_type = "Cuboid" + dom_params = {"l1": 0.0, "r1": 3.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0} + domain_class = getattr(domains, dom_type) + domain = domain_class(**dom_params) + + if tesselation: + Np = None + ppb = 160 + loading_params = LoadingParameters(ppb=ppb, loading="tesselation") + else: + Np = 160000 + ppb = None + loading_params = LoadingParameters(Np=Np, ppb=ppb, seed=1607) + + # background + background = ConstantVelocity(n=1.5, density_profile="constant") + background.domain = domain + + # perturbation + if bc_x in ("periodic", "fixed"): + fun_exact = lambda e1, e2, e3: 1.5 - np.sin(2 * np.pi * e1) + pert = {"n": perturbations.ModesSin(ls=(1,), amps=(-1e-0,))} + elif bc_x == "mirror": + fun_exact = lambda e1, e2, e3: 1.5 - np.cos(2 * np.pi * e1) + pert = {"n": perturbations.ModesCos(ls=(1,), amps=(-1e-0,))} + + # exact solution + eta1 = np.linspace(0, 1.0, eval_pts) # add offset for non-periodic boundary conditions, TODO: implement Neumann + eta2 = np.array([0.0]) + eta3 = np.array([0.0]) + ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + exact_eval = fun_exact(ee1, ee2, ee3) + + # boundary conditions + boundary_params = BoundaryParameters(bc_sph=(bc_x, "periodic", "periodic")) + + # loop + h_vec = [((2**k) * 10**-3 * 0.25) for k in range(2, 12)] + err_vec = [] + for h1 in h_vec: + particles = ParticlesSPH( + comm_world=comm, + loading_params=loading_params, + boundary_params=boundary_params, + boxes_per_dim=boxes_per_dim, + bufsize=1.0, + domain=domain, + background=background, + perturbations=pert, + n_as_volume_form=True, + verbose=False, + ) + + particles.draw_markers(sort=False, verbose=False) + particles.mpi_sort_markers() + particles.initialize_weights() + h2 = 1 / boxes_per_dim[1] + h3 = 1 / boxes_per_dim[2] + + test_eval = particles.eval_density(ee1, ee2, ee3, h1=h1, h2=h2, h3=h3) + all_eval = np.zeros_like(test_eval) + + comm.Allreduce(test_eval, all_eval, op=MPI.SUM) + + if show_plot and comm.Get_rank() == 0: + plt.figure() + plt.plot(ee1.squeeze(), exact_eval.squeeze(), label="exact") + plt.plot(ee1.squeeze(), all_eval.squeeze(), "--.", label="eval_sph") + plt.title(f"{h1 = }") + # plt.savefig(f"fun_{h1}.png") + + # error in max-norm + diff = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + + print(f"{h1 = }, {diff = }") + + if tesselation and h1 < 0.256: + assert diff < 0.036 + + err_vec += [diff] + + if tesselation: + fit = np.polyfit(np.log(h_vec[1:5]), np.log(err_vec[1:5]), 1) + else: + fit = np.polyfit(np.log(h_vec[:-2]), np.log(err_vec[:-2]), 1) + + if show_plot and comm.Get_rank() == 0: plt.figure(figsize=(12, 8)) - plt.plot(ee1.squeeze(), fun_exact(ee1, ee2, ee3).squeeze(), label="exact") - plt.plot(ee1.squeeze(), all_eval.squeeze(), "--.", label="eval_sph") + plt.loglog(h_vec, err_vec, label="Convergence") + plt.loglog(h_vec, np.exp(fit[1]) * np.array(h_vec) ** (fit[0]), "--", label=f"fit with slope {fit[0]}") plt.legend() + plt.show() + # plt.savefig("Convergence_SPH") + + if comm.Get_rank() == 0: + print(f"\n{bc_x = }, {eval_pts = }, {tesselation = }, {fit[0] = }") + if not tesselation: + assert np.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate + + +@pytest.mark.parametrize("boxes_per_dim", [(12, 1, 1)]) +@pytest.mark.parametrize("bc_x", ["periodic", "fixed", "mirror"]) +@pytest.mark.parametrize("eval_pts", [11, 16]) +@pytest.mark.parametrize("tesselation", [False, True]) +def test_evaluation_mc_Np_and_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselation, show_plot=False): + comm = MPI.COMM_WORLD + + # DOMAIN object + dom_type = "Cuboid" + dom_params = {"l1": 0.0, "r1": 3.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0} + domain_class = getattr(domains, dom_type) + domain = domain_class(**dom_params) + + if tesselation: + ppbs = [4, 8, 16, 32, 64] + Nps = [None] * len(ppbs) + else: + Nps = [(2**k) * 10**3 for k in range(-2, 9)] + ppbs = [None] * len(Nps) + + # background + background = ConstantVelocity(n=1.5, density_profile="constant") + background.domain = domain + + # perturbation + if bc_x in ("periodic", "fixed"): + fun_exact = lambda e1, e2, e3: 1.5 - np.sin(2 * np.pi * e1) + pert = {"n": perturbations.ModesSin(ls=(1,), amps=(-1e-0,))} + elif bc_x == "mirror": + fun_exact = lambda e1, e2, e3: 1.5 - np.cos(2 * np.pi * e1) + pert = {"n": perturbations.ModesCos(ls=(1,), amps=(-1e-0,))} + + # exact solution + eta1 = np.linspace(0, 1.0, eval_pts) + eta2 = np.array([0.0]) + eta3 = np.array([0.0]) + ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + exact_eval = fun_exact(ee1, ee2, ee3) + + # boundary conditions + boundary_params = BoundaryParameters(bc_sph=(bc_x, "periodic", "periodic")) + + h_arr = [((2**k) * 10**-3 * 0.25) for k in range(2, 12)] + err_vec = [] + for h in h_arr: + err_vec += [[]] + for Np, ppb in zip(Nps, ppbs): + if tesselation: + loading_params = LoadingParameters(ppb=ppb, loading="tesselation") + else: + loading_params = LoadingParameters(Np=Np, seed=1607) + + particles = ParticlesSPH( + comm_world=comm, + loading_params=loading_params, + boundary_params=boundary_params, + boxes_per_dim=boxes_per_dim, + bufsize=1.0, + domain=domain, + background=background, + perturbations=pert, + n_as_volume_form=True, + verbose=False, + ) + + particles.draw_markers(sort=False, verbose=False) + particles.mpi_sort_markers() + particles.initialize_weights() + + h2 = 1 / boxes_per_dim[1] + h3 = 1 / boxes_per_dim[2] + + test_eval = particles.eval_density(ee1, ee2, ee3, h1=h, h2=h2, h3=h3) + all_eval = np.zeros_like(test_eval) + comm.Allreduce(test_eval, all_eval, op=MPI.SUM) + + # error in max-norm + diff = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + err_vec[-1] += [diff] + + if comm.Get_rank() == 0: + print(f"{Np = }, {ppb = }, {diff = }") + # if show_plot: + # plt.figure() + # plt.plot(ee1.squeeze(), fun_exact(ee1, ee2, ee3).squeeze(), label="exact") + # plt.plot(ee1.squeeze(), all_eval.squeeze(), "--.", label="eval_sph") + # plt.title(f"{h = }, {Np = }") + # # plt.savefig(f"fun_h{h}_N{Np}_ppb{ppb}.png") + + err_vec = np.array(err_vec) + err_min = np.min(err_vec) + + if show_plot and comm.Get_rank() == 0: + if tesselation: + h_mesh, n_mesh = np.meshgrid(np.log10(h_arr), np.log10(ppbs), indexing="ij") + if not tesselation: + h_mesh, n_mesh = np.meshgrid(np.log10(h_arr), np.log10(Nps), indexing="ij") + plt.figure(figsize=(6, 6)) + plt.pcolor(h_mesh, n_mesh, np.log10(err_vec), shading="auto") + plt.title("Error") + plt.colorbar(label="log10(error)") + plt.xlabel("log10(h)") + plt.ylabel("log10(particles)") + + min_indices = np.argmin(err_vec, axis=0) + min_h_values = [] + for mi in min_indices: + min_h_values += [np.log10(h_arr[mi])] + if tesselation: + log_particles = np.log10(ppbs) + else: + log_particles = np.log10(Nps) + plt.plot(min_h_values, log_particles, "r-", label="Min error h for each Np", linewidth=2) + plt.legend() + # plt.savefig("SPH_conv_in_h_and_N.png") + + plt.show() + + if comm.Get_rank() == 0: + print(f"\n{tesselation = }, {bc_x = }, {err_min = }") + + if tesselation: + if bc_x == "periodic": + assert np.min(err_vec) < 7.7e-5 + elif bc_x == "fixed": + assert err_min < 7.7e-5 + else: + assert err_min < 7.7e-5 + else: + if bc_x in ("periodic", "fixed"): + assert err_min < 0.0089 + else: + assert err_min < 0.021 + + +@pytest.mark.parametrize("boxes_per_dim", [(24, 24, 1)]) +@pytest.mark.parametrize("bc_x", ["periodic", "fixed", "mirror"]) +@pytest.mark.parametrize("bc_y", ["periodic", "fixed", "mirror"]) +@pytest.mark.parametrize("tesselation", [False, True]) +def test_evaluation_SPH_Np_convergence_2d(boxes_per_dim, bc_x, bc_y, tesselation, show_plot=False): + comm = MPI.COMM_WORLD + + # DOMAIN object + dom_type = "Cuboid" + + Lx = 1.0 + Ly = 1.0 + dom_params = {"l1": 0.0, "r1": Lx, "l2": 0.0, "r2": Ly, "l3": 0.0, "r3": 1.0} + domain_class = getattr(domains, dom_type) + domain = domain_class(**dom_params) + + if tesselation: + ppbs = [4, 8, 16, 32, 64, 200] + Nps = [None] * len(ppbs) + else: + Nps = [(2**k) * 10**3 for k in range(-2, 9)] + ppbs = [None] * len(Nps) + + # background + background = ConstantVelocity(n=1.5, density_profile="constant") + background.domain = domain + + # perturbation + if bc_x in ("periodic", "fixed"): + if bc_y in ("periodic", "fixed"): + fun_exact = lambda x, y, z: 1.5 - np.sin(2 * np.pi / Lx * x) * np.sin(2 * np.pi / Ly * y) + pert = {"n": perturbations.ModesSinSin(ls=(1,), ms=(1,), amps=(-1e-0,))} + elif bc_y == "mirror": + fun_exact = lambda x, y, z: 1.5 - np.sin(2 * np.pi / Lx * x) * np.cos(2 * np.pi / Ly * y) + pert = {"n": perturbations.ModesSinCos(ls=(1,), ms=(1,), amps=(-1e-0,))} + + elif bc_x == "mirror": + if bc_y in ("periodic", "fixed"): + fun_exact = lambda x, y, z: 1.5 - np.cos(2 * np.pi / Lx * x) * np.sin(2 * np.pi / Ly * y) + pert = {"n": perturbations.ModesCosSin(ls=(1,), ms=(1,), amps=(-1e-0,))} + elif bc_y == "mirror": + fun_exact = lambda x, y, z: 1.5 - np.cos(2 * np.pi / Lx * x) * np.cos(2 * np.pi / Ly * y) + pert = {"n": perturbations.ModesCosCos(ls=(1,), ms=(1,), amps=(-1e-0,))} + + # exact solution + eta1 = np.linspace(0, 1.0, 41) + eta2 = np.linspace(0, 1.0, 86) + eta3 = np.array([0.0]) + ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + x, y, z = domain(eta1, eta2, eta3) + exact_eval = fun_exact(x, y, z) + + # boundary conditions + boundary_params = BoundaryParameters(bc_sph=(bc_x, bc_y, "periodic")) + + err_vec = [] + for Np, ppb in zip(Nps, ppbs): + if tesselation: + loading_params = LoadingParameters(ppb=ppb, loading="tesselation") + else: + loading_params = LoadingParameters(Np=Np, seed=1607) + + particles = ParticlesSPH( + comm_world=comm, + loading_params=loading_params, + boundary_params=boundary_params, + boxes_per_dim=boxes_per_dim, + bufsize=1.0, + box_bufsize=4.0, + domain=domain, + background=background, + perturbations=pert, + n_as_volume_form=True, + verbose=False, + ) + + if comm.Get_rank() == 0: + print(f"{particles.domain_array}") + + particles.draw_markers(sort=False, verbose=False) + particles.mpi_sort_markers() + particles.initialize_weights() + h1 = 1 / boxes_per_dim[0] + h2 = 1 / boxes_per_dim[1] + h3 = 1 / boxes_per_dim[2] + + test_eval = particles.eval_density(ee1, ee2, ee3, h1=h1, h2=h2, h3=h3, kernel_type="gaussian_2d") + all_eval = np.zeros_like(test_eval) + + comm.Allreduce(test_eval, all_eval, op=MPI.SUM) + + # error in max-norm + diff = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + err_vec += [diff] + + if tesselation: + assert diff < 0.06 + + if comm.Get_rank() == 0: + print(f"{Np = }, {ppb = }, {diff = }") + if show_plot: + fig, ax = plt.subplots() + d = ax.pcolor(ee1.squeeze(), ee2.squeeze(), all_eval.squeeze(), label="eval_sph", vmin=1.0, vmax=2.0) + fig.colorbar(d, ax=ax, label="2d_SPH") + ax.set_xlabel("ee1") + ax.set_ylabel("ee2") + ax.set_title(f"{Np}_{ppb = }") + # fig.savefig(f"2d_sph_{Np}_{ppb}.png") + + if tesselation: + fit = np.polyfit(np.log(ppbs), np.log(err_vec), 1) + xvec = ppbs + else: + fit = np.polyfit(np.log(Nps), np.log(err_vec), 1) + xvec = Nps + + if show_plot and comm.Get_rank() == 0: + plt.figure(figsize=(12, 8)) + plt.loglog(xvec, err_vec, label="Convergence") + plt.loglog(xvec, np.exp(fit[1]) * np.array(xvec) ** (fit[0]), "--", label=f"fit with slope {fit[0]}") + plt.legend() plt.show() + # plt.savefig(f"Convergence_SPH_{tesselation=}") + + if comm.Get_rank() == 0: + print(f"\n{bc_x = }, {tesselation = }, {fit[0] = }") - # print(f'{fun_exact(ee1, ee2, ee3) = }') - # print(f'{comm.Get_rank() = }, {all_eval = }') - # print(f'{np.max(np.abs(all_eval - fun_exact(ee1, ee2, ee3))) = }') - assert np.all(np.abs(all_eval - fun_exact(ee1, ee2, ee3)) < 0.017) + if not tesselation: + assert np.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate if __name__ == "__main__": - # test_evaluation_mc(40000, "periodic", show_plot=True) - test_evaluation_tesselation((8, 1, 1), 4, "periodic", show_plot=True) + test_sph_evaluation_1d( + (24, 1, 1), + "trigonometric_1d", + # "gaussian_1d", + 1, + # "periodic", + "mirror", + 16, + tesselation=False, + show_plot=True, + ) + + # test_sph_evaluation_2d( + # (12, 12, 1), + # # "trigonometric_2d", + # "gaussian_2d", + # 1, + # "periodic", + # "periodic", + # 16, + # show_plot=True + # ) + + # test_sph_evaluation_3d( + # (12, 8, 8), + # # "trigonometric_2d", + # "gaussian_3d", + # 2, + # "periodic", + # "periodic", + # "periodic", + # 11, + # show_plot=True + # ) + + # for nb in range(4, 25): + # print(f"\n{nb = }") + # test_evaluation_SPH_Np_convergence_1d((12,1,1), "fixed", eval_pts=16, tesselation=False, show_plot=True) + # test_evaluation_SPH_h_convergence_1d((12,1,1), "periodic", eval_pts=16, tesselation=True, show_plot=True) + # test_evaluation_mc_Np_and_h_convergence_1d((12,1,1),"mirror", eval_pts=16, tesselation = False, show_plot=True) + # test_evaluation_SPH_Np_convergence_2d((24, 24, 1), "periodic", "periodic", tesselation=True, show_plot=True) + # test_evaluation_SPH_Np_convergence_2d((24, 24, 1), "periodic", "fixed", tesselation=True, show_plot=True) + # test_evaluation_SPH_Np_convergence_2d((32, 32, 1), "fixed", "periodic", tesselation=True, show_plot=True) + # test_evaluation_SPH_Np_convergence_2d((32, 32, 1), "fixed", "fixed", tesselation=True, show_plot=True) + # test_evaluation_SPH_Np_convergence_2d((32, 32, 1), "mirror", "mirror", tesselation=True, show_plot=True) diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index e66ab8753..117f0dcfa 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -172,6 +172,26 @@ def __init__( self.ranges = ranges +class KernelDensityPlot: + """SPH density plot in configuration space. + + Parameters + ---------- + pts_e1, pts_e2, pts_e3 : int + Number of evaluation points in each direction. + """ + + def __init__( + self, + pts_e1: int = 11, + pts_e2: int = 11, + pts_e3: int = 11, + ): + self.pts_e1 = pts_e1 + self.pts_e2 = pts_e2 + self.pts_e3 = pts_e3 + + def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): """ This function is for getting kinetic energy of the case when canonical momentum is used, rather than velocity diff --git a/src/struphy/post_processing/likwid/plot_time_traces.py b/src/struphy/post_processing/likwid/plot_time_traces.py index f959a18a6..818d96935 100644 --- a/src/struphy/post_processing/likwid/plot_time_traces.py +++ b/src/struphy/post_processing/likwid/plot_time_traces.py @@ -1,16 +1,43 @@ import os import pickle +import re -import matplotlib import matplotlib.pyplot as plt import numpy as np +import plotly.io as pio + +# pio.kaleido.scope.mathjax = None +import struphy.post_processing.likwid.maxplotlylib as mply + + +def glob_to_regex(pat: str) -> str: + # Escape all regex metachars, then convert \* → .* and \? → . + esc = re.escape(pat) + return "^" + esc.replace(r"\*", ".*").replace(r"\?", ".") + "$" + + +def plot_region(region_name, groups_include=["*"], groups_skip=[]): + # skips first + for pat in groups_skip: + rx = glob_to_regex(pat) + if re.fullmatch(rx, region_name): + return False + + # includes next + for pat in groups_include: + rx = glob_to_regex(pat) + if re.fullmatch(rx, region_name): + return True + + return False def plot_time_vs_duration( - path, + paths, output_path, groups_include=["*"], groups_skip=[], + show_spans=False, ): """ Plot start times versus durations for all profiling regions from all MPI ranks, @@ -21,38 +48,43 @@ def plot_time_vs_duration( path: str Path to the file where profiling data is saved. """ - with open(path, "rb") as file: - profiling_data = pickle.load(file) - # Create a color map for each unique region - unique_regions = set(region_name for rank_data in profiling_data["rank_data"].values() for region_name in rank_data) - color_map = {region_name: plt.cm.tab10(i % 10) for i, region_name in enumerate(unique_regions)} + if isinstance(paths, str): + paths = [paths] + plt.figure(figsize=(10, 6)) + for path in paths: + print(f"{path = }") + with open(path, "rb") as file: + profiling_data = pickle.load(file) - # Iterate through each rank's data - for rank_name, rank_data in profiling_data["rank_data"].items(): - for region_name, info in rank_data.items(): - skip_region = False - for reg in ["main", "setup", "initialize"]: - if reg in region_name or reg == region_name: - skip_region = True - if skip_region: - continue - start_times = info["start_times"] - durations = info["durations"] + # Create a color map for each unique region + unique_regions = set( + region_name for rank_data in profiling_data["rank_data"].values() for region_name in rank_data + ) + color_map = {region_name: plt.cm.tab10(i % 10) for i, region_name in enumerate(unique_regions)} - # Use the color from the color_map for each region - color = color_map[region_name] - plt.plot( - start_times, durations, "x-", color=color, label=f"{region_name}" if rank_name == "rank_0" else None - ) - xmax = max(start_times) - x = 0 - while x < xmax: - xa = x + info["config"]["sample_duration"] - xb = xa + (info["config"]["sample_interval"] - info["config"]["sample_duration"]) - plt.axvspan(xa, xb, alpha=0.5, color="red", zorder=1) - x += info["config"]["sample_interval"] + # Iterate through each rank's data + for rank_name, rank_data in profiling_data["rank_data"].items(): + for region_name, info in rank_data.items(): + if plot_region(region_name=region_name, groups_include=groups_include, groups_skip=groups_skip): + start_times = info["start_times"] + durations = info["durations"] + + # Use the color from the color_map for each region + color = color_map[region_name] + label = f"{region_name}" if rank_name == "rank_0" else None + if len(paths) > 1: + label += f" ({path.split('/')[-2]})" + plt.plot(start_times, durations, "x-", color=color, label=label) + xmax = max(start_times) + x = 0 + if show_spans: + while x < xmax: + xa = x + info["config"]["sample_duration"] + xb = xa + (info["config"]["sample_interval"] - info["config"]["sample_duration"]) + plt.axvspan(xa, xb, alpha=0.5, color="red", zorder=1) + x += info["config"]["sample_interval"] plt.title("Time vs. Duration for Profiling Regions") plt.xlabel("Start Time (s)") @@ -66,28 +98,80 @@ def plot_time_vs_duration( print(f"Saved time trace to:{figure_path}") -def plot_gantt_chart( +def plot_avg_duration_bar_chart( path, output_path, groups_include=["*"], groups_skip=[], ): - """ - Plot Gantt chart of profiling regions from all MPI ranks using a grouped bar plot, - where bars are grouped by region and stacked for different ranks, with each rank having a specific color. - - Parameters - ---------- - path: str - Path to the file where profiling data is saved. - """ - # Load the profiling data from the specified path with open(path, "rb") as file: profiling_data = pickle.load(file) - plt.figure(figsize=(12, 8)) + region_durations = {} + + # Gather all durations per region across all ranks + for rank_data in profiling_data["rank_data"].values(): + for region_name, info in rank_data.items(): + if any(skip in region_name for skip in groups_skip): + continue + if groups_include != ["*"] and not any(inc in region_name for inc in groups_include): + continue + durations = info["durations"] + region_durations.setdefault(region_name, []).extend(durations) + + # Compute statistics per region + regions = sorted(region_durations.keys()) + avg_durations = [np.mean(region_durations[r]) for r in regions] + min_durations = [np.min(region_durations[r]) for r in regions] + max_durations = [np.max(region_durations[r]) for r in regions] + yerr = [ + [avg - min_ for avg, min_ in zip(avg_durations, min_durations)], + [max_ - avg for avg, max_ in zip(avg_durations, max_durations)], + ] + + # Plot bar chart with error bars (min-max spans) + plt.figure(figsize=(12, 6)) + x = np.arange(len(regions)) + plt.bar(x, avg_durations, yerr=yerr, capsize=5, color="skyblue", edgecolor="k") + plt.yscale("log") + plt.xticks(x, regions, rotation=45, ha="right") + plt.ylabel("Duration (s)") + plt.title("Average Duration per Profiling Region (with Min-Max Span)") + plt.grid(True, linestyle="--", alpha=0.5) + plt.tight_layout() + + # Save the figure + figure_path = os.path.join(output_path, "avg_duration_per_region.pdf") + plt.savefig(figure_path) + print(f"Saved average duration bar chart to: {figure_path}") + + +import plotly.graph_objects as go + + +def plot_region(region_name, groups_include=["*"], groups_skip=[]): + from fnmatch import fnmatch + + for pattern in groups_skip: + if fnmatch(region_name, pattern): + return False + for pattern in groups_include: + if fnmatch(region_name, pattern): + return True + return False + + +def plot_gantt_chart_plotly( + path: str, + output_path: str, + groups_include: list = ["*"], + groups_skip: list = [], + show: bool = False, +): + # print(f'Parsing {path}...') + with open(path, "rb") as file: + profiling_data = pickle.load(file) - # Collect unique region names and their earliest start times region_start_times = {} for rank_data in profiling_data["rank_data"].values(): for region_name, info in rank_data.items(): @@ -95,28 +179,16 @@ def plot_gantt_chart( if region_name not in region_start_times or first_start_time < region_start_times[region_name]: region_start_times[region_name] = first_start_time - # Sort region names by their earliest start time region_names = sorted(region_start_times, key=region_start_times.get) - - # Generate a color map for each rank rank_names = list(profiling_data["rank_data"].keys()) num_ranks = len(rank_names) - # color_map = matplotlib.cm.get_cmap('tab10') # Use 'tab10' colormap - color_map = plt.get_cmap("tab10") - # Parameters for spacing - bar_height = 0.1 # Height of each bar for a rank - rank_spacing = 0.1 # Vertical spacing between bars for different ranks - region_spacing = 0.5 # Vertical spacing between groups for different regions - - y_positions = [] # To store y positions for labels - region_labels = [] # To store labels for the y-axis - current_y = 0 # Initial y position - - # Iterate through each region in the sorted order + + bars = [] + for region_idx, region_name in enumerate(region_names): - start_y = current_y # Starting y position for the first rank in this region + if not plot_region(region_name, groups_include, groups_skip): + continue - # Plot bars for each rank for the current region for rank_idx, (rank_name, rank_data) in enumerate(profiling_data["rank_data"].items()): if region_name in rank_data: info = rank_data[region_name] @@ -124,43 +196,172 @@ def plot_gantt_chart( end_times = info["end_times"] durations = end_times - start_times - # Calculate the y position for this bar - y_position = current_y + rank_idx * rank_spacing - # Plot each call as a bar with a specific color for the rank - plt.barh( - y=y_position, - width=durations, - left=start_times, - color=color_map(rank_idx / num_ranks), - alpha=0.6, - height=bar_height, - label=f"{rank_name}" if region_idx == 0 else None, - ) - - # Calculate the middle y position for this region's label - middle_y = start_y + (rank_idx * rank_spacing) / 2 - y_positions.append(middle_y) - region_labels.append(region_name) - - # Move to the next y position for the next region, adding region spacing - current_y += (rank_idx + 1) * rank_spacing + region_spacing - - # Customize the plot - plt.xlim(left=-1) - plt.yticks(y_positions, region_labels) # Label the y-axis with region names - plt.xlabel("Elapsed time (s)") - plt.ylabel("Profiling Regions") - plt.title("Gantt chart of profiling regions (colored by MPI rank)") - if num_ranks < 10: - plt.legend(title="MPI Ranks", loc="upper left") # Add legend for MPI ranks - plt.grid(visible=True, linestyle="--", alpha=0.5, axis="x") # Grid only on x-axis - plt.tight_layout() - # plt.show() + for i in range(len(start_times)): + bars.append( + dict( + Task=region_name, + Rank=rank_name, + Start=start_times[i], + Finish=end_times[i], + Duration=durations[i], + ) + ) - # Save the plot as a PDF file - figure_path = os.path.join(output_path, "gantt_chart.pdf") - plt.savefig(figure_path) - print(f"Saved gantt chart to:{figure_path}") + if len(bars) == 0: + print("No regions matched the filter.") + return + + # Create a color map per rank + rank_color_map = {rank: f"hsl({360 * i / max(1, num_ranks)}, 70%, 50%)" for i, rank in enumerate(rank_names)} + + # Create plotly figure + fig = go.Figure() + for bar in bars: + fig.add_trace( + go.Bar( + x=[bar["Duration"]], + y=[bar["Task"]], + base=[bar["Start"]], + orientation="h", + name=bar["Rank"], + marker_color=rank_color_map[bar["Rank"]], + hovertemplate=f"Rank: {bar['Rank']}
Start: {bar['Start']:.3f}s
Duration: {bar['Duration']:.3f}s", + ) + ) + + fig.update_layout( + barmode="stack", + # title="Gantt Chart of Profiling Regions", + xaxis_title="Elapsed Time (s)", + yaxis_title="Profiling Regions", + height=600 + 20 * len(region_names), + # legend_title="MPI Ranks", + margin=dict(t=0, b=0, l=0, r=0), + showlegend=False, + ) + + mply.format_axes(fig) + mply.format_font(fig) + mply.format_grid(fig) + mply.format_size(fig, width=1600, height=800) + + # Save the plot as HTML + figure_path = os.path.join(output_path, "gantt_chart_plotly.html") + figure_path_pdf = os.path.join(output_path, "gantt_chart_plotly.pdf") + + if show: + fig.show() + fig.write_html(figure_path) + + # fig.write_image(figure_path_pdf) + print(f"Saved interactive gantt chart to: {figure_path}") + + +def plot_gantt_chart( + paths, + output_path, + groups_include=["*"], + groups_skip=[], + backend="matplotlib", +): + """ + Plot Gantt chart of profiling regions from all MPI ranks using a grouped bar plot, + where bars are grouped by region and stacked for different ranks, with each rank having a specific color. + + Parameters + ---------- + path: str + Path to the file where profiling data is saved. + """ + assert backend in ["matplotlib", "plotly"], "backend must be either matplotlib or plotly" + + if isinstance(paths, str): + paths = [paths] + + for path in paths: + # Load the profiling data from the specified path + with open(path, "rb") as file: + profiling_data = pickle.load(file) + + plt.figure(figsize=(12, 8)) + # color_map = matplotlib.cm.get_cmap('tab10') # Use 'tab10' colormap + color_map = plt.get_cmap("tab10") + + # Collect unique region names and their earliest start times + region_start_times = {} + for rank_data in profiling_data["rank_data"].values(): + for region_name, info in rank_data.items(): + first_start_time = np.min(info["start_times"]) + if region_name not in region_start_times or first_start_time < region_start_times[region_name]: + region_start_times[region_name] = first_start_time + + # Sort region names by their earliest start time + region_names = sorted(region_start_times, key=region_start_times.get) + + # Generate a color map for each rank + rank_names = list(profiling_data["rank_data"].keys()) + num_ranks = len(rank_names) + + # Parameters for spacing + bar_height = 0.1 # Height of each bar for a rank + rank_spacing = 0.1 # Vertical spacing between bars for different ranks + region_spacing = 0.5 # Vertical spacing between groups for different regions + + y_positions = [] # To store y positions for labels + region_labels = [] # To store labels for the y-axis + current_y = 0 # Initial y position + + # Iterate through each region in the sorted order + for region_idx, region_name in enumerate(region_names): + if not plot_region(region_name=region_name, groups_include=groups_include, groups_skip=groups_skip): + continue + start_y = current_y # Starting y position for the first rank in this region + + # Plot bars for each rank for the current region + for rank_idx, (rank_name, rank_data) in enumerate(profiling_data["rank_data"].items()): + if region_name in rank_data: + info = rank_data[region_name] + start_times = info["start_times"] + end_times = info["end_times"] + durations = end_times - start_times + + # Calculate the y position for this bar + y_position = current_y + rank_idx * rank_spacing + # Plot each call as a bar with a specific color for the rank + plt.barh( + y=y_position, + width=durations, + left=start_times, + color=color_map(rank_idx / num_ranks), + alpha=0.6, + height=bar_height, + label=f"{rank_name}" if region_idx == 0 else None, + ) + + # Calculate the middle y position for this region's label + middle_y = start_y + (rank_idx * rank_spacing) / 2 + y_positions.append(middle_y) + region_labels.append(region_name) + + # Move to the next y position for the next region, adding region spacing + current_y += (rank_idx + 1) * rank_spacing + region_spacing + + # Customize the plot + plt.xlim(left=-1) + plt.yticks(y_positions, region_labels) # Label the y-axis with region names + plt.xlabel("Elapsed time (s)") + plt.ylabel("Profiling Regions") + # plt.title("Gantt chart of profiling regions") + if num_ranks < 10: + plt.legend(title="MPI Ranks", loc="upper left") # Add legend for MPI ranks + plt.grid(visible=True, linestyle="--", alpha=0.5, axis="x") # Grid only on x-axis + plt.tight_layout() + # plt.show() + + # Save the plot as a PDF file + figure_path = os.path.join(output_path, "gantt_chart.pdf") + plt.savefig(figure_path) + print(f"Saved gantt chart to:{figure_path}") if __name__ == "__main__": @@ -174,16 +375,40 @@ def plot_gantt_chart( # Parse command-line arguments parser = argparse.ArgumentParser(description="Plot profiling time trace.") + # parser.add_argument( + # "--path", + # type=str, + # default=os.path.join(o_path, "sim_1", "profiling_time_trace.pkl"), + # help="Path to the profiling data file (default: o_path from struphy state)", + # ) + + parser.add_argument( + "--simulations", + nargs="+", + default=["sim_1"], + help="One or more simulations to compare (e.g. --simulations sim_1 sim_2)", + ) + + parser.add_argument( + "--groups", + nargs="+", + default=["*"], + help="One or more groups to include (e.g. --groups model.integrate PushEta PushVxB)", + ) + parser.add_argument( - "--path", - type=str, - default=os.path.join(o_path, "sim_1", "profiling_time_trace.pkl"), - help="Path to the profiling data file (default: o_path from struphy state)", + "--groups-skip", + nargs="+", + default=[], + help="One or more groups to skip (e.g. --groups-skip model.integrate PushEta PushVxB)", ) args = parser.parse_args() - path = os.path.abspath(args.path) # Convert to absolute path + # path = os.path.abspath(args.path) # Convert to absolute path + # simulations = parser.simulations + + paths = [os.path.join(o_path, simulation, "profiling_time_trace.pkl") for simulation in args.simulations] # Plot the time trace - plot_time_vs_duration(path, output_path=o_path) - plot_gantt_chart(path, output_path=o_path) + plot_time_vs_duration(paths=paths, output_path=o_path, groups_include=args.groups, groups_skip=args.groups_skip) + plot_gantt_chart(paths=paths, output_path=o_path, groups_include=args.groups, groups_skip=args.groups_skip) diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 39ac730ac..0c7893acb 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -17,7 +17,7 @@ from struphy.kinetic_background import maxwellians from struphy.kinetic_background.base import KineticBackground from struphy.models.base import StruphyModel, setup_derham -from struphy.models.species import KineticSpecies +from struphy.models.species import ParticleSpecies from struphy.models.variables import PICVariable from struphy.topology.grids import TensorProductGrid @@ -648,7 +648,7 @@ def post_process_f( Whether to compute the kinetic background values and add them to the binning data. This is used if non-standard weights are binned. """ - # get # of MPI processes from meta.txt file + # get # of MPI processes from meta file with open(os.path.join(path_in, "meta.yml"), "r") as f: meta = yaml.load(f, Loader=yaml.FullLoader) nproc = meta["MPI processes"] @@ -735,7 +735,7 @@ def post_process_f( # maxw_params=maxw_params, # ) - spec: KineticSpecies = getattr(params_in.model, species) + spec: ParticleSpecies = getattr(params_in.model, species) var: PICVariable = spec.var f_bckgr: KineticBackground = var.backgrounds @@ -800,7 +800,13 @@ def post_process_f( file.close() -def post_process_n_sph(path_in, path_out, species, step=1, compute_bckgr=False): +def post_process_n_sph( + path_in, + params_in: ParamsIn, + path_out, + species, + step=1, +): """Computes and saves the density n of saved sph data during a simulation. Parameters @@ -808,6 +814,9 @@ def post_process_n_sph(path_in, path_out, species, step=1, compute_bckgr=False): path_in : str Absolute path of simulation output folder. + params_in : ParamsIn + Simulation parameters. + path_out : str Absolute path of where to store the .txt files. Will be in path_out/orbits. @@ -816,21 +825,11 @@ def post_process_n_sph(path_in, path_out, species, step=1, compute_bckgr=False): step : int, optional Whether to do post-processing at every time step (step=1, default), every second time step (step=2), etc. - - compute_bckgr : bool - Whehter to compute the kinetic background values and add them to the binning data. - This is used if non-standard weights are binned. """ - - # get model name and # of MPI processes from meta.txt file - with open(os.path.join(path_in, "meta.txt"), "r") as f: - lines = f.readlines() - - nproc = lines[4].split()[-1] - - # load parameters - with open(os.path.join(path_in, "parameters.yml"), "r") as f: - params = yaml.load(f, Loader=yaml.FullLoader) + # get model name and # of MPI processes from meta file + with open(os.path.join(path_in, "meta.yml"), "r") as f: + meta = yaml.load(f, Loader=yaml.FullLoader) + nproc = meta["MPI processes"] # open hdf5 files files = [ diff --git a/src/struphy/post_processing/pproc_struphy.py b/src/struphy/post_processing/pproc_struphy.py index 06f64c37f..8d26c8c8d 100644 --- a/src/struphy/post_processing/pproc_struphy.py +++ b/src/struphy/post_processing/pproc_struphy.py @@ -20,7 +20,6 @@ def main( guiding_center: bool = False, classify: bool = False, no_vtk: bool = False, - time_trace: bool = False, ): """Post-processing of finished Struphy runs. @@ -61,14 +60,6 @@ def main( shutil.rmtree(path_pproc) os.mkdir(path_pproc) - if time_trace: - from struphy.post_processing.likwid.plot_time_traces import plot_gantt_chart, plot_time_vs_duration - - path_time_trace = os.path.join(path, "profiling_time_trace.pkl") - plot_time_vs_duration(path_time_trace, output_path=path_pproc) - plot_gantt_chart(path_time_trace, output_path=path_pproc) - return - # check for fields and kinetic data in hdf5 file that need post processing file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") @@ -239,8 +230,6 @@ def main( parser.add_argument("--no-vtk", help="whether vtk files creation should be skipped", action="store_true") - parser.add_argument("--time-trace", help="whether to plot the time trace", action="store_true") - args = parser.parse_args() main( @@ -251,5 +240,4 @@ def main( guiding_center=args.guiding_center, classify=args.classify, no_vtk=args.no_vtk, - time_trace=args.time_trace, ) diff --git a/src/struphy/profiling/profiling.py b/src/struphy/profiling/profiling.py index b84da9a9d..160002c15 100644 --- a/src/struphy/profiling/profiling.py +++ b/src/struphy/profiling/profiling.py @@ -163,7 +163,7 @@ def save_to_pickle(cls, file_path): comm = MPI.COMM_WORLD rank = comm.Get_rank() - size = comm.Get_size() + # size = comm.Get_size() # Prepare the data to be gathered local_data = {} @@ -209,7 +209,10 @@ def save_to_pickle(cls, file_path): if rank == 0: # Combine the data from all processes - combined_data = {"config": None, "rank_data": {f"rank_{i}": data for i, data in enumerate(all_data)}} + combined_data = { + "config": None, + "rank_data": {f"rank_{i}": data for i, data in enumerate(all_data)}, + } # Add the likwid data if likwid_data: @@ -267,22 +270,28 @@ def __init__(self, region_name, time_trace=False): self._region_name = self.config.simulation_label + region_name self._time_trace = time_trace self._ncalls = 0 - self._start_times = [] - self._end_times = [] - self._durations = [] + self._start_times = np.empty(1, dtype=float) + self._end_times = np.empty(1, dtype=float) + self._durations = np.empty(1, dtype=float) self._started = False def __enter__(self): + if self._ncalls == len(self._start_times): + self._start_times = np.append(self._start_times, np.zeros_like(self._start_times)) + self._end_times = np.append(self._end_times, np.zeros_like(self._end_times)) + self._durations = np.append(self._durations, np.zeros_like(self._durations)) + if self.config.likwid: self._pylikwid().markerstartregion(self.region_name) - self._ncalls += 1 - if self._time_trace: self._start_time = MPI.Wtime() - if self._start_time % self.config.sample_interval < self.config.sample_duration or self._ncalls == 1: - self._start_times.append(self._start_time) + if self._start_time % self.config.sample_interval < self.config.sample_duration or self._ncalls == 0: + self._start_times[self._ncalls] = self._start_time self._started = True + + self._ncalls += 1 + return self def __exit__(self, exc_type, exc_value, traceback): @@ -290,8 +299,8 @@ def __exit__(self, exc_type, exc_value, traceback): self._pylikwid().markerstopregion(self.region_name) if self._time_trace and self.started: end_time = MPI.Wtime() - self._end_times.append(end_time) - self._durations.append(end_time - self._start_time) + self._end_times[self._ncalls - 1] = end_time + self._durations[self._ncalls - 1] = end_time - self._start_time self._started = False def _pylikwid(self): diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 71a36358e..a41170a34 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -9,9 +9,11 @@ import numpy as np import scipy as sc from line_profiler import profile +from matplotlib import pyplot as plt from mpi4py import MPI from numpy import zeros -from psydac.linalg.basic import IdentityOperator, ZeroOperator +from psydac.api.essential_bc import apply_essential_bc_stencil +from psydac.linalg.basic import ComposedLinearOperator, IdentityOperator, ZeroOperator from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace from psydac.linalg.solvers import inverse from psydac.linalg.stencil import StencilVector @@ -19,8 +21,14 @@ import struphy.feec.utilities as util from struphy.examples.restelli2018 import callables from struphy.feec import preconditioner -from struphy.feec.basis_projection_ops import BasisProjectionOperator, BasisProjectionOperatorLocal, CoordinateProjector -from struphy.feec.mass import WeightedMassOperator +from struphy.feec.basis_projection_ops import ( + BasisProjectionOperator, + BasisProjectionOperatorLocal, + BasisProjectionOperators, + CoordinateProjector, +) +from struphy.feec.linear_operators import BoundaryOperator +from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators from struphy.feec.preconditioner import MassMatrixPreconditioner from struphy.feec.projectors import L2Projector from struphy.feec.psydac_derham import Derham, SplineFunction @@ -7572,6 +7580,7 @@ def options(default=False): dct["method_to_solve"] = "DirectNPInverse" dct["preconditioner"] = False dct["spectralanalysis"] = False + dct["lifting"] = False dct["dimension"] = "2D" dct["1D_dt"] = 0.001 if default: @@ -7600,6 +7609,7 @@ def __init__( method_to_solve: str = options(default=True)["method_to_solve"], preconditioner: bool = options(default=True)["preconditioner"], spectralanalysis: bool = options(default=True)["spectralanalysis"], + lifting: bool = options(default=False)["lifting"], dimension: str = options(default=True)["dimension"], D1_dt: float = options(default=True)["1D_dt"], ): @@ -7626,35 +7636,46 @@ def __init__( self._preconditioner = preconditioner self._dimension = dimension self._spectralanalysis = spectralanalysis + self._lifting = lifting - if self._variant == "GMRES": - self._M2 = getattr(self.mass_ops, "M2") - self._M3 = getattr(self.mass_ops, "M3") - self._M2B = -getattr(self.mass_ops, "M2B") - # Define block matrix [[A BT], [B 0]] (without time step size dt in the diagonals) - _A11 = ( - self._M2 - - self._M2B / self._eps_norm - + self._nu - * ( - self.derham.div.T @ self._M3 @ self.derham.div - + self.basis_ops.S21.T @ self.derham.curl.T @ self._M2 @ self.derham.curl @ self.basis_ops.S21 - ) + # Lifting for nontrivial boundary conditions + # derham had boundary conditions in eta1 direction, the following is in space Hdiv_0 + if self._lifting: + self.derhamv0 = Derham( + self.derham.Nel, + self.derham.p, + self.derham.spl_kind, + domain=self.domain, + dirichlet_bc=[[True, True], [False, False], [False, False]], ) - _A12 = None - _A21 = _A12 - _A22 = ( - -self._stab_sigma * IdentityOperator(_A11.domain) - + self._M2B / self._eps_norm - + self._nu_e - * ( - self.derham.div.T @ self._M3 @ self.derham.div - + self.basis_ops.S21.T @ self.derham.curl.T @ self._M2 @ self.derham.curl @ self.basis_ops.S21 - ) + + self._mass_opsv0 = WeightedMassOperators( + self.derhamv0, + self.domain, + verbose=solver["verbose"], + eq_mhd=self.mass_ops.weights["eq_mhd"], + ) + self._basis_opsv0 = BasisProjectionOperators( + self.derhamv0, + self.domain, + verbose=solver["verbose"], + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) + else: + self.derhamnumpy = Derham( + self.derham.Nel, + self.derham.p, + self.derham.spl_kind, + domain=self.domain, + # dirichlet_bc=self.derham.dirichlet_bc, + # nquads = self.derham._nquads, + # nq_pr = self.derham._nq_pr, + # comm = MPI.COMM_SELF, # self.derham._comm, + # polar_ck= self.derham._polar_ck, + # local_projectors=self.derham.with_local_projectors ) - _B1 = -self._M3 @ self.derham.div - _B2 = self._M3 @ self.derham.div + # get forceterms for according dimension if self._dimension in ["2D", "1D"]: ### Manufactured solution ### _forceterm_logical = lambda e1, e2, e3: 0 * e1 @@ -7664,7 +7685,8 @@ def __init__( b0=self._B0, nu=self._nu, dimension=self._dimension, - epsilon=self._stab_sigma, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, dt=D1_dt, ) _funy = getattr(callables, "ManufacturedSolutionForceterm")( @@ -7673,7 +7695,8 @@ def __init__( b0=self._B0, nu=self._nu, dimension=self._dimension, - epsilon=self._stab_sigma, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, dt=D1_dt, ) _funelectronsx = getattr(callables, "ManufacturedSolutionForceterm")( @@ -7682,7 +7705,8 @@ def __init__( b0=self._B0, nu_e=self._nu_e, dimension=self._dimension, - epsilon=self._stab_sigma, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, dt=D1_dt, ) _funelectronsy = getattr(callables, "ManufacturedSolutionForceterm")( @@ -7691,7 +7715,8 @@ def __init__( b0=self._B0, nu_e=self._nu_e, dimension=self._dimension, - epsilon=self._stab_sigma, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, dt=D1_dt, ) @@ -7713,8 +7738,8 @@ def __init__( forcetermelectrons_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain ) l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - self._F1 = l2_proj.get_dofs([funx, funy, _forceterm_logical]) - self._F2 = l2_proj.get_dofs([fun_electronsx, fun_electronsy, _forceterm_logical]) + self._F1 = l2_proj([funx, funy, _forceterm_logical]) + self._F2 = l2_proj([fun_electronsx, fun_electronsy, _forceterm_logical]) elif self._dimension == "Restelli": ### Restelli ### @@ -7727,6 +7752,7 @@ def __init__( Bp=self._Bp, alpha=self._alpha, beta=self._beta, + eps=self._eps_norm, ) _funelectrons = getattr(callables, "RestelliForcingTerm")( B0=self._B0, @@ -7735,6 +7761,7 @@ def __init__( Bp=self._Bp, alpha=self._alpha, beta=self._beta, + eps=self._eps_norm, ) # get callable(s) for specified init type @@ -7742,19 +7769,188 @@ def __init__( forcetermelectrons_class = [_forceterm_logical, _forceterm_logical, _funelectrons] # pullback callable - fun_pb = TransformedPformComponent( + fun_pb_1 = TransformedPformComponent( + forceterm_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + ) + fun_pb_2 = TransformedPformComponent( + forceterm_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + ) + fun_pb_3 = TransformedPformComponent( forceterm_class, fun_basis="physical", out_form="2", comp=2, domain=self.domain ) - fun_electrons_pb = TransformedPformComponent( + fun_electrons_pb_1 = TransformedPformComponent( + forcetermelectrons_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + ) + fun_electrons_pb_2 = TransformedPformComponent( + forcetermelectrons_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + ) + fun_electrons_pb_3 = TransformedPformComponent( forcetermelectrons_class, fun_basis="physical", out_form="2", comp=2, domain=self.domain ) - l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - self._F1 = l2_proj.get_dofs([_forceterm_logical, _forceterm_logical, fun_pb]) - self._F2 = l2_proj.get_dofs([_forceterm_logical, _forceterm_logical, fun_electrons_pb]) + if self._lifting: + l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) + else: + l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + self._F1 = l2_proj([fun_pb_1, fun_pb_2, fun_pb_3], apply_bc=self._lifting) + self._F2 = l2_proj([fun_electrons_pb_1, fun_electrons_pb_2, fun_electrons_pb_3], apply_bc=self._lifting) ### End Restelli ### + elif self._dimension == "Tokamak": + ### Tokamak geometry curl-free manufactured solution ### + + _forceterm_logical = lambda e1, e2, e3: 0 * e1 + _funx = getattr(callables, "ManufacturedSolutionForceterm")( + species="Ions", + comp="0", + b0=self._B0, + nu=self._nu, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funy = getattr(callables, "ManufacturedSolutionForceterm")( + species="Ions", + comp="1", + b0=self._B0, + nu=self._nu, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funz = getattr(callables, "ManufacturedSolutionForceterm")( + species="Ions", + comp="2", + b0=self._B0, + nu=self._nu, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funelectronsx = getattr(callables, "ManufacturedSolutionForceterm")( + species="Electrons", + comp="0", + b0=self._B0, + nu_e=self._nu_e, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funelectronsy = getattr(callables, "ManufacturedSolutionForceterm")( + species="Electrons", + comp="1", + b0=self._B0, + nu_e=self._nu_e, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + _funelectronsz = getattr(callables, "ManufacturedSolutionForceterm")( + species="Electrons", + comp="2", + b0=self._B0, + nu_e=self._nu_e, + dimension=self._dimension, + stab_sigma=self._stab_sigma, + eps=self._eps_norm, + dt=D1_dt, + a=self._a, + Bp=self._Bp, + alpha=self._alpha, + beta=self._beta, + ) + + # get callable(s) for specified init type + forceterm_class = [_funx, _funy, _funz] + forcetermelectrons_class = [_funelectronsx, _funelectronsy, _funelectronsz] + + # pullback callable + fun_pb_1 = TransformedPformComponent( + forceterm_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + ) + fun_pb_2 = TransformedPformComponent( + forceterm_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + ) + fun_pb_3 = TransformedPformComponent( + forceterm_class, fun_basis="physical", out_form="2", comp=2, domain=self.domain + ) + fun_electrons_pb_1 = TransformedPformComponent( + forcetermelectrons_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + ) + fun_electrons_pb_2 = TransformedPformComponent( + forcetermelectrons_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + ) + fun_electrons_pb_3 = TransformedPformComponent( + forcetermelectrons_class, fun_basis="physical", out_form="2", comp=2, domain=self.domain + ) + if self._lifting: + l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) + else: + l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) + self._F1 = l2_proj([fun_pb_1, fun_pb_2, fun_pb_3], apply_bc=self._lifting) + self._F2 = l2_proj([fun_electrons_pb_1, fun_electrons_pb_2, fun_electrons_pb_3], apply_bc=self._lifting) + + ### End Tokamak geometry manufactured solution ### + if self._variant == "GMRES": + if self._lifting: + self._M2 = getattr(self._mass_opsv0, "M2") + self._M3 = getattr(self._mass_opsv0, "M3") + self._M2B = -getattr(self._mass_opsv0, "M2B") + self._div = self.derhamv0.div + self._curl = self.derhamv0.curl + self._S21 = self._basis_opsv0.S21 + else: + self._M2 = getattr(self.mass_ops, "M2") + self._M3 = getattr(self.mass_ops, "M3") + self._M2B = -getattr(self.mass_ops, "M2B") + self._div = self.derham.div + self._curl = self.derham.curl + self._S21 = self.basis_ops.S21 + + # Define block matrix [[A BT], [B 0]] (without time step size dt in the diagonals) + _A11 = ( + self._M2 + - self._M2B / self._eps_norm + + self._nu + * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + ) + _A12 = None + _A21 = _A12 + _A22 = ( + -self._stab_sigma * IdentityOperator(_A11.domain) + + self._M2B / self._eps_norm + + self._nu_e + * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + ) + _B1 = -self._M3 @ self._div + _B2 = self._M3 @ self._div + if _A12 is not None: assert _A11.codomain == _A12.codomain if _A21 is not None: @@ -7779,45 +7975,163 @@ def __init__( elif self._variant == "Uzawa": # Numpy - fun = [] - for m in range(3): - fun += [[]] - for n in range(3): - fun[-1] += [ - lambda e1, e2, e3, m=m, n=n: self.basis_ops.G(e1, e2, e3)[:, :, :, m, n] - / self.basis_ops.sqrt_g(e1, e2, e3), - ] - self._S21 = None - if self.derham.with_local_projectors: - self._S21 = BasisProjectionOperatorLocal( - self.derham._Ploc["1"], self.derham.Vh_fem["2"], fun, transposed=False - ) - self.derhamnumpy = Derham( - self.derham.Nel, - self.derham.p, - self.derham.spl_kind, - domain=self.domain, - ) - if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - self._M2np = self.mass_ops.M2._mat.toarray() - self._M3np = self.mass_ops.M3._mat.toarray() - self._Dnp = self.derhamnumpy.div.toarray() - self._Cnp = self.derhamnumpy.curl.toarray() - if self._S21 is not None: - self._Hodgenp = self._S21.toarray - else: - self._Hodgenp = self.basis_ops.S21.toarray_struphy() # self.basis_ops.S21.toarray - self._M2Bnp = -self.mass_ops.M2B._mat.toarray() - elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - self._M2np = self.mass_ops.M2._mat.tosparse() - self._M3np = self.mass_ops.M3._mat.tosparse() - self._Dnp = self.derhamnumpy.div.tosparse() - self._Cnp = self.derhamnumpy.curl.tosparse() - if self._S21 is not None: - self._Hodgenp = self._S21.tosparse - else: - self._Hodgenp = self.basis_ops.S21.toarray_struphy(is_sparse=True) - self._M2Bnp = -self.mass_ops.M2B._mat.tosparse() + if self._lifting: + fun = [] + for m in range(3): + fun += [[]] + for n in range(3): + fun[-1] += [ + lambda e1, e2, e3, m=m, n=n: self._basis_opsv0.G(e1, e2, e3)[:, :, :, m, n] + / self._basis_opsv0.sqrt_g(e1, e2, e3), + ] + self._S21 = None + if self.derhamv0.with_local_projectors: + self._S21 = BasisProjectionOperatorLocal( + self.derhamv0._Ploc["1"], self.derhamv0.Vh_fem["2"], fun, transposed=False + ) + + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + Vbc = self._mass_opsv0.M2._V_boundary_op.toarray_struphy() + Wbc = self._mass_opsv0.M2._W_boundary_op.toarray_struphy() + M2_mat = self._mass_opsv0.M2._mat.toarray() + self._M2np = Wbc @ M2_mat @ Vbc.T + Vbc = self._mass_opsv0.M3._V_boundary_op.toarray_struphy() + Wbc = self._mass_opsv0.M3._W_boundary_op.toarray_struphy() + M3_mat = self._mass_opsv0.M3._mat.toarray() + self._M3np = Wbc @ M3_mat @ Vbc.T + if isinstance(self.derhamv0.div, ComposedLinearOperator): + for mult in self.derhamv0.div.multiplicants: + if isinstance(mult, BlockLinearOperator): + if hasattr(self, "_Dnp"): + self._Dnp = self._Dnp @ mult.toarray() + else: + self._Dnp = mult.toarray() + # print(f"{type(mult.toarray())=}") #with_pads = True + elif isinstance(mult, BoundaryOperator): + if hasattr(self, "_Dnp"): + self._Dnp = self._Dnp @ mult.T.toarray_struphy() + else: + self._Dnp = mult.toarray_struphy() + elif isinstance(self.derhamv0.div, BlockLinearOperator): + self._Dnp = self.derhamv0.div.toarray() + if isinstance(self.derhamv0.curl, ComposedLinearOperator): + for mult in self.derhamv0.curl.multiplicants: + if isinstance(mult, BlockLinearOperator): + if hasattr(self, "_Cnp"): + self._Cnp = self._Cnp @ mult.toarray() + else: + self._Cnp = mult.toarray() + elif isinstance(mult, BoundaryOperator): + if hasattr(self, "_Cnp"): + self._Cnp = self._Cnp @ mult.T.toarray_struphy() + else: + self._Cnp = mult.toarray_struphy() + elif isinstance(self.derhamv0.curl, BlockLinearOperator): + self._Dnp = self.derhamv0.curl.toarray() + + if self._S21 is not None: + self._Hodgenp = self._S21.toarray + else: + self._Hodgenp = self._basis_opsv0.S21.toarray_struphy() # self.basis_ops.S21.toarray + Vbc = self._mass_opsv0.M2B._V_boundary_op.toarray_struphy() + Wbc = self._mass_opsv0.M2B._W_boundary_op.toarray_struphy() + M2B_mat = -self._mass_opsv0.M2B._mat.toarray() # - sign because of the definition of M2B + self._M2Bnp = Wbc @ M2B_mat @ Vbc.T + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + Vbc = self._mass_opsv0.M2._V_boundary_op.toarray_struphy(is_sparse=True) + Wbc = self._mass_opsv0.M2._W_boundary_op.toarray_struphy(is_sparse=True) + M2_mat = self._mass_opsv0.M2._mat.tosparse() + self._M2np = Wbc @ M2_mat @ Vbc.T + Vbc = self._mass_opsv0.M3._V_boundary_op.toarray_struphy(is_sparse=True) + Wbc = self._mass_opsv0.M3._W_boundary_op.toarray_struphy(is_sparse=True) + M3_mat = self._mass_opsv0.M3._mat.tosparse() + self._M3np = Wbc @ M3_mat @ Vbc.T + if self._S21 is not None: + self._Hodgenp = self._S21.tosparse + else: + self._Hodgenp = self._basis_opsv0.S21.toarray_struphy(is_sparse=True) + Vbc = self._mass_opsv0.M2B._V_boundary_op.toarray_struphy(is_sparse=True) + Wbc = self._mass_opsv0.M2B._W_boundary_op.toarray_struphy(is_sparse=True) + M2B_mat = self._mass_opsv0.M2B._mat.tosparse() + self._M2Bnp = -Wbc @ M2B_mat @ Vbc.T # - sign because of the definition of M2B + + if isinstance(self.derhamv0.div, ComposedLinearOperator): + for mult in self.derhamv0.div.multiplicants: + if isinstance(mult, BlockLinearOperator): + if hasattr(self, "_Dnp"): + self._Dnp = self._Dnp @ mult.tosparse() + else: + self._Dnp = mult.tosparse() + elif isinstance(mult, BoundaryOperator): + if hasattr(self, "_Dnp"): + self._Dnp = self._Dnp @ mult.toarray_struphy(is_sparse=True) + else: + self._Dnp = mult.toarray_struphy(is_sparse=True) + elif isinstance(self.derhamv0.div, BlockLinearOperator): + self._Dnp = self.derhamv0.div.tosparse() + + if isinstance(self.derhamv0.curl, ComposedLinearOperator): + for mult in self.derhamv0.curl.multiplicants: + if isinstance(mult, BlockLinearOperator): + if hasattr(self, "_Cnp"): + self._Cnp = self._Cnp @ mult.tosparse() + else: + self._Cnp = mult.tosparse() + elif isinstance(mult, BoundaryOperator): + if hasattr(self, "_Cnp"): + self._Cnp = self._Cnp @ mult.toarray_struphy(is_sparse=True) + else: + self._Cnp = mult.toarray_struphy(is_sparse=True) + elif isinstance(self.derhamv0.curl, BlockLinearOperator): + self._Dnp = self.derhamv0.curl.tosparse() + + else: # no lifting, use original Derham + fun = [] + for m in range(3): + fun += [[]] + for n in range(3): + fun[-1] += [ + lambda e1, e2, e3, m=m, n=n: self.basis_ops.G(e1, e2, e3)[:, :, :, m, n] + / self.basis_ops.sqrt_g(e1, e2, e3), + ] + self._S21 = None + if self.derham.with_local_projectors: + self._S21 = BasisProjectionOperatorLocal( + self.derham._Ploc["1"], self.derham.Vh_fem["2"], fun, transposed=False + ) + + if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + Vbc = self.mass_ops.M2._V_boundary_op.toarray_struphy() + Wbc = self.mass_ops.M2._W_boundary_op.toarray_struphy() + M2_mat = self.mass_ops.M2._mat.toarray() + self._M2np = Wbc @ M2_mat @ Vbc.T + Vbc = self.mass_ops.M3._V_boundary_op.toarray_struphy() + Wbc = self.mass_ops.M3._W_boundary_op.toarray_struphy() + M3_mat = self.mass_ops.M3._mat.toarray() + self._M3np = Wbc @ M3_mat @ Vbc.T + self._Dnp = self.derhamnumpy.div.toarray() + self._Cnp = self.derhamnumpy.curl.toarray() + + if self._S21 is not None: + self._Hodgenp = self._S21.toarray + else: + self._Hodgenp = self.basis_ops.S21.toarray_struphy() + Vbc = self.mass_ops.M2B._V_boundary_op.toarray_struphy() + Wbc = self.mass_ops.M2B._W_boundary_op.toarray_struphy() + M2B_mat = -self.mass_ops.M2B._mat.toarray() + self._M2Bnp = Wbc @ M2B_mat @ Vbc.T + elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + self._M2np = self.mass_ops.M2.tosparse + self._M3np = self.mass_ops.M3.tosparse + if self._S21 is not None: + self._Hodgenp = self._S21.tosparse + else: + self._Hodgenp = self.basis_ops.S21.toarray_struphy(is_sparse=True) + self._M2Bnp = -self.mass_ops.M2B.tosparse + + self._Dnp = self.derhamnumpy.div.tosparse() + self._Cnp = self.derhamnumpy.curl.tosparse() + self._A11np_notimedependency = ( self._nu * ( @@ -7829,6 +8143,7 @@ def __init__( A11np = self._M2np + self._A11np_notimedependency if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + A11np += self._stab_sigma * np.identity(A11np.shape[0]) self.A22np = ( self._stab_sigma * np.identity(A11np.shape[0]) + self._nu_e @@ -7842,6 +8157,7 @@ def __init__( np.identity(self.A22np.shape[0]) * self._stab_sigma ) # + self._nu_e * (self._Dnp.T @ self._M3np @ self._Dnp) elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + A11np += self._stab_sigma * sc.sparse.eye(A11np.shape[0], format="csr") self.A22np = ( self._stab_sigma * sc.sparse.eye(A11np.shape[0], format="csr") + self._nu_e @@ -7851,11 +8167,12 @@ def __init__( ) + self._M2Bnp / self._eps_norm ) - self._A22prenp = sc.sparse.eye(self.A22np.shape[0], format="csr") + self._A22prenp = self._stab_sigma * sc.sparse.eye(self.A22np.shape[0], format="csr") B1np = -self._M3np @ self._Dnp B2np = self._M3np @ self._Dnp self._B1np = B1np + self._B2np = B2np self._F1np = self._F1.toarray() self._F2np = self._F2.toarray() _Anp = [A11np, self.A22np] @@ -7881,15 +8198,16 @@ def __init__( elif self._variant == "Uzawa": self._solver_UzawaNumpy = SaddlePointSolver( + Apre=_Anppre, A=_Anp, B=_Bnp, F=_Fnp, - Apre=_Anppre, method_to_solve=self._method_to_solve, preconditioner=self._preconditioner, spectralanalysis=spectralanalysis, tol=solver["tol"], max_iter=solver["maxiter"], + verbose=solver["verbose"], ) def __call__(self, dt): @@ -7898,65 +8216,68 @@ def __call__(self, dt): uenfeec = self.feec_vars[1] phinfeec = self.feec_vars[2] - # print(f"{min(unfeec.toarray()) =}") - # print(f"{max(unfeec.toarray()) =}") - # print(f"{min(self._F1.toarray()) =}") - # print(f"{max(self._F1.toarray()) =}") - # f1test = self.derham.create_spline_function('f1', 'Hdiv') - # f1test.vector = self._F1 - # utest = self.derham.create_spline_function('u', 'Hdiv') - # utest.vector = self._feec_vars[0] - - # print(f'{type(self._F1)=}') - # print(f'{type(self.feec_vars[0])=}') - - # x1 = np.linspace(0.,1.,100) - # x2 = np.linspace(0.,1.,100) - # x3 = np.linspace(0.,1.,2) - - # f1call = f1test(x1,x2,x3) - # ucall = utest(x1,x2,x3) - - # from matplotlib import pyplot as plt - # plt.figure(figsize=(8, 6)) - # plt.subplot(3,2,1) - # plt.contourf(f1call[2][:,:,0]) - # plt.colorbar() - # plt.subplot(3,2,2) - # plt.plot(self._F2.toarray()) - # plt.subplot(3,2,3) - # plt.contourf(ucall[0][:,:,0]) - # plt.colorbar() - # plt.subplot(3,2,4) - # plt.plot(uenfeec.toarray()) - - # plt.savefig("Forceterm.png") - # exit() - if self._variant == "GMRES": - # Define block matrix [[A BT], [B 0]] + if self._lifting: + phinfeeccopy = self.derhamv0.create_spline_function("phi", space_id="L2") + phinfeeccopy.vector = phinfeec + # unfeec in space Hdiv, u0 in space Hdiv_0 + unfeeccopy = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u0 = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u_prime = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u0.vector = uenfeec + unfeeccopy.vector = uenfeec + apply_essential_bc_stencil(u0.vector[0], axis=0, ext=-1, order=0) + apply_essential_bc_stencil(u0.vector[0], axis=0, ext=1, order=0) + u_prime.vector = unfeeccopy.vector - u0.vector + + uenfeeccopy = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue0 = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue_prime = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue0.vector = uenfeec + uenfeeccopy.vector = uenfeec + apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=-1, order=0) + apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=1, order=0) + ue_prime.vector = uenfeeccopy.vector - ue0.vector + _A11 = ( self._M2 / dt - self._M2B / self._eps_norm + self._nu - * ( - self.derham.div.T @ self._M3 @ self.derham.div - + self.basis_ops.S21.T @ self.derham.curl.T @ self._M2 @ self.derham.curl @ self.basis_ops.S21 - ) + * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) ) _A12 = None _A21 = _A12 _A22 = ( self._nu_e - * ( - self.derham.div.T @ self._M3 @ self.derham.div - + self.basis_ops.S21.T @ self.derham.curl.T @ self._M2 @ self.derham.curl @ self.basis_ops.S21 - ) + * (self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21) + self._M2B / self._eps_norm - self._stab_sigma * IdentityOperator(_A11.domain) ) - _B1 = -self._M3 @ self.derham.div - _B2 = self._M3 @ self.derham.div + + if self._lifting: + _A11prime = -self._M2B / self._eps_norm + self._nu * ( + self.derhamv0.div.T @ self._M3 @ self.derhamv0.div + + self._basis_opsv0.S21.T + @ self.derhamv0.curl.T + @ self._M2 + @ self.derhamv0.curl + @ self._basis_opsv0.S21 + ) + _A22prime = ( + self._nu_e + * ( + self.derhamv0.div.T @ self._M3 @ self.derhamv0.div + + self._basis_opsv0.S21.T + @ self.derhamv0.curl.T + @ self._M2 + @ self.derhamv0.curl + @ self._basis_opsv0.S21 + ) + + self._M2B / self._eps_norm + - self._stab_sigma * IdentityOperator(_A11.domain) + ) + _B1 = -self._M3 @ self._div + _B2 = self._M3 @ self._div if _A12 is not None: assert _A11.codomain == _A12.codomain @@ -7974,7 +8295,16 @@ def __call__(self, dt): _A = BlockLinearOperator(self._block_domainA, self._block_codomainA, blocks=_blocksA) _blocksB = [[_B1, _B2]] _B = BlockLinearOperator(self._block_domainB, self._block_codomainB, blocks=_blocksB) - _blocksF = [self._F1 + self._M2.dot(unfeec, out=self._untemp) / dt, self._F2] + if self._lifting: + _blocksF = [ + self._M2.dot(self._F1) + self._M2.dot(u0.vector) / dt - _A11prime.dot(u_prime.vector), + self._M2.dot(self._F2) - _A22prime.dot(ue_prime.vector), + ] + else: + _blocksF = [ + self._M2.dot(self._F1) + self._M2.dot(unfeec) / dt, + self._M2.dot(self._F2), + ] _F = BlockVector(self._block_domainA, blocks=_blocksF) # Imported solver @@ -7982,43 +8312,93 @@ def __call__(self, dt): self._solver_GMRES.B = _B self._solver_GMRES.F = _F - ( - _sol1, - _sol2, - info, - ) = self._solver_GMRES(unfeec, uenfeec, phinfeec) - un = _sol1[0] - uen = _sol1[1] - phin = _sol2 - + if self._lifting: + ( + _sol1, + _sol2, + info, + ) = self._solver_GMRES(u0.vector, ue0.vector, phinfeec) + un = _sol1[0] + u_prime.vector + uen = _sol1[1] + ue_prime.vector + phin = _sol2 + else: + ( + _sol1, + _sol2, + info, + ) = self._solver_GMRES(unfeec, uenfeec, phinfeec) + un = _sol1[0] + uen = _sol1[1] + phin = _sol2 # write new coeffs into self.feec_vars max_du, max_due, max_dphi = self.feec_vars_update(un, uen, phin) elif self._variant == "Uzawa": - print(f"{self._eps_norm = }") # Numpy A11np = self._M2np / dt + self._A11np_notimedependency if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): + A11np += self._stab_sigma * np.identity(A11np.shape[0]) _A22prenp = self._A22prenp A22np = self.A22np elif self._method_to_solve in ("SparseSolver", "ScipySparse"): + A11np += self._stab_sigma * sc.sparse.eye(A11np.shape[0], format="csr") _A22prenp = self._A22prenp A22np = self.A22np # _Anp[1] and _Anppre[1] remain unchanged _Anp = [A11np, A22np] if self._preconditioner == True: - _A11prenp = self._M2np / dt + self._A11prenp_notimedependency + _A11prenp = self._M2np / dt # + self._A11prenp_notimedependency _Anppre = [_A11prenp, _A22prenp] - _F1np = self._F1np + 1.0 / dt * self._M2np.dot(unfeec.toarray()) - _Fnp = [_F1np, self._F2np] + + if self._lifting: + # unfeec in space Hdiv, u0 in space Hdiv_0 + unfeeccopy = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u0 = self.derhamv0.create_spline_function("u", space_id="Hdiv") + u_prime = self.derham.create_spline_function("u", space_id="Hdiv") + u0.vector = unfeec + unfeeccopy.vector = unfeec + apply_essential_bc_stencil(u0.vector[0], axis=0, ext=-1, order=0) + apply_essential_bc_stencil(u0.vector[0], axis=0, ext=1, order=0) + u_prime.vector = unfeeccopy.vector - u0.vector + + uenfeeccopy = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue0 = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue_prime = self.derhamv0.create_spline_function("ue", space_id="Hdiv") + ue0.vector = uenfeec + uenfeeccopy.vector = uenfeec + apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=-1, order=0) + apply_essential_bc_stencil(ue0.vector[0], axis=0, ext=1, order=0) + ue_prime.vector = uenfeeccopy.vector - ue0.vector + + _F1np = ( + self._M2np @ self._F1np + + 1.0 / dt * self._M2np.dot(u0.vector.toarray()) + - self._A11np_notimedependency.dot(u_prime.vector.toarray()) + ) + _F2np = self._M2np @ self._F2np - self.A22np.dot(ue_prime.vector.toarray()) + _Fnp = [_F1np, _F2np] + else: + _F1np = self._M2np @ self._F1np + 1.0 / dt * self._M2np.dot(unfeec.toarray()) + _F2np = self._M2np @ self._F2np + _Fnp = [_F1np, _F2np] if self.rank == 0: - self._solver_UzawaNumpy.A = _Anp if self._preconditioner == True: self._solver_UzawaNumpy.Apre = _Anppre + self._solver_UzawaNumpy.A = _Anp self._solver_UzawaNumpy.F = _Fnp - un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy(unfeec, uenfeec, phinfeec) + if self._lifting: + un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( + u0.vector, ue0.vector, phinfeec + ) + + un += u_prime.vector.toarray() + uen += ue_prime.vector.toarray() + else: + un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( + unfeec, uenfeec, phinfeec + ) dimlist = [[shp - 2 * pi for shp, pi in zip(unfeec[i][:].shape, self.derham.p)] for i in range(3)] dimphi = [shp - 2 * pi for shp, pi in zip(phinfeec[:].shape, self.derham.p)] diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index fc2aa2de9..a5596827d 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -17,7 +17,7 @@ from struphy.fields_background.equils import set_defaults from struphy.io.options import OptsMPIsort, check_option from struphy.io.setup import descend_options_dict -from struphy.models.variables import FEECVariable, PICVariable +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.ode.utils import ButcherTableau from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.pic.base import Particles @@ -48,15 +48,15 @@ class PushEta(Propagator): class Variables: def __init__(self): - self._var: PICVariable = None + self._var: PICVariable | SPHVariable = None @property - def var(self) -> PICVariable: + def var(self) -> PICVariable | SPHVariable: return self._var @var.setter def var(self, new): - assert isinstance(new, PICVariable) + assert isinstance(new, PICVariable | SPHVariable) self._var = new def __init__(self): From 0f8f5e22c77fc34ff81bc419a0088cdd7817f6c4 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 29 Sep 2025 09:59:14 +0000 Subject: [PATCH 124/292] Port sph tutorial to 318 --- src/struphy/fields_background/generic.py | 8 + src/struphy/io/options.py | 23 +- src/struphy/main.py | 44 +- src/struphy/models/base.py | 182 ++-- src/struphy/models/fluid.py | 144 +-- src/struphy/models/tests/test_models.py | 1 + src/struphy/models/toy.py | 31 +- src/struphy/models/variables.py | 135 +-- src/struphy/pic/base.py | 11 +- src/struphy/pic/utilities.py | 58 +- .../propagators/propagators_markers.py | 237 +++-- tutorials/tutorial_01_parameter_files.ipynb | 24 +- tutorials/tutorial_02_test_particles.ipynb | 22 +- ...l_03_smoothed_particle_hydrodynamics.ipynb | 983 ++++++++++++++++++ 14 files changed, 1527 insertions(+), 376 deletions(-) create mode 100644 tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb diff --git a/src/struphy/fields_background/generic.py b/src/struphy/fields_background/generic.py index f309f1eef..0b82d7e17 100644 --- a/src/struphy/fields_background/generic.py +++ b/src/struphy/fields_background/generic.py @@ -1,3 +1,5 @@ +import copy + from struphy.fields_background.base import ( CartesianFluidEquilibrium, CartesianFluidEquilibriumWithB, @@ -17,6 +19,9 @@ def __init__( p_xyz: callable = None, n_xyz: callable = None, ): + # use params setter + self.params = copy.deepcopy(locals()) + if u_xyz is None: u_xyz = lambda x, y, z: (0.0 * x, 0.0 * x, 0.0 * x) else: @@ -57,6 +62,9 @@ def __init__( b_xyz: callable = None, gradB_xyz: callable = None, ): + # use params setter + self.params = copy.deepcopy(locals()) + super().__init__(u_xyz=u_xyz, p_xyz=p_xyz, n_xyz=n_xyz) if b_xyz is None: diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 601f228ef..13c1ae134 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -32,10 +32,31 @@ OptsPICSpace = Literal["Particles6D", "DeltaFParticles6D", "Particles5D", "Particles3D"] OptsMarkerBC = Literal["periodic", "reflect"] OptsRecontructBC = Literal["periodic", "mirror", "fixed"] -OptsLoading = Literal["pseudo_random", "sobol_standard", "sobol_antithetic", "external", "restart", "tesselation"] +OptsLoading = Literal[ + "pseudo_random", + "sobol_standard", + "sobol_antithetic", + "external", + "restart", + "tesselation", +] OptsSpatialLoading = Literal["uniform", "disc"] OptsMPIsort = Literal["each", "last", None] +# sph +OptsKernel = Literal[ + "trigonometric_1d", + "gaussian_1d", + "linear_1d", + "trigonometric_2d", + "gaussian_2d", + "linear_2d", + "trigonometric_3d", + "gaussian_3d", + "linear_isotropic_3d", + "linear_3d", +] + ## Option classes diff --git a/src/struphy/main.py b/src/struphy/main.py index 74bdd33a8..5955f16f8 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -612,10 +612,10 @@ def pproc( if exist_kinetic["n_sph"]: post_process_n_sph( path, + params_in, path_kinetics_species, species, step, - compute_bckgr=compute_bckgr, ) @@ -633,7 +633,7 @@ def __init__(self, path: str): self._orbits = {} self._f = {} self._spline_values = {} - self.sph_species = {} + self._n_sph = {} self.grids_log: list[np.ndarray] = None self.grids_phy: list[np.ndarray] = None self.t_grid: np.ndarray = None @@ -653,6 +653,11 @@ def spline_values(self) -> dict[str, dict[str, np.ndarray]]: """Keys: species name. Values: dicts of variable names with values being 3d arrays on the grid.""" return self._spline_values + @property + def n_sph(self) -> dict[str, dict[str, dict[str, np.ndarray]]]: + """Keys: species name. Values: dicts of view names ('view_0' etc.) holding dicts of corresponding np.arrays for plotting.""" + return self._n_sph + @property def Nt(self) -> dict[str, int]: """Number of available time points (snap shots) for each species.""" @@ -752,6 +757,7 @@ def load_data(path: str) -> SimData: for folder in sub_folders: path_dat = os.path.join(path_spec, folder) sub_wlk = os.walk(path_dat) + if "orbits" in folder: files = next(sub_wlk)[2] Nt = len(files) // 2 @@ -765,6 +771,7 @@ def load_data(path: str) -> SimData: simdata._orbits[spec] = np.zeros((Nt, *tmp.shape), dtype=float) simdata._orbits[spec][step] = tmp n += 1 + elif "distribution_function" in folder: simdata._f[spec] = {} slices = next(sub_wlk)[1] @@ -779,18 +786,25 @@ def load_data(path: str) -> SimData: tmp = np.load(os.path.join(path_dat, sli, file)) # print(f"{name = }") simdata._f[spec][sli][name] = tmp + + elif "n_sph" in folder: + simdata._n_sph[spec] = {} + slices = next(sub_wlk)[1] + # print(f"{slices = }") + for sli in slices: + simdata._n_sph[spec][sli] = {} + # print(f"{sli = }") + files = next(sub_wlk)[2] + # print(f"{files = }") + for file in files: + name = file.split(".")[0] + tmp = np.load(os.path.join(path_dat, sli, file)) + # print(f"{name = }") + simdata._n_sph[spec][sli][name] = tmp + else: print(f"{folder = }") raise NotImplementedError - # # simdata.pic_species[spec][folder] = {} - # tmp = {} - # for file in files: - # # print(f"{file = }") - # if ".npy" in file: - # var = file.split(".")[0] - # tmp[var] = np.load(os.path.join(path_dat, file)) - # # sort dict - # simdata.pic_species[spec][folder] = dict(sorted(tmp.items())) print("\nThe following data has been loaded:") print(f"{simdata.time_grid_size = }") @@ -810,7 +824,13 @@ def load_data(path: str) -> SimData: print(f" {kk}") for kkk, vvv in vv.items(): print(f" {kkk}") - print(f"simdata.sph_species:") + print(f"\nsimdata.n_sph:") + for k, v in simdata.n_sph.items(): + print(f" {k}") + for kk, vv in v.items(): + print(f" {kk}") + for kkk, vvv in vv.items(): + print(f" {kkk}") return simdata diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 479c5b63a..780c3efff 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -3,7 +3,7 @@ import os from abc import ABCMeta, abstractmethod from functools import reduce -from typing import Callable +from textwrap import indent import numpy as np import yaml @@ -511,7 +511,7 @@ def update_scalar(self, name, value=None): # Perform MPI operations based on the compute flags if "sum_world" in compute_operations: - self.comm_world.Allreduce( + MPI.COMM_WORLD.Allreduce( MPI.IN_PLACE, value_array, op=MPI.SUM, @@ -698,35 +698,14 @@ def update_markers_to_be_saved(self): obj = var.particles assert isinstance(obj, Particles) - # allocate array for saving markers if not present - if not hasattr(self, "_n_markers_saved"): - n_markers = species.n_markers - - if isinstance(n_markers, float): - if n_markers > 1.0: - self._n_markers_saved = int(n_markers) - else: - self._n_markers_saved = int(obj.n_mks_global * n_markers) - else: - self._n_markers_saved = n_markers - - assert self._n_markers_saved <= obj.Np, ( - f"The number of markers for which data should be stored (={self._n_markers_saved}) murst be <= than the total number of markers (={obj.Np})" - ) - if self._n_markers_saved > 0: - var.particle_data["markers"] = np.zeros( - (self._n_markers_saved, obj.markers.shape[1]), - dtype=float, - ) - - if self._n_markers_saved > 0: + if var.n_to_save > 0: markers_on_proc = np.logical_and( obj.markers[:, -1] >= 0.0, - obj.markers[:, -1] < self._n_markers_saved, + obj.markers[:, -1] < var.n_to_save, ) n_markers_on_proc = np.count_nonzero(markers_on_proc) - var.particle_data["markers"][:] = -1.0 - var.particle_data["markers"][:n_markers_on_proc] = obj.markers[markers_on_proc] + var.saved_markers[:] = -1.0 + var.saved_markers[:n_markers_on_proc] = obj.markers[markers_on_proc] @profile def update_distr_functions(self): @@ -749,20 +728,21 @@ def update_distr_functions(self): str_dn = f"d{i + 1}" dim_to_int[str_dn] = 3 + obj.vdim + 3 + i - if species.binning_plots: - for slice_i, edges in var.particle_data["bin_edges"].items(): - comps = slice_i.split("_") - components = [False] * (3 + obj.vdim + 3 + obj.n_cols_diagnostics) + for bin_plot in species.binning_plots: + comps = bin_plot.slice.split("_") + components = [False] * (3 + obj.vdim + 3 + obj.n_cols_diagnostics) - for comp in comps: - components[dim_to_int[comp]] = True + for comp in comps: + components[dim_to_int[comp]] = True - f_slice, df_slice = obj.binning(components, edges) + edges = bin_plot.bin_edges + divide_by_jac = bin_plot.divide_by_jac + f_slice, df_slice = obj.binning(components, edges, divide_by_jac=divide_by_jac) - var.particle_data["f"][slice_i][:] = f_slice - var.particle_data["df"][slice_i][:] = df_slice + bin_plot.f[:] = f_slice + bin_plot.df[:] = df_slice - if species.kernel_density_plots: + for kd_plot in species.kernel_density_plots: h1 = 1 / obj.boxes_per_dim[0] h2 = 1 / obj.boxes_per_dim[1] h3 = 1 / obj.boxes_per_dim[2] @@ -773,16 +753,16 @@ def update_distr_functions(self): else: kernel_type = "gaussian_" + str(ndim) + "d" - for i, pts in enumerate(var.particle_data["plot_pts"]): - n_sph = obj.eval_density( - *pts, - h1=h1, - h2=h2, - h3=h3, - kernel_type=kernel_type, - fast=True, - ) - var.particle_data["n_sph"][i][:] = n_sph + pts = kd_plot.plot_pts + n_sph = obj.eval_density( + *pts, + h1=h1, + h2=h2, + h3=h3, + kernel_type=kernel_type, + fast=True, + ) + kd_plot.n_sph[:] = n_sph def print_scalar_quantities(self): """ @@ -1105,44 +1085,39 @@ def initialize_data_output(self, data: DataContainer, size): key_spec = os.path.join("kinetic", name) key_spec_restart = os.path.join("restart", name) - data.add_data({key_spec_restart: obj._markers}) + # restart data + data.add_data({key_spec_restart: obj.markers}) - # TODO: particle_data should be a KineticData object, not a dict - for key1, val1 in var.particle_data.items(): - key_dat = os.path.join(key_spec, key1) + # marker data + key_mks = os.path.join(key_spec, "markers") + data.add_data({key_mks: var.saved_markers}) - if key1 == "bin_edges": - continue - elif key1 == "f" or key1 == "df": - assert isinstance(val1, dict) - for key2, val2 in val1.items(): - key_f = os.path.join(key_dat, key2) - data.add_data({key_f: val2}) - - dims = (len(key2) - 2) // 3 + 1 - for dim in range(dims): - data.file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = ( - var.particle_data["bin_edges"][key2][dim][:-1] - + ( - var.particle_data["bin_edges"][key2][dim][1] - - var.particle_data["bin_edges"][key2][dim][0] - ) - / 2 - ) - # case of "n_sph" - elif isinstance(val1, list): - for i, v1 in enumerate(val1): - key_n = os.path.join(key_dat, "view_", str(i)) - data.add_data({key_n: v1}) - # save 1d point values, not meshgrids, because attrs size is limited - eta1 = var.particle_data["plot_pts"][i][0][:, 0, 0] - eta2 = var.particle_data["plot_pts"][i][1][0, :, 0] - eta3 = var.particle_data["plot_pts"][i][2][0, 0, :] - data.file[key_n].attrs["eta1"] = eta1 - data.file[key_n].attrs["eta2"] = eta2 - data.file[key_n].attrs["eta3"] = eta3 - else: - data.add_data({key_dat: val1}) + # binning plot data + for bin_plot in species.binning_plots: + key_f = os.path.join(key_spec, "f", bin_plot.slice) + key_df = os.path.join(key_spec, "df", bin_plot.slice) + + data.add_data({key_f: bin_plot.f}) + data.add_data({key_df: bin_plot.df}) + + for dim, be in enumerate(bin_plot.bin_edges): + data.file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = be[:-1] + (be[1] - be[0]) / 2 + + for i, kd_plot in enumerate(species.kernel_density_plots): + key_n = os.path.join(key_spec, "n_sph", f"view_{i}") + + data.add_data({key_n: kd_plot.n_sph}) + # save 1d point values, not meshgrids, because attrs size is limited + eta1 = kd_plot.plot_pts[0][:, 0, 0] + eta2 = kd_plot.plot_pts[1][0, :, 0] + eta3 = kd_plot.plot_pts[2][0, 0, :] + data.file[key_n].attrs["eta1"] = eta1 + data.file[key_n].attrs["eta2"] = eta2 + data.file[key_n].attrs["eta3"] = eta3 + + # TODO: maybe add other data + # else: + # data.add_data({key_dat: val1}) # keys to be saved at each time step and only at end (restart) save_keys_all = [] @@ -1329,7 +1304,13 @@ def generate_default_parameter_file( particle_params += f"\nloading_params = LoadingParameters()\n" particle_params += f"weights_params = WeightsParameters()\n" particle_params += f"boundary_params = BoundaryParameters()\n" - particle_params += f"model.{sn}.set_markers(loading_params=loading_params, weights_params=weights_params, boundary_params=boundary_params)\n" + particle_params += f"model.{sn}.set_markers(loading_params=loading_params,\n" + txt = f"weights_params=weights_params,\n" + particle_params += indent(txt, " " * len(f"model.{sn}.set_markers(")) + txt = f"boundary_params=boundary_params,\n" + particle_params += indent(txt, " " * len(f"model.{sn}.set_markers(")) + txt = f")\n" + particle_params += indent(txt, " " * len(f"model.{sn}.set_markers(")) particle_params += f"model.{sn}.set_sorting_boxes()\n" particle_params += f"model.{sn}.set_save_data()\n" @@ -1385,7 +1366,12 @@ def generate_default_parameter_file( file.write("from struphy.kinetic_background import maxwellians\n") file.write( - "from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters, BinningPlot\n" + "from struphy.pic.utilities import (LoadingParameters,\n\ + WeightsParameters,\n\ + BoundaryParameters,\n\ + BinningPlot,\n\ + KernelDensityPlot,\n\ + )\n" ) file.write("from struphy import main\n") @@ -1434,7 +1420,7 @@ def generate_default_parameter_file( for prop in self.propagators.__dict__: file.write(f"model.propagators.{prop}.options = model.propagators.{prop}.Options()\n") - file.write("\n# background and initial conditions\n") + file.write("\n# background, perturbations and initial conditions\n") if has_feec: file.write(init_bckgr_feec) file.write(init_pert_feec) @@ -1451,24 +1437,24 @@ def generate_default_parameter_file( file.write('\nif __name__ == "__main__":\n') file.write(" # start run\n") file.write( - " main.run(model, \n\ - params_path=__file__, \n\ - env=env, \n\ - base_units=base_units, \n\ - time_opts=time_opts, \n\ - domain=domain, \n\ - equil=equil, \n\ - grid=grid, \n\ - derham_opts=derham_opts, \n\ - verbose=verbose, \n\ - )" + " main.run(model,\n\ + params_path=__file__,\n\ + env=env,\n\ + base_units=base_units,\n\ + time_opts=time_opts,\n\ + domain=domain,\n\ + equil=equil,\n\ + grid=grid,\n\ + derham_opts=derham_opts,\n\ + verbose=verbose,\n\ + )" ) file.close() print( f"\nDefault parameter file for '{self.__class__.__name__}' has been created in the cwd ({path}).\n\ -You can now launch a simlaltion with 'python params_{self.__class__.__name__}.py'" +You can now launch a simulation with 'python params_{self.__class__.__name__}.py'" ) return path diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 6b1047541..26e6314ee 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -2177,8 +2177,8 @@ def diagnostics_dct(): __diagnostics__ = diagnostics_dct() -class IsothermalEulerSPH(StruphyModel): - r"""Isothermal Euler equations discretized with smoothed particle hydrodynamics (SPH). +class EulerSPH(StruphyModel): + r"""Euler equations discretized with smoothed particle hydrodynamics (SPH). :ref:`normalization`: @@ -2198,105 +2198,107 @@ class IsothermalEulerSPH(StruphyModel): \partial_t S + \mathbf u \cdot \nabla S &= 0\,, \end{align} - where :math:`S` denotes the entropy per unit mass and the internal energy per unit mass is + where :math:`S` denotes the entropy per unit mass. + The internal energy per unit mass can be defined in two ways: .. math:: - \mathcal U(\rho, S) = \kappa(S) \log \rho\,. + \mathrm{"isothermal:"}\qquad &\mathcal U(\rho, S) = \kappa(S) \log \rho\,. + + \mathrm{"polytropic:"}\qquad &\mathcal U(\rho, S) = \kappa(S) \frac{\rho^{\gamma - 1}}{\gamma - 1}\,. :ref:`propagators` (called in sequence): 1. :class:`~struphy.propagators.propagators_markers.PushEta` 2. :class:`~struphy.propagators.propagators_markers.PushVxB` 3. :class:`~struphy.propagators.propagators_markers.PushVinSPHpressure` - - :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species - dct["kinetic"]["euler_fluid"] = "ParticlesSPH" - return dct + class EulerFluid(ParticleSpecies): + def __init__(self): + self.var = SPHVariable() + self.init_variables() - @staticmethod - def bulk_species(): - return "euler_fluid" + ## propagators - @staticmethod - def velocity_scale(): - return "thermal" + class Propagators: + def __init__(self, with_B0: bool = True): + self.push_eta = propagators_markers.PushEta() + if with_B0: + self.push_vxb = propagators_markers.PushVxB() + self.push_sph_p = propagators_markers.PushVinSPHpressure() - # @staticmethod - # def diagnostics_dct(): - # dct = {} - # dct["projected_density"] = "L2" - # return dct + ## abstract methods - @staticmethod - def propagators_dct(): - return { - propagators_markers.PushEta: ["euler_fluid"], - # propagators_markers.PushVxB: ["euler_fluid"], - propagators_markers.PushVinSPHpressure: ["euler_fluid"], - } + def __init__(self, with_B0: bool = True): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + self.with_B0 = with_B0 - def __init__(self, params, comm, clone_config=None): - super().__init__(params, comm=comm, clone_config=clone_config) + # 1. instantiate all species + self.euler_fluid = self.EulerFluid() - # prelim - _p = params["kinetic"]["euler_fluid"] - algo_eta = _p["options"]["PushEta"]["algo"] - # algo_vxb = _p["options"]["PushVxB"]["algo"] - kernel_type = _p["options"]["PushVinSPHpressure"]["kernel_type"] - algo_sph = _p["options"]["PushVinSPHpressure"]["algo"] - gravity = _p["options"]["PushVinSPHpressure"]["gravity"] - thermodynamics = _p["options"]["PushVinSPHpressure"]["thermodynamics"] + # 2. instantiate all propagators + self.propagators = self.Propagators(with_B0=with_B0) - # magnetic field - # self._b_eq = self.projected_equil.b2 + # 3. assign variables to propagators + self.propagators.push_eta.variables.var = self.euler_fluid.var + if with_B0: + self.propagators.push_vxb.variables.ions = self.euler_fluid.var + self.propagators.push_sph_p.variables.fluid = self.euler_fluid.var - # set keyword arguments for propagators - self._kwargs[propagators_markers.PushEta] = { - "algo": algo_eta, - # "density_field": self.pointer["projected_density"], - } + # define scalars for update_scalar_quantities + self.add_scalar("en_kin", compute="from_sph", variable=self.euler_fluid.var) - # self._kwargs[propagators_markers.PushVxB] = { - # "algo": algo_vxb, - # "kappa": 1.0, - # "b2": self._b_eq, - # "b2_add": None, - # } - - self._kwargs[propagators_markers.PushVinSPHpressure] = { - "kernel_type": kernel_type, - "algo": algo_sph, - "gravity": gravity, - "thermodynamics": thermodynamics, - } + @property + def bulk_species(self): + return self.euler_fluid - # Initialize propagators used in splitting substeps - self.init_propagators() + @property + def velocity_scale(self): + return "thermal" - # Scalar variables to be saved during simulation - self.add_scalar("en_kin", compute="from_sph", species="euler_fluid") + def allocate_helpers(self): + pass + + # @staticmethod + # def diagnostics_dct(): + # dct = {} + # dct["projected_density"] = "L2" + # return dct def update_scalar_quantities(self): - valid_markers = self.pointer["euler_fluid"].markers_wo_holes_and_ghost + particles = self.euler_fluid.var.particles + valid_markers = particles.markers_wo_holes_and_ghost en_kin = valid_markers[:, 6].dot( valid_markers[:, 3] ** 2 + valid_markers[:, 4] ** 2 + valid_markers[:, 5] ** 2 - ) / (2.0 * self.pointer["euler_fluid"].Np) + ) / (2.0 * particles.Np) self.update_scalar("en_kin", en_kin) + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "push_vxb.Options" in line: + new_file += ["if model.with_B0:\n"] + new_file += [" " + line] + elif "set_save_data" in line: + new_file += ["\nkd_plot = KernelDensityPlot()\n"] + new_file += ["model.euler_fluid.set_save_data(kernel_density_plots=(kd_plot,))\n"] + elif "base_units = BaseUnits" in line: + new_file += ["base_units = BaseUnits(kBT=1.0)\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class HasegawaWakatani(StruphyModel): r"""Hasegawa-Wakatani equations in 2D. diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index af5c28410..24aa086f1 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -27,6 +27,7 @@ fluid_models = [ "LinearMHD", + "EulerSPH", ] # for name, obj in inspect.getmembers(fluid): # if inspect.isclass(obj) and "models.fluid" in obj.__module__: diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 1a8d3562e..ee263db44 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1050,7 +1050,9 @@ class PressureLessSPH(StruphyModel): &\partial_t \rho + \nabla \cdot ( \rho \mathbf u ) = 0 \,, \\[4mm] - &\partial_t (\rho \mathbf u) + \nabla \cdot (\rho \mathbf u \otimes \mathbf u) = 0 \,. + &\partial_t (\rho \mathbf u) + \nabla \cdot (\rho \mathbf u \otimes \mathbf u) = - \nabla \phi_0 \,, + + where :math:`\phi_0` is a static external potential. :ref:`propagators` (called in sequence): @@ -1071,6 +1073,7 @@ def __init__(self): class Propagators: def __init__(self): self.push_eta = propagators_markers.PushEta() + self.push_v = propagators_markers.PushVinEfield() ## abstract methods @@ -1086,6 +1089,7 @@ def __init__(self): # 3. assign variables to propagators self.propagators.push_eta.variables.var = self.cold_fluid.var + self.propagators.push_v.variables.var = self.cold_fluid.var # define scalars for update_scalar_quantities self.add_scalar("en_kin", compute="from_particles", variable=self.cold_fluid.var) @@ -1109,14 +1113,29 @@ def allocate_helpers(self): def update_scalar_quantities(self): particles = self.cold_fluid.var.particles - en_kin = particles.markers_wo_holes_and_ghost[:, 6].dot( - particles.markers_wo_holes_and_ghost[:, 3] ** 2 - + particles.markers_wo_holes_and_ghost[:, 4] ** 2 - + particles.markers_wo_holes_and_ghost[:, 5] ** 2 - ) / (2.0 * particles.Np) + valid_parts = particles.markers_wo_holes_and_ghost + en_kin = valid_parts[:, 6].dot(valid_parts[:, 3] ** 2 + valid_parts[:, 4] ** 2 + valid_parts[:, 5] ** 2) / ( + 2.0 * particles.Np + ) self.update_scalar("en_kin", en_kin) + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "push_v.Options" in line: + new_file += ["phi = equil.p0\n"] + new_file += ["model.propagators.push_v.options = model.propagators.push_v.Options(phi=phi)\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class TwoFluidQuasiNeutralToy(StruphyModel): r"""Linearized, quasi-neutral two-fluid model with zero electron inertia. diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 01482dc55..101e0d497 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -145,7 +145,6 @@ class PICVariable(Variable): def __init__(self, space: OptsPICSpace = "Particles6D"): check_option(space, OptsPICSpace) self._space = space - self._particle_data = {} @property def space(self): @@ -155,10 +154,6 @@ def space(self): def particles(self) -> Particles: return self._particles - @property - def particle_data(self) -> dict: - return self._particle_data - @property def species(self) -> ParticleSpecies: if not hasattr(self, "_species"): @@ -248,53 +243,36 @@ def allocate( self.particles.draw_markers(sort=sort, verbose=verbose) self.particles.initialize_weights() - # for storing the binned distribution function - self.particle_data["bin_edges"] = {} - self.particle_data["f"] = {} - self.particle_data["df"] = {} - - for bin_plot in self.species.binning_plots: - sli = bin_plot.slice - n_bins = bin_plot.n_bins - ranges = bin_plot.ranges + # allocate array for saving markers if not present + n_markers = self.species.n_markers + if isinstance(n_markers, float): + if n_markers > 1.0: + self._n_to_save = int(n_markers) + else: + self._n_to_save = int(self.particles.n_mks_global * n_markers) + else: + self._n_to_save = n_markers - assert ((len(sli) - 2) / 3).is_integer(), f"Binning coordinates must be separated by '_', but reads {sli}." - assert len(sli.split("_")) == len(ranges) == len(n_bins), ( - f"Number of slices names ({len(sli.split('_'))}), number of bins ({len(n_bins)}), and number of ranges ({len(ranges)}) are inconsistent with each other!\n\n" - ) - self.particle_data["bin_edges"][sli] = [] - dims = (len(sli) - 2) // 3 + 1 - for j in range(dims): - self.particle_data["bin_edges"][sli] += [ - np.linspace( - ranges[j][0], - ranges[j][1], - n_bins[j] + 1, - ), - ] - self.particle_data["f"][sli] = np.zeros(n_bins, dtype=float) - self.particle_data["df"][sli] = np.zeros(n_bins, dtype=float) - - # for storing an sph evaluation of the density n - self.particle_data["n_sph"] = [] - self.particle_data["plot_pts"] = [] - - for kd_plot in self.species.kernel_density_plots: - eta1 = np.linspace(0.0, 1.0, kd_plot.pts_e1) - eta2 = np.linspace(0.0, 1.0, kd_plot.pts_e2) - eta3 = np.linspace(0.0, 1.0, kd_plot.pts_e3) - ee1, ee2, ee3 = np.meshgrid( - eta1, - eta2, - eta3, - indexing="ij", + assert self._n_to_save <= self.particles.Np, ( + f"The number of markers for which data should be stored (={self._n_to_save}) murst be <= than the total number of markers (={obj.Np})" + ) + if self._n_to_save > 0: + self._saved_markers = np.zeros( + (self._n_to_save, self.particles.markers.shape[1]), + dtype=float, ) - self.particle_data["plot_pts"] += [(ee1, ee2, ee3)] - self.particle_data["n_sph"] += [np.zeros(ee1.shape, dtype=float)] # other data (wave-particle power exchange, etc.) # TODO + @property + def n_to_save(self) -> int: + return self._n_to_save + + @property + def saved_markers(self) -> np.ndarray: + return self._saved_markers + class SPHVariable(Variable): def __init__(self): @@ -406,49 +384,32 @@ def allocate( self.particles.draw_markers(sort=sort, verbose=verbose) self.particles.initialize_weights() - # for storing the binned distribution function - self.particle_data["bin_edges"] = {} - self.particle_data["f"] = {} - self.particle_data["df"] = {} - - for bin_plot in self.species.binning_plots: - sli = bin_plot.slice - n_bins = bin_plot.n_bins - ranges = bin_plot.ranges + # allocate array for saving markers if not present + n_markers = self.species.n_markers + if isinstance(n_markers, float): + if n_markers > 1.0: + self._n_to_save = int(n_markers) + else: + self._n_to_save = int(self.particles.n_mks_global * n_markers) + else: + self._n_to_save = n_markers - assert ((len(sli) - 2) / 3).is_integer(), f"Binning coordinates must be separated by '_', but reads {sli}." - assert len(sli.split("_")) == len(ranges) == len(n_bins), ( - f"Number of slices names ({len(sli.split('_'))}), number of bins ({len(n_bins)}), and number of ranges ({len(ranges)}) are inconsistent with each other!\n\n" - ) - self.particle_data["bin_edges"][sli] = [] - dims = (len(sli) - 2) // 3 + 1 - for j in range(dims): - self.particle_data["bin_edges"][sli] += [ - np.linspace( - ranges[j][0], - ranges[j][1], - n_bins[j] + 1, - ), - ] - self.particle_data["f"][sli] = np.zeros(n_bins, dtype=float) - self.particle_data["df"][sli] = np.zeros(n_bins, dtype=float) - - # for storing an sph evaluation of the density n - self.particle_data["n_sph"] = [] - self.particle_data["plot_pts"] = [] - - for kd_plot in self.species.kernel_density_plots: - eta1 = np.linspace(0.0, 1.0, kd_plot.pts_e1) - eta2 = np.linspace(0.0, 1.0, kd_plot.pts_e2) - eta3 = np.linspace(0.0, 1.0, kd_plot.pts_e3) - ee1, ee2, ee3 = np.meshgrid( - eta1, - eta2, - eta3, - indexing="ij", + assert self._n_to_save <= self.particles.Np, ( + f"The number of markers for which data should be stored (={self._n_to_save}) murst be <= than the total number of markers (={obj.Np})" + ) + if self._n_to_save > 0: + self._saved_markers = np.zeros( + (self._n_to_save, self.particles.markers.shape[1]), + dtype=float, ) - self.particle_data["plot_pts"] += [(ee1, ee2, ee3)] - self.particle_data["n_sph"] += [np.zeros(ee1.shape, dtype=float)] # other data (wave-particle power exchange, etc.) # TODO + + @property + def n_to_save(self) -> int: + return self._n_to_save + + @property + def saved_markers(self) -> np.ndarray: + return self._saved_markers diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 4408f2593..91c41437a 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1876,16 +1876,21 @@ def reset_marker_ids(self): self.marker_ids = first_marker_id + np.arange(self.n_mks_loc, dtype=int) @profile - def binning(self, components, bin_edges, divide_by_jac=True): + def binning( + self, + components: tuple[bool], + bin_edges: tuple[np.ndarray], + divide_by_jac: bool = True, + ): r"""Computes full-f and delta-f distribution functions via marker binning in logical space. Numpy's histogramdd is used, following the algorithm outlined in :ref:`binning`. Parameters ---------- - components : list[bool] + components : tuple[bool] List of length 3 + vdim; an entry is True if the direction in phase space is to be binned. - bin_edges : list[array] + bin_edges : tuple[array] List of bin edges (resolution) having the length of True entries in components. divide_by_jac : boll diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index 117f0dcfa..99b597c9b 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -151,8 +151,11 @@ class BinningPlot: n_bins : int | tuple[int] Number of bins for each coordinate. - ranges : tuple[int] | tuple[tuple[int]]= (0.0, 1.0) + ranges : tuple[int] | tuple[tuple[int]] = (0.0, 1.0) Binning range (as an interval in R) for each coordinate. + + divide_by_jac : bool + Whether to divide by the Jacobian determinant (volume-to-0-form). """ def __init__( @@ -160,16 +163,45 @@ def __init__( slice: str = "e1", n_bins: int | tuple[int] = 128, ranges: tuple[float] | tuple[tuple[float]] = (0.0, 1.0), + divide_by_jac: bool = True, ): - self.slice = slice - if isinstance(n_bins, int): n_bins = (n_bins,) - self.n_bins = n_bins if not isinstance(ranges[0], tuple): ranges = (ranges,) + + assert ((len(slice) - 2) / 3).is_integer(), f"Binning coordinates must be separated by '_', but reads {slice}." + assert len(slice.split("_")) == len(ranges) == len(n_bins), ( + f"Number of slices names ({len(slice.split('_'))}), number of bins ({len(n_bins)}), and number of ranges ({len(ranges)}) are inconsistent with each other!\n\n" + ) + self.slice = slice + self.n_bins = n_bins self.ranges = ranges + self.divide_by_jac = divide_by_jac + + # computations and allocations + self._bin_edges = [] + for nb, rng in zip(n_bins, ranges): + self._bin_edges += [np.linspace(rng[0], rng[1], nb + 1)] + self._bin_edges = tuple(self.bin_edges) + + self._f = np.zeros(n_bins, dtype=float) + self._df = np.zeros(n_bins, dtype=float) + + @property + def bin_edges(self) -> tuple: + return self._bin_edges + + @property + def f(self) -> np.ndarray: + """The binned distribution function (full-f).""" + return self._f + + @property + def df(self) -> np.ndarray: + """The binned distribution function minus the background (delta-f).""" + return self._df class KernelDensityPlot: @@ -187,9 +219,21 @@ def __init__( pts_e2: int = 11, pts_e3: int = 11, ): - self.pts_e1 = pts_e1 - self.pts_e2 = pts_e2 - self.pts_e3 = pts_e3 + e1 = np.linspace(0.0, 1.0, pts_e1) + e2 = np.linspace(0.0, 1.0, pts_e2) + e3 = np.linspace(0.0, 1.0, pts_e3) + ee1, ee2, ee3 = np.meshgrid(e1, e2, e3, indexing="ij") + self._plot_pts = (ee1, ee2, ee3) + self._n_sph = np.zeros(ee1.shape, dtype=float) + + @property + def plot_pts(self) -> tuple: + return self._plot_pts + + @property + def n_sph(self) -> np.ndarray: + """The evaluated density.""" + return self._n_sph def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index a5596827d..e046bd442 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -2,7 +2,7 @@ import copy from dataclasses import dataclass -from typing import Literal, get_args +from typing import Callable, Literal, get_args import numpy as np from line_profiler import profile @@ -15,7 +15,11 @@ from struphy.feec.mass import WeightedMassOperators from struphy.fields_background.base import MHDequilibrium from struphy.fields_background.equils import set_defaults -from struphy.io.options import OptsMPIsort, check_option +from struphy.io.options import ( + OptsKernel, + OptsMPIsort, + check_option, +) from struphy.io.setup import descend_options_dict from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.ode.utils import ButcherTableau @@ -143,16 +147,16 @@ class PushVxB(Propagator): class Variables: def __init__(self): - self._ions: PICVariable = None + self._ions: PICVariable | SPHVariable = None @property - def ions(self) -> PICVariable: + def ions(self) -> PICVariable | SPHVariable: return self._ions @ions.setter def ions(self, new): - assert isinstance(new, PICVariable) - assert new.space == "Particles6D" + assert isinstance(new, PICVariable | SPHVariable) + assert new.space in ("Particles6D", "ParticlesSPH") self._ions = new def __init__(self): @@ -259,57 +263,107 @@ class PushVinEfield(Propagator): .. math:: - \frac{\text{d} \mathbf{v}_p}{\text{d} t} = \kappa \, \mathbf{E}(\mathbf{x}_p) \,, + \frac{\text{d} \mathbf{v}_p}{\text{d} t} = \frac{1}{\varepsilon} \, \mathbf{E}(\mathbf{x}_p) \,, - where :math:`\kappa \in \mathbb R` is a constant and in logical coordinates, given by :math:`\mathbf x = F(\boldsymbol \eta)`: + where :math:`\varepsilon \in \mathbb R` is a constant. In logical coordinates, given by :math:`\mathbf x = F(\boldsymbol \eta)`: .. math:: - \frac{\text{d} \mathbf{v}_p}{\text{d} t} = \kappa \, DF^{-\top} \hat{\mathbf E}^1(\boldsymbol \eta_p) \,, + \frac{\text{d} \mathbf{v}_p}{\text{d} t} = \frac{1}{\varepsilon} \, DF^{-\top} \hat{\mathbf E}^1(\boldsymbol \eta_p) \,, - which is solved analytically. + which is solved analytically. :math:`\mathbf E` can optionally be defined + through a potential, :math:`\mathbf E = - \nabla \phi`. """ - @staticmethod - def options(): - pass + class Variables: + def __init__(self): + self._var: PICVariable | SPHVariable = None - def __init__( - self, - particles: Particles6D, - *, - e_field: BlockVector | PolarVector, - kappa: float = 1.0, - ): - super().__init__(particles) + @property + def var(self) -> PICVariable | SPHVariable: + return self._var - self.kappa = kappa + @var.setter + def var(self, new): + assert isinstance(new, PICVariable | SPHVariable) + assert new.space in ("Particles6D", "ParticlesSPH") + self._var = new - assert isinstance(e_field, (BlockVector, PolarVector)) - self._e_field = e_field + def __init__(self): + self.variables = self.Variables() - # instantiate Pusher - args_kernel = ( - self.derham.args_derham, - self._e_field[0]._data, - self._e_field[1]._data, - self._e_field[2]._data, - self.kappa, - ) + @dataclass + class Options: + # propagator options + e_field: FEECVariable | tuple[Callable] = None + phi: FEECVariable | Callable = None - self._pusher = Pusher( - particles, - pusher_kernels.push_v_with_efield, - args_kernel, - self.domain.args_domain, - alpha_in_kernel=1.0, - ) + def __post_init__(self): + # checks + if self.e_field is not None: + assert isinstance(self.e_field, tuple[Callable]) or self.e_field.space == "Hcurl" + else: + if self.phi is not None: + assert isinstance(self.phi, Callable) or self.phi.space == "H1" + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + # scaling factor + self._epsilon = self.variables.var.species.equation_params.epsilon + + self._e_field = None + + if self.options.e_field is not None: + if isinstance(self.options.e_field, tuple[Callable]): + self._e_field = self.derham.P["1"](self.options.e_field) + else: + self._e_field = self.options.e_field.spline.vector + + if self.options.phi is not None: + if isinstance(self.options.phi, Callable): + _phi = self.derham.P["0"](self.options.phi) + else: + _phi = self.options.phi.spline.vector + self._e_field = self.derham.grad.dot(_phi) + self._e_field.update_ghost_regions() # very important, we will move it inside grad + self._e_field *= -1.0 + + if self._e_field is not None: + # instantiate Pusher + args_kernel = ( + self.derham.args_derham, + self._e_field[0]._data, + self._e_field[1]._data, + self._e_field[2]._data, + 1.0 / self._epsilon, + ) + + self._pusher = Pusher( + self.variables.var.particles, + pusher_kernels.push_v_with_efield, + args_kernel, + self.domain.args_domain, + alpha_in_kernel=1.0, + ) def __call__(self, dt): - """ - TODO - """ - self._pusher(dt) + if self._e_field is not None: + self._pusher(dt) class PushEtaPC(Propagator): @@ -1613,35 +1667,63 @@ class PushVinSPHpressure(Propagator): * Explicit from :class:`~struphy.ode.utils.ButcherTableau` """ - @staticmethod - def options(default=False): - dct = {} - dct["kernel_type"] = list(Particles.ker_dct()) - dct["algo"] = [ - "forward_euler", - ] # "heun2", "rk2", "heun3", "rk4"] - dct["gravity"] = (0.0, 0.0, 0.0) - dct["thermodynamics"] = ["isothermal", "polytropic"] - if default: - dct = descend_options_dict(dct, []) - return dct + class Variables: + def __init__(self): + self._fluid: SPHVariable = None - def __init__( - self, - particles: ParticlesSPH, - *, - kernel_type: str = "gaussian_2d", - kernel_width: tuple = None, - algo: str = options(default=True)["algo"], # TODO: implement other algos than forward Euler - gravity: tuple = options(default=True)["gravity"], - thermodynamics: str = options(default=True)["thermodynamics"], - ): - # base class constructor call - super().__init__(particles) + @property + def fluid(self) -> SPHVariable: + return self._fluid + + @fluid.setter + def fluid(self, new): + assert isinstance(new, SPHVariable) + assert new.space == "ParticlesSPH" + self._fluid = new + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsAlgo = Literal["forward_euler"] + OptsThermo = Literal["isothermal", "polytropic"] + # propagator options + kernel_type: OptsKernel = "gaussian_2d" + kernel_width: tuple = None + algo: OptsAlgo = "forward_euler" + gravity: tuple = (0.0, 0.0, 0.0) + thermodynamics: OptsThermo = "isothermal" + + def __post_init__(self): + # checks + check_option(self.kernel_type, OptsKernel) + check_option(self.algo, self.OptsAlgo) + check_option(self.thermodynamics, self.OptsThermo) + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): # init kernel for evaluating density etc. before each time step. init_kernel = eval_kernels_gc.sph_pressure_coeffs + particles = self.variables.fluid.particles + first_free_idx = particles.args_markers.first_free_idx comps = (0, 1, 2) @@ -1649,12 +1731,12 @@ def __init__( neighbours = particles.sorting_boxes.neighbours holes = particles.holes periodic = [bci == "periodic" for bci in particles.bc] - kernel_nr = particles.ker_dct()[kernel_type] + kernel_nr = particles.ker_dct()[self.options.kernel_type] - if kernel_width is None: - kernel_width = tuple([1 / ni for ni in self.particles[0].boxes_per_dim]) + if self.options.kernel_width is None: + self.options.kernel_width = tuple([1 / ni for ni in particles.boxes_per_dim]) else: - assert all([hi <= 1 / ni for hi, ni in zip(kernel_width, self.particles[0].boxes_per_dim)]) + assert all([hi <= 1 / ni for hi, ni in zip(self.options.kernel_width, particles.boxes_per_dim)]) # init kernel args_init = ( @@ -1663,7 +1745,7 @@ def __init__( holes, *periodic, kernel_nr, - *kernel_width, + *self.options.kernel_width, ) self.add_init_kernel( @@ -1674,12 +1756,12 @@ def __init__( ) # pusher kernel - if thermodynamics == "isothermal": + if self.options.thermodynamics == "isothermal": kernel = pusher_kernels.push_v_sph_pressure - elif thermodynamics == "polytropic": + elif self.options.thermodynamics == "polytropic": kernel = pusher_kernels.push_v_sph_pressure_ideal_gas - gravity = np.array(gravity, dtype=float) + gravity = np.array(self.options.gravity, dtype=float) args_kernel = ( boxes, @@ -1687,7 +1769,7 @@ def __init__( holes, *periodic, kernel_nr, - *kernel_width, + *self.options.kernel_width, gravity, ) @@ -1701,6 +1783,7 @@ def __init__( init_kernels=self.init_kernels, ) + @profile def __call__(self, dt): - self.particles[0].put_particles_in_boxes() + self.variables.fluid.particles.put_particles_in_boxes() self._pusher(dt) diff --git a/tutorials/tutorial_01_parameter_files.ipynb b/tutorials/tutorial_01_parameter_files.ipynb index 13a6c1839..f05cbb734 100644 --- a/tutorials/tutorial_01_parameter_files.ipynb +++ b/tutorials/tutorial_01_parameter_files.ipynb @@ -7,12 +7,12 @@ "source": [ "# 1 - Parameters and `struphy.main`\n", "\n", - "Struphy is based on \"models\". Each model is a parallelized simulation code for a certain physical model, described py a set partial differntial equations (PDEs).\n", + "Struphy is a collection of \"models\". Each model is a parallelized simulation code for a physics model described py a set partial differntial equations (PDEs).\n", "Check the documentation for a [list of currently available models](https://struphy.pages.mpcdf.de/struphy/sections/models.html).\n", - "A model is launched through its parameter file, where all simulation parameters can be specified by the user.\n", + "A model is launched through its parameter file, where all simulation parameters can be specified.\n", "\n", "Model parameter files are Python scripts (.py) that can be executed with the Python interpreter.\n", - "For each `MODEL`, the default parameter file can be generated from the console via\n", + "For each `MODEL`, the default parameter file can be generated from the console:\n", "\n", "```\n", "struphy params MODEL\n", @@ -24,7 +24,7 @@ "python3 params_MODEL.py\n", "```\n", "\n", - "The user can modify the parameter file to launch a specific simulation.\n", + "One can modify the parameter file to launch a specific simulation.\n", "For example, the parameter file of the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.Vlasov) can be generated from\n", "\n", "```\n", @@ -57,12 +57,16 @@ "from struphy.io.options import FieldsBackground\n", "from struphy.initial import perturbations\n", "from struphy.kinetic_background import maxwellians\n", - "from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters\n", + "from struphy.pic.utilities import (LoadingParameters, \n", + " WeightsParameters, \n", + " BoundaryParameters,\n", + " BinningPlot,\n", + " KernelDensityPlot,\n", + " )\n", "from struphy import main\n", "\n", "# import model, set verbosity\n", - "from struphy.models.toy import Vlasov\n", - "verbose = True" + "from struphy.models.toy import Vlasov" ] }, { @@ -71,7 +75,7 @@ "metadata": {}, "source": [ "All parameter files import the modules listed above, even though some of them might not be needed in a specific model. \n", - "The last import imports the model itself. The `verbose` flag controls the screen output during the simulation run.\n", + "The last import imports the model itself.\n", "\n", "## Part 2: Generic options\n", "\n", @@ -232,7 +236,7 @@ "source": [ "## Part 6: `main.run`\n", "\n", - "In the final part of the parameter file, the `main.run` command is invoked. This command will allocate memory and run the specified simulation. The run command is not executed when the parameter file is imported in another Python script." + "In the final part of the parameter file, the `main.run` command is invoked. This command will allocate memory and run the specified simulation. The run command is not executed when the parameter file is imported in another Python script. The `verbose` flag controls the screen output during the simulation run." ] }, { @@ -242,6 +246,8 @@ "metadata": {}, "outputs": [], "source": [ + "verbose = True\n", + "\n", "main.run(model, \n", " params_path=None, \n", " env=env, \n", diff --git a/tutorials/tutorial_02_test_particles.ipynb b/tutorials/tutorial_02_test_particles.ipynb index 3c4b05b56..c65d01cbd 100644 --- a/tutorials/tutorial_02_test_particles.ipynb +++ b/tutorials/tutorial_02_test_particles.ipynb @@ -45,12 +45,16 @@ "from struphy.io.options import FieldsBackground\n", "from struphy.initial import perturbations\n", "from struphy.kinetic_background import maxwellians\n", - "from struphy.pic.utilities import LoadingParameters, WeightsParameters, BoundaryParameters\n", + "from struphy.pic.utilities import (LoadingParameters, \n", + " WeightsParameters, \n", + " BoundaryParameters,\n", + " BinningPlot,\n", + " KernelDensityPlot,\n", + " )\n", "from struphy import main\n", "\n", "# import model, set verbosity\n", - "from struphy.models.toy import Vlasov\n", - "verbose = False" + "from struphy.models.toy import Vlasov" ] }, { @@ -58,8 +62,6 @@ "id": "2", "metadata": {}, "source": [ - "Note that we set `verbose = False` which will be passed to `main.run` to supress screen output during the simulation.\n", - "\n", "We shall create two simulations, which we store in different output folders, defined throught the environment variables:" ] }, @@ -258,6 +260,8 @@ "metadata": {}, "outputs": [], "source": [ + "verbose = False\n", + "\n", "main.run(model, \n", " params_path=None, \n", " env=env, \n", @@ -455,6 +459,8 @@ "metadata": {}, "outputs": [], "source": [ + "verbose = False\n", + "\n", "main.run(model, \n", " params_path=None, \n", " env=env, \n", @@ -653,6 +659,8 @@ "outputs": [], "source": [ "# run\n", + "verbose = False\n", + "\n", "main.run(model, \n", " params_path=None, \n", " env=env, \n", @@ -1015,6 +1023,8 @@ "source": [ "time_opts = Time(dt=0.2, Tend=3000, split_algo=\"Strang\")\n", "\n", + "verbose = False\n", + "\n", "main.run(model, \n", " params_path=None, \n", " env=env, \n", @@ -1202,6 +1212,8 @@ "source": [ "time_opts = Time(dt=0.1, Tend=100, split_algo=\"Strang\")\n", "\n", + "verbose = False\n", + "\n", "main.run(model, \n", " params_path=None, \n", " env=env, \n", diff --git a/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb b/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb new file mode 100644 index 000000000..ffe72d95c --- /dev/null +++ b/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb @@ -0,0 +1,983 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d44f9387", + "metadata": {}, + "source": [ + "# 3 - Smoothed-particle hydrodynamics\n", + "\n", + "Struphy provides several models for smoothed-particle hydrodynamics (SPH). In this tutorial, we shall exlore\n", + "\n", + "1. Pressure-less fluid flow in a Beltrami force field (model [PressureLessSPH](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_toy.html#struphy.models.toy.PressureLessSPH))\n", + "2. Gas expansion with SPH (model [EulerSPH](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_fluid.html#struphy.models.fluid.EulerSPH))\n", + "\n", + "The parameter files discussed below can be obtained with the console commands\n", + "\n", + "```\n", + "struphy params PressureLessSPH\n", + "struphy params EulerSPH \n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "29c3fc2b", + "metadata": {}, + "source": [ + "## Pressure-less fluid flow in Beltrami force field\n", + "\n", + "Let $\\mathbf u_0(\\mathbf x)$ stand for an initial flow velocity field. Moreover,\n", + "let $\\Omega \\subset \\mathbb R^3$ be a box (cuboid). We search for trajectories $(\\mathbf x_p, \\mathbf v_p): [0,T] \\to \\Omega \\times \\mathbb R^3$, $p = 0, \\ldots, N-1$ that satisfy\n", + "\n", + "$$\n", + "\\begin{align}\n", + " \\dot{\\mathbf x}_p &= \\mathbf v_p\\,,\\qquad && \\mathbf x_p(0) = \\mathbf x_{p0}\\,,\n", + " \\\\[2mm]\n", + " \\dot{\\mathbf v}_p &= -\\nabla p(\\mathbf x_p) \\qquad && \\mathbf v_p(0) = \\mathbf u_0(\\mathbf x_p(0))\\,,\n", + " \\end{align}\n", + "$$\n", + "\n", + "where $p \\in H^1(\\Omega)$ is some given function. In what follows we shall set $p$ to give a Beltrami force field to guide the fluid particles.\n", + "\n", + "We start with the generic imports:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bb4344bf", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.io.options import EnvironmentOptions, BaseUnits, Time\n", + "from struphy.geometry import domains\n", + "from struphy.fields_background import equils\n", + "from struphy.topology import grids\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.io.options import FieldsBackground\n", + "from struphy.initial import perturbations\n", + "from struphy.kinetic_background import maxwellians\n", + "from struphy.pic.utilities import (LoadingParameters,\n", + " WeightsParameters,\n", + " BoundaryParameters,\n", + " BinningPlot,\n", + " KernelDensityPlot,\n", + " )\n", + "from struphy import main\n", + "\n", + "# import model, set verbosity\n", + "from struphy.models.toy import PressureLessSPH" + ] + }, + { + "cell_type": "markdown", + "id": "0f9a7269", + "metadata": {}, + "source": [ + "We shall use the Strang splitting algorithm for the propagators, and simulate in Cartesian geometry (Cuboid):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9e7a378a", + "metadata": {}, + "outputs": [], + "source": [ + "# environment options\n", + "env = EnvironmentOptions()\n", + "\n", + "# units\n", + "base_units = BaseUnits()\n", + "\n", + "# time stepping\n", + "time_opts = Time(dt=0.02, Tend=4, split_algo=\"Strang\")\n", + "\n", + "# geometry\n", + "l1 = -.5\n", + "r1 = .5\n", + "l2 = -.5\n", + "r2 = .5\n", + "l3 = 0.\n", + "r3 = 1.\n", + "domain = domains.Cuboid(l1=l1, r1=r1, l2=l2, r2=r2, l3=l3, r3=r3)" + ] + }, + { + "cell_type": "markdown", + "id": "3b88acc5", + "metadata": {}, + "source": [ + "The Beltrami flow can be specified as a lambda function and then passed to `GenericCartesianFluidEquilibrium`. This fluid equilibirum will be set as initial condition below." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11c45bd8", + "metadata": {}, + "outputs": [], + "source": [ + "# construct Beltrami flow\n", + "import numpy as np\n", + "\n", + "def u_fun(x, y, z):\n", + " ux = -np.cos(np.pi*x)*np.sin(np.pi*y)\n", + " uy = np.sin(np.pi*x)*np.cos(np.pi*y)\n", + " uz = 0 * x \n", + " return ux, uy, uz\n", + "\n", + "p_fun = lambda x, y, z: 0.5*(np.sin(np.pi*x)**2 + np.sin(np.pi*y)**2)\n", + "n_fun = lambda x, y, z: 1. + 0*x\n", + "\n", + "# put the functions in a generic equilibirum container\n", + "from struphy.fields_background.generic import GenericCartesianFluidEquilibrium\n", + "bel_flow = GenericCartesianFluidEquilibrium(u_xyz=u_fun, p_xyz=p_fun, n_xyz=n_fun)" + ] + }, + { + "cell_type": "markdown", + "id": "45d672b4", + "metadata": {}, + "source": [ + "We will also need a grid and Derham complex for projecting the fluid equilibirum onto a spline basis:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d166e75d", + "metadata": {}, + "outputs": [], + "source": [ + "# fluid equilibrium (can be used as part of initial conditions)\n", + "equil = None\n", + "\n", + "# grid\n", + "grid = grids.TensorProductGrid(Nel=(64, 64, 1))\n", + "\n", + "# derham options\n", + "derham_opts = DerhamOptions(p=(3, 3, 1), spl_kind=(False, False, True))" + ] + }, + { + "cell_type": "markdown", + "id": "9e7ecca1", + "metadata": {}, + "source": [ + "In the next step, the light-weight instance of the model is created. We set the parameter `epsilon` to 1.0, appearing in the propagator [PushVinEfield](https://struphy.pages.mpcdf.de/struphy/sections/subsections/propagators_markers.html#struphy.propagators.propagators_markers.PushVinEfield). Moreover, we launch with 1000 particles and save 100 % percent of them through `n_markers=1.0`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0d22885b", + "metadata": {}, + "outputs": [], + "source": [ + "# light-weight model instance\n", + "model = PressureLessSPH()\n", + "\n", + "# species parameters\n", + "model.cold_fluid.set_phys_params(epsilon=1.0)\n", + "\n", + "loading_params = LoadingParameters(Np=1000)\n", + "weights_params = WeightsParameters()\n", + "boundary_params = BoundaryParameters(bc=('reflect', 'reflect', 'periodic'))\n", + "model.cold_fluid.set_markers(loading_params=loading_params,\n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params,\n", + " )\n", + "model.cold_fluid.set_sorting_boxes(boxes_per_dim=(1, 1, 1))\n", + "model.cold_fluid.set_save_data(n_markers=1.0)" + ] + }, + { + "cell_type": "markdown", + "id": "3f577b80", + "metadata": {}, + "source": [ + "We now set the propagator options. Here, it is important to pass the pressure from the Beltrami flow as auxiliary field to `push_v`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "263332e0", + "metadata": {}, + "outputs": [], + "source": [ + "# propagator options\n", + "from struphy.ode.utils import ButcherTableau\n", + "butcher = ButcherTableau(algo=\"forward_euler\")\n", + "model.propagators.push_eta.options = model.propagators.push_eta.Options(butcher=butcher)\n", + "\n", + "phi = bel_flow.p0\n", + "model.propagators.push_v.options = model.propagators.push_v.Options(phi=phi)" + ] + }, + { + "cell_type": "markdown", + "id": "a2fa6f8b", + "metadata": {}, + "source": [ + "The initial condition of the species is defined by a background function, plus possible perturbations, which we ignore here. The background of the speceis is set to the Beltrami flow:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d658ceb0", + "metadata": {}, + "outputs": [], + "source": [ + "# background, perturbations and initial conditions\n", + "model.cold_fluid.var.add_background(bel_flow)" + ] + }, + { + "cell_type": "markdown", + "id": "d6504849", + "metadata": {}, + "source": [ + "Let us start the run, post-process the raw data, load the post-processed data and plot the particle trajectories:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "be28bb9a", + "metadata": {}, + "outputs": [], + "source": [ + "verbose = False\n", + "\n", + "main.run(model,\n", + " params_path=None,\n", + " env=env,\n", + " base_units=base_units,\n", + " time_opts=time_opts,\n", + " domain=domain,\n", + " equil=equil,\n", + " grid=grid,\n", + " derham_opts=derham_opts,\n", + " verbose=verbose,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0eadf863", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "\n", + "main.pproc(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "de74bc5d", + "metadata": {}, + "outputs": [], + "source": [ + "simdata = main.load_data(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1f18caa7", + "metadata": {}, + "outputs": [], + "source": [ + "from matplotlib import pyplot as plt\n", + "plt.figure(figsize=(12, 28))\n", + "\n", + "orbits = simdata.orbits[\"cold_fluid\"]\n", + "\n", + "coloring = np.select([orbits[0, :, 0]<=-0.2, \n", + " np.abs(orbits[0, :, 0]) < +0.2, \n", + " orbits[0, :, 0] >= 0.2],\n", + " [-1.0, 0.0, +1.0])\n", + "\n", + "dt = time_opts.dt\n", + "Nt = simdata.time_grid_size - 1\n", + "interval = Nt/20\n", + "plot_ct = 0\n", + "for i in range(Nt):\n", + " if i % interval == 0:\n", + " print(f'{i = }')\n", + " plot_ct += 1\n", + " plt.subplot(5, 2, plot_ct)\n", + " ax = plt.gca() \n", + " plt.scatter(orbits[i, :, 0], orbits[i, :, 1], c=coloring)\n", + " plt.axis('square')\n", + " plt.title('n0_scatter')\n", + " plt.xlim(l1, r1)\n", + " plt.ylim(l2, r2)\n", + " plt.colorbar()\n", + " plt.title(f'Gas at t={i*dt}')\n", + " if plot_ct == 10:\n", + " break" + ] + }, + { + "cell_type": "markdown", + "id": "d3eafa60", + "metadata": {}, + "source": [ + "Let us perform another simulation, similar to the previous one. We will save the results in the folder `sim_2`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9561df94", + "metadata": {}, + "outputs": [], + "source": [ + "# environment options\n", + "env = EnvironmentOptions(sim_folder=\"sim_2\")" + ] + }, + { + "cell_type": "markdown", + "id": "34422982", + "metadata": {}, + "source": [ + "This time, we shall draw the markers on a regular grid obtained from a tesselation of the domain: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fe9a414c", + "metadata": {}, + "outputs": [], + "source": [ + "# light-weight model instance\n", + "model = PressureLessSPH()\n", + "\n", + "# species parameters\n", + "model.cold_fluid.set_phys_params(epsilon=1.0)\n", + "\n", + "loading_params = LoadingParameters(ppb=4, loading=\"tesselation\")\n", + "weights_params = WeightsParameters()\n", + "boundary_params = BoundaryParameters(bc=('reflect', 'reflect', 'periodic'))\n", + "model.cold_fluid.set_markers(loading_params=loading_params,\n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params,\n", + " bufsize=0.5\n", + " )\n", + "model.cold_fluid.set_sorting_boxes(boxes_per_dim=(16, 16, 1))\n", + "model.cold_fluid.set_save_data(n_markers=1.0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d3354fe7", + "metadata": {}, + "outputs": [], + "source": [ + "# propagator options\n", + "from struphy.ode.utils import ButcherTableau\n", + "butcher = ButcherTableau(algo=\"forward_euler\")\n", + "model.propagators.push_eta.options = model.propagators.push_eta.Options(butcher=butcher)\n", + "\n", + "phi = bel_flow.p0\n", + "model.propagators.push_v.options = model.propagators.push_v.Options(phi=phi)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "69ada6ed", + "metadata": {}, + "outputs": [], + "source": [ + "# background, perturbations and initial conditions\n", + "model.cold_fluid.var.add_background(bel_flow)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b2bd98c9", + "metadata": {}, + "outputs": [], + "source": [ + "main.run(model,\n", + " params_path=None,\n", + " env=env,\n", + " base_units=base_units,\n", + " time_opts=time_opts,\n", + " domain=domain,\n", + " equil=equil,\n", + " grid=grid,\n", + " derham_opts=derham_opts,\n", + " verbose=verbose,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "98659a03", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "path = os.path.join(os.getcwd(), \"sim_2\")\n", + "\n", + "main.pproc(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b78574d6", + "metadata": {}, + "outputs": [], + "source": [ + "simdata = main.load_data(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "05d19960", + "metadata": {}, + "outputs": [], + "source": [ + "from matplotlib import pyplot as plt\n", + "plt.figure(figsize=(12, 28))\n", + "\n", + "orbits = simdata.orbits[\"cold_fluid\"]\n", + "\n", + "coloring = np.select([orbits[0, :, 0]<=-0.2, \n", + " np.abs(orbits[0, :, 0]) < +0.2, \n", + " orbits[0, :, 0] >= 0.2],\n", + " [-1.0, 0.0, +1.0])\n", + "\n", + "dt = time_opts.dt\n", + "Nt = simdata.time_grid_size - 1\n", + "interval = Nt/20\n", + "plot_ct = 0\n", + "for i in range(Nt):\n", + " if i % interval == 0:\n", + " print(f'{i = }')\n", + " plot_ct += 1\n", + " plt.subplot(5, 2, plot_ct)\n", + " ax = plt.gca() \n", + " plt.scatter(orbits[i, :, 0], orbits[i, :, 1], c=coloring)\n", + " plt.axis('square')\n", + " plt.title('n0_scatter')\n", + " plt.xlim(l1, r1)\n", + " plt.ylim(l2, r2)\n", + " plt.colorbar()\n", + " plt.title(f'Gas at t={i*dt}')\n", + " if plot_ct == 10:\n", + " break" + ] + }, + { + "cell_type": "markdown", + "id": "99087247", + "metadata": {}, + "source": [ + "## Gas expansion\n", + "\n", + "We use SPH to solve Euler's equations (model [EulerSPH](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_fluid.html#struphy.models.fluid.EulerSPH)):\n", + "\n", + "$$\n", + "\\begin{align}\n", + " \\partial_t \\rho + \\nabla \\cdot (\\rho \\mathbf u) &= 0\\,,\n", + " \\\\[2mm]\n", + " \\rho(\\partial_t \\mathbf u + \\mathbf u \\cdot \\nabla \\mathbf u) &= - \\nabla \\left(\\rho^2 \\frac{\\partial \\mathcal U(\\rho, S)}{\\partial \\rho} \\right)\\,,\n", + " \\\\[2mm]\n", + " \\partial_t S + \\mathbf u \\cdot \\nabla S &= 0\\,,\n", + " \\end{align}\n", + "$$\n", + "\n", + "where $S$ denotes the entropy per unit mass and the internal energy per unit mass is \n", + "\n", + "$$\n", + "\\mathcal U(\\rho, S) = \\kappa(S) \\log \\rho\\,.\n", + "$$\n", + "\n", + "The SPH discretization leads to ODEs for $N$ particles indexed by $p$,\n", + "\n", + "$$\n", + "\\begin{align}\n", + " \\dot{\\mathbf x}_p &= \\mathbf v_p\\,,\\qquad && \\mathbf x_p(0) = \\mathbf x_{p0}\\,,\n", + " \\\\[2mm]\n", + " \\dot{\\mathbf v}_p &= -\\kappa_{p}(0) \\sum_{i=1}^N w_i \\left(\\frac{1}{\\rho^{N,h}(\\mathbf x_p)} + \\frac{1}{\\rho^{N,h}(\\mathbf x_i)} \\right) \\nabla W_h(\\mathbf x_p - \\mathbf x_i) \\qquad && \\mathbf v_p(0) = \\mathbf u(\\mathbf x_p(0))\\,,\n", + " \\end{align}\n", + "$$\n", + "\n", + "where the smoothed density reads\n", + "\n", + "$$\n", + " \\rho^{N,h}(\\mathbf x) = \\sum_{j=1}^N w_j W_h(\\mathbf x - \\mathbf x_j)\\,,\n", + "$$\n", + "\n", + "with weights $w_p = const.$ and where $W_h(\\mathbf x)$ is a suitable smoothing kernel.\n", + "The velocity update is performed with the Propagator [PushVinSPHpressure](https://struphy.pages.mpcdf.de/struphy/sections/subsections/propagators_markers.html#struphy.propagators.propagators_markers.PushVinSPHpressure).\n", + "\n", + "We shall now compute a gas expansion in 2d (nonlinear example). First, check out some of the smoothing kernels available for SPH evaluations:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a6f19621", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.pic.sph_smoothing_kernels import linear_uni, trigonometric_uni, gaussian_uni\n", + "\n", + "x = np.linspace(-1, 1, 200)\n", + "out1 = np.zeros_like(x)\n", + "out2 = np.zeros_like(x)\n", + "out3 = np.zeros_like(x)\n", + "\n", + "for i, xi in enumerate(x):\n", + " out1[i] = trigonometric_uni(xi, 1.)\n", + " out2[i] = gaussian_uni(xi, 1.)\n", + " out3[i] = linear_uni(xi, 1.)\n", + "plt.plot(x, out1, label=\"trigonometric\")\n", + "plt.plot(x, out2, label=\"gaussian\")\n", + "plt.plot(x, out3, label = \"linear\")\n", + "plt.title('Some smoothing kernels')\n", + "plt.legend()" + ] + }, + { + "cell_type": "markdown", + "id": "8245c936", + "metadata": {}, + "source": [ + "We start with the generic imports and also import the model `EulerSPH`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "77c93c55", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.io.options import EnvironmentOptions, BaseUnits, Time\n", + "from struphy.geometry import domains\n", + "from struphy.fields_background import equils\n", + "from struphy.topology import grids\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.io.options import FieldsBackground\n", + "from struphy.initial import perturbations\n", + "from struphy.kinetic_background import maxwellians\n", + "from struphy.pic.utilities import (LoadingParameters,\n", + " WeightsParameters,\n", + " BoundaryParameters,\n", + " BinningPlot,\n", + " KernelDensityPlot,\n", + " )\n", + "from struphy import main\n", + "\n", + "# import model, set verbosity\n", + "from struphy.models.fluid import EulerSPH" + ] + }, + { + "cell_type": "markdown", + "id": "8fa5cb3c", + "metadata": {}, + "source": [ + "Here, it is important to set the base unit `kBT` in order to derive the velocity unit:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "405a42f0", + "metadata": {}, + "outputs": [], + "source": [ + "# environment options\n", + "env = EnvironmentOptions()\n", + "\n", + "# units\n", + "base_units = BaseUnits(kBT=1.0)\n", + "\n", + "# time stepping\n", + "time_opts = Time(dt=0.04, Tend=1.6, split_algo=\"Strang\")\n", + "\n", + "# geometry\n", + "l1 = -3.0\n", + "r1 = 3.0\n", + "l2 = -3.0\n", + "r2 = 3.0\n", + "l3 = 0.\n", + "r3 = 1.\n", + "domain = domains.Cuboid(l1=l1, r1=r1, l2=l2, r2=r2, l3=l3, r3=r3)" + ] + }, + { + "cell_type": "markdown", + "id": "a2a23116", + "metadata": {}, + "source": [ + "As background, which goes into the initial condition below, we define a Gaussian blob in the xy-plane:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c56515fb", + "metadata": {}, + "outputs": [], + "source": [ + "# gaussian initial blob\n", + "from struphy.fields_background.generic import GenericCartesianFluidEquilibrium\n", + "import numpy as np\n", + "T_h = 0.2\n", + "gamma = 5/3\n", + "n_fun = lambda x, y, z: np.exp(-(x**2 + y**2)/T_h) / 35\n", + "\n", + "blob = GenericCartesianFluidEquilibrium(n_xyz=n_fun)" + ] + }, + { + "cell_type": "markdown", + "id": "814c12cb", + "metadata": {}, + "source": [ + "We also need a grid and Derham complex for projecting the fluid background on a spline basis:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8fffd408", + "metadata": {}, + "outputs": [], + "source": [ + "# fluid equilibrium (can be used as part of initial conditions)\n", + "equil = None\n", + "\n", + "# grid\n", + "grid = grids.TensorProductGrid(Nel=(64, 64, 1))\n", + "\n", + "# derham options\n", + "derham_opts = DerhamOptions(p=(3, 3, 1), spl_kind=(False, False, True))" + ] + }, + { + "cell_type": "markdown", + "id": "35bc5741", + "metadata": {}, + "source": [ + "Now we create the light-weight instance of `EulerSPH`, without the optional propagator for particles in a magnetic background field. Note as well that we shall reject particles whose weight is below a certain threshold in order to save computing time:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd88cb8c", + "metadata": {}, + "outputs": [], + "source": [ + "# light-weight model instance\n", + "model = EulerSPH(with_B0=False)\n", + "\n", + "# species parameters\n", + "model.euler_fluid.set_phys_params()\n", + "\n", + "loading_params = LoadingParameters(ppb=400)\n", + "weights_params = WeightsParameters(reject_weights=True, threshold=3e-3)\n", + "boundary_params = BoundaryParameters()\n", + "model.euler_fluid.set_markers(loading_params=loading_params,\n", + " weights_params=weights_params,\n", + " boundary_params=boundary_params,\n", + " )\n", + "nx = 16\n", + "ny = 16\n", + "model.euler_fluid.set_sorting_boxes(boxes_per_dim=(nx, ny, 1))" + ] + }, + { + "cell_type": "markdown", + "id": "6f5818f5", + "metadata": {}, + "source": [ + "For visualization of the result, we want to save a binning plot and a kernel density plot (sph evaluation of the fluid density):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3be20f57", + "metadata": {}, + "outputs": [], + "source": [ + "bin_plot = BinningPlot(slice=\"e1_e2\", \n", + " n_bins=(64, 64), \n", + " ranges=((0.0, 1.0), (0.0, 1.0)), \n", + " divide_by_jac=False,\n", + " )\n", + "pts_e1 = 100\n", + "pts_e2 = 90\n", + "kd_plot = KernelDensityPlot(pts_e1=pts_e1, pts_e2=pts_e2, pts_e3=1)\n", + "model.euler_fluid.set_save_data(n_markers=1.0,\n", + " binning_plots=(bin_plot,),\n", + " kernel_density_plots=(kd_plot,),\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "f1f36fc7", + "metadata": {}, + "source": [ + "We choose `gaussian_2d` as the smoothing kernel for sph evaluations during the pressure step:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "57b7a966", + "metadata": {}, + "outputs": [], + "source": [ + "# propagator options\n", + "from struphy.ode.utils import ButcherTableau\n", + "butcher = ButcherTableau(algo=\"forward_euler\")\n", + "model.propagators.push_eta.options = model.propagators.push_eta.Options(butcher=butcher)\n", + "\n", + "model.propagators.push_sph_p.options = model.propagators.push_sph_p.Options(kernel_type=\"gaussian_2d\")" + ] + }, + { + "cell_type": "markdown", + "id": "cab5de8f", + "metadata": {}, + "source": [ + "Now we set the initial condition and run the simulation. The time steps take long at the beginning, but get faster towards the end, when particles are spread out over the domain. The simulation launched in the console will be a lot faster than in the notebook, especially when using MPI, or compiling with GPU." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "33f8717c", + "metadata": {}, + "outputs": [], + "source": [ + "# background, perturbations and initial conditions\n", + "model.euler_fluid.var.add_background(blob)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18ed3d41", + "metadata": {}, + "outputs": [], + "source": [ + "verbose = True\n", + "\n", + "main.run(model,\n", + " params_path=None,\n", + " env=env,\n", + " base_units=base_units,\n", + " time_opts=time_opts,\n", + " domain=domain,\n", + " equil=equil,\n", + " grid=grid,\n", + " derham_opts=derham_opts,\n", + " verbose=verbose,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ed420df4", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "path = os.path.join(os.getcwd(), \"sim_1\")\n", + "main.pproc(path)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3d1e2242", + "metadata": {}, + "outputs": [], + "source": [ + "simdata = main.load_data(path)" + ] + }, + { + "cell_type": "markdown", + "id": "db6498c1", + "metadata": {}, + "source": [ + "The above output of the `simdata` object tells us where to find the post-processed simulation data:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14d59054", + "metadata": {}, + "outputs": [], + "source": [ + "# analytical functions\n", + "n_xyz = blob.n_xyz\n", + "n3 = blob.n3\n", + "\n", + "# grids\n", + "x = np.linspace(l1, r1, pts_e1)\n", + "y = np.linspace(l2, r2, pts_e2)\n", + "xx, yy = np.meshgrid(x, y, indexing=\"ij\")\n", + "ee1, ee2, ee3 = simdata.n_sph[\"euler_fluid\"][\"view_0\"][\"grid_n_sph\"]\n", + "eta1 = ee1[:, 0, 0]\n", + "eta2 = ee2[0, :, 0]\n", + "bc_x = simdata.f[\"euler_fluid\"][\"e1_e2\"][\"grid_e1\"]\n", + "bc_y = simdata.f[\"euler_fluid\"][\"e1_e2\"][\"grid_e2\"]\n", + "\n", + "# markers\n", + "orbits = simdata.orbits[\"euler_fluid\"]\n", + "positions = orbits[0, :, :3]\n", + "weights = orbits[0, :, 6]\n", + "\n", + "# binning and sph eval\n", + "n_sph = simdata.n_sph[\"euler_fluid\"][\"view_0\"][\"n_sph\"][0]\n", + "f_bin = simdata.f[\"euler_fluid\"][\"e1_e2\"][\"f_binned\"][0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d3b3befa", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt \n", + "plt.figure(figsize=(12, 15))\n", + "\n", + "# plots\n", + "plt.subplot(3, 2, 1)\n", + "plt.pcolor(xx, yy, n_fun(xx, yy, 0))\n", + "plt.axis('square')\n", + "plt.title('n_xyz initial')\n", + "plt.colorbar()\n", + "\n", + "plt.subplot(3, 2, 2)\n", + "plt.pcolor(eta1, eta2, n3(eta1, eta2, 0, squeeze_out=True).T)\n", + "plt.axis('square')\n", + "plt.title('$\\hat{n}^{\\t{vol}}$ initial (volume form)')\n", + "plt.colorbar()\n", + "\n", + "make_scatter = True\n", + "if make_scatter:\n", + " plt.subplot(3, 2, 3)\n", + " ax = plt.gca()\n", + " ax.set_xticks(np.linspace(l1, r1, nx + 1))\n", + " ax.set_yticks(np.linspace(l2, r2, ny + 1))\n", + " plt.tick_params(labelbottom = False) \n", + " coloring = weights\n", + " plt.scatter(positions[:, 0], positions[:, 1], c=coloring, s=.25)\n", + " plt.grid(c='k')\n", + " plt.axis('square')\n", + " plt.title('$\\hat{n}^{\\t{vol}}$ initial scatter (random)')\n", + " plt.xlim(l1, r1)\n", + " plt.ylim(l2, r2)\n", + " plt.colorbar()\n", + "\n", + "plt.subplot(3, 2, 4)\n", + "ax = plt.gca()\n", + "ax.set_xticks(np.linspace(0, 1, nx + 1))\n", + "ax.set_yticks(np.linspace(0, 1., ny + 1))\n", + "plt.tick_params(labelbottom = False) \n", + "plt.pcolor(ee1[:,:,0], ee2[:,:,0], n_sph[:,:,0])\n", + "plt.grid()\n", + "plt.axis('square')\n", + "plt.title(f'n_sph initial (random)')\n", + "plt.colorbar()\n", + "\n", + "plt.subplot(3, 2, 5)\n", + "ax = plt.gca()\n", + "plt.pcolor(bc_x, bc_y, f_bin)\n", + "plt.axis('square')\n", + "plt.title(f'n_binned initial (random)')\n", + "plt.colorbar()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d86a7691", + "metadata": {}, + "outputs": [], + "source": [ + "dt = time_opts.dt\n", + "Nt = simdata.time_grid_size - 1\n", + "\n", + "positions = orbits[:, :, :3]\n", + "\n", + "interval = Nt/10\n", + "plot_ct = 0\n", + "\n", + "plt.figure(figsize=(12, 24))\n", + "for i in range(Nt):\n", + " if i % interval == 0:\n", + " print(f'{i = }')\n", + " plot_ct += 1\n", + " plt.subplot(4, 2, plot_ct)\n", + " ax = plt.gca() \n", + " coloring = weights\n", + " plt.scatter(positions[i, :, 0], positions[i, :, 1], c=coloring, s=.25)\n", + " plt.axis('square')\n", + " plt.title('n0_scatter')\n", + " plt.xlim(l1, r1)\n", + " plt.ylim(l2, r2)\n", + " plt.colorbar()\n", + " plt.title(f'Gas at t={i*dt}')\n", + " if plot_ct == 8:\n", + " break" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "env", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From d3d46ca56aa4fcac0a69f644c6bfcec5c095bbd9 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 1 Oct 2025 08:45:10 +0000 Subject: [PATCH 125/292] Recent updates from devel into 318 --- src/struphy/console/format.py | 81 +++- src/struphy/console/main.py | 2 +- src/struphy/io/inp/parameters.yml | 4 +- src/struphy/models/__init__.py | 74 ++++ src/struphy/pic/particles.py | 2 +- src/struphy/pic/pushing/eval_kernels_gc.py | 415 ++++++++++++++++++ src/struphy/pic/pushing/pusher_kernels.py | 140 ++++++ src/struphy/pic/tests/test_sorting.py | 20 + src/struphy/propagators/__init__.py | 95 ++++ .../propagators/propagators_markers.py | 131 ++++++ .../tests/test_gyrokinetic_poisson.py | 2 +- src/struphy/propagators/tests/test_poisson.py | 2 +- 12 files changed, 958 insertions(+), 10 deletions(-) diff --git a/src/struphy/console/format.py b/src/struphy/console/format.py index 86af155d4..7db28daed 100644 --- a/src/struphy/console/format.py +++ b/src/struphy/console/format.py @@ -97,6 +97,10 @@ PASS_GREEN = f"{GREEN_COLOR}PASS{BLACK_COLOR}" +MODELS_INIT_PATH = os.path.join(LIBPATH, "models/__init__.py") +PROPAGATORS_INIT_PATH = os.path.join(LIBPATH, "propagators/__init__.py") + + def check_omp_flags(file_path, verbose=False): """Checks if a file contains incorrect OpenMP-like flags (`# $`). @@ -412,7 +416,9 @@ def get_python_files(input_type, path=None): python_files = parse_path(LIBPATH) elif input_type == "path": + print(path) if os.path.isfile(path): + print("isfile") python_files = [path] else: python_files = parse_path(path) @@ -464,8 +470,9 @@ def get_python_files(input_type, path=None): print("No Python files found to check.") return [] - python_files = [file for file in python_files if not re.search(r"__\w+__", file)] - + python_files = [ + f for f in python_files if not (re.match(r"^__\w+__\.py$", os.path.basename(f)) and "__init__.py" not in f) + ] return python_files @@ -669,7 +676,7 @@ def struphy_format(config, verbose, yes=False): config : dict Configuration dictionary containing the following keys: - input_type : str, optional - The type of files to format ('all', 'path', 'staged', or 'branch'). Defaults to 'all'. + The type of files to format ('all', 'path', 'staged', 'branch', or '__init__.py'). Defaults to 'all'. - path : str, optional Directory or file path where files will be formatted. - linters : list @@ -693,8 +700,24 @@ def struphy_format(config, verbose, yes=False): if input_type is None and path is not None: input_type = "path" - python_files = get_python_files(input_type, path) + if input_type == "__init__.py": + print(f"Rewriting {PROPAGATORS_INIT_PATH}") + propagators_init = construct_propagators_init_file() + with open(PROPAGATORS_INIT_PATH, "w") as f: + f.write(propagators_init) + + print(f"Rewriting {MODELS_INIT_PATH}") + models_init = construct_models_init_file() + with open(MODELS_INIT_PATH, "w") as f: + f.write(models_init) + + python_files = [PROPAGATORS_INIT_PATH, MODELS_INIT_PATH] + input_type = "path" + else: + python_files = get_python_files(input_type, path) + if len(python_files) == 0: + print("No Python files to format.") sys.exit(0) confirm_formatting(python_files, linters, yes) @@ -1302,3 +1325,53 @@ def parse_json_file_to_html(json_file_path, html_output_path): print(f"Error: Failed to parse JSON file. {e}") except Exception as e: print(f"An unexpected error occurred: {e}") + + +def construct_models_init_file() -> str: + """ + Constructs the content for the __init__.py file for the models module. + + Returns: + str: The content for the __init__.py file as a string. + """ + import struphy.models.fluid as fluid + import struphy.models.hybrid as hybrid + import struphy.models.kinetic as kinetic + import struphy.models.toy as toy + from struphy.models.base import StruphyModel + + models_init = "" + + model_names = [] + for model_type in [toy, fluid, hybrid, kinetic]: + for _, cls in model_type.__dict__.items(): + if isinstance(cls, type) and issubclass(cls, StruphyModel) and cls != StruphyModel: + model_names.append(cls.__name__) + models_init += f"from {model_type.__name__} import {cls.__name__}\n" + models_init += "\n\n" + models_init += f"__all__ = {model_names}\n" + return models_init + + +def construct_propagators_init_file() -> str: + """ + Constructs the content for the __init__.py file for the propagators module. + + Returns: + str: The content for the __init__.py file as a string. + """ + import struphy.propagators.propagators_coupling as propagators_coupling + import struphy.propagators.propagators_fields as propagators_fields + import struphy.propagators.propagators_markers as propagators_markers + from struphy.propagators.base import Propagator + + propagators_init = "" + propagators_names = [] + for model_type in [propagators_coupling, propagators_fields, propagators_markers]: + for _, cls in model_type.__dict__.items(): + if isinstance(cls, type) and issubclass(cls, Propagator) and cls != Propagator: + propagators_names.append(cls.__name__) + propagators_init += f"from {model_type.__name__} import {cls.__name__}\n" + propagators_init += "\n\n" + propagators_init += f"__all__ = {propagators_names}\n" + return propagators_init diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index ee3f46c41..19832a830 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -1013,7 +1013,7 @@ def add_parser_format(subparsers): subparser.add_argument( "input_type", type=str, - choices=["all", "staged", "branch"], + choices=["all", "staged", "branch", "__init__.py"], nargs="?", # optional help="specify the files to process", ) diff --git a/src/struphy/io/inp/parameters.yml b/src/struphy/io/inp/parameters.yml index 7a16d3d8f..73fbf48ae 100644 --- a/src/struphy/io/inp/parameters.yml +++ b/src/struphy/io/inp/parameters.yml @@ -86,7 +86,7 @@ kinetic : ppc : null # number of markers per 3d grid cell ppb : null # number of markers per sorting box bc : [periodic, periodic, periodic] # marker boundary conditions: remove, reflect or periodic - bufsize : 1. # MPI send/receive buffer + bufsize : 2. # MPI send/receive buffer loading : pseudo_random # particle loading mechanism loading_params : seed : 1234 # seed for random number generator @@ -97,7 +97,7 @@ kinetic : weights : reject_weights : False threshold : 0.0 - boxes_per_dim: [16, 16, 1] # for particle sorting and sph evaluations via smoothing kernels + boxes_per_dim: [12, 8, 1] # for particle sorting and sph evaluations via smoothing kernels dims_mask: [True, True, True] # which directions to use in domain decomposition (for meshless methods, otherwise given under grid) save_data : n_markers : 3 # number of markers (int) to be saved during simulation; if float and <= 1.0, it is a perecentage of Np diff --git a/src/struphy/models/__init__.py b/src/struphy/models/__init__.py index e69de29bb..0c7aef8c7 100644 --- a/src/struphy/models/__init__.py +++ b/src/struphy/models/__init__.py @@ -0,0 +1,74 @@ +from struphy.models.fluid import ( + ColdPlasma, + HasegawaWakatani, + EulerSPH, + LinearExtendedMHDuniform, + LinearMHD, + ViscoresistiveDeltafMHD, + ViscoresistiveDeltafMHD_with_q, + ViscoresistiveLinearMHD, + ViscoresistiveLinearMHD_with_q, + ViscoresistiveMHD, + ViscoresistiveMHD_with_p, + ViscoresistiveMHD_with_q, + ViscousFluid, +) +from struphy.models.hybrid import ColdPlasmaVlasov, LinearMHDDriftkineticCC, LinearMHDVlasovCC, LinearMHDVlasovPC +from struphy.models.kinetic import ( + DriftKineticElectrostaticAdiabatic, + LinearVlasovAmpereOneSpecies, + LinearVlasovMaxwellOneSpecies, + VlasovAmpereOneSpecies, + VlasovMaxwellOneSpecies, +) +from struphy.models.toy import ( + DeterministicParticleDiffusion, + GuidingCenter, + Maxwell, + Poisson, + PressureLessSPH, + RandomParticleDiffusion, + ShearAlfven, + TwoFluidQuasiNeutralToy, + VariationalBarotropicFluid, + VariationalCompressibleFluid, + VariationalPressurelessFluid, + Vlasov, +) + +__all__ = [ + "Maxwell", + "Vlasov", + "GuidingCenter", + "ShearAlfven", + "VariationalPressurelessFluid", + "VariationalBarotropicFluid", + "VariationalCompressibleFluid", + "Poisson", + "DeterministicParticleDiffusion", + "RandomParticleDiffusion", + "PressureLessSPH", + "TwoFluidQuasiNeutralToy", + "LinearMHD", + "LinearExtendedMHDuniform", + "ColdPlasma", + "ViscoresistiveMHD", + "ViscousFluid", + "ViscoresistiveMHD_with_p", + "ViscoresistiveLinearMHD", + "ViscoresistiveDeltafMHD", + "ViscoresistiveMHD_with_q", + "ViscoresistiveLinearMHD_with_q", + "ViscoresistiveDeltafMHD_with_q", + "IsothermalEulerSPH", + "HasegawaWakatani", + "LinearMHDVlasovCC", + "LinearMHDVlasovPC", + "LinearMHDDriftkineticCC", + "ColdPlasmaVlasov", + "VlasovAmpereOneSpecies", + "VlasovMaxwellOneSpecies", + "LinearVlasovAmpereOneSpecies", + "LinearVlasovMaxwellOneSpecies", + "DriftKineticElectrostaticAdiabatic", +] diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 1252a530e..35df68d04 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -802,7 +802,7 @@ def __init__( # default number of diagnostics and auxiliary columns self._n_cols_diagnostics = kwargs.pop("n_cols_diagn", 0) - self._n_cols_aux = kwargs.pop("n_cols_aux", 5) + self._n_cols_aux = kwargs.pop("n_cols_aux", 24) clone_config = kwargs.get("clone_config", None) assert clone_config is None, "SPH can only be launched with --nclones 1" diff --git a/src/struphy/pic/pushing/eval_kernels_gc.py b/src/struphy/pic/pushing/eval_kernels_gc.py index ed4fb66c4..5e4dfe9a8 100644 --- a/src/struphy/pic/pushing/eval_kernels_gc.py +++ b/src/struphy/pic/pushing/eval_kernels_gc.py @@ -595,3 +595,418 @@ def sph_isotherm_kappa( continue markers[ip, first_diagnostic_idx] = 1.0 + + +@stack_array("eta_k", "eta_n", "eta", "grad_H", "e_field") +def sph_mean_velocity_coeffs( + alpha: "float[:]", + column_nr: int, + comps: "int[:]", + args_markers: "MarkerArguments", + args_domain: "DomainArguments", + boxes: "int[:, :]", + neighbours: "int[:, :]", + holes: "bool[:]", + periodic1: "bool", + periodic2: "bool", + periodic3: "bool", + kernel_type: "int", + h1: "float", + h2: "float", + h3: "float", +): + r"""Evaluate the :math:`\boldsymbol \eta`-gradient of the Hamiltonian + + .. math:: + + H(\mathbf Z_p) = H(\boldsymbol \eta_p, v_{\parallel,p}) = \varepsilon \frac{v_{\parallel,p}^2}{2} + + \varepsilon \mu |\hat \mathbf B| (\boldsymbol \eta_p) + \hat \phi(\boldsymbol \eta_p)\,, + + that is + + .. math:: + + \hat \nabla H(\mathbf Z_p) = \varepsilon \mu \hat \nabla |\hat \mathbf B| (\boldsymbol \eta_p) + + \hat \nabla \hat \phi(\boldsymbol \eta_p)\,, + + where the evaluation point is the weighted average + :math:`Z_{p,i} = \alpha_i Z_{p,i}^{n+1,k} + (1 - \alpha_i) Z_{p,i}^n`, + for :math:`i=1,2,3,4`. Markers must be sorted according to the evaluation point + :math:`\boldsymbol \eta_p` beforehand. + + The components specified in ``comps`` are save at ``column_nr:column_nr + len(comps)`` + in markers array for each particle. + """ + + gamma = 5 / 3 + + # get marker arguments + markers = args_markers.markers + n_markers = args_markers.n_markers + n_cols = shape(markers)[1] + Np = args_markers.Np + vdim = args_markers.vdim + weight_idx = args_markers.weight_idx + valid_mks = args_markers.valid_mks + + for ip in range(n_markers): + # only do something if particle is a "true" particle + if not valid_mks[ip]: + continue + + eta1 = markers[ip, 0] + eta2 = markers[ip, 1] + eta3 = markers[ip, 2] + loc_box = int(markers[ip, n_cols - 2]) + n_at_eta = sph_eval_kernels.boxed_based_kernel( + args_markers, + eta1, + eta2, + eta3, + loc_box, + boxes, + neighbours, + holes, + periodic1, + periodic2, + periodic3, + weight_idx, + kernel_type, + h1, + h2, + h3, + ) + weight = markers[ip, weight_idx] + velocities = markers[ip, 3:6] + # save + markers[ip, column_nr] = weight / n_at_eta * velocities[0] + markers[ip, column_nr + 1] = weight / n_at_eta * velocities[1] + markers[ip, column_nr + 2] = weight / n_at_eta * velocities[2] + + +@stack_array("eta_k", "eta_n", "eta", "grad_H", "e_field") +def sph_mean_velocity( + alpha: "float[:]", + column_nr: int, + comps: "int[:]", + args_markers: "MarkerArguments", + args_domain: "DomainArguments", + boxes: "int[:, :]", + neighbours: "int[:, :]", + holes: "bool[:]", + periodic1: "bool", + periodic2: "bool", + periodic3: "bool", + kernel_type: "int", + h1: "float", + h2: "float", + h3: "float", +): + r"""Evaluate the :math:`\boldsymbol \eta`-gradient of the Hamiltonian + + .. math:: + + H(\mathbf Z_p) = H(\boldsymbol \eta_p, v_{\parallel,p}) = \varepsilon \frac{v_{\parallel,p}^2}{2} + + \varepsilon \mu |\hat \mathbf B| (\boldsymbol \eta_p) + \hat \phi(\boldsymbol \eta_p)\,, + + that is + + .. math:: + + \hat \nabla H(\mathbf Z_p) = \varepsilon \mu \hat \nabla |\hat \mathbf B| (\boldsymbol \eta_p) + + \hat \nabla \hat \phi(\boldsymbol \eta_p)\,, + + where the evaluation point is the weighted average + :math:`Z_{p,i} = \alpha_i Z_{p,i}^{n+1,k} + (1 - \alpha_i) Z_{p,i}^n`, + for :math:`i=1,2,3,4`. Markers must be sorted according to the evaluation point + :math:`\boldsymbol \eta_p` beforehand. + + The components specified in ``comps`` are save at ``column_nr:column_nr + len(comps)`` + in markers array for each particle. + """ + + gamma = 5 / 3 + + # get marker arguments + markers = args_markers.markers + n_markers = args_markers.n_markers + n_cols = shape(markers)[1] + Np = args_markers.Np + vdim = args_markers.vdim + weight_idx = args_markers.weight_idx + first_free_idx = args_markers.first_free_idx + valid_mks = args_markers.valid_mks + + for ip in range(n_markers): + # only do something if particle is a "true" particle + if not valid_mks[ip]: + continue + + eta1 = markers[ip, 0] + eta2 = markers[ip, 1] + eta3 = markers[ip, 2] + loc_box = int(markers[ip, n_cols - 2]) + v1_at_eta = sph_eval_kernels.boxed_based_kernel( + args_markers, + eta1, + eta2, + eta3, + loc_box, + boxes, + neighbours, + holes, + periodic1, + periodic2, + periodic3, + first_free_idx, + kernel_type, + h1, + h2, + h3, + ) + + v2_at_eta = sph_eval_kernels.boxed_based_kernel( + args_markers, + eta1, + eta2, + eta3, + loc_box, + boxes, + neighbours, + holes, + periodic1, + periodic2, + periodic3, + first_free_idx + 1, + kernel_type, + h1, + h2, + h3, + ) + + v3_at_eta = sph_eval_kernels.boxed_based_kernel( + args_markers, + eta1, + eta2, + eta3, + loc_box, + boxes, + neighbours, + holes, + periodic1, + periodic2, + periodic3, + first_free_idx + 2, + kernel_type, + h1, + h2, + h3, + ) + # save + markers[ip, column_nr] = v1_at_eta + markers[ip, column_nr + 1] = v2_at_eta + markers[ip, column_nr + 2] = v3_at_eta + + +@stack_array("eta_k", "eta_n", "eta", "grad_H", "e_field") +def sph_grad_mean_velocity( + alpha: "float[:]", + column_nr: int, + comps: "int[:]", + args_markers: "MarkerArguments", + args_domain: "DomainArguments", + boxes: "int[:, :]", + neighbours: "int[:, :]", + holes: "bool[:]", + periodic1: "bool", + periodic2: "bool", + periodic3: "bool", + kernel_type: "int", + h1: "float", + h2: "float", + h3: "float", +): + r"""Evaluate the :math:`\boldsymbol \eta`-gradient of the Hamiltonian + + .. math:: + + H(\mathbf Z_p) = H(\boldsymbol \eta_p, v_{\parallel,p}) = \varepsilon \frac{v_{\parallel,p}^2}{2} + + \varepsilon \mu |\hat \mathbf B| (\boldsymbol \eta_p) + \hat \phi(\boldsymbol \eta_p)\,, + + that is + + .. math:: + + \hat \nabla H(\mathbf Z_p) = \varepsilon \mu \hat \nabla |\hat \mathbf B| (\boldsymbol \eta_p) + + \hat \nabla \hat \phi(\boldsymbol \eta_p)\,, + + where the evaluation point is the weighted average + :math:`Z_{p,i} = \alpha_i Z_{p,i}^{n+1,k} + (1 - \alpha_i) Z_{p,i}^n`, + for :math:`i=1,2,3,4`. Markers must be sorted according to the evaluation point + :math:`\boldsymbol \eta_p` beforehand. + + The components specified in ``comps`` are save at ``column_nr:column_nr + len(comps)`` + in markers array for each particle. + """ + + gamma = 5 / 3 + + # get marker arguments + markers = args_markers.markers + n_markers = args_markers.n_markers + n_cols = shape(markers)[1] + Np = args_markers.Np + vdim = args_markers.vdim + weight_idx = args_markers.weight_idx + first_free_idx = args_markers.first_free_idx + valid_mks = args_markers.valid_mks + + grad_v_at_eta = zeros((3, 3), dtype=float) + for ip in range(n_markers): + # only do something if particle is a "true" particle + if not valid_mks[ip]: + continue + + eta1 = markers[ip, 0] + eta2 = markers[ip, 1] + eta3 = markers[ip, 2] + loc_box = int(markers[ip, n_cols - 2]) + for j in range(3): + for k in range(3): + grad_v_at_eta[j, k] = sph_eval_kernels.boxed_based_kernel( + args_markers, + eta1, + eta2, + eta3, + loc_box, + boxes, + neighbours, + holes, + periodic1, + periodic2, + periodic3, + first_free_idx + j, + kernel_type + 1 + k, + h1, + h2, + h3, + ) + + # save + markers[ip, column_nr + 3 * j + k] = grad_v_at_eta[j, k] + + +@stack_array("eta_k", "eta_n", "eta", "grad_H", "e_field") +def sph_viscosity_tensor( + alpha: "float[:]", + column_nr: int, + comps: "int[:]", + args_markers: "MarkerArguments", + args_domain: "DomainArguments", + boxes: "int[:, :]", + neighbours: "int[:, :]", + holes: "bool[:]", + periodic1: "bool", + periodic2: "bool", + periodic3: "bool", + kernel_type: "int", + h1: "float", + h2: "float", + h3: "float", +): + r"""Evaluate the :math:`\boldsymbol \eta`-gradient of the Hamiltonian + + .. math:: + + H(\mathbf Z_p) = H(\boldsymbol \eta_p, v_{\parallel,p}) = \varepsilon \frac{v_{\parallel,p}^2}{2} + + \varepsilon \mu |\hat \mathbf B| (\boldsymbol \eta_p) + \hat \phi(\boldsymbol \eta_p)\,, + + that is + + .. math:: + + \hat \nabla H(\mathbf Z_p) = \varepsilon \mu \hat \nabla |\hat \mathbf B| (\boldsymbol \eta_p) + + \hat \nabla \hat \phi(\boldsymbol \eta_p)\,, + + where the evaluation point is the weighted average + :math:`Z_{p,i} = \alpha_i Z_{p,i}^{n+1,k} + (1 - \alpha_i) Z_{p,i}^n`, + for :math:`i=1,2,3,4`. Markers must be sorted according to the evaluation point + :math:`\boldsymbol \eta_p` beforehand. + + The components specified in ``comps`` are save at ``column_nr:column_nr + len(comps)`` + in markers array for each particle. + """ + + gamma = 5 / 3 + + # get marker arguments + markers = args_markers.markers + n_markers = args_markers.n_markers + n_cols = shape(markers)[1] + Np = args_markers.Np + vdim = args_markers.vdim + weight_idx = args_markers.weight_idx + first_free_idx = args_markers.first_free_idx + valid_mks = args_markers.valid_mks + + grad_v_at_eta = zeros((3, 3), dtype=float) + d_dev = zeros((3, 3), dtype=float) + for ip in range(n_markers): + # only do something if particle is a "true" particle + if not valid_mks[ip]: + continue + + eta1 = markers[ip, 0] + eta2 = markers[ip, 1] + eta3 = markers[ip, 2] + loc_box = int(markers[ip, n_cols - 2]) + n_at_eta = sph_eval_kernels.boxed_based_kernel( + args_markers, + eta1, + eta2, + eta3, + loc_box, + boxes, + neighbours, + holes, + periodic1, + periodic2, + periodic3, + weight_idx, + kernel_type, + h1, + h2, + h3, + ) + weight = markers[ip, weight_idx] + for j in range(3): + for k in range(3): + grad_v_at_eta[j, k] = sph_eval_kernels.boxed_based_kernel( + args_markers, + eta1, + eta2, + eta3, + loc_box, + boxes, + neighbours, + holes, + periodic1, + periodic2, + periodic3, + first_free_idx + j, + kernel_type + 1 + k, + h1, + h2, + h3, + ) + + mu = 0.7 + d = 0.5 * (grad_v_at_eta + grad_v_at_eta.T) + trace_d = d[0, 0] + d[1, 1] + d[2, 2] + d_dev[0, 0] = d[0, 0] - (trace_d / 3.0) + d_dev[1, 1] = d[1, 1] - (trace_d / 3.0) + d_dev[2, 2] = d[2, 2] - (trace_d / 3.0) + d_dev *= 2 * mu * weight / n_at_eta + for j in range(3): + for k in range(3): + markers[ip, column_nr + 3 * j + k] = d_dev[j, k] diff --git a/src/struphy/pic/pushing/pusher_kernels.py b/src/struphy/pic/pushing/pusher_kernels.py index 42d1284ea..429e1c722 100644 --- a/src/struphy/pic/pushing/pusher_kernels.py +++ b/src/struphy/pic/pushing/pusher_kernels.py @@ -3445,3 +3445,143 @@ def push_v_sph_pressure_ideal_gas( markers[ip, 3:6] -= dt * (grad_u_cart - gravity) # -- removed omp: #$ omp end parallel + + +@stack_array("grad_u", "grad_u_cart", "tmp1", "dfinv", "dfinvT") +def push_v_viscosity( + dt: float, + stage: int, + args_markers: "MarkerArguments", + args_domain: "DomainArguments", + boxes: "int[:,:]", + neighbours: "int[:, :]", + holes: "bool[:]", + periodic1: "bool", + periodic2: "bool", + periodic3: "bool", + kernel_type: "int", + h1: "float", + h2: "float", + h3: "float", +): + r"""Updates particle velocities as + + .. math:: + + \frac{\mathbf v^{n+1} - \mathbf v^n}{\Delta t} = \kappa_p \sum_{q} w_p\,w_q \left( \frac{1}{\rho^{N,h}(\boldsymbol \eta_p)} + \frac{1}{\rho^{N,h}(\boldsymbol \eta_q)} \right) G^{-1}\nabla W_h(\boldsymbol \eta_p - \boldsymbol \eta_q) \,, + + where :math:`G^{-1}` denotes the inverse metric tensor, and with the smoothed density + + .. math:: + + \rho^{N,h}(\boldsymbol \eta_p) = \frac 1N \sum_q w_q \, W_h(\boldsymbol \eta_p - \boldsymbol \eta_q)\,, + + where :math:`W_h(\boldsymbol \eta)` is a smoothing kernel from :mod:`~struphy.pic.sph_smoothing_kernels`. + + Parameters + ---------- + boxes : 2d array + Box array of the sorting boxes structure. + + neighbours : 2d array + Array containing the 27 neighbouring boxes of each box. + + holes : bool + 1D array of length markers.shape[0]. True if markers[i] is a hole. + + periodic1, periodic2, periodic3 : bool + True if periodic in that dimension. + + kernel_type : int + Number of the smoothing kernel. + + h1, h2, h3 : float + Kernel width in respective dimension. + + gravity: np.ndarray + Constant gravitational force as 3-vector. + """ + # allocate arrays + grad_u = zeros(3, dtype=float) + grad_u_cart = zeros(3, dtype=float) + tmp1 = zeros((3, 3), dtype=float) + dfinv = zeros((3, 3), dtype=float) + dfinvT = zeros((3, 3), dtype=float) + + # get marker arguments + markers = args_markers.markers + n_markers = args_markers.n_markers + Np = args_markers.Np + weight_idx = args_markers.weight_idx + first_free_idx = args_markers.first_free_idx + valid_mks = args_markers.valid_mks + n_cols = shape(markers)[1] + f_visc = zeros(3, dtype=float) + f_visc_cart = zeros(3, dtype=float) + + # -- removed omp: #$ omp parallel private(ip, eta1, eta2, eta3, dfinv) + # -- removed omp: #$ omp for + for ip in range(n_markers): + if not valid_mks[ip]: + continue + + eta1 = markers[ip, 0] + eta2 = markers[ip, 1] + eta3 = markers[ip, 2] + kappa = 1.0 # markers[ip, first_diagnostics_idx] + # n_at_eta = markers[ip, first_free_idx] + loc_box = int(markers[ip, n_cols - 2]) + + for j in range(3): # row of viscosity tensor + for k in range(3): # column = derivative direction + coeff_idx = first_free_idx + 3 * j + k + 15 + + # if k == 0: + # deriv_type = kernel_type + 1 + # use_component = True + # elif k == 1 and kernel_type >= 340: + # deriv_type = kernel_type + 2 + # use_component = True + # elif k == 2 and kernel_type >= 670: + # deriv_type = kernel_type + 3 + # use_component = True + # else: + # use_component = False + + # if use_component: + f_visc[j] += sph_eval_kernels.boxed_based_kernel( + args_markers, + eta1, + eta2, + eta3, + loc_box, + boxes, + neighbours, + holes, + periodic1, + periodic2, + periodic3, + coeff_idx, + kernel_type + 1 + k, + h1, + h2, + h3, + ) + + # push to Cartesian coordinates + evaluation_kernels.df_inv( + eta1, + eta2, + eta3, + args_domain, + tmp1, + False, + dfinv, + ) + linalg_kernels.transpose(dfinv, dfinvT) + linalg_kernels.matrix_vector(dfinvT, f_visc, f_visc_cart) + + # update velocities + markers[ip, 3:6] -= dt * (f_visc_cart) + + # -- removed omp: #$ omp end parallel diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index ce404d3e2..85f8f2935 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -10,6 +10,26 @@ from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters +@pytest.mark.parametrize("nx", [8, 70]) +@pytest.mark.parametrize("ny", [16, 80]) +@pytest.mark.parametrize("nz", [32, 90]) +@pytest.mark.parametrize("algo", ["fortran_ordering", "c_ordering"]) +def test_flattening(nx, ny, nz, algo): + from struphy.pic.sorting_kernels import flatten_index, unflatten_index + + n1s = np.array(np.random.rand(10) * (nx + 1), dtype=int) + n2s = np.array(np.random.rand(10) * (ny + 1), dtype=int) + n3s = np.array(np.random.rand(10) * (nz + 1), dtype=int) + for n1 in n1s: + for n2 in n2s: + for n3 in n3s: + n_glob = flatten_index(int(n1), int(n2), int(n3), nx, ny, nz, algo) + n1n, n2n, n3n = unflatten_index(n_glob, nx, ny, nz, algo) + assert n1n == n1 + assert n2n == n2 + assert n3n == n3 + + @pytest.mark.parametrize("nx", [8, 70]) @pytest.mark.parametrize("ny", [16, 80]) @pytest.mark.parametrize("nz", [32, 90]) diff --git a/src/struphy/propagators/__init__.py b/src/struphy/propagators/__init__.py index e69de29bb..3f61b5c9a 100644 --- a/src/struphy/propagators/__init__.py +++ b/src/struphy/propagators/__init__.py @@ -0,0 +1,95 @@ +from struphy.propagators.propagators_coupling import ( + CurrentCoupling5DCurlb, + CurrentCoupling5DGradB, + CurrentCoupling6DCurrent, + EfieldWeights, + PressureCoupling6D, + VlasovAmpere, +) +from struphy.propagators.propagators_fields import ( + AdiabaticPhi, + CurrentCoupling5DDensity, + CurrentCoupling6DDensity, + FaradayExtended, + Hall, + HasegawaWakatani, + ImplicitDiffusion, + JxBCold, + Magnetosonic, + MagnetosonicCurrentCoupling5D, + MagnetosonicUniform, + Maxwell, + OhmCold, + Poisson, + ShearAlfven, + ShearAlfvenB1, + ShearAlfvenCurrentCoupling5D, + TimeDependentSource, + TwoFluidQuasiNeutralFull, + VariationalDensityEvolve, + VariationalEntropyEvolve, + VariationalMagFieldEvolve, + VariationalMomentumAdvection, + VariationalPBEvolve, + VariationalQBEvolve, + VariationalResistivity, + VariationalViscosity, +) +from struphy.propagators.propagators_markers import ( + PushDeterministicDiffusion, + PushEta, + PushEtaPC, + PushGuidingCenterBxEstar, + PushGuidingCenterParallel, + PushRandomDiffusion, + PushVinEfield, + PushVinSPHpressure, + PushVxB, + StepStaticEfield, +) + +__all__ = [ + "VlasovAmpere", + "EfieldWeights", + "PressureCoupling6D", + "CurrentCoupling6DCurrent", + "CurrentCoupling5DCurlb", + "CurrentCoupling5DGradB", + "Maxwell", + "OhmCold", + "JxBCold", + "ShearAlfven", + "ShearAlfvenB1", + "Hall", + "Magnetosonic", + "MagnetosonicUniform", + "FaradayExtended", + "CurrentCoupling6DDensity", + "ShearAlfvenCurrentCoupling5D", + "MagnetosonicCurrentCoupling5D", + "CurrentCoupling5DDensity", + "ImplicitDiffusion", + "Poisson", + "VariationalMomentumAdvection", + "VariationalDensityEvolve", + "VariationalEntropyEvolve", + "VariationalMagFieldEvolve", + "VariationalPBEvolve", + "VariationalQBEvolve", + "VariationalViscosity", + "VariationalResistivity", + "TimeDependentSource", + "AdiabaticPhi", + "HasegawaWakatani", + "TwoFluidQuasiNeutralFull", + "PushEta", + "PushVxB", + "PushVinEfield", + "PushEtaPC", + "PushGuidingCenterBxEstar", + "PushGuidingCenterParallel", + "StepStaticEfield", + "PushDeterministicDiffusion", + "PushRandomDiffusion", + "PushVinSPHpressure", +] diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index e046bd442..a890e5d27 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -1787,3 +1787,134 @@ def allocate(self): def __call__(self, dt): self.variables.fluid.particles.put_particles_in_boxes() self._pusher(dt) + + +class PushVinViscousPotential(Propagator): + r"""For each marker :math:`p`, solves + + .. math:: + + \frac{\textnormal d \mathbf v_p(t)}{\textnormal d t} = \kappa_p \sum_{i=1}^N w_i \left( \frac{1}{\rho^{N,h}(\boldsymbol \eta_p)} + \frac{1}{\rho^{N,h}(\boldsymbol \eta_i)} \right) DF^{-\top}\nabla W_h(\boldsymbol \eta_p - \boldsymbol \eta_i) \,, + + where :math:`DF^{-\top}` denotes the inverse transpose Jacobian, and with the smoothed density + + .. math:: + + \rho^{N,h}(\boldsymbol \eta) = \frac 1N \sum_{j=1}^N w_j \, W_h(\boldsymbol \eta - \boldsymbol \eta_j)\,, + + where :math:`W_h(\boldsymbol \eta)` is a smoothing kernel from :mod:`~struphy.pic.sph_smoothing_kernels`. + Time stepping: + + * Explicit from :class:`~struphy.ode.utils.ButcherTableau` + """ + + @staticmethod + def options(default=False): + dct = {} + dct["kernel_type"] = list(Particles.ker_dct()) + dct["kernel_width"] = None + dct["algo"] = [ + "forward_euler", + ] # "heun2", "rk2", "heun3", "rk4"] + if default: + dct = descend_options_dict(dct, []) + return dct + + def __init__( + self, + particles: ParticlesSPH, + *, + kernel_type: str = "gaussian_2d", + kernel_width: tuple = None, + algo: str = options(default=True)["algo"], # TODO: implement other algos than forward Euler + ): + # base class constructor call + super().__init__(particles) + + # init kernel for evaluating density etc. before each time step. + init_kernel_1 = eval_kernels_gc.sph_mean_velocity_coeffs + first_free_idx = particles.args_markers.first_free_idx + comps = (0, 1, 2) + + init_kernel_2 = eval_kernels_gc.sph_mean_velocity + # first_free_idx = particles.args_markers.first_free_idx + # comps = (0, 1, 2) + + init_kernel_3 = eval_kernels_gc.sph_grad_mean_velocity + comps_tensor = (0, 1, 2, 3, 4, 5, 6, 7, 8) + + init_kernel_4 = eval_kernels_gc.sph_viscosity_tensor + + boxes = particles.sorting_boxes.boxes + neighbours = particles.sorting_boxes.neighbours + holes = particles.holes + periodic = [bci == "periodic" for bci in particles.bc] + kernel_nr = particles.ker_dct()[kernel_type] + + if kernel_width is None: + kernel_width = tuple([1 / ni for ni in self.particles[0].boxes_per_dim]) + else: + assert all([hi <= 1 / ni for hi, ni in zip(kernel_width, self.particles[0].boxes_per_dim)]) + + # init kernel + args_init = ( + boxes, + neighbours, + holes, + *periodic, + kernel_nr, + *kernel_width, + ) + + self.add_init_kernel( + init_kernel_1, + first_free_idx, + comps, + args_init, + ) + + self.add_init_kernel( + init_kernel_2, + first_free_idx + 3, # +3 so that the previous one is not overwritten + comps, + args_init, + ) + + self.add_init_kernel( + init_kernel_3, + first_free_idx + 6, # +3 so that the previous one is not overwritten + comps_tensor, + args_init, + ) + + self.add_init_kernel( + init_kernel_4, + first_free_idx + 15, + comps_tensor, + args_init, + ) + + kernel = pusher_kernels.push_v_viscosity + + args_kernel = ( + boxes, + neighbours, + holes, + *periodic, + kernel_nr, + *kernel_width, + ) + + # the Pusher class wraps around all kernels + self._pusher = Pusher( + particles, + kernel, + args_kernel, + self.domain.args_domain, + alpha_in_kernel=0.0, + init_kernels=self.init_kernels, + ) + + def __call__(self, dt): + self.particles[0].put_particles_in_boxes() + self._pusher(dt) diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index d5cfc2554..4d2d7087e 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -9,8 +9,8 @@ from struphy.geometry import domains from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable +from struphy.propagators import ImplicitDiffusion from struphy.propagators.base import Propagator -from struphy.propagators.propagators_fields import ImplicitDiffusion comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 2460968b6..5a1de11ba 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -9,8 +9,8 @@ from struphy.geometry import domains from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable +from struphy.propagators import ImplicitDiffusion from struphy.propagators.base import Propagator -from struphy.propagators.propagators_fields import ImplicitDiffusion comm = MPI.COMM_WORLD rank = comm.Get_rank() From 649eb929cbbad58103292e4e8b87149ef35c6967 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 1 Oct 2025 13:52:16 +0000 Subject: [PATCH 126/292] Updated __init__.py files in 318 --- src/struphy/models/__init__.py | 4 ++-- src/struphy/propagators/__init__.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/struphy/models/__init__.py b/src/struphy/models/__init__.py index 0c7aef8c7..3878287f4 100644 --- a/src/struphy/models/__init__.py +++ b/src/struphy/models/__init__.py @@ -1,7 +1,7 @@ from struphy.models.fluid import ( ColdPlasma, - HasegawaWakatani, EulerSPH, + HasegawaWakatani, LinearExtendedMHDuniform, LinearMHD, ViscoresistiveDeltafMHD, @@ -60,7 +60,7 @@ "ViscoresistiveMHD_with_q", "ViscoresistiveLinearMHD_with_q", "ViscoresistiveDeltafMHD_with_q", - "IsothermalEulerSPH", + "EulerSPH", "HasegawaWakatani", "LinearMHDVlasovCC", "LinearMHDVlasovPC", diff --git a/src/struphy/propagators/__init__.py b/src/struphy/propagators/__init__.py index 3f61b5c9a..2dbef2b10 100644 --- a/src/struphy/propagators/__init__.py +++ b/src/struphy/propagators/__init__.py @@ -44,6 +44,7 @@ PushRandomDiffusion, PushVinEfield, PushVinSPHpressure, + PushVinViscousPotential, PushVxB, StepStaticEfield, ) @@ -92,4 +93,5 @@ "PushDeterministicDiffusion", "PushRandomDiffusion", "PushVinSPHpressure", + "PushVinViscousPotential", ] From 1952c26ad82aa2f107571962cfda3a9b820ed2c1 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 6 Oct 2025 06:36:09 +0000 Subject: [PATCH 127/292] Add SPH and VlasovAmpere verification tests --- src/struphy/main.py | 17 +- src/struphy/models/base.py | 89 ++++++---- src/struphy/models/kinetic.py | 5 +- src/struphy/models/species.py | 9 +- .../models/tests/test_verif_EulerSPH.py | 166 +++++++++++++++++ .../test_verif_VlasovAmpereOneSpecies.py | 168 ++++++++++++++++++ src/struphy/pic/base.py | 3 +- src/struphy/pic/utilities.py | 6 +- .../propagators/propagators_markers.py | 1 + src/struphy/topology/grids.py | 2 +- 10 files changed, 400 insertions(+), 66 deletions(-) create mode 100644 src/struphy/models/tests/test_verif_EulerSPH.py create mode 100644 src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py diff --git a/src/struphy/main.py b/src/struphy/main.py index 5955f16f8..9592db6bd 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -191,22 +191,7 @@ def run( # domain and fluid background model.setup_domain_and_equil(domain, equil) - # default grid - if grid is None: - Nel = (16, 16, 16) - if rank == 0: - print(f"\nNo grid specified - using TensorProductGrid with {Nel = }.") - grid = grids.TensorProductGrid(Nel=Nel) - - # allocate derham-related objects - if derham_opts is None: - p = (3, 3, 3) - spl_kind = (False, False, False) - if rank == 0: - print( - f"\nNo Derham options specified - creating Derham with {p = } and {spl_kind = } for projecting equilibrium." - ) - derham_opts = DerhamOptions(p=p, spl_kind=spl_kind) + # feec model.allocate_feec(grid, derham_opts) # equation paramters diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 780c3efff..458b1cfba 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -171,54 +171,67 @@ def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): else: derham_comm = self.clone_config.sub_comm - self._derham = setup_derham( - grid, - derham_opts, - comm=derham_comm, - domain=self.domain, - verbose=self.verbose, - ) + if grid is None or derham_opts is None: + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\n{grid = }, {derham_opts = }: no Derham object set up.") + self._derham = None + else: + self._derham = setup_derham( + grid, + derham_opts, + comm=derham_comm, + domain=self.domain, + verbose=self.verbose, + ) # create weighted mass operators - self._mass_ops = WeightedMassOperators( - self.derham, - self.domain, - verbose=self.verbose, - eq_mhd=self.equil, - ) - - # create projected equilibrium - if isinstance(self.equil, MHDequilibrium): - self._projected_equil = ProjectedMHDequilibrium( - self.equil, - self.derham, - ) - elif isinstance(self.equil, FluidEquilibriumWithB): - self._projected_equil = ProjectedFluidEquilibriumWithB( - self.equil, - self.derham, - ) - elif isinstance(self.equil, FluidEquilibrium): - self._projected_equil = ProjectedFluidEquilibrium( - self.equil, + if self.derham is None: + self._mass_ops = None + else: + self._mass_ops = WeightedMassOperators( self.derham, + self.domain, + verbose=self.verbose, + eq_mhd=self.equil, ) - else: + + # create projected equilibrium + if self.derham is None: self._projected_equil = None + else: + if isinstance(self.equil, MHDequilibrium): + self._projected_equil = ProjectedMHDequilibrium( + self.equil, + self.derham, + ) + elif isinstance(self.equil, FluidEquilibriumWithB): + self._projected_equil = ProjectedFluidEquilibriumWithB( + self.equil, + self.derham, + ) + elif isinstance(self.equil, FluidEquilibrium): + self._projected_equil = ProjectedFluidEquilibrium( + self.equil, + self.derham, + ) + else: + self._projected_equil = None def allocate_propagators(self): # set propagators base class attributes (then available to all propagators) Propagator.derham = self.derham Propagator.domain = self.domain - if self.derham is not None: - Propagator.mass_ops = self.mass_ops + Propagator.mass_ops = self.mass_ops + if self.derham is None: + Propagator.basis_ops = None + else: Propagator.basis_ops = BasisProjectionOperators( self.derham, self.domain, verbose=self.verbose, eq_mhd=self.equil, ) - Propagator.projected_equil = self.projected_equil + Propagator.projected_equil = self.projected_equil assert len(self.prop_list) > 0, "No propagators in this model, check the model class." for prop in self.prop_list: @@ -1394,12 +1407,12 @@ def generate_default_parameter_file( file.write("\n# fluid equilibrium (can be used as part of initial conditions)\n") file.write("equil = equils.HomogenSlab()\n") - if has_feec: - grid = "grid = grids.TensorProductGrid()\n" - derham = "derham_opts = DerhamOptions()\n" - else: - grid = "grid = None\n" - derham = "derham_opts = None\n" + # if has_feec: + grid = "grid = grids.TensorProductGrid()\n" + derham = "derham_opts = DerhamOptions()\n" + # else: + # grid = "grid = None\n" + # derham = "derham_opts = None\n" file.write("\n# grid\n") file.write(grid) diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 4f5dd89bd..0d6274d50 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -207,10 +207,7 @@ def allocate_propagators(self): alpha = self.kinetic_ions.equation_params.alpha epsilon = self.kinetic_ions.equation_params.epsilon - l2_proj = L2Projector(space_id="H1", mass_ops=self.mass_ops) - rho_coeffs = l2_proj.solve(charge_accum.vectors[0]) - - self.initial_poisson.options.rho = alpha**2 / epsilon * rho_coeffs + self.initial_poisson.options.rho = alpha**2 / epsilon * charge_accum.vectors[0] self.initial_poisson.allocate() # Solve with dt=1. and compute electric field diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 0c1a2311e..d70de2c0e 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -95,19 +95,22 @@ def __init__( self.alpha = om_p / om_c else: self.alpha = alpha - warnings.warn(f"Override equation parameter {self.alpha = }") + if MPI.COMM_WORLD.Get_rank() == 0: + warnings.warn(f"Override equation parameter {self.alpha = }") if epsilon is None: self.epsilon = 1.0 / (om_c * units.t) else: self.epsilon = epsilon - warnings.warn(f"Override equation parameter {self.epsilon = }") + if MPI.COMM_WORLD.Get_rank() == 0: + warnings.warn(f"Override equation parameter {self.epsilon = }") if kappa is None: self.kappa = om_p * units.t else: self.kappa = kappa - warnings.warn(f"Override equation parameter {self.kappa = }") + if MPI.COMM_WORLD.Get_rank() == 0: + warnings.warn(f"Override equation parameter {self.kappa = }") if verbose and MPI.COMM_WORLD.Get_rank() == 0: print(f"\nSet normalization parameters for species {species.__class__.__name__}:") diff --git a/src/struphy/models/tests/test_verif_EulerSPH.py b/src/struphy/models/tests/test_verif_EulerSPH.py new file mode 100644 index 000000000..fa83945e6 --- /dev/null +++ b/src/struphy/models/tests/test_verif_EulerSPH.py @@ -0,0 +1,166 @@ +import os + +import numpy as np +import pytest +from matplotlib import pyplot as plt +from matplotlib.ticker import FormatStrFormatter +from mpi4py import MPI + +from struphy import main +from struphy.fields_background import equils +from struphy.geometry import domains +from struphy.initial import perturbations +from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time +from struphy.kinetic_background import maxwellians +from struphy.pic.utilities import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, +) +from struphy.topology import grids + +test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") + + +@pytest.mark.parametrize("nx", [12, 24]) +@pytest.mark.parametrize("plot_pts", [11, 32]) +def test_soundwave_1d(nx: int, plot_pts: int, do_plot: bool = False): + """Verification test for SPH discretization of isthermal Euler equations. + A standing sound wave with c_s=1 traveserses the domain once. + """ + # import model + from struphy.models.fluid import EulerSPH + + # environment options + out_folders = os.path.join(test_folder, "EulerSPH") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="soundwave_1d") + + # units + base_units = BaseUnits(kBT=1.0) + + # time stepping + time_opts = Time(dt=0.03125, Tend=2.5, split_algo="Strang") + + # geometry + r1 = 2.5 + domain = domains.Cuboid(r1=r1) + + # fluid equilibrium (can be used as part of initial conditions) + equil = None + + # grid + grid = None + + # derham options + derham_opts = None + + # light-weight model instance + model = EulerSPH(with_B0=False) + + # species parameters + model.euler_fluid.set_phys_params() + + loading_params = LoadingParameters(ppb=8, loading="tesselation") + weights_params = WeightsParameters() + boundary_params = BoundaryParameters() + model.euler_fluid.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + ) + model.euler_fluid.set_sorting_boxes( + boxes_per_dim=(nx, 1, 1), + dims_maks=(True, False, False), + ) + + bin_plot = BinningPlot(slice="e1", n_bins=(32,), ranges=(0.0, 1.0)) + kd_plot = KernelDensityPlot(pts_e1=plot_pts, pts_e2=1) + model.euler_fluid.set_save_data( + binning_plots=(bin_plot,), + kernel_density_plots=(kd_plot,), + ) + + # propagator options + from struphy.ode.utils import ButcherTableau + + butcher = ButcherTableau(algo="forward_euler") + model.propagators.push_eta.options = model.propagators.push_eta.Options(butcher=butcher) + if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + model.propagators.push_sph_p.options = model.propagators.push_sph_p.Options(kernel_type="gaussian_1d") + + # background, perturbations and initial conditions + background = equils.ConstantVelocity() + model.euler_fluid.var.add_background(background) + perturbation = perturbations.ModesSin(ls=(1,), amps=(1.0e-2,)) + model.euler_fluid.var.add_perturbation(del_n=perturbation) + + # start run + main.run( + model, + params_path=None, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=True, + ) + + # post processing + if MPI.COMM_WORLD.Get_rank() == 0: + main.pproc(env.path_out) + + # diagnostics + simdata = main.load_data(env.path_out) + + ee1, ee2, ee3 = simdata.n_sph["euler_fluid"]["view_0"]["grid_n_sph"] + n_sph = simdata.n_sph["euler_fluid"]["view_0"]["n_sph"] + + if do_plot: + ppb = 8 + dt = time_opts.dt + end_time = time_opts.Tend + Nt = int(end_time // dt) + x = ee1 * r1 + + plt.figure(figsize=(10, 8)) + interval = Nt / 10 + plot_ct = 0 + for i in range(0, Nt + 1): + if i % interval == 0: + print(f"{i = }") + plot_ct += 1 + ax = plt.gca() + + if plot_ct <= 6: + style = "-" + else: + style = "." + plt.plot(x.squeeze(), n_sph[i, :, 0, 0], style, label=f"time={i * dt:4.2f}") + plt.xlim(0, 2.5) + plt.legend() + ax.set_xticks(np.linspace(0, 2.5, nx + 1)) + ax.xaxis.set_major_formatter(FormatStrFormatter("%.2f")) + plt.grid(c="k") + plt.xlabel("x") + plt.ylabel(r"$\rho$") + + plt.title(f"standing sound wave ($c_s = 1$) for {nx = } and {ppb = }") + if plot_ct == 11: + break + + plt.show() + + error = np.max(np.abs(n_sph[0] - n_sph[-1])) + print(f"SPH sound wave {error = }.") + assert error < 6e-4 + print("Assertion passed.") + + +if __name__ == "__main__": + test_soundwave_1d(nx=12, plot_pts=11, do_plot=True) diff --git a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py new file mode 100644 index 000000000..d2420294e --- /dev/null +++ b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py @@ -0,0 +1,168 @@ +import os + +import h5py +import numpy as np +import pytest +from matplotlib import pyplot as plt +from matplotlib.ticker import FormatStrFormatter +from mpi4py import MPI + +from struphy import main +from struphy.fields_background import equils +from struphy.geometry import domains +from struphy.initial import perturbations +from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time +from struphy.kinetic_background import maxwellians +from struphy.pic.utilities import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, +) +from struphy.topology import grids + +test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") + + +def test_weak_Landau(do_plot: bool = False): + """Verification test for weak Landau damping. + The computed damping rate is compared to the analytical rate. + """ + # import model + from struphy.models.kinetic import VlasovAmpereOneSpecies + + # environment options + out_folders = os.path.join(test_folder, "VlasovAmpereOneSpecies") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="weak_Landau") + + # units + base_units = BaseUnits() + + # time stepping + time_opts = Time(dt=0.05, Tend=15) + + # geometry + r1 = 12.56 + domain = domains.Cuboid(r1=r1) + + # fluid equilibrium (can be used as part of initial conditions) + equil = None + + # grid + grid = grids.TensorProductGrid(Nel=(32, 1, 1)) + + # derham options + derham_opts = DerhamOptions(p=(3, 1, 1)) + + # light-weight model instance + model = VlasovAmpereOneSpecies(with_B0=False) + + # species parameters + model.kinetic_ions.set_phys_params(alpha=1.0, epsilon=-1.0) + + ppc = 1000 + loading_params = LoadingParameters(ppc=ppc, seed=1234) + weights_params = WeightsParameters(control_variate=True) + boundary_params = BoundaryParameters() + model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + bufsize=0.4, + ) + model.kinetic_ions.set_sorting_boxes(boxes_per_dim=(16, 1, 1), do_sort=True) + + binplot = BinningPlot(slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-5.0, 5.0))) + model.kinetic_ions.set_save_data(binning_plots=(binplot,)) + + # propagator options + model.propagators.push_eta.options = model.propagators.push_eta.Options() + if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options() + model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0") + + # background and initial conditions + background = maxwellians.Maxwellian3D(n=(1.0, None)) + model.kinetic_ions.var.add_background(background) + + # if .add_initial_condition is not called, the background is the initial condition + perturbation = perturbations.ModesCos(ls=(1,), amps=(1e-3,)) + init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) + model.kinetic_ions.var.add_initial_condition(init) + + # start run + main.run( + model, + params_path=None, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=False, + ) + + # post processing not needed for scalar data + + # exat solution + gamma = -0.1533 + + def E_exact(t): + eps = 0.001 + k = 0.5 + r = 0.3677 + omega = 1.4156 + phi = 0.5362 + return 2 * eps**2 * np.pi / k**2 * r**2 * np.exp(2 * gamma * t) * np.cos(omega * t - phi) ** 2 + + # get parameters + dt = time_opts.dt + algo = time_opts.split_algo + Nel = grid.Nel + p = derham_opts.p + + # get scalar data + if MPI.COMM_WORLD.Get_rank() == 0: + pa_data = os.path.join(env.path_out, "data") + with h5py.File(os.path.join(pa_data, "data_proc0.hdf5"), "r") as f: + time = f["time"]["value"][()] + E = f["scalar"]["en_E"][()] + logE = np.log10(E) + + # find where time derivative of E is zero + dEdt = (np.roll(logE, -1) - np.roll(logE, 1))[1:-1] / (2.0 * dt) + zeros = dEdt * np.roll(dEdt, -1) < 0.0 + maxima_inds = np.logical_and(zeros, dEdt > 0.0) + maxima = logE[1:-1][maxima_inds] + t_maxima = time[1:-1][maxima_inds] + + # plot + if do_plot: + plt.figure(figsize=(18, 12)) + plt.plot(time, logE, label="numerical") + plt.plot(time, np.log10(E_exact(time)), label="exact") + plt.legend() + plt.title(f"{dt=}, {algo=}, {Nel=}, {p=}, {ppc=}") + plt.xlabel("time [m/c]") + plt.plot(t_maxima[:5], maxima[:5], "r") + plt.plot(t_maxima[:5], maxima[:5], "or", markersize=10) + plt.ylim([-10, -4]) + + plt.show() + + # linear fit + linfit = np.polyfit(t_maxima[:5], maxima[:5], 1) + gamma_num = linfit[0] + + # assert + rel_error = np.abs(gamma_num - gamma) / np.abs(gamma) + assert rel_error < 0.22, f"Assertion for weak Landau damping failed: {gamma_num = } vs. {gamma = }." + print(f"Assertion for weak Landau damping passed ({rel_error = }).") + + +if __name__ == "__main__": + test_weak_Landau(do_plot=True) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 91c41437a..82f709907 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1676,7 +1676,8 @@ def draw_markers( if self._initialized_sorting and sort: if self.mpi_rank == 0 and verbose: print("Sorting the markers after initial draw") - self.mpi_sort_markers() + if self.mpi_comm is not None: + self.mpi_sort_markers() self.do_sort() @profile diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index 99b597c9b..7414d0714 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -215,9 +215,9 @@ class KernelDensityPlot: def __init__( self, - pts_e1: int = 11, - pts_e2: int = 11, - pts_e3: int = 11, + pts_e1: int = 16, + pts_e2: int = 16, + pts_e3: int = 1, ): e1 = np.linspace(0.0, 1.0, pts_e1) e2 = np.linspace(0.0, 1.0, pts_e2) diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index a890e5d27..3d89b3d94 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -193,6 +193,7 @@ def options(self, new): def allocate(self): # scaling factor self._epsilon = self.variables.ions.species.equation_params.epsilon + assert self.derham is not None, f"{self.__class__.__name__} needs a Derham object." # TODO: treat PolarVector as well, but polar splines are being reworked at the moment if self.projected_equil is not None: diff --git a/src/struphy/topology/grids.py b/src/struphy/topology/grids.py index 0593bce05..b26887326 100644 --- a/src/struphy/topology/grids.py +++ b/src/struphy/topology/grids.py @@ -17,5 +17,5 @@ class TensorProductGrid: If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed. """ - Nel: tuple = (16, 1, 1) + Nel: tuple = (24, 10, 1) mpi_dims_mask: tuple = (True, True, True) From 797f3d80d65f4dac82731c5aa9943671a90aec37 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 7 Oct 2025 14:48:40 +0000 Subject: [PATCH 128/292] Resolve "Port some toy models to 318" --- src/struphy/console/main.py | 2 +- src/struphy/console/test.py | 2 +- src/struphy/io/options.py | 4 +- src/struphy/kinetic_background/maxwellians.py | 45 +- src/struphy/main.py | 34 +- src/struphy/models/__init__.py | 146 ++--- src/struphy/models/base.py | 14 +- src/struphy/models/kinetic.py | 5 +- src/struphy/models/tests/test_models.py | 11 +- .../models/tests/test_verif_Maxwell.py | 10 +- .../models/tests/test_verif_Poisson.py | 150 +++++ src/struphy/models/toy.py | 423 +++++++------- src/struphy/pic/accumulation/accum_kernels.py | 4 +- src/struphy/pic/particles.py | 10 +- src/struphy/pic/tests/test_accum_vec_H1.py | 2 +- src/struphy/propagators/__init__.py | 192 +++---- src/struphy/propagators/propagators_fields.py | 515 +++++++++--------- .../propagators/propagators_markers.py | 154 ++++-- .../tests/test_gyrokinetic_poisson.py | 76 +-- src/struphy/propagators/tests/test_poisson.py | 283 ++++++++-- 20 files changed, 1244 insertions(+), 838 deletions(-) create mode 100644 src/struphy/models/tests/test_verif_Poisson.py diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index 19832a830..94c8f3220 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -954,7 +954,7 @@ def add_parser_test(subparsers, list_models): type=int, metavar="N", help="set number of MPI processes used in tests (default=2))", - default=2, + default=1, ) parser_test.add_argument( diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 09a0872ac..938571d75 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -4,7 +4,7 @@ def struphy_test( group: str, *, - mpi: int = 2, + mpi: int = 1, with_desc: bool = False, vrbose: bool = False, show_plots: bool = False, diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 13c1ae134..1e035a8ec 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -25,8 +25,10 @@ # solvers OptsSymmSolver = Literal["pcg", "cg"] -OptsGenSolver = Literal["pbicgstab", "bicgstab"] +OptsGenSolver = Literal["pbicgstab", "bicgstab", "GMRES"] OptsMassPrecond = Literal["MassMatrixPreconditioner", None] +OptsSaddlePointSolver = Literal["Uzawa", "GMRES"] +OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] # markers OptsPICSpace = Literal["Particles6D", "DeltaFParticles6D", "Particles5D", "Particles3D"] diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index 5e240a1b7..eeffa34ff 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -655,20 +655,28 @@ def default_maxw_params(cls): def __init__( self, - maxw_params: dict = None, - pert_params: dict = None, + n: tuple[float | Callable, Perturbation] = (1.0, None), + u1: tuple[float | Callable, Perturbation] = (0.0, None), + u2: tuple[float | Callable, Perturbation] = (0.0, None), + u3: tuple[float | Callable, Perturbation] = (0.0, None), equil: FluidEquilibriumWithB = None, ): - super().__init__( - maxw_params=maxw_params, - pert_params=pert_params, - equil=equil, - ) + self._maxw_params = {} + self._maxw_params["n"] = n + self._maxw_params["u1"] = u1 + self._maxw_params["u2"] = u2 + self._maxw_params["u3"] = u3 + self._maxw_params["vth1"] = (0.0, None) + self._maxw_params["vth2"] = (0.0, None) + self._maxw_params["vth3"] = (0.0, None) + + self.check_maxw_params() + + self._equil = equil - # make sure temperatures are zero - self._maxw_params["vth1"] = 0.0 - self._maxw_params["vth2"] = 0.0 - self._maxw_params["vth3"] = 0.0 + @property + def maxw_params(self): + return self._maxw_params @property def coords(self): @@ -691,6 +699,10 @@ def volume_form(self): return False @property + def equil(self) -> FluidEquilibriumWithB: + """Fluid background with B-field.""" + return self._equil + def velocity_jacobian_det(self, eta1, eta2, eta3, *v): """Jacobian determinant of the velocity coordinate transformation.""" return 1.0 @@ -718,14 +730,3 @@ def vth(self, eta1, eta2, eta3): def __call__(self, eta1, eta2, eta3): return self.n(eta1, eta2, eta3) - - @property - def add_perturbation(self) -> bool: - if not hasattr(self, "_add_perturbation"): - self._add_perturbation = True - return self._add_perturbation - - @add_perturbation.setter - def add_perturbation(self, new): - assert isinstance(new, bool) - self._add_perturbation = new diff --git a/src/struphy/main.py b/src/struphy/main.py index 9592db6bd..4d1d4d951 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -327,17 +327,17 @@ def run( print(message, end="\n") print() + # update time and index (round time to 10 decimals for a clean time grid!) + time_state["value"][0] = round(time_state["value"][0] + dt, 10) + time_state["value_sec"][0] = round(time_state["value_sec"][0] + dt * model.units.t, 10) + time_state["index"][0] += 1 + # perform one time step dt t0 = time.time() with ProfileManager.profile_region("model.integrate"): model.integrate(dt, split_algo) t1 = time.time() - # update time and index (round time to 10 decimals for a clean time grid!) - time_state["value"][0] = round(time_state["value"][0] + dt, 10) - time_state["value_sec"][0] = round(time_state["value_sec"][0] + dt * model.units.t, 10) - time_state["index"][0] += 1 - run_time_now = (time.time() - start_simulation) / 60 # update diagnostics data and save data @@ -670,18 +670,6 @@ def Nattr(self) -> dict[str, int]: self._Nattr[spec] = orbs.shape[2] return self._Nattr - @property - def spline_grid_resolution(self): - if self.grids_log is not None: - res = [x.size for x in self.grids_log] - else: - res = None - return res - - @property - def time_grid_size(self): - return self.t_grid.size - def load_data(path: str) -> SimData: """Load data generated during post-processing. @@ -792,8 +780,16 @@ def load_data(path: str) -> SimData: raise NotImplementedError print("\nThe following data has been loaded:") - print(f"{simdata.time_grid_size = }") - print(f"{simdata.spline_grid_resolution = }") + print(f"\ngrids:") + print(f"{simdata.t_grid.shape = }") + if simdata.grids_log is not None: + print(f"{simdata.grids_log[0].shape = }") + print(f"{simdata.grids_log[1].shape = }") + print(f"{simdata.grids_log[2].shape = }") + if simdata.grids_phy is not None: + print(f"{simdata.grids_phy[0].shape = }") + print(f"{simdata.grids_phy[1].shape = }") + print(f"{simdata.grids_phy[2].shape = }") print(f"\nsimdata.spline_values:") for k, v in simdata.spline_values.items(): print(f" {k}") diff --git a/src/struphy/models/__init__.py b/src/struphy/models/__init__.py index 3878287f4..0153467f6 100644 --- a/src/struphy/models/__init__.py +++ b/src/struphy/models/__init__.py @@ -1,74 +1,74 @@ -from struphy.models.fluid import ( - ColdPlasma, - EulerSPH, - HasegawaWakatani, - LinearExtendedMHDuniform, - LinearMHD, - ViscoresistiveDeltafMHD, - ViscoresistiveDeltafMHD_with_q, - ViscoresistiveLinearMHD, - ViscoresistiveLinearMHD_with_q, - ViscoresistiveMHD, - ViscoresistiveMHD_with_p, - ViscoresistiveMHD_with_q, - ViscousFluid, -) -from struphy.models.hybrid import ColdPlasmaVlasov, LinearMHDDriftkineticCC, LinearMHDVlasovCC, LinearMHDVlasovPC -from struphy.models.kinetic import ( - DriftKineticElectrostaticAdiabatic, - LinearVlasovAmpereOneSpecies, - LinearVlasovMaxwellOneSpecies, - VlasovAmpereOneSpecies, - VlasovMaxwellOneSpecies, -) -from struphy.models.toy import ( - DeterministicParticleDiffusion, - GuidingCenter, - Maxwell, - Poisson, - PressureLessSPH, - RandomParticleDiffusion, - ShearAlfven, - TwoFluidQuasiNeutralToy, - VariationalBarotropicFluid, - VariationalCompressibleFluid, - VariationalPressurelessFluid, - Vlasov, -) +# from struphy.models.fluid import ( +# ColdPlasma, +# EulerSPH, +# HasegawaWakatani, +# LinearExtendedMHDuniform, +# LinearMHD, +# ViscoresistiveDeltafMHD, +# ViscoresistiveDeltafMHD_with_q, +# ViscoresistiveLinearMHD, +# ViscoresistiveLinearMHD_with_q, +# ViscoresistiveMHD, +# ViscoresistiveMHD_with_p, +# ViscoresistiveMHD_with_q, +# ViscousFluid, +# ) +# from struphy.models.hybrid import ColdPlasmaVlasov, LinearMHDDriftkineticCC, LinearMHDVlasovCC, LinearMHDVlasovPC +# from struphy.models.kinetic import ( +# DriftKineticElectrostaticAdiabatic, +# LinearVlasovAmpereOneSpecies, +# LinearVlasovMaxwellOneSpecies, +# VlasovAmpereOneSpecies, +# VlasovMaxwellOneSpecies, +# ) +# from struphy.models.toy import ( +# DeterministicParticleDiffusion, +# GuidingCenter, +# Maxwell, +# Poisson, +# PressureLessSPH, +# RandomParticleDiffusion, +# ShearAlfven, +# TwoFluidQuasiNeutralToy, +# VariationalBarotropicFluid, +# VariationalCompressibleFluid, +# VariationalPressurelessFluid, +# Vlasov, +# ) -__all__ = [ - "Maxwell", - "Vlasov", - "GuidingCenter", - "ShearAlfven", - "VariationalPressurelessFluid", - "VariationalBarotropicFluid", - "VariationalCompressibleFluid", - "Poisson", - "DeterministicParticleDiffusion", - "RandomParticleDiffusion", - "PressureLessSPH", - "TwoFluidQuasiNeutralToy", - "LinearMHD", - "LinearExtendedMHDuniform", - "ColdPlasma", - "ViscoresistiveMHD", - "ViscousFluid", - "ViscoresistiveMHD_with_p", - "ViscoresistiveLinearMHD", - "ViscoresistiveDeltafMHD", - "ViscoresistiveMHD_with_q", - "ViscoresistiveLinearMHD_with_q", - "ViscoresistiveDeltafMHD_with_q", - "EulerSPH", - "HasegawaWakatani", - "LinearMHDVlasovCC", - "LinearMHDVlasovPC", - "LinearMHDDriftkineticCC", - "ColdPlasmaVlasov", - "VlasovAmpereOneSpecies", - "VlasovMaxwellOneSpecies", - "LinearVlasovAmpereOneSpecies", - "LinearVlasovMaxwellOneSpecies", - "DriftKineticElectrostaticAdiabatic", -] +# __all__ = [ +# "Maxwell", +# "Vlasov", +# "GuidingCenter", +# "ShearAlfven", +# "VariationalPressurelessFluid", +# "VariationalBarotropicFluid", +# "VariationalCompressibleFluid", +# "Poisson", +# "DeterministicParticleDiffusion", +# "RandomParticleDiffusion", +# "PressureLessSPH", +# "TwoFluidQuasiNeutralToy", +# "LinearMHD", +# "LinearExtendedMHDuniform", +# "ColdPlasma", +# "ViscoresistiveMHD", +# "ViscousFluid", +# "ViscoresistiveMHD_with_p", +# "ViscoresistiveLinearMHD", +# "ViscoresistiveDeltafMHD", +# "ViscoresistiveMHD_with_q", +# "ViscoresistiveLinearMHD_with_q", +# "ViscoresistiveDeltafMHD_with_q", +# "EulerSPH", +# "HasegawaWakatani", +# "LinearMHDVlasovCC", +# "LinearMHDVlasovPC", +# "LinearMHDDriftkineticCC", +# "ColdPlasmaVlasov", +# "VlasovAmpereOneSpecies", +# "VlasovMaxwellOneSpecies", +# "LinearVlasovAmpereOneSpecies", +# "LinearVlasovMaxwellOneSpecies", +# "DriftKineticElectrostaticAdiabatic", +# ] diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 458b1cfba..acc198e9f 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -344,6 +344,8 @@ def kwargs(self): @property def scalar_quantities(self): """A dictionary of scalar quantities to be saved during the simulation.""" + if not hasattr(self, "_scalar_quantities"): + self._scalar_quantities = {} return self._scalar_quantities @property @@ -1350,7 +1352,7 @@ def generate_default_parameter_file( init_bckgr_pic += f"maxwellian_2 = maxwellians.Maxwellian3D(n=(0.1, None))\n" init_pert_pic += f"maxwellian_1pt = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n" init_pert_pic += f"init = maxwellian_1pt + maxwellian_2\n" - init_pert_pic += f"model.kinetic_ions.var.add_initial_condition(init)\n" + init_pert_pic += f"model.{sn}.{vn}.add_initial_condition(init)\n" elif "5D" in var.space: init_bckgr_pic = f"maxwellian_1 = maxwellians.GyroMaxwellian2D(n=(1.0, None), equil=equil)\n" init_bckgr_pic += f"maxwellian_2 = maxwellians.GyroMaxwellian2D(n=(0.1, None), equil=equil)\n" @@ -1358,7 +1360,13 @@ def generate_default_parameter_file( f"maxwellian_1pt = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), equil=equil)\n" ) init_pert_pic += f"init = maxwellian_1pt + maxwellian_2\n" - init_pert_pic += f"model.kinetic_ions.var.add_initial_condition(init)\n" + init_pert_pic += f"model.{sn}.{vn}.add_initial_condition(init)\n" + if "3D" in var.space: + init_bckgr_pic = f"maxwellian_1 = maxwellians.ColdPlasma(n=(1.0, None))\n" + init_bckgr_pic += f"maxwellian_2 = maxwellians.ColdPlasma(n=(0.1, None))\n" + init_pert_pic += f"maxwellian_1pt = maxwellians.ColdPlasma(n=(1.0, perturbation))\n" + init_pert_pic += f"init = maxwellian_1pt + maxwellian_2\n" + init_pert_pic += f"model.{sn}.{vn}.add_initial_condition(init)\n" init_bckgr_pic += f"background = maxwellian_1 + maxwellian_2\n" init_bckgr_pic += f"model.{sn}.{vn}.add_background(background)\n" @@ -1390,7 +1398,6 @@ def generate_default_parameter_file( file.write("\n# import model, set verbosity\n") file.write(f"from {self.__module__} import {self.__class__.__name__}\n") - file.write("verbose = True\n") file.write("\n# environment options\n") file.write("env = EnvironmentOptions()\n") @@ -1449,6 +1456,7 @@ def generate_default_parameter_file( file.write('\nif __name__ == "__main__":\n') file.write(" # start run\n") + file.write(" verbose = True\n\n") file.write( " main.run(model,\n\ params_path=__file__,\n\ diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 0d6274d50..f75ac4b8c 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -199,15 +199,14 @@ def allocate_propagators(self): self.domain.args_domain, ) - charge_accum(particles.vdim) - # another sanity check: compute FE coeffs of density # charge_accum.show_accumulated_spline_field(self.mass_ops) alpha = self.kinetic_ions.equation_params.alpha epsilon = self.kinetic_ions.equation_params.epsilon - self.initial_poisson.options.rho = alpha**2 / epsilon * charge_accum.vectors[0] + self.initial_poisson.options.rho = charge_accum + self.initial_poisson.options.rho_coeffs = alpha**2 / epsilon self.initial_poisson.allocate() # Solve with dt=1. and compute electric field diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index 24aa086f1..b4462af54 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -18,6 +18,10 @@ "Vlasov", "GuidingCenter", "PressureLessSPH", + "Poisson", + "DeterministicParticleDiffusion", + "RandomParticleDiffusion", + "TwoFluidQuasiNeutralToy", ] # for name, obj in inspect.getmembers(toy): # if inspect.isclass(obj) and "models.toy" in obj.__module__: @@ -61,12 +65,17 @@ def call_test(model_name: str, module: ModuleType = None, verbose=True): if rank == 0: print(f"\n*** Testing '{model_name}':") + # exceptions + if model_name == "TwoFluidQuasiNeutralToy" and MPI.COMM_WORLD.Get_size() > 1: + print(f"WARNING: Model {model_name} cannot be tested for {MPI.COMM_WORLD.Get_size() = }") + return + if module is None: submods = [toy, fluid, kinetic, hybrid] for submod in submods: try: model = getattr(submod, model_name)() - except: + except AttributeError: continue else: diff --git a/src/struphy/models/tests/test_verif_Maxwell.py b/src/struphy/models/tests/test_verif_Maxwell.py index f191a93c4..9068eccb8 100644 --- a/src/struphy/models/tests/test_verif_Maxwell.py +++ b/src/struphy/models/tests/test_verif_Maxwell.py @@ -13,6 +13,7 @@ from struphy.initial import perturbations from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time from struphy.kinetic_background import maxwellians +from struphy.models.toy import Maxwell from struphy.topology import grids test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") @@ -21,11 +22,6 @@ @pytest.mark.mpi(min_size=3) @pytest.mark.parametrize("algo", ["implicit", "explicit"]) def test_light_wave_1d(algo: str, do_plot: bool = False): - # import model, set verbosity - from struphy.models.toy import Maxwell - - verbose = True - # environment options out_folders = os.path.join(test_folder, "Maxwell") env = EnvironmentOptions(out_folders=out_folders, sim_folder="light_wave_1d") @@ -58,7 +54,9 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0, seed=123)) model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=1, seed=123)) - # # start run + # start run + verbose = True + main.run( model, params_path=None, diff --git a/src/struphy/models/tests/test_verif_Poisson.py b/src/struphy/models/tests/test_verif_Poisson.py new file mode 100644 index 000000000..2abda9dec --- /dev/null +++ b/src/struphy/models/tests/test_verif_Poisson.py @@ -0,0 +1,150 @@ +import os + +import numpy as np +import pytest +from matplotlib import pyplot as plt +from mpi4py import MPI + +from struphy import main +from struphy.fields_background import equils +from struphy.geometry import domains +from struphy.initial import perturbations +from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time +from struphy.kinetic_background import maxwellians +from struphy.models.toy import Poisson +from struphy.pic.utilities import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, +) +from struphy.topology import grids + +test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") + + +def test_poisson_1d(do_plot=False): + # environment options + out_folders = os.path.join(test_folder, "Poisson") + env = EnvironmentOptions(out_folders=out_folders, sim_folder="time_source_1d") + + # units + base_units = BaseUnits() + + # time stepping + time_opts = Time(dt=0.1, Tend=2.0) + + # geometry + l1 = -5.0 + r1 = 5.0 + l2 = -5.0 + r2 = 5.0 + l3 = -6.0 + r3 = 6.0 + domain = domains.Cuboid( + l1=l1, + r1=r1, + ) # l2=l2, r2=r2, l3=l3, r3=r3) + + # fluid equilibrium (can be used as part of initial conditions) + equil = None + + # grid + grid = grids.TensorProductGrid(Nel=(48, 1, 1)) + + # derham options + derham_opts = DerhamOptions() + + # light-weight model instance + model = Poisson() + + # propagator options + omega = 2 * np.pi + model.propagators.source.options = model.propagators.source.Options(omega=omega) + model.propagators.poisson.options = model.propagators.poisson.Options(rho=model.em_fields.source) + + # background, perturbations and initial conditions + l = 2 + amp = 1e-1 + pert = perturbations.ModesCos(ls=(l,), amps=(amp,)) + model.em_fields.source.add_perturbation(pert) + + # analytical solution + Lx = r1 - l1 + rhs_exact = lambda e1, e2, e3, t: amp * np.cos(l * 2 * np.pi / Lx * e1) * np.cos(omega * t) + phi_exact = ( + lambda e1, e2, e3, t: amp / (l * 2 * np.pi / Lx) ** 2 * np.cos(l * 2 * np.pi / Lx * e1) * np.cos(omega * t) + ) + + # start run + verbose = True + + main.run( + model, + params_path=None, + env=env, + base_units=base_units, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + verbose=verbose, + ) + + # post processing + if MPI.COMM_WORLD.Get_rank() == 0: + main.pproc(env.path_out) + + # diagnostics + if MPI.COMM_WORLD.Get_rank() == 0: + simdata = main.load_data(env.path_out) + + phi = simdata.spline_values["em_fields"]["phi_log"] + source = simdata.spline_values["em_fields"]["source_log"] + x = simdata.grids_phy[0][:, 0, 0] + y = simdata.grids_phy[1][0, :, 0] + z = simdata.grids_phy[2][0, 0, :] + time = simdata.t_grid + + interval = 2 + c = 0 + if do_plot: + fig = plt.figure(figsize=(12, 40)) + + err = 0.0 + for i, t in enumerate(phi): + phi_h = phi[t][0][:, 0, 0] + phi_e = phi_exact(x, 0, 0, t) + new_err = np.abs(np.max(phi_h - phi_e)) / (amp / (l * 2 * np.pi / Lx) ** 2) + if new_err > err: + err = new_err + + if do_plot and i % interval == 0: + plt.subplot(5, 2, 2 * c + 1) + plt.plot(x, phi_h, label="phi") + plt.plot(x, phi_e, "r--", label="exact") + plt.title(f"phi at {t = }") + plt.ylim(-amp / (l * 2 * np.pi / Lx) ** 2, amp / (l * 2 * np.pi / Lx) ** 2) + plt.legend() + + plt.subplot(5, 2, 2 * c + 2) + plt.plot(x, source[t][0][:, 0, 0], label="rhs") + plt.plot(x, rhs_exact(x, 0, 0, t), "r--", label="exact") + plt.title(f"source at {t = }") + plt.ylim(-amp, amp) + plt.legend() + + c += 1 + if c > 4: + break + + plt.show() + print(f"{err = }") + assert err < 0.0057 + + +if __name__ == "__main__": + # test_light_wave_1d(algo="explicit", do_plot=True) + test_poisson_1d(do_plot=False) diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index ee263db44..f91cd4c81 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -804,65 +804,89 @@ class Poisson(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species - dct["em_fields"]["phi"] = "H1" - dct["em_fields"]["source"] = "H1" - return dct + class EMFields(FieldSpecies): + def __init__(self): + self.phi = FEECVariable(space="H1") + self.source = FEECVariable(space="H1") + self.init_variables() - @staticmethod - def bulk_species(): - return None + ## propagators - @staticmethod - def velocity_scale(): - return None + class Propagators: + def __init__(self): + self.source = propagators_fields.TimeDependentSource() + self.poisson = propagators_fields.Poisson() - @staticmethod - def propagators_dct(): - return { - propagators_fields.TimeDependentSource: ["source"], - propagators_fields.ImplicitDiffusion: ["phi"], - } + ## abstract methods - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - def __init__(self, params, comm, clone_config=None): - super().__init__(params, comm=comm, clone_config=clone_config) + # 1. instantiate all species + self.em_fields = self.EMFields() - # extract necessary parameters - model_params = params["em_fields"]["options"]["ImplicitDiffusion"]["model"] - solver_params = params["em_fields"]["options"]["ImplicitDiffusion"]["solver"] - omega = params["em_fields"]["options"]["TimeDependentSource"]["omega"] - hfun = params["em_fields"]["options"]["TimeDependentSource"]["hfun"] + # 2. instantiate all propagators + self.propagators = self.Propagators() - # set keyword arguments for propagators - self._kwargs[propagators_fields.TimeDependentSource] = { - "omega": omega, - "hfun": hfun, - } + # 3. assign variables to propagators + self.propagators.source.variables.source = self.em_fields.source + self.propagators.poisson.variables.phi = self.em_fields.phi - self._kwargs[propagators_fields.ImplicitDiffusion] = { - "sigma_1": model_params["sigma_1"], - "stab_mat": model_params["stab_mat"], - "diffusion_mat": model_params["diffusion_mat"], - "rho": self.pointer["source"], - "solver": solver_params, - } + @property + def bulk_species(self): + return None - # Initialize propagators used in splitting substeps - self.init_propagators() + @property + def velocity_scale(self): + return None + + def allocate_helpers(self): + pass def update_scalar_quantities(self): pass + def allocate_propagators(self): + """Solve initial Poisson equation. + + :meta private: + """ + + # initialize fields and particles + super().allocate_propagators() + + # # use setter to assign source + # self.propagators.poisson.rho = self.mass_ops.M0.dot(self.em_fields.source.spline.vector) + + # Solve with dt=1. and compute electric field + if MPI.COMM_WORLD.Get_rank() == 0: + print("\nSolving initial Poisson problem...") + + self.propagators.poisson(1.0) + + if MPI.COMM_WORLD.Get_rank() == 0: + print("Done.") + + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "poisson.Options" in line: + new_file += [ + "model.propagators.poisson.options = model.propagators.poisson.Options(rho=model.em_fields.source)\n" + ] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class DeterministicParticleDiffusion(StruphyModel): r"""Diffusion equation discretized with a deterministic particle method; @@ -890,64 +914,49 @@ class DeterministicParticleDiffusion(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["kinetic"]["species1"] = "Particles3D" - return dct + ## species - @staticmethod - def bulk_species(): - return "species1" + class Hydrogen(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="Particles3D") + self.init_variables() - @staticmethod - def velocity_scale(): - return None + ## propagators - @staticmethod - def propagators_dct(): - return {propagators_markers.PushDeterministicDiffusion: ["species1"]} + class Propagators: + def __init__(self): + self.det_diff = propagators_markers.PushDeterministicDiffusion() - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + ## abstract methods - def __init__(self, params, comm, clone_config=None): - super().__init__(params, comm=comm, clone_config=clone_config) + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - from mpi4py.MPI import IN_PLACE, SUM + # 1. instantiate all species + self.hydrogen = self.Hydrogen() - # prelim - params = self.kinetic["species1"]["params"] - algo = params["options"]["PushDeterministicDiffusion"]["algo"] - diffusion_coefficient = params["options"]["PushDeterministicDiffusion"]["diffusion_coefficient"] + # 2. instantiate all propagators + self.propagators = self.Propagators() - # # project magnetic background - # self._b_eq = self.derham.P['2']([self.equil.b2_1, - # self.equil.b2_2, - # self.equil.b2_3]) + # 3. assign variables to propagators + self.propagators.det_diff.variables.var = self.hydrogen.var - # set keyword arguments for propagators - self._kwargs[propagators_markers.PushDeterministicDiffusion] = { - "algo": algo, - "bc_type": params["markers"]["bc"], - "diffusion_coefficient": diffusion_coefficient, - } + # define scalars for update_scalar_quantities + # self.add_scalar("electric energy") + # self.add_scalar("magnetic energy") + # self.add_scalar("total energy") - # Initialize propagators used in splitting substeps - self.init_propagators() + @property + def bulk_species(self): + return self.hydrogen - # Scalar variables to be saved during simulation - self.add_scalar("en_f") + @property + def velocity_scale(self): + return None - # MPI operations needed for scalar variables - self._mpi_sum = SUM - self._mpi_in_place = IN_PLACE - self._tmp = np.empty(1, dtype=float) + def allocate_helpers(self): + pass def update_scalar_quantities(self): pass @@ -978,64 +987,49 @@ class RandomParticleDiffusion(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["kinetic"]["species1"] = "Particles3D" - return dct + ## species - @staticmethod - def bulk_species(): - return "species1" + class Hydrogen(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="Particles3D") + self.init_variables() - @staticmethod - def velocity_scale(): - return None + ## propagators - @staticmethod - def propagators_dct(): - return {propagators_markers.PushRandomDiffusion: ["species1"]} + class Propagators: + def __init__(self): + self.rand_diff = propagators_markers.PushRandomDiffusion() - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + ## abstract methods - def __init__(self, params, comm, clone_config=None): - super().__init__(params, comm=comm, clone_config=clone_config) + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - from mpi4py.MPI import IN_PLACE, SUM + # 1. instantiate all species + self.hydrogen = self.Hydrogen() - # prelim - species1_params = self.kinetic["species1"]["params"] - algo = species1_params["options"]["PushRandomDiffusion"]["algo"] - diffusion_coefficient = species1_params["options"]["PushRandomDiffusion"]["diffusion_coefficient"] + # 2. instantiate all propagators + self.propagators = self.Propagators() - # # project magnetic background - # self._b_eq = self.derham.P['2']([self.equil.b2_1, - # self.equil.b2_2, - # self.equil.b2_3]) + # 3. assign variables to propagators + self.propagators.rand_diff.variables.var = self.hydrogen.var - # set keyword arguments for propagators - self._kwargs[propagators_markers.PushRandomDiffusion] = { - "algo": algo, - "bc_type": species1_params["markers"]["bc"], - "diffusion_coefficient": diffusion_coefficient, - } + # define scalars for update_scalar_quantities + # self.add_scalar("electric energy") + # self.add_scalar("magnetic energy") + # self.add_scalar("total energy") - # Initialize propagators used in splitting substeps - self.init_propagators() + @property + def bulk_species(self): + return self.hydrogen - # Scalar variables to be saved during simulation - self.add_scalar("en_f") + @property + def velocity_scale(self): + return None - # MPI operations needed for scalar variables - self._mpi_sum = SUM - self._mpi_in_place = IN_PLACE - self._tmp = np.empty(1, dtype=float) + def allocate_helpers(self): + pass def update_scalar_quantities(self): pass @@ -1175,116 +1169,75 @@ class TwoFluidQuasiNeutralToy(StruphyModel): in plasma physics, Journal of Computational Physics 2018. """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species - dct["em_fields"]["potential"] = "L2" - dct["fluid"]["ions"] = { - "u": "Hdiv", - } - dct["fluid"]["electrons"] = { - "u": "Hdiv", - } - return dct + class EMfields(FieldSpecies): + def __init__(self): + self.phi = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def bulk_species(): - return "ions" + class Ions(FluidSpecies): + def __init__(self): + self.u = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def velocity_scale(): - return "thermal" + class Electrons(FluidSpecies): + def __init__(self): + self.u = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def propagators_dct(): - return {propagators_fields.TwoFluidQuasiNeutralFull: ["ions_u", "electrons_u", "potential"]} + ## propagators - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + class Propagators: + def __init__(self): + self.qn_full = propagators_fields.TwoFluidQuasiNeutralFull() - # add special options - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["fluid", "electrons"], - option=propagators_fields.TwoFluidQuasiNeutralFull, - dct=dct, - ) - return dct + ## abstract methods - def __init__(self, params, comm, clone_config=None): - super().__init__(params, comm=comm, clone_config=clone_config) + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # get species paramaters - electrons_params = params["fluid"]["electrons"] + # 1. instantiate all species + self.em_fields = self.EMfields() + self.ions = self.Ions() + self.electrons = self.Electrons() - # Get coupling strength - if electrons_params["options"]["TwoFluidQuasiNeutralFull"]["override_eq_params"]: - self._epsilon = electrons_params["options"]["TwoFluidQuasiNeutralFull"]["eps_norm"] - print( - f"\n!!! Override equation parameters: {self._epsilon = }.", - ) - else: - self._epsilon = self.equation_params["electrons"]["epsilon"] + # 2. instantiate all propagators + self.propagators = self.Propagators() - # extract necessary parameters - stokes_solver = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["solver"] - stokes_nu = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["nu"] - stokes_nu_e = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["nu_e"] - stokes_a = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["a"] - stokes_R0 = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["R0"] - stokes_B0 = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["B0"] - stokes_Bp = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["Bp"] - stokes_alpha = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["alpha"] - stokes_beta = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["beta"] - stokes_sigma = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["stab_sigma"] - stokes_variant = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["variant"] - stokes_method_to_solve = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["method_to_solve"] - stokes_preconditioner = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["preconditioner"] - stokes_spectralanalysis = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"][ - "spectralanalysis" - ] - stokes_lifting = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["lifting"] - stokes_dimension = params["fluid"]["electrons"]["options"]["TwoFluidQuasiNeutralFull"]["dimension"] - stokes_1D_dt = params["time"]["dt"] - - # Check MPI size to ensure only one MPI process - size = comm.Get_size() - if size != 1 and stokes_variant == "Uzawa": - if comm.Get_rank() == 0: - print(f"Error: TwoFluidQuasiNeutralToy only runs with one MPI process.") - return # Early return to stop execution for multiple MPI processes + # 3. assign variables to propagators + self.propagators.qn_full.variables.u = self.ions.u + self.propagators.qn_full.variables.ue = self.electrons.u + self.propagators.qn_full.variables.phi = self.em_fields.phi - # set keyword arguments for propagators - self._kwargs[propagators_fields.TwoFluidQuasiNeutralFull] = { - "solver": stokes_solver, - "nu": stokes_nu, - "nu_e": stokes_nu_e, - "eps_norm": self._epsilon, - "a": stokes_a, - "R0": stokes_R0, - "B0": stokes_B0, - "Bp": stokes_Bp, - "alpha": stokes_alpha, - "beta": stokes_beta, - "stab_sigma": stokes_sigma, - "variant": stokes_variant, - "method_to_solve": stokes_method_to_solve, - "preconditioner": stokes_preconditioner, - "spectralanalysis": stokes_spectralanalysis, - "dimension": stokes_dimension, - "D1_dt": stokes_1D_dt, - "lifting": stokes_lifting, - } + # define scalars for update_scalar_quantities - # Initialize propagators used in splitting substeps - self.init_propagators() + @property + def bulk_species(self): + return self.ions + + @property + def velocity_scale(self): + return "thermal" + + def allocate_helpers(self): + pass def update_scalar_quantities(self): pass + + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "BaseUnits()" in line: + new_file += ["base_units = BaseUnits(kBT=1.0)\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) diff --git a/src/struphy/pic/accumulation/accum_kernels.py b/src/struphy/pic/accumulation/accum_kernels.py index 8d3c2923b..20d807a53 100644 --- a/src/struphy/pic/accumulation/accum_kernels.py +++ b/src/struphy/pic/accumulation/accum_kernels.py @@ -33,7 +33,6 @@ def charge_density_0form( args_derham: "DerhamArguments", args_domain: "DomainArguments", vec: "float[:,:,:]", - vdim: "int", ): r""" Kernel for :class:`~struphy.pic.accumulation.particles_to_grid.AccumulatorVector` into V0 with the filling @@ -45,6 +44,7 @@ def charge_density_0form( markers = args_markers.markers Np = args_markers.Np + weight_idx = args_markers.weight_idx # -- removed omp: #$ omp parallel private (ip, eta1, eta2, eta3, filling) # -- removed omp: #$ omp for reduction ( + :vec) @@ -59,7 +59,7 @@ def charge_density_0form( eta3 = markers[ip, 2] # filling = w_p/N - filling = markers[ip, 3 + vdim] / Np + filling = markers[ip, weight_idx] / Np particle_to_mat_kernels.vec_fill_b_v0( args_derham, diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 35df68d04..3489e701a 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -647,8 +647,8 @@ class Particles3D(Particles): """ @classmethod - def default_bckgr_params(cls): - return {"ColdPlasma": {}} + def default_background(cls): + return maxwellians.ColdPlasma() def __init__( self, @@ -656,8 +656,10 @@ def __init__( ): kwargs["type"] = "full_f" - if "bckgr_params" not in kwargs: - kwargs["bckgr_params"] = self.default_bckgr_params() + if "background" not in kwargs: + kwargs["background"] = self.default_background() + elif kwargs["background"] is None: + kwargs["background"] = self.default_background() # default number of diagnostics and auxiliary columns self._n_cols_diagnostics = kwargs.pop("n_cols_diagn", 0) diff --git a/src/struphy/pic/tests/test_accum_vec_H1.py b/src/struphy/pic/tests/test_accum_vec_H1.py index 3f2ce4923..67e838fc3 100644 --- a/src/struphy/pic/tests/test_accum_vec_H1.py +++ b/src/struphy/pic/tests/test_accum_vec_H1.py @@ -132,7 +132,7 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): domain.args_domain, ) - acc(particles.vdim) + acc() # sum all MC integrals _sum_within_clone = np.empty(1, dtype=float) diff --git a/src/struphy/propagators/__init__.py b/src/struphy/propagators/__init__.py index 2dbef2b10..d8f1bd3bd 100644 --- a/src/struphy/propagators/__init__.py +++ b/src/struphy/propagators/__init__.py @@ -1,97 +1,97 @@ -from struphy.propagators.propagators_coupling import ( - CurrentCoupling5DCurlb, - CurrentCoupling5DGradB, - CurrentCoupling6DCurrent, - EfieldWeights, - PressureCoupling6D, - VlasovAmpere, -) -from struphy.propagators.propagators_fields import ( - AdiabaticPhi, - CurrentCoupling5DDensity, - CurrentCoupling6DDensity, - FaradayExtended, - Hall, - HasegawaWakatani, - ImplicitDiffusion, - JxBCold, - Magnetosonic, - MagnetosonicCurrentCoupling5D, - MagnetosonicUniform, - Maxwell, - OhmCold, - Poisson, - ShearAlfven, - ShearAlfvenB1, - ShearAlfvenCurrentCoupling5D, - TimeDependentSource, - TwoFluidQuasiNeutralFull, - VariationalDensityEvolve, - VariationalEntropyEvolve, - VariationalMagFieldEvolve, - VariationalMomentumAdvection, - VariationalPBEvolve, - VariationalQBEvolve, - VariationalResistivity, - VariationalViscosity, -) -from struphy.propagators.propagators_markers import ( - PushDeterministicDiffusion, - PushEta, - PushEtaPC, - PushGuidingCenterBxEstar, - PushGuidingCenterParallel, - PushRandomDiffusion, - PushVinEfield, - PushVinSPHpressure, - PushVinViscousPotential, - PushVxB, - StepStaticEfield, -) +# from struphy.propagators.propagators_coupling import ( +# CurrentCoupling5DCurlb, +# CurrentCoupling5DGradB, +# CurrentCoupling6DCurrent, +# EfieldWeights, +# PressureCoupling6D, +# VlasovAmpere, +# ) +# from struphy.propagators.propagators_fields import ( +# AdiabaticPhi, +# CurrentCoupling5DDensity, +# CurrentCoupling6DDensity, +# FaradayExtended, +# Hall, +# HasegawaWakatani, +# ImplicitDiffusion, +# JxBCold, +# Magnetosonic, +# MagnetosonicCurrentCoupling5D, +# MagnetosonicUniform, +# Maxwell, +# OhmCold, +# Poisson, +# ShearAlfven, +# ShearAlfvenB1, +# ShearAlfvenCurrentCoupling5D, +# TimeDependentSource, +# TwoFluidQuasiNeutralFull, +# VariationalDensityEvolve, +# VariationalEntropyEvolve, +# VariationalMagFieldEvolve, +# VariationalMomentumAdvection, +# VariationalPBEvolve, +# VariationalQBEvolve, +# VariationalResistivity, +# VariationalViscosity, +# ) +# from struphy.propagators.propagators_markers import ( +# PushDeterministicDiffusion, +# PushEta, +# PushEtaPC, +# PushGuidingCenterBxEstar, +# PushGuidingCenterParallel, +# PushRandomDiffusion, +# PushVinEfield, +# PushVinSPHpressure, +# PushVinViscousPotential, +# PushVxB, +# StepStaticEfield, +# ) -__all__ = [ - "VlasovAmpere", - "EfieldWeights", - "PressureCoupling6D", - "CurrentCoupling6DCurrent", - "CurrentCoupling5DCurlb", - "CurrentCoupling5DGradB", - "Maxwell", - "OhmCold", - "JxBCold", - "ShearAlfven", - "ShearAlfvenB1", - "Hall", - "Magnetosonic", - "MagnetosonicUniform", - "FaradayExtended", - "CurrentCoupling6DDensity", - "ShearAlfvenCurrentCoupling5D", - "MagnetosonicCurrentCoupling5D", - "CurrentCoupling5DDensity", - "ImplicitDiffusion", - "Poisson", - "VariationalMomentumAdvection", - "VariationalDensityEvolve", - "VariationalEntropyEvolve", - "VariationalMagFieldEvolve", - "VariationalPBEvolve", - "VariationalQBEvolve", - "VariationalViscosity", - "VariationalResistivity", - "TimeDependentSource", - "AdiabaticPhi", - "HasegawaWakatani", - "TwoFluidQuasiNeutralFull", - "PushEta", - "PushVxB", - "PushVinEfield", - "PushEtaPC", - "PushGuidingCenterBxEstar", - "PushGuidingCenterParallel", - "StepStaticEfield", - "PushDeterministicDiffusion", - "PushRandomDiffusion", - "PushVinSPHpressure", - "PushVinViscousPotential", -] +# __all__ = [ +# "VlasovAmpere", +# "EfieldWeights", +# "PressureCoupling6D", +# "CurrentCoupling6DCurrent", +# "CurrentCoupling5DCurlb", +# "CurrentCoupling5DGradB", +# "Maxwell", +# "OhmCold", +# "JxBCold", +# "ShearAlfven", +# "ShearAlfvenB1", +# "Hall", +# "Magnetosonic", +# "MagnetosonicUniform", +# "FaradayExtended", +# "CurrentCoupling6DDensity", +# "ShearAlfvenCurrentCoupling5D", +# "MagnetosonicCurrentCoupling5D", +# "CurrentCoupling5DDensity", +# "ImplicitDiffusion", +# "Poisson", +# "VariationalMomentumAdvection", +# "VariationalDensityEvolve", +# "VariationalEntropyEvolve", +# "VariationalMagFieldEvolve", +# "VariationalPBEvolve", +# "VariationalQBEvolve", +# "VariationalViscosity", +# "VariationalResistivity", +# "TimeDependentSource", +# "AdiabaticPhi", +# "HasegawaWakatani", +# "TwoFluidQuasiNeutralFull", +# "PushEta", +# "PushVxB", +# "PushVinEfield", +# "PushEtaPC", +# "PushGuidingCenterBxEstar", +# "PushGuidingCenterParallel", +# "StepStaticEfield", +# "PushDeterministicDiffusion", +# "PushRandomDiffusion", +# "PushVinSPHpressure", +# "PushVinViscousPotential", +# ] diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index a41170a34..68f5fa1db 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -1,10 +1,9 @@ "Only FEEC variables are updated." import copy -from collections.abc import Callable from copy import deepcopy from dataclasses import dataclass -from typing import Literal, get_args +from typing import Callable, Literal, get_args import numpy as np import scipy as sc @@ -41,7 +40,15 @@ from struphy.fields_background.equils import set_defaults from struphy.geometry.utilities import TransformedPformComponent from struphy.initial import perturbations -from struphy.io.options import OptsGenSolver, OptsMassPrecond, OptsSymmSolver, OptsVecSpace, check_option +from struphy.io.options import ( + OptsDirectSolver, + OptsGenSolver, + OptsMassPrecond, + OptsSaddlePointSolver, + OptsSymmSolver, + OptsVecSpace, + check_option, +) from struphy.io.setup import descend_options_dict from struphy.kinetic_background.base import Maxwellian from struphy.kinetic_background.maxwellians import GyroMaxwellian2D, Maxwellian3D @@ -2649,35 +2656,6 @@ class ImplicitDiffusion(Propagator): * :math:`\sigma_1=\sigma_2=0` and :math:`\sigma_3 = \Delta t`: **Poisson solver** with a given charge density :math:`\sum_i\rho_i`. * :math:`\sigma_2=0` and :math:`\sigma_1 = \sigma_3 = \Delta t` : Poisson with **adiabatic electrons**. * :math:`\sigma_1=\sigma_2=1` and :math:`\sigma_3 = 0`: **Implicit heat equation**. - - Parameters - ---------- - phi : StencilVector - FE coefficients of the solution as a discrete 0-form. - - sigma_1, sigma_2, sigma_3 : float | int - Equation parameters. - - divide_by_dt : bool - Whether to divide the sigmas by dt during __call__. - - stab_mat : str - Name of the matrix :math:`M^0_{n_0}`. - - diffusion_mat : str - Name of the matrix :math:`M^1_{D_0}`. - - rho : StencilVector or tuple or list - (List of) right-hand side FE coefficients of a 0-form (optional, can be set with a setter later). - Can be either a) StencilVector or b) 2-tuple, or a list of those. - In case b) the first tuple entry must be :class:`~struphy.pic.accumulation.particles_to_grid.AccumulatorVector`, - and the second entry must be :class:`~struphy.pic.base.Particles`. - - x0 : StencilVector - Initial guess for the iterative solver (optional, can be set with a setter later). - - solver : dict - Parameters for the iterative solver (see ``__init__`` for details). """ class Variables: @@ -2709,7 +2687,8 @@ class Options: divide_by_dt: bool = False stab_mat: OptsStabMat = "M0" diffusion_mat: OptsDiffusionMat = "M1" - rho: StencilVector | tuple | list | Callable = None + rho: FEECVariable | Callable | tuple[AccumulatorVector, Particles] | list = None + rho_coeffs: float | list = None x0: StencilVector = None solver: OptsSymmSolver = "pcg" precond: OptsMassPrecond = "MassMatrixPreconditioner" @@ -2758,30 +2737,39 @@ def allocate(self): phi = self.variables.phi.spline.vector # collect rhs - rho = self.options.rho + def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: + """Perform preliminary operations on rho to comute the rhs and return the result.""" + if rho is None: + rhs = phi.space.zeros() + elif isinstance(rho, FEECVariable): + assert rho.space == "H1" + rhs = rho + elif isinstance(rho, AccumulatorVector): + rhs = rho + elif isinstance(rho, Callable): + rhs = L2Projector("H1", self.mass_ops).get_dofs(rho, apply_bc=True) + else: + raise TypeError(f"{type(rho) = } is not accepted.") - if rho is None: - self._rho = [phi.space.zeros()] + return rhs + + rho = self.options.rho + if isinstance(rho, list): + self._sources = [] + for r in rho: + self._sources += [verify_rhs(r)] else: - if isinstance(rho, list): - for r in rho: - if isinstance(r, tuple): - assert isinstance(r[0], AccumulatorVector) - assert isinstance(r[1], Particles) - # assert r.space_id == 'H1' - else: - assert r.space == phi.space - elif isinstance(rho, tuple): - assert isinstance(rho[0], AccumulatorVector) - assert isinstance(rho[1], Particles) - # assert rho[0].space_id == 'H1' - rho = [rho] - elif isinstance(rho, Callable): - rho = [rho()] + self._sources = [verify_rhs(rho)] + + # coeffs of rhs + if self.options.rho_coeffs is not None: + if isinstance(self.options.rho_coeffs, (list, tuple)): + self._coeffs = self.options.rho_coeffs else: - assert rho.space == phi.space - rho = [rho] - self._rho = rho + self._coeffs = [self.options.rho_coeffs] + assert len(self._coeffs) == len(self._sources) + else: + self._coeffs = [1.0 for src in self.sources] # initial guess and solver params self._x0 = self.options.x0 @@ -2827,65 +2815,41 @@ def allocate(self): self._tmp = phi.space.zeros() self._rhs = phi.space.zeros() self._rhs2 = phi.space.zeros() + self._tmp_src = phi.space.zeros() @property - def rho(self): + def sources(self) -> list[StencilVector | FEECVariable | AccumulatorVector]: """ - (List of) right-hand side FE coefficients of a 0-form. - The list entries can be either a) StencilVectors or b) 2-tuples; - in the latter case, the first tuple entry must be :class:`~struphy.pic.accumulation.particles_to_grid.AccumulatorVector`, - and the second entry must be :class:`~struphy.pic.base.Particles`. + Right-hand side of the equation (sources). """ - return self._rho + return self._sources - @rho.setter - def rho(self, value): - """In-place setter for StencilVector/PolarVector. - If rho is a list, len(value) msut be len(rho) and value can contain None. + @property + def coeffs(self) -> list[float]: """ - if isinstance(value, list): - assert len(value) == len(self.rho) - for i, (val, r) in enumerate(zip(value, self.rho)): - if val is None: - continue - elif isinstance(val, tuple): - assert isinstance(val[0], AccumulatorVector) - assert isinstance(val[1], Particles) - assert isinstance(r, tuple) - self._rho[i] = val - else: - assert val.space == r.space - r[:] = val[:] - elif isinstance(ValueError, tuple): - assert isinstance(value[0], AccumulatorVector) - assert isinstance(value[1], Particles) - assert len(self.rho) == 1 - # assert rho[0].space_id == 'H1' - self._rho[0] = value - else: - assert value.space == self.derham.Vh["0"] - assert len(self.rho) == 1 - self._rho[0][:] = value[:] + Same length as self.sources. Coefficients multiplied with sources before solve (default is 1.0). + """ + return self._coeffs @property def x0(self): """ psydac.linalg.stencil.StencilVector or struphy.polar.basic.PolarVector. First guess of the iterative solver. """ - return self._x0 + return self.options.x0 @x0.setter - def x0(self, value): + def x0(self, value: StencilVector): """In-place setter for StencilVector/PolarVector. First guess of the iterative solver.""" assert value.space == self.derham.Vh["0"] assert value.space.symbolic_space == "H1", ( f"Right-hand side must be in H1, but is in {value.space.symbolic_space}." ) - if self._x0 is None: - self._x0 = value + if self.options.x0 is None: + self.options.x0 = value else: - self._x0[:] = value[:] + self.options.x0[:] = value[:] @profile def __call__(self, dt): @@ -2905,12 +2869,15 @@ def __call__(self, dt): rhs *= sig_2 self._rhs2 *= 0.0 - for rho in self._rho: - if isinstance(rho, tuple): - rho[0]() # accumulate - self._rhs2 += sig_3 * rho[0].vectors[0] - else: - self._rhs2 += sig_3 * rho + for src, coeff in zip(self.sources, self.coeffs): + if isinstance(src, StencilVector): + self._rhs2 += sig_3 * coeff * src + elif isinstance(src, FEECVariable): + v = src.spline.vector + self._rhs2 += sig_3 * coeff * self.mass_ops.M0.dot(v, out=self._tmp_src) + elif isinstance(src, AccumulatorVector): + src() # accumulate + self._rhs2 += sig_3 * coeff * src.vectors[0] rhs += self._rhs2 @@ -2979,7 +2946,8 @@ class Options: # propagator options stab_eps: float = 0.0 stab_mat: OptsStabMat = "Id" - rho: StencilVector | tuple | list | Callable = None + rho: FEECVariable | Callable | tuple[AccumulatorVector, Particles] | list = None + rho_coeffs: float | list = None x0: StencilVector = None solver: OptsSymmSolver = "pcg" precond: OptsMassPrecond = "MassMatrixPreconditioner" @@ -7053,48 +7021,74 @@ class TimeDependentSource(Propagator): * :math:`h(\omega t) = \sin(\omega t)` """ - @staticmethod - def options(default=False): - dct = {} - dct["omega"] = 1.0 - dct["hfun"] = ["cos", "sin"] - if default: - dct = descend_options_dict(dct, []) - return dct + class Variables: + def __init__(self): + self._source: FEECVariable = None - def __init__( - self, - c: StencilVector, - *, - omega: float = options()["omega"], - hfun: str = options(default=True)["hfun"], - ): - super().__init__(c) + @property + def source(self) -> FEECVariable: + return self._source + + @source.setter + def source(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._source = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsTimeSource = Literal["cos", "sin"] + # propagator options + omega: float = 2.0 * np.pi + hfun: OptsTimeSource = "cos" + + def __post_init__(self): + # checks + check_option(self.hfun, self.OptsTimeSource) + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new - if hfun == "cos": + @profile + def allocate(self): + if self.options.hfun == "cos": def hfun(t): - return np.cos(omega * t) - elif hfun == "sin": + return np.cos(self.options.omega * t) + elif self.options.hfun == "sin": def hfun(t): - return np.sin(omega * t) + return np.sin(self.options.omega * t) else: - raise NotImplementedError(f"{hfun = } not implemented.") + raise NotImplementedError(f"{self.options.hfun = } not implemented.") self._hfun = hfun + self._c0 = self.variables.source.spline.vector.copy() + @profile def __call__(self, dt): - print(f"{self.time_state[0] = }") - if self.time_state[0] == 0.0: - self._c0 = self.feec_vars[0].copy() - print("Initial source coeffs set.") - # new coeffs cn1 = self._c0 * self._hfun(self.time_state[0]) # write new coeffs into self.feec_vars - max_dc = self.feec_vars_update(cn1) + # max_dc = self.feec_vars_update(cn1) + self.update_feec_variables(source=cn1) class AdiabaticPhi(Propagator): @@ -7546,97 +7540,122 @@ class TwoFluidQuasiNeutralFull(Propagator): :ref:`time_discret`: fully implicit. """ - def allocate(self): - pass + class Variables: + def __init__(self): + self._u: FEECVariable = None + self._ue: FEECVariable = None + self._phi: FEECVariable = None - def set_options(self, **kwargs): - pass + @property + def u(self) -> FEECVariable: + return self._u - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("gmres", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["nu"] = 1.0 - dct["nu_e"] = 0.01 - dct["override_eq_params"] = [False, {"epsilon": 1.0}] - dct["eps_norm"] = 1.0 - dct["a"] = 1.0 - dct["R0"] = 1.0 - dct["B0"] = 10.0 - dct["Bp"] = 12.5 - dct["alpha"] = 0.1 - dct["beta"] = 1.0 - dct["stab_sigma"] = 0.00001 - dct["variant"] = "Uzawa" - dct["method_to_solve"] = "DirectNPInverse" - dct["preconditioner"] = False - dct["spectralanalysis"] = False - dct["lifting"] = False - dct["dimension"] = "2D" - dct["1D_dt"] = 0.001 - if default: - dct = descend_options_dict(dct, []) + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._u = new - return dct + @property + def ue(self) -> FEECVariable: + return self._ue - def __init__( - self, - u: BlockVector, - ue: BlockVector, - phi: BlockVector, - *, - nu: float = options(default=True)["nu"], - nu_e: float = options(default=True)["nu_e"], - eps_norm: float = options(default=True)["eps_norm"], - solver: dict = options(default=True)["solver"], - a: float = options(default=True)["a"], - R0: float = options(default=True)["R0"], - B0: float = options(default=True)["B0"], - Bp: float = options(default=True)["Bp"], - alpha: float = options(default=True)["alpha"], - beta: float = options(default=True)["beta"], - stab_sigma: float = options(default=True)["stab_sigma"], - variant: str = options(default=True)["variant"], - method_to_solve: str = options(default=True)["method_to_solve"], - preconditioner: bool = options(default=True)["preconditioner"], - spectralanalysis: bool = options(default=True)["spectralanalysis"], - lifting: bool = options(default=False)["lifting"], - dimension: str = options(default=True)["dimension"], - D1_dt: float = options(default=True)["1D_dt"], - ): - super().__init__(u, ue, phi) + @ue.setter + def ue(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._ue = new - self._info = solver["info"] + @property + def phi(self) -> FEECVariable: + return self._phi + + @phi.setter + def phi(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._phi = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsDimension = Literal["1D", "2D", "Restelli", "Tokamak"] + # propagator options + nu: float = 1.0 + nu_e: float = 0.01 + eps_norm: float = 1.0 + solver: OptsGenSolver = "GMRES" + solver_params: SolverParameters = None + a: float = 1.0 + R0: float = 1.0 + B0: float = 10.0 + Bp: float = 12.0 + alpha: float = 0.1 + beta: float = 1.0 + stab_sigma: float = 1e-5 + variant: OptsSaddlePointSolver = "Uzawa" + method_to_solve: OptsDirectSolver = "DirectNPInverse" + preconditioner: bool = False + spectralanalysis: bool = False + lifting: bool = False + dimension: OptsDimension = "2D" + D1_dt: float = 1e-3 + + def __post_init__(self): + # checks + check_option(self.solver, OptsGenSolver) + check_option(self.variant, OptsSaddlePointSolver) + check_option(self.method_to_solve, OptsDirectSolver) + check_option(self.dimension, self.OptsDimension) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._info = self.options.solver_params.info if self.derham.comm is not None: self._rank = self.derham.comm.Get_rank() else: self._rank = 0 - self._nu = nu - self._nu_e = nu_e - self._eps_norm = eps_norm - self._a = a - self._R0 = R0 - self._B0 = B0 - self._Bp = Bp - self._alpha = alpha - self._beta = beta - self._stab_sigma = stab_sigma - self._variant = variant - self._method_to_solve = method_to_solve - self._preconditioner = preconditioner - self._dimension = dimension - self._spectralanalysis = spectralanalysis - self._lifting = lifting + self._nu = self.options.nu + self._nu_e = self.options.nu_e + self._eps_norm = self.options.eps_norm + self._a = self.options.a + self._R0 = self.options.R0 + self._B0 = self.options.B0 + self._Bp = self.options.Bp + self._alpha = self.options.alpha + self._beta = self.options.beta + self._stab_sigma = self.options.stab_sigma + self._variant = self.options.variant + self._method_to_solve = self.options.method_to_solve + self._preconditioner = self.options.preconditioner + self._dimension = self.options.dimension + self._spectralanalysis = self.options.spectralanalysis + self._lifting = self.options.lifting + + solver_params = self.options.solver_params # Lifting for nontrivial boundary conditions # derham had boundary conditions in eta1 direction, the following is in space Hdiv_0 @@ -7652,13 +7671,13 @@ def __init__( self._mass_opsv0 = WeightedMassOperators( self.derhamv0, self.domain, - verbose=solver["verbose"], + verbose=solver_params.verbose, eq_mhd=self.mass_ops.weights["eq_mhd"], ) self._basis_opsv0 = BasisProjectionOperators( self.derhamv0, self.domain, - verbose=solver["verbose"], + verbose=solver_params.verbose, eq_mhd=self.basis_ops.weights["eq_mhd"], ) else: @@ -7687,7 +7706,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, ) _funy = getattr(callables, "ManufacturedSolutionForceterm")( species="Ions", @@ -7697,7 +7716,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, ) _funelectronsx = getattr(callables, "ManufacturedSolutionForceterm")( species="Electrons", @@ -7707,7 +7726,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, ) _funelectronsy = getattr(callables, "ManufacturedSolutionForceterm")( species="Electrons", @@ -7717,7 +7736,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, ) # get callable(s) for specified init type @@ -7726,16 +7745,16 @@ def __init__( # pullback callable funx = TransformedPformComponent( - forceterm_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + forceterm_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain ) funy = TransformedPformComponent( - forceterm_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + forceterm_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain ) fun_electronsx = TransformedPformComponent( - forcetermelectrons_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain ) fun_electronsy = TransformedPformComponent( - forcetermelectrons_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain ) l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) self._F1 = l2_proj([funx, funy, _forceterm_logical]) @@ -7770,22 +7789,22 @@ def __init__( # pullback callable fun_pb_1 = TransformedPformComponent( - forceterm_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + forceterm_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain ) fun_pb_2 = TransformedPformComponent( - forceterm_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + forceterm_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain ) fun_pb_3 = TransformedPformComponent( - forceterm_class, fun_basis="physical", out_form="2", comp=2, domain=self.domain + forceterm_class, given_in_basis="physical", out_form="2", comp=2, domain=self.domain ) fun_electrons_pb_1 = TransformedPformComponent( - forcetermelectrons_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain ) fun_electrons_pb_2 = TransformedPformComponent( - forcetermelectrons_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain ) fun_electrons_pb_3 = TransformedPformComponent( - forcetermelectrons_class, fun_basis="physical", out_form="2", comp=2, domain=self.domain + forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=2, domain=self.domain ) if self._lifting: l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) @@ -7808,7 +7827,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, a=self._a, Bp=self._Bp, alpha=self._alpha, @@ -7822,7 +7841,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, a=self._a, Bp=self._Bp, alpha=self._alpha, @@ -7836,7 +7855,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, a=self._a, Bp=self._Bp, alpha=self._alpha, @@ -7850,7 +7869,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, a=self._a, Bp=self._Bp, alpha=self._alpha, @@ -7864,7 +7883,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, a=self._a, Bp=self._Bp, alpha=self._alpha, @@ -7878,7 +7897,7 @@ def __init__( dimension=self._dimension, stab_sigma=self._stab_sigma, eps=self._eps_norm, - dt=D1_dt, + dt=self.options.D1_dt, a=self._a, Bp=self._Bp, alpha=self._alpha, @@ -7891,22 +7910,22 @@ def __init__( # pullback callable fun_pb_1 = TransformedPformComponent( - forceterm_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + forceterm_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain ) fun_pb_2 = TransformedPformComponent( - forceterm_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + forceterm_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain ) fun_pb_3 = TransformedPformComponent( - forceterm_class, fun_basis="physical", out_form="2", comp=2, domain=self.domain + forceterm_class, given_in_basis="physical", out_form="2", comp=2, domain=self.domain ) fun_electrons_pb_1 = TransformedPformComponent( - forcetermelectrons_class, fun_basis="physical", out_form="2", comp=0, domain=self.domain + forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain ) fun_electrons_pb_2 = TransformedPformComponent( - forcetermelectrons_class, fun_basis="physical", out_form="2", comp=1, domain=self.domain + forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain ) fun_electrons_pb_3 = TransformedPformComponent( - forcetermelectrons_class, fun_basis="physical", out_form="2", comp=2, domain=self.domain + forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=2, domain=self.domain ) if self._lifting: l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) @@ -8187,10 +8206,10 @@ def __init__( A=_A, B=_B, F=_F, - solver_name=solver["type"][0], - tol=solver["tol"], - max_iter=solver["maxiter"], - verbose=solver["verbose"], + solver_name=self.options.solver, + tol=self.options.solver_params.tol, + max_iter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, pc=None, ) # Allocate memory for call @@ -8204,17 +8223,17 @@ def __init__( F=_Fnp, method_to_solve=self._method_to_solve, preconditioner=self._preconditioner, - spectralanalysis=spectralanalysis, - tol=solver["tol"], - max_iter=solver["maxiter"], - verbose=solver["verbose"], + spectralanalysis=self.options.spectralanalysis, + tol=self.options.solver_params.tol, + max_iter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, ) def __call__(self, dt): # current variables - unfeec = self.feec_vars[0] - uenfeec = self.feec_vars[1] - phinfeec = self.feec_vars[2] + unfeec = self.variables.u.spline.vector + uenfeec = self.variables.ue.spline.vector + phinfeec = self.variables.phi.spline.vector if self._variant == "GMRES": if self._lifting: @@ -8331,7 +8350,7 @@ def __call__(self, dt): uen = _sol1[1] phin = _sol2 # write new coeffs into self.feec_vars - max_du, max_due, max_dphi = self.feec_vars_update(un, uen, phin) + max_du, max_due, max_dphi = self.update_feec_variables(u=un, ue=uen, phi=phin) elif self._variant == "Uzawa": # Numpy @@ -8430,7 +8449,7 @@ def __call__(self, dt): print(f"TwoFluidQuasiNeutralFull is only running on one MPI.") # write new coeffs into self.feec_vars - max_du, max_due, max_dphi = self.feec_vars_update(u_temp, ue_temp, phi_temp) + max_du, max_due, max_dphi = self.update_feec_variables(u=u_temp, ue=ue_temp, phi=phi_temp) if self._info and self._rank == 0: print("Status for TwoFluidQuasiNeutralFull:", info["success"]) diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 3d89b3d94..ecd00fa88 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -24,6 +24,7 @@ from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.ode.utils import ButcherTableau from struphy.pic.accumulation import accum_kernels, accum_kernels_gc +from struphy.pic.accumulation.particles_to_grid import AccumulatorVector from struphy.pic.base import Particles from struphy.pic.particles import Particles3D, Particles5D, Particles6D, ParticlesSPH from struphy.pic.pushing import eval_kernels_gc, pusher_kernels, pusher_kernels_gc @@ -1470,40 +1471,66 @@ class PushDeterministicDiffusion(Propagator): * Explicit from :class:`~struphy.ode.utils.ButcherTableau` """ - @staticmethod - def options(default=False): - dct = {} - dct["algo"] = ["rk4", "forward_euler", "heun2", "rk2", "heun3"] - dct["diffusion_coefficient"] = 1.0 - if default: - dct = descend_options_dict(dct, []) - return dct + class Variables: + def __init__(self): + self._var: PICVariable = None - def __init__( - self, - particles: Particles3D, - *, - algo: str = options(default=True)["algo"], - bc_type: list = ["periodic", "periodic", "periodic"], - diffusion_coefficient: float = options()["diffusion_coefficient"], - ): - from struphy.pic.accumulation.particles_to_grid import AccumulatorVector + @property + def var(self) -> PICVariable: + return self._var - super().__init__(particles) + @var.setter + def var(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles3D" + self._var = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + butcher: ButcherTableau = None + bc_type: tuple = ("periodic", "periodic", "periodic") + diff_coeff: float = 1.0 - self._bc_type = bc_type - self._diffusion = diffusion_coefficient + def __post_init__(self): + # defaults + if self.butcher is None: + self.butcher = ButcherTableau() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._bc_type = self.options.bc_type + self._diffusion = self.options.diff_coeff self._tmp = self.derham.Vh["1"].zeros() # choose algorithm - self._butcher = ButcherTableau(algo) + self._butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: import numpy as np self._butcher._a = np.diag(self._butcher.a, k=-1) self._butcher._a = np.array(list(self._butcher.a) + [0.0]) + particles = self.variables.var.particles + self._u_on_grid = AccumulatorVector( particles, "H1", @@ -1539,9 +1566,10 @@ def __call__(self, dt): """ TODO """ + particles = self.variables.var.particles # accumulate - self._u_on_grid(self.particles[0].vdim) + self._u_on_grid() # take gradient pi_u = self._u_on_grid.vectors[0] @@ -1552,8 +1580,8 @@ def __call__(self, dt): self._pusher(dt) # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() + if particles.control_variate: + particles.update_weights() class PushRandomDiffusion(Propagator): @@ -1576,31 +1604,59 @@ class PushRandomDiffusion(Propagator): * ``forward_euler`` (1st order) """ - @staticmethod - def options(default=False): - dct = {} - dct["algo"] = ["forward_euler"] - dct["diffusion_coefficient"] = 1.0 - if default: - dct = descend_options_dict(dct, []) - return dct + class Variables: + def __init__(self): + self._var: PICVariable = None - def __init__( - self, - particles: Particles3D, - algo: str = options(default=True)["algo"], - bc_type: list = ["periodic", "periodic", "periodic"], - diffusion_coefficient: float = options()["diffusion_coefficient"], - ): - super().__init__(particles) + @property + def var(self) -> PICVariable: + return self._var - self._bc_type = bc_type - self._diffusion = diffusion_coefficient + @var.setter + def var(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles3D" + self._var = new - self._noise = array(self.particles[0].markers[:, :3]) + def __init__(self): + self.variables = self.Variables() - # choose algorithm - self._butcher = ButcherTableau("forward_euler") + @dataclass + class Options: + butcher: ButcherTableau = None + bc_type: tuple = ("periodic", "periodic", "periodic") + diff_coeff: float = 1.0 + + def __post_init__(self): + # defaults + if self.butcher is None: + self.butcher = ButcherTableau() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._bc_type = self.options.bc_type + self._diffusion = self.options.diff_coeff + + particles = self.variables.var.particles + + self._noise = array(particles.markers[:, :3]) + + self._butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: import numpy as np @@ -1635,18 +1691,20 @@ def __call__(self, dt): TODO """ + particles = self.variables.var.particles + self._noise[:] = random.multivariate_normal( self._mean, self._cov, - len(self.particles[0].markers), + len(particles.markers), ) # push markers self._pusher(dt) # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() + if particles.control_variate: + particles.update_weights() class PushVinSPHpressure(Propagator): diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index 4d2d7087e..2cac1a195 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -7,10 +7,11 @@ from struphy.feec.projectors import L2Projector from struphy.feec.psydac_derham import Derham from struphy.geometry import domains +from struphy.geometry.base import Domain from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable -from struphy.propagators import ImplicitDiffusion from struphy.propagators.base import Propagator +from struphy.propagators.propagators_fields import ImplicitDiffusion comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -27,7 +28,8 @@ ["Orthogonal", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 3.0}], ], ) -def test_poisson_M1perp_1d(direction, bc_type, mapping, show_plot=False): +@pytest.mark.parametrize("projected_rhs", [False, True]) +def test_poisson_M1perp_1d(direction, bc_type, mapping, projected_rhs, show_plot=False): """ Test the convergence of Poisson solver with M1perp diffusion matrix in 1D by means of manufactured solutions. @@ -38,7 +40,7 @@ def test_poisson_M1perp_1d(direction, bc_type, mapping, show_plot=False): dom_params = mapping[1] domain_class = getattr(domains, dom_type) - domain = domain_class(**dom_params) + domain: Domain = domain_class(**dom_params) if dom_type == "Cuboid": Lx = dom_params["r1"] - dom_params["l1"] @@ -132,10 +134,16 @@ def rho1_xyz(x, y, z): Propagator.mass_ops = mass_ops # pullbacks of right-hand side - def rho1(e1, e2, e3): - return domain.pull(rho1_xyz, e1, e2, e3, kind="0", squeeze_out=True) - - rho_vec = L2Projector("H1", mass_ops).get_dofs(rho1, apply_bc=True) + def rho_pulled(e1, e2, e3): + return domain.pull(rho1_xyz, e1, e2, e3, kind="0", squeeze_out=False) + + # define how to pass rho + if projected_rhs: + rho = FEECVariable(space="H1") + rho.allocate(derham=derham, domain=domain) + rho.spline.vector = derham.P["0"](rho_pulled) + else: + rho = rho_pulled # create Poisson solver solver_params = SolverParameters( @@ -158,7 +166,7 @@ def rho1(e1, e2, e3): sigma_3=1.0, divide_by_dt=True, diffusion_mat="M1perp", - rho=rho_vec, + rho=rho, solver="pcg", precond="MassMatrixPreconditioner", solver_params=solver_params, @@ -198,7 +206,7 @@ def rho1(e1, e2, e3): m, _ = np.polyfit(np.log(Nels), np.log(errors), deg=1) print(f"For {pi = }, solution converges in {direction=} with rate {-m = } ") - assert -m > (pi + 1 - 0.06) + assert -m > (pi + 1 - 0.07) # Plot convergence in 1D if show_plot: @@ -234,7 +242,8 @@ def rho1(e1, e2, e3): ["Orthogonal", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 1.0}], ], ) -def test_poisson_M1perp_2d(Nel, p, bc_type, mapping, show_plot=False): +@pytest.mark.parametrize("projected_rhs", [False, True]) +def test_poisson_M1perp_2d(Nel, p, bc_type, mapping, projected_rhs, show_plot=False): """ Test the Poisson solver with M1perp diffusion matrix by means of manufactured solutions in 2D . @@ -245,7 +254,7 @@ def test_poisson_M1perp_2d(Nel, p, bc_type, mapping, show_plot=False): dom_params = mapping[1] domain_class = getattr(domains, dom_type) - domain = domain_class(**dom_params) + domain: Domain = domain_class(**dom_params) if dom_type == "Cuboid": Lx = dom_params["r1"] - dom_params["l1"] @@ -326,16 +335,24 @@ def rho1_xyz(x, y, z): e3 = np.linspace(0.0, 1.0, 1) # pullbacks of right-hand side - def rho1(e1, e2, e3): - return domain.pull(rho1_xyz, e1, e2, e3, kind="0", squeeze_out=True) + def rho1_pulled(e1, e2, e3): + return domain.pull(rho1_xyz, e1, e2, e3, kind="0", squeeze_out=False) - def rho2(e1, e2, e3): - return domain.pull(rho2_xyz, e1, e2, e3, kind="0", squeeze_out=True) + def rho2_pulled(e1, e2, e3): + return domain.pull(rho2_xyz, e1, e2, e3, kind="0", squeeze_out=False) - # discrete right-hand sides - l2_proj = L2Projector("H1", mass_ops) - rho_vec1 = l2_proj.get_dofs(rho1, apply_bc=True) - rho_vec2 = l2_proj.get_dofs(rho2, apply_bc=True) + # how to pass right-hand sides + if projected_rhs: + rho1 = FEECVariable(space="H1") + rho1.allocate(derham=derham, domain=domain) + rho1.spline.vector = derham.P["0"](rho1_pulled) + + rho2 = FEECVariable(space="H1") + rho2.allocate(derham=derham, domain=domain) + rho2.spline.vector = derham.P["0"](rho2_pulled) + else: + rho1 = rho1_pulled + rho2 = rho2_pulled # Create Poisson solvers solver_params = SolverParameters( @@ -358,7 +375,7 @@ def rho2(e1, e2, e3): sigma_3=1.0, divide_by_dt=True, diffusion_mat="M1perp", - rho=rho_vec1, + rho=rho1, solver="pcg", precond="MassMatrixPreconditioner", solver_params=solver_params, @@ -378,7 +395,7 @@ def rho2(e1, e2, e3): sigma_3=1.0, divide_by_dt=True, diffusion_mat="M1perp", - rho=rho_vec2, + rho=rho2, solver="pcg", precond="MassMatrixPreconditioner", solver_params=solver_params, @@ -431,7 +448,7 @@ def rho2(e1, e2, e3): plt.show() assert error1 < 0.0044 - assert error2 < 0.021 + assert error2 < 0.023 @pytest.mark.skip(reason="Not clear if the 2.5d strategy is sound.") @@ -459,7 +476,7 @@ def test_poisson_M1perp_3d_compare_2p5d(Nel, p, mapping, show_plot=False): dom_params = mapping[1] domain_class = getattr(domains, dom_type) - domain = domain_class(**dom_params) + domain: Domain = domain_class(**dom_params) # boundary conditions spl_kind = [False, True, True] @@ -515,7 +532,7 @@ def rho(e1, e2, e3): sigma_3=1.0, divide_by_dt=True, diffusion_mat="M1perp", - rho=rho_vec, + rho=rho, solver="pcg", precond="MassMatrixPreconditioner", solver_params=solver_params, @@ -539,7 +556,6 @@ def rho(e1, e2, e3): _phi_small = FEECVariable(space="H1") _phi_small.allocate(derham=derham, domain=domain) - rhs = derham.create_spline_function("rhs", "H1") poisson_solver_2p5d = ImplicitDiffusion() poisson_solver_2p5d.variables.phi = _phi_small @@ -550,7 +566,7 @@ def rho(e1, e2, e3): sigma_3=1.0, divide_by_dt=True, diffusion_mat="M1perp", - rho=rhs.vector, + rho=rho, solver="pcg", precond="MassMatrixPreconditioner", solver_params=solver_params, @@ -569,8 +585,6 @@ def rho(e1, e2, e3): t0 = time() t_inner = 0.0 for n in range(s[2], e[2] + 1): - # scale the rhs with Nel[2] !! - rhs.vector[s[0] : e[0] + 1, s[1] : e[1] + 1, 0] = rho_vec[s[0] : e[0] + 1, s[1] : e[1] + 1, n] * Nel[2] t0i = time() poisson_solver_2p5d(dt) t1i = time() @@ -637,7 +651,7 @@ def rho(e1, e2, e3): # mapping = ['Orthogonal', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 1.}] # test_poisson_M1perp_2d(Nel, p, bc_type, mapping, show_plot=True) - Nel = [64, 64, 16] - p = [2, 2, 1] - mapping = ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}] + # Nel = [64, 64, 16] + # p = [2, 2, 1] + # mapping = ["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}] # test_poisson_M1perp_3d_compare_2p5d(Nel, p, mapping, show_plot=True) diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 5a1de11ba..812df0593 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -7,10 +7,22 @@ from struphy.feec.projectors import L2Projector from struphy.feec.psydac_derham import Derham from struphy.geometry import domains +from struphy.geometry.base import Domain +from struphy.initial import perturbations +from struphy.kinetic_background.maxwellians import Maxwellian3D from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable -from struphy.propagators import ImplicitDiffusion +from struphy.pic.accumulation.accum_kernels import charge_density_0form +from struphy.pic.accumulation.particles_to_grid import AccumulatorVector +from struphy.pic.particles import Particles6D +from struphy.pic.utilities import ( + BinningPlot, + BoundaryParameters, + LoadingParameters, + WeightsParameters, +) from struphy.propagators.base import Propagator +from struphy.propagators.propagators_fields import ImplicitDiffusion, Poisson comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -27,7 +39,14 @@ ["Orthogonal", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 3.0}], ], ) -def test_poisson_1d(direction, bc_type, mapping, show_plot=False): +@pytest.mark.parametrize("projected_rhs", [False, True]) +def test_poisson_1d( + direction: int, + bc_type: str, + mapping: list[str, dict], + projected_rhs: bool, + show_plot: bool = False, +): """ Test the convergence of Poisson solver in 1D by means of manufactured solutions. """ @@ -37,7 +56,7 @@ def test_poisson_1d(direction, bc_type, mapping, show_plot=False): dom_params = mapping[1] domain_class = getattr(domains, dom_type) - domain = domain_class(**dom_params) + domain: Domain = domain_class(**dom_params) if dom_type == "Cuboid": Lx = dom_params["r1"] - dom_params["l1"] @@ -155,10 +174,16 @@ def rho1_xyz(x, y, z): Propagator.mass_ops = mass_ops # pullbacks of right-hand side - def rho1(e1, e2, e3): - return domain.pull(rho1_xyz, e1, e2, e3, kind="0", squeeze_out=True) - - rho_vec = L2Projector("H1", mass_ops).get_dofs(rho1, apply_bc=True) + def rho_pulled(e1, e2, e3): + return domain.pull(rho1_xyz, e1, e2, e3, kind="0", squeeze_out=False) + + # define how to pass rho + if projected_rhs: + rho = FEECVariable(space="H1") + rho.allocate(derham=derham, domain=domain) + rho.spline.vector = derham.P["0"](rho_pulled) + else: + rho = rho_pulled # create Poisson solver solver_params = SolverParameters( @@ -172,14 +197,14 @@ def rho1(e1, e2, e3): _phi = FEECVariable(space="H1") _phi.allocate(derham=derham, domain=domain) - poisson_solver = ImplicitDiffusion() + poisson_solver = Poisson() poisson_solver.variables.phi = _phi poisson_solver.options = poisson_solver.Options( - sigma_1=1e-12, - sigma_2=0.0, - sigma_3=1.0, - rho=rho_vec, + stab_eps=1e-12, + # sigma_2=0.0, + # sigma_3=1.0, + rho=rho, solver="pcg", precond="MassMatrixPreconditioner", solver_params=solver_params, @@ -223,7 +248,7 @@ def rho1(e1, e2, e3): m, _ = np.polyfit(np.log(Nels), np.log(errors), deg=1) print(f"For {pi = }, solution converges in {direction=} with rate {-m = } ") - assert -m > (pi + 1 - 0.06) + assert -m > (pi + 1 - 0.07) # Plot convergence in 1D if show_plot: @@ -248,6 +273,161 @@ def rho1(e1, e2, e3): plt.show() +@pytest.mark.parametrize( + "mapping", + [ + ["Cuboid", {"l1": 0.0, "r1": 4.0, "l2": 0.0, "r2": 2.0, "l3": 0.0, "r3": 3.0}], + # ["Orthogonal", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 3.0}], + ], +) +def test_poisson_accum_1d(mapping, do_plot=False): + """Pass accumulators as rhs.""" + # create domain object + dom_type = mapping[0] + dom_params = mapping[1] + + domain_class = getattr(domains, dom_type) + domain: Domain = domain_class(**dom_params) + + if dom_type == "Cuboid": + Lx = dom_params["r1"] - dom_params["l1"] + else: + Lx = dom_params["Lx"] + + # create derham object + Nel = (16, 1, 1) + p = (2, 1, 1) + spl_kind = (True, True, True) + derham = Derham(Nel, p, spl_kind, comm=comm) + + # mass matrices + mass_ops = WeightedMassOperators(derham, domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + # 6D particle object + domain_array = derham.domain_array + nprocs = derham.domain_decomposition.nprocs + domain_decomp = (domain_array, nprocs) + + lp = LoadingParameters(ppc=4000, seed=765) + wp = WeightsParameters(control_variate=True) + bp = BoundaryParameters() + + backgr = Maxwellian3D(n=(1.0, None)) + l = 1 + amp = 1e-1 + pert = perturbations.ModesCos(ls=(l,), amps=(amp,)) + maxw = Maxwellian3D(n=(1.0, pert)) + + pert_exact = lambda x, y, z: amp * np.cos(l * 2 * np.pi / Lx * x) + phi_exact = lambda x, y, z: amp / (l * 2 * np.pi / Lx) ** 2 * np.cos(l * 2 * np.pi / Lx * x) + e_exact = lambda x, y, z: amp / (l * 2 * np.pi / Lx) * np.sin(l * 2 * np.pi / Lx * x) + + particles = Particles6D( + comm_world=comm, + domain_decomp=domain_decomp, + loading_params=lp, + weights_params=wp, + boundary_params=bp, + domain=domain, + background=backgr, + initial_condition=maxw, + ) + particles.draw_markers() + particles.initialize_weights() + + # particle to grid coupling + kernel = charge_density_0form + accum = AccumulatorVector(particles, "H1", kernel, mass_ops, domain.args_domain) + # accum() + # if do_plot: + # accum.show_accumulated_spline_field(mass_ops) + + rho = accum + + # create Poisson solver + solver_params = SolverParameters( + tol=1.0e-13, + maxiter=3000, + info=True, + verbose=False, + recycle=False, + ) + + _phi = FEECVariable(space="H1") + _phi.allocate(derham=derham, domain=domain) + + poisson_solver = Poisson() + poisson_solver.variables.phi = _phi + + poisson_solver.options = poisson_solver.Options( + stab_eps=1e-6, + # sigma_2=0.0, + # sigma_3=1.0, + rho=rho, + solver="pcg", + precond="MassMatrixPreconditioner", + solver_params=solver_params, + ) + + poisson_solver.allocate() + + # Solve Poisson (call propagator with dt=1.) + dt = 1.0 + poisson_solver(dt) + + # push numerical solution and compare + e1 = np.linspace(0.0, 1.0, 50) + e2 = 0.0 + e3 = 0.0 + + num_values = domain.push(_phi.spline, e1, e2, e3, kind="0") + x, y, z = domain(e1, e2, e3) + pert_values = pert_exact(x, y, z) + analytic_values = phi_exact(x, y, z) + e_values = e_exact(x, y, z) + + _e = FEECVariable(space="Hcurl") + _e.allocate(derham=derham, domain=domain) + derham.grad.dot(-_phi.spline.vector, out=_e.spline.vector) + num_values_e = domain.push(_e.spline, e1, e2, e3, kind="1") + + if do_plot: + field = derham.create_spline_function("accum_field", "H1") + field.vector = accum.vectors[0] + accum_values = field(e1, e2, e3) + + plt.figure(figsize=(18, 12)) + plt.subplot(1, 3, 1) + plt.plot(x[:, 0, 0], num_values[:, 0, 0], "ob", label="numerical") + plt.plot(x[:, 0, 0], analytic_values[:, 0, 0], "r--", label="exact") + plt.xlabel("x") + plt.title("phi") + plt.legend() + plt.subplot(1, 3, 2) + plt.plot(x[:, 0, 0], accum_values[:, 0, 0], "ob", label="numerical, without L2-proj") + plt.plot(x[:, 0, 0], pert_values[:, 0, 0], "r--", label="exact") + plt.xlabel("x") + plt.title("rhs") + plt.legend() + plt.subplot(1, 3, 3) + plt.plot(x[:, 0, 0], num_values_e[0][:, 0, 0], "ob", label="numerical") + plt.plot(x[:, 0, 0], e_values[:, 0, 0], "r--", label="exact") + plt.xlabel("x") + plt.title("e_field") + plt.legend() + + plt.show() + + error = np.max(np.abs(num_values_e[0][:, 0, 0] - e_values[:, 0, 0])) / np.max(np.abs(e_values[:, 0, 0])) + print(f"{error=}") + + assert error < 0.0086 + + @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize("Nel", [[64, 64, 1]]) @pytest.mark.parametrize("p", [[1, 1, 1], [2, 2, 1]]) @@ -259,7 +439,8 @@ def rho1(e1, e2, e3): ["Colella", {"Lx": 4.0, "Ly": 2.0, "alpha": 0.1, "Lz": 1.0}], ], ) -def test_poisson_2d(Nel, p, bc_type, mapping, show_plot=False): +@pytest.mark.parametrize("projected_rhs", [False, True]) +def test_poisson_2d(Nel, p, bc_type, mapping, projected_rhs, show_plot=False): """ Test the Poisson solver by means of manufactured solutions in 2D . """ @@ -269,7 +450,7 @@ def test_poisson_2d(Nel, p, bc_type, mapping, show_plot=False): dom_params = mapping[1] domain_class = getattr(domains, dom_type) - domain = domain_class(**dom_params) + domain: Domain = domain_class(**dom_params) if dom_type == "Cuboid": Lx = dom_params["r1"] - dom_params["l1"] @@ -350,16 +531,24 @@ def rho1_xyz(x, y, z): e3 = np.linspace(0.0, 1.0, 1) # pullbacks of right-hand side - def rho1(e1, e2, e3): - return domain.pull(rho1_xyz, e1, e2, e3, kind="0", squeeze_out=True) + def rho1_pulled(e1, e2, e3): + return domain.pull(rho1_xyz, e1, e2, e3, kind="0", squeeze_out=False) - def rho2(e1, e2, e3): - return domain.pull(rho2_xyz, e1, e2, e3, kind="0", squeeze_out=True) + def rho2_pulled(e1, e2, e3): + return domain.pull(rho2_xyz, e1, e2, e3, kind="0", squeeze_out=False) - # discrete right-hand sides - l2_proj = L2Projector("H1", mass_ops) - rho_vec1 = l2_proj.get_dofs(rho1, apply_bc=True) - rho_vec2 = l2_proj.get_dofs(rho2, apply_bc=True) + # how to pass right-hand sides + if projected_rhs: + rho1 = FEECVariable(space="H1") + rho1.allocate(derham=derham, domain=domain) + rho1.spline.vector = derham.P["0"](rho1_pulled) + + rho2 = FEECVariable(space="H1") + rho2.allocate(derham=derham, domain=domain) + rho2.spline.vector = derham.P["0"](rho2_pulled) + else: + rho1 = rho1_pulled + rho2 = rho2_pulled # Create Poisson solvers solver_params = SolverParameters( @@ -373,14 +562,14 @@ def rho2(e1, e2, e3): _phi1 = FEECVariable(space="H1") _phi1.allocate(derham=derham, domain=domain) - poisson_solver1 = ImplicitDiffusion() + poisson_solver1 = Poisson() poisson_solver1.variables.phi = _phi1 poisson_solver1.options = poisson_solver1.Options( - sigma_1=1e-8, - sigma_2=0.0, - sigma_3=1.0, - rho=rho_vec1, + stab_eps=1e-8, + # sigma_2=0.0, + # sigma_3=1.0, + rho=rho1, solver="pcg", precond="MassMatrixPreconditioner", solver_params=solver_params, @@ -389,21 +578,27 @@ def rho2(e1, e2, e3): poisson_solver1.allocate() # _phi1 = derham.create_spline_function("test1", "H1") - # poisson_solver1 = ImplicitDiffusion( + # poisson_solver1 = Poisson( # _phi1.vector, sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, rho=rho_vec1, solver=solver_params # ) _phi2 = FEECVariable(space="H1") _phi2.allocate(derham=derham, domain=domain) - poisson_solver2 = ImplicitDiffusion() + poisson_solver2 = Poisson() poisson_solver2.variables.phi = _phi2 + stab_eps = 1e-8 + err_lim = 0.03 + if bc_type == "neumann" and dom_type == "Colella": + stab_eps = 1e-4 + err_lim = 0.046 + poisson_solver2.options = poisson_solver2.Options( - sigma_1=1e-8, - sigma_2=0.0, - sigma_3=1.0, - rho=rho_vec2, + stab_eps=stab_eps, + # sigma_2=0.0, + # sigma_3=1.0, + rho=rho2, solver="pcg", precond="MassMatrixPreconditioner", solver_params=solver_params, @@ -412,7 +607,7 @@ def rho2(e1, e2, e3): poisson_solver2.allocate() # _phi2 = derham.create_spline_function("test2", "H1") - # poisson_solver2 = ImplicitDiffusion( + # poisson_solver2 = Poisson( # _phi2.vector, sigma_1=1e-8, sigma_2=0.0, sigma_3=1.0, rho=rho_vec2, solver=solver_params # ) @@ -463,21 +658,23 @@ def rho2(e1, e2, e3): if p[0] == 1 and bc_type == "neumann" and mapping[0] == "Colella": pass else: - assert error1 < 0.0044 - assert error2 < 0.021 + assert error1 < 0.0053 + assert error2 < err_lim if __name__ == "__main__": - direction = 2 - bc_type = "dirichlet" + # direction = 0 + # bc_type = "dirichlet" mapping = ["Cuboid", {"l1": 0.0, "r1": 4.0, "l2": 0.0, "r2": 2.0, "l3": 0.0, "r3": 3.0}] # mapping = ['Orthogonal', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 3.}] - test_poisson_1d(direction, bc_type, mapping, show_plot=True) + # test_poisson_1d(direction, bc_type, mapping, projected_rhs=True, show_plot=True) # Nel = [64, 64, 1] # p = [2, 2, 1] # bc_type = 'neumann' - # #mapping = ['Cuboid', {'l1': 0., 'r1': 4., 'l2': 0., 'r2': 2., 'l3': 0., 'r3': 3.}] - # #mapping = ['Orthogonal', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 1.}] + # # mapping = ['Cuboid', {'l1': 0., 'r1': 4., 'l2': 0., 'r2': 2., 'l3': 0., 'r3': 3.}] + # # mapping = ['Orthogonal', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 1.}] # mapping = ['Colella', {'Lx': 4., 'Ly': 2., 'alpha': .1, 'Lz': 1.}] - # test_poisson_2d(Nel, p, bc_type, mapping, show_plot=True) + # test_poisson_2d(Nel, p, bc_type, mapping, projected_rhs=True, show_plot=True) + + test_poisson_accum_1d(mapping, do_plot=True) From b5754f06f63ed027f3a16434741788d289e665ad Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 8 Oct 2025 14:05:33 +0200 Subject: [PATCH 129/292] remove double --check-file --- src/struphy/console/main.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index 6718cadf8..94c8f3220 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -713,13 +713,6 @@ def add_parser_params(subparsers, list_models, model_message): help="check if the parameters in the .yml file are valid", ) - parser_params.add_argument( - "--check-file", - type=str, - metavar="FILE", - help="check if the parameters in the .yml file are valid", - ) - parser_params.add_argument( "-y", "--yes", From 9f5701dd5fb482f221adecde58818185d1d05429 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 8 Oct 2025 12:15:21 +0000 Subject: [PATCH 130/292] Update shearalfven to new format --- src/struphy/console/params.py | 3 +- src/struphy/models/tests/test_models.py | 1 + src/struphy/models/toy.py | 92 +++++++++++-------------- 3 files changed, 44 insertions(+), 52 deletions(-) diff --git a/src/struphy/console/params.py b/src/struphy/console/params.py index bbe33ace3..1e5ed7d5f 100644 --- a/src/struphy/console/params.py +++ b/src/struphy/console/params.py @@ -45,4 +45,5 @@ def struphy_params(model_name: str, params_path: str, yes: bool = False, check_f else: prompt = not yes - model.generate_default_parameter_file(path=params_path, prompt=prompt) + print(f"Generating default parameter file for {model_class}.") + model_class().generate_default_parameter_file(path=params_path, prompt=prompt) diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index b4462af54..0fd193a58 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -19,6 +19,7 @@ "GuidingCenter", "PressureLessSPH", "Poisson", + "ShearAlfven", "DeterministicParticleDiffusion", "RandomParticleDiffusion", "TwoFluidQuasiNeutralToy", diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index f91cd4c81..62f3b4722 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -307,43 +307,30 @@ class ShearAlfven(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["em_fields"]["b2"] = "Hdiv" - dct["fluid"]["mhd"] = {"u2": "Hdiv"} - return dct - - @staticmethod - def bulk_species(): - return "mhd" - - @staticmethod - def velocity_scale(): - return "alfvén" - - @staticmethod - def propagators_dct(): - return {propagators_fields.ShearAlfven: ["mhd_u2", "b2"]} + ## species + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + class MHD(FluidSpecies): + def __init__(self): + self.velocity = FEECVariable(space="Hdiv") + self.init_variables() - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + class Propagators: + def __init__(self) -> None: + self.shear_alf = propagators_fields.ShearAlfven() - from struphy.polar.basic import PolarVector + @property + def bulk_species(self): + return self.mhd - # extract necessary parameters - alfven_solver = params["fluid"]["mhd"]["options"]["ShearAlfven"]["solver"] - alfven_algo = params["fluid"]["mhd"]["options"]["ShearAlfven"]["algo"] + @property + def velocity_scale(self): + return "alfvén" + def allocate_helpers(self): # project background magnetic field (2-form) and pressure (3-form) self._b_eq = self.derham.P["2"]( [ @@ -353,21 +340,26 @@ def __init__(self, params, comm, clone_config=None): ] ) - # set keyword arguments for propagators - self._kwargs[propagators_fields.ShearAlfven] = { - "u_space": "Hdiv", - "solver": alfven_solver, - "algo": alfven_algo, - } + # temporary vectors for scalar quantities + self._tmp_b1 = self.derham.Vh["2"].zeros() + self._tmp_b2 = self.derham.Vh["2"].zeros() - # Initialize propagators used in splitting substeps - self.init_propagators() + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.shear_alf.variables.u = self.mhd.velocity + self.propagators.shear_alf.variables.b = self.em_fields.b_field # Scalar variables to be saved during simulation - # self.add_scalar('en_U') - # self.add_scalar('en_B') - # self.add_scalar('en_B_eq') - # self.add_scalar('en_B_tot') self.add_scalar("en_tot") self.add_scalar("en_U", compute="from_field") @@ -376,14 +368,12 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("en_B_tot", compute="from_field") self.add_scalar("en_tot2", summands=["en_U", "en_B", "en_B_eq"]) - # temporary vectors for scalar quantities - self._tmp_b1 = self.derham.Vh["2"].zeros() - self._tmp_b2 = self.derham.Vh["2"].zeros() - def update_scalar_quantities(self): # perturbed fields - en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.pointer["mhd_u2"], self.pointer["mhd_u2"]) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b2"], self.pointer["b2"]) + en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.mhd.velocity.spline.vector, self.mhd.velocity.spline.vector) + en_B = 0.5 * self.mass_ops.M2.dot_inner( + self.em_fields.b_field.spline.vector, self.em_fields.b_field.spline.vector + ) self.update_scalar("en_U", en_U) self.update_scalar("en_B", en_B) @@ -396,7 +386,7 @@ def update_scalar_quantities(self): # total magnetic field self._b_eq.copy(out=self._tmp_b1) - self._tmp_b1 += self.pointer["b2"] + self._tmp_b1 += self.em_fields.b_field.spline.vector self.mass_ops.M2.dot(self._tmp_b1, apply_bc=False, out=self._tmp_b2) en_Btot = self._tmp_b1.inner(self._tmp_b2) / 2 From fd2a463d496bca31a619361beb8883c4fddde5a7 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 8 Oct 2025 14:44:32 +0200 Subject: [PATCH 131/292] test hook --- src/struphy/pic/particles.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 72e25d205..9382b7518 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -18,6 +18,7 @@ from struphy.pic.base import Particles + class Particles6D(Particles): """ A class for initializing particles in models that use the full 6D phase space. From 4c77173f252d99ca86519934d34ebac15f2080e8 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 8 Oct 2025 14:46:10 +0200 Subject: [PATCH 132/292] update pre-commit yaml --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index af670854e..9a0819492 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 # Use the latest stable version + rev: v6.0.0 # Use the latest stable version hooks: - id: check-added-large-files # Prevent giant files from being committed. args: ["--maxkb=1000"] From 16ae1464b4e99d7c195bac5b304be990ae925cc6 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 8 Oct 2025 14:47:44 +0200 Subject: [PATCH 133/292] test pre-commit hooks --- src/struphy/pic/particles.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 9382b7518..f393e8ed6 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -17,7 +17,8 @@ from struphy.pic import utilities_kernels from struphy.pic.base import Particles - +<<<<<<< HEAD +>>>>>>> devel class Particles6D(Particles): """ From b1343def4fa45690ef5b6be4905298c67395e6fd Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 8 Oct 2025 14:50:30 +0200 Subject: [PATCH 134/292] try hook option --assume-in-merge --- .pre-commit-config.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9a0819492..3f3107d09 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -5,6 +5,7 @@ repos: - id: check-added-large-files # Prevent giant files from being committed. args: ["--maxkb=1000"] - id: check-merge-conflict # Check for files that contain merge conflict strings. + args: ["--assume-in-merge"] - id: check-toml # Attempts to load all TOML files to verify syntax. - id: check-yaml # Attempts to load all yaml files to verify syntax. args: ["--unsafe"] From fc4225997a63e81717aff1eb04d745cb5729d4d0 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 8 Oct 2025 14:59:10 +0200 Subject: [PATCH 135/292] test hook From efc0fc1089979f45f369b4581da1fba44529cf68 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 8 Oct 2025 15:06:16 +0200 Subject: [PATCH 136/292] remove merge conflicts --- src/struphy/pic/particles.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index f393e8ed6..3489e701a 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -2,12 +2,8 @@ import numpy as np -<<<<<<< HEAD from struphy.fields_background import equils from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB -======= -from struphy.fields_background.base import FluidEquilibriumWithB ->>>>>>> devel from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB from struphy.geometry.base import Domain from struphy.geometry.utilities import TransformedPformComponent @@ -17,8 +13,6 @@ from struphy.pic import utilities_kernels from struphy.pic.base import Particles -<<<<<<< HEAD ->>>>>>> devel class Particles6D(Particles): """ From b75d347bd4d3db94e8e1dbee381d13e941f4c336 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 8 Oct 2025 13:10:23 +0000 Subject: [PATCH 137/292] Resolve "Port remaining toy models" --- src/struphy/feec/mass.py | 94 ++++ src/struphy/io/options.py | 1 + src/struphy/linear_algebra/solver.py | 14 + src/struphy/models/tests/test_models.py | 19 +- src/struphy/models/toy.py | 416 ++++++++-------- src/struphy/propagators/propagators_fields.py | 459 +++++++++++------- 6 files changed, 593 insertions(+), 410 deletions(-) diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 9d5531e98..e0443228e 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -1,4 +1,5 @@ import inspect +from copy import deepcopy import numpy as np from mpi4py import MPI @@ -7,6 +8,7 @@ from psydac.fem.vector import VectorFemSpace from psydac.linalg.basic import IdentityOperator, LinearOperator, Vector from psydac.linalg.block import BlockLinearOperator, BlockVector +from psydac.linalg.solvers import inverse from psydac.linalg.stencil import StencilDiagonalMatrix, StencilMatrix, StencilVector from struphy.feec import mass_kernels @@ -728,6 +730,12 @@ def M1gyro(self): return self._M1gyro + @property + def WMM(self): + if not hasattr(self, "_WMM"): + self._WMM = self.H1vecMassMatrix_density(self.derham, self, self.domain) + return self._WMM + ####################################### # Wrapper around WeightedMassOperator # ####################################### @@ -1024,6 +1032,92 @@ def _operate(self, f1, f2, op, e1, e2, e3): return out + ####################################### + # Aux classes (to be removed in TODO) # + ####################################### + class H1vecMassMatrix_density: + """Wrapper around a Weighted mass operator from H1vec to H1vec whose weights are given by a 3 form""" + + def __init__(self, derham, mass_ops, domain): + self._massop = mass_ops.create_weighted_mass("H1vec", "H1vec") + self.field = derham.create_spline_function("field", "L2") + + integration_grid = [grid_1d.flatten() for grid_1d in derham.quad_grid_pts["0"]] + + self.integration_grid_spans, self.integration_grid_bn, self.integration_grid_bd = ( + derham.prepare_eval_tp_fixed( + integration_grid, + ) + ) + + grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) + self._f_values = np.zeros(grid_shape, dtype=float) + + metric = domain.metric(*integration_grid) + self._mass_metric_term = deepcopy(metric) + self._full_term_mass = deepcopy(metric) + + @property + def massop( + self, + ): + """The WeightedMassOperator""" + return self._massop + + @property + def inv( + self, + ): + """The inverse WeightedMassOperator""" + if not hasattr(self, "_inv"): + self._create_inv() + return self._inv + + def update_weight(self, coeffs): + """Update the weighted mass matrix operator""" + + self.field.vector = coeffs + f_values = self.field.eval_tp_fixed_loc( + self.integration_grid_spans, + self.integration_grid_bd, + out=self._f_values, + ) + for i in range(3): + for j in range(3): + self._full_term_mass[i, j] = f_values * self._mass_metric_term[i, j] + + self._massop.assemble( + [ + [self._full_term_mass[0, 0], self._full_term_mass[0, 1], self._full_term_mass[0, 2]], + [ + self._full_term_mass[1, 0], + self._full_term_mass[ + 1, + 1, + ], + self._full_term_mass[1, 2], + ], + [self._full_term_mass[2, 0], self._full_term_mass[2, 1], self._full_term_mass[2, 2]], + ], + verbose=False, + ) + + if hasattr(self, "_inv") and self.inv._options["pc"] is not None: + self.inv._options["pc"].update_mass_operator(self.massop) + + def _create_inv(self, type="pcg", tol=1e-16, maxiter=500, verbose=False): + """Inverse the weighted mass matrix, preconditioner must be set outside + via self._inv._options['pc'] = ...""" + self._inv = inverse( + self.massop, + type, + pc=None, + tol=tol, + maxiter=maxiter, + verbose=verbose, + recycle=True, + ) + class WeightedMassOperatorsOldForTesting: r""" diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 1e035a8ec..a98e2cd5c 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -29,6 +29,7 @@ OptsMassPrecond = Literal["MassMatrixPreconditioner", None] OptsSaddlePointSolver = Literal["Uzawa", "GMRES"] OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] +OptsNonlinearSolver = Literal["Picard", "Newton"] # markers OptsPICSpace = Literal["Particles6D", "DeltaFParticles6D", "Particles5D", "Particles3D"] diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py index 0e36cfcaf..f14a98915 100644 --- a/src/struphy/linear_algebra/solver.py +++ b/src/struphy/linear_algebra/solver.py @@ -1,5 +1,7 @@ from dataclasses import dataclass +from struphy.io.options import OptsNonlinearSolver + @dataclass class SolverParameters: @@ -10,3 +12,15 @@ class SolverParameters: info: bool = False verbose: bool = False recycle: bool = True + + +@dataclass +class NonlinearSolverParameters: + """Parameters for psydac solvers.""" + + tol: float = 1e-8 + maxiter: int = 100 + info: bool = False + verbose: bool = False + type: OptsNonlinearSolver = "Picard" + linearize: bool = False diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index 0fd193a58..30771a2de 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -1,3 +1,4 @@ +import inspect import os from types import ModuleType @@ -13,20 +14,10 @@ rank = MPI.COMM_WORLD.Get_rank() # available models -toy_models = [ - "Maxwell", - "Vlasov", - "GuidingCenter", - "PressureLessSPH", - "Poisson", - "ShearAlfven", - "DeterministicParticleDiffusion", - "RandomParticleDiffusion", - "TwoFluidQuasiNeutralToy", -] -# for name, obj in inspect.getmembers(toy): -# if inspect.isclass(obj) and "models.toy" in obj.__module__: -# toy_models += [name] +toy_models = [] +for name, obj in inspect.getmembers(toy): + if inspect.isclass(obj) and "models.toy" in obj.__module__: + toy_models += [name] if rank == 0: print(f"\n{toy_models = }") diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 62f3b4722..32ef5331b 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -3,6 +3,8 @@ import numpy as np from mpi4py import MPI +from struphy.feec.projectors import L2Projector +from struphy.feec.variational_utilities import InternalEnergyEvaluator from struphy.models.base import StruphyModel from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable @@ -419,77 +421,74 @@ class VariationalPressurelessFluid(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["fluid"]["fluid"] = {"rho3": "L2", "uv": "H1vec"} - return dct + ## species - @staticmethod - def bulk_species(): - return "fluid" + class Fluid(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + ## propagators - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["fluid_rho3", "fluid_uv"], - propagators_fields.VariationalMomentumAdvection: ["fluid_uv"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.mass import WeightedMassOperator - from struphy.feec.variational_utilities import H1vecMassMatrix_density - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - # Initialize mass matrix - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_momentum = params["fluid"]["fluid"]["options"]["VariationalMomentumAdvection"]["lin_solver"] - nonlin_solver_momentum = params["fluid"]["fluid"]["options"]["VariationalMomentumAdvection"]["nonlin_solver"] - lin_solver_density = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - - gamma = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": "pressureless", - "gamma": gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - } - - self._kwargs[propagators_fields.VariationalMomentumAdvection] = { - "mass_ops": self.WMM, - "lin_solver": lin_solver_momentum, - "nonlin_solver": nonlin_solver_momentum, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() + class Propagators: + def __init__(self): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_mom = propagators_fields.VariationalMomentumAdvection() - # Scalar variables to be saved during simulation + ## abstract methods + + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.fluid = self.Fluid() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.fluid.density + self.propagators.variat_dens.variables.u = self.fluid.velocity + self.propagators.variat_mom.variables.u = self.fluid.velocity + + # define scalars for update_scalar_quantities self.add_scalar("en_U") + @property + def bulk_species(self): + return self.fluid + + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + pass + def update_scalar_quantities(self): - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["fluid_uv"], self.pointer["fluid_uv"]) + u = self.fluid.velocity.spline.vector + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='pressureless')\n" + ] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class VariationalBarotropicFluid(StruphyModel): r"""Barotropic fluid equations discretized with a variational method. @@ -518,84 +517,84 @@ class VariationalBarotropicFluid(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["fluid"]["fluid"] = {"rho3": "L2", "uv": "H1vec"} - return dct + ## species - @staticmethod - def bulk_species(): - return "fluid" + class Fluid(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + ## propagators - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["fluid_rho3", "fluid_uv"], - propagators_fields.VariationalMomentumAdvection: ["fluid_uv"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.variational_utilities import H1vecMassMatrix_density - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - # Initialize mass matrix - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_momentum = params["fluid"]["fluid"]["options"]["VariationalMomentumAdvection"]["lin_solver"] - nonlin_solver_momentum = params["fluid"]["fluid"]["options"]["VariationalMomentumAdvection"]["nonlin_solver"] - lin_solver_density = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - - gamma = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": "barotropic", - "gamma": gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - } - - self._kwargs[propagators_fields.VariationalMomentumAdvection] = { - "mass_ops": self.WMM, - "lin_solver": lin_solver_momentum, - "nonlin_solver": nonlin_solver_momentum, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() + class Propagators: + def __init__(self): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_mom = propagators_fields.VariationalMomentumAdvection() - # Scalar variables to be saved during simulation + ## abstract methods + + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.fluid = self.Fluid() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.fluid.density + self.propagators.variat_dens.variables.u = self.fluid.velocity + self.propagators.variat_mom.variables.u = self.fluid.velocity + + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_thermo") self.add_scalar("en_tot") + @property + def bulk_species(self): + return self.fluid + + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + pass + def update_scalar_quantities(self): - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["fluid_uv"], self.pointer["fluid_uv"]) + rho = self.fluid.density.spline.vector + u = self.fluid.velocity.spline.vector + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) - en_thermo = 0.5 * self.mass_ops.M3.dot_inner(self.pointer["fluid_rho3"], self.pointer["fluid_rho3"]) + en_thermo = 0.5 * self.mass_ops.M3.dot_inner(rho, rho) self.update_scalar("en_thermo", en_thermo) en_tot = en_U + en_thermo self.update_scalar("en_tot", en_tot) + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='barotropic')\n" + ] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class VariationalCompressibleFluid(StruphyModel): r"""Fully compressible fluid equations discretized with a variational method. @@ -627,96 +626,56 @@ class VariationalCompressibleFluid(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["fluid"]["fluid"] = {"rho3": "L2", "s3": "L2", "uv": "H1vec"} - return dct + ## species - @staticmethod - def bulk_species(): - return "fluid" + class Fluid(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.entropy = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + ## propagators - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["fluid_rho3", "fluid_uv"], - propagators_fields.VariationalMomentumAdvection: ["fluid_uv"], - propagators_fields.VariationalEntropyEvolve: ["fluid_s3", "fluid_uv"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import H1vecMassMatrix_density - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - # Initialize mass matrix - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_momentum = params["fluid"]["fluid"]["options"]["VariationalMomentumAdvection"]["lin_solver"] - nonlin_solver_momentum = params["fluid"]["fluid"]["options"]["VariationalMomentumAdvection"]["nonlin_solver"] - lin_solver_density = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - lin_solver_entropy = params["fluid"]["fluid"]["options"]["VariationalEntropyEvolve"]["lin_solver"] - nonlin_solver_entropy = params["fluid"]["fluid"]["options"]["VariationalEntropyEvolve"]["nonlin_solver"] - - self._gamma = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - model = "full" - - from struphy.feec.variational_utilities import InternalEnergyEvaluator - - self._energy_evaluator = InternalEnergyEvaluator(self.derham, self._gamma) - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": model, - "s": self.pointer["fluid_s3"], - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - "energy_evaluator": self._energy_evaluator, - } - - self._kwargs[propagators_fields.VariationalMomentumAdvection] = { - "mass_ops": self.WMM, - "lin_solver": lin_solver_momentum, - "nonlin_solver": nonlin_solver_momentum, - } - - self._kwargs[propagators_fields.VariationalEntropyEvolve] = { - "model": model, - "rho": self.pointer["fluid_rho3"], - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_entropy, - "nonlin_solver": nonlin_solver_entropy, - "energy_evaluator": self._energy_evaluator, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() + class Propagators: + def __init__(self): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_mom = propagators_fields.VariationalMomentumAdvection() + self.variat_ent = propagators_fields.VariationalEntropyEvolve() - # Scalar variables to be saved during simulation + ## abstract methods + + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.fluid = self.Fluid() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.fluid.density + self.propagators.variat_dens.variables.u = self.fluid.velocity + self.propagators.variat_mom.variables.u = self.fluid.velocity + self.propagators.variat_ent.variables.s = self.fluid.entropy + self.propagators.variat_ent.variables.u = self.fluid.velocity + + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_thermo") self.add_scalar("en_tot") - # temporary vectors for scalar quantities + @property + def bulk_species(self): + return self.fluid + + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): projV3 = L2Projector("L2", self._mass_ops) def f(e1, e2, e3): @@ -725,8 +684,13 @@ def f(e1, e2, e3): f = np.vectorize(f) self._integrator = projV3(f) + self._energy_evaluator = InternalEnergyEvaluator(self.derham, self.propagators.variat_ent.options.gamma) + def update_scalar_quantities(self): - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["fluid_uv"], self.pointer["fluid_uv"]) + rho = self.fluid.density.spline.vector + u = self.fluid.velocity.spline.vector + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) en_thermo = self.update_thermo_energy() @@ -734,15 +698,45 @@ def update_scalar_quantities(self): en_tot = en_U + en_thermo self.update_scalar("en_tot", en_tot) + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full',\n" + ] + new_file += [ + " s=model.fluid.entropy)\n" + ] + elif "variat_ent.Options" in line: + new_file += [ + "model.propagators.variat_ent.options = model.propagators.variat_ent.Options(model='full',\n" + ] + new_file += [ + " rho=model.fluid.density)\n" + ] + elif "entropy.add_background" in line: + new_file += ["model.fluid.density.add_background(FieldsBackground())\n"] + new_file += [line] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + def update_thermo_energy(self): """Reuse tmp used in VariationalEntropyEvolve to compute the thermodynamical energy. :meta private: """ - en_prop = self._propagators[2] + en_prop = self.propagators.variat_ent - self._energy_evaluator.sf.vector = self.pointer["fluid_s3"] - self._energy_evaluator.rhof.vector = self.pointer["fluid_rho3"] + self._energy_evaluator.sf.vector = self.fluid.entropy.spline.vector + self._energy_evaluator.rhof.vector = self.fluid.density.spline.vector sf_values = self._energy_evaluator.sf.eval_tp_fixed_loc( self._energy_evaluator.integration_grid_spans, self._energy_evaluator.integration_grid_bd, @@ -763,7 +757,7 @@ def update_thermo_energy(self): def __ener(self, rho, s): """Themodynamical energy as a function of rho and s, usign the perfect gaz hypothesis E(rho, s) = rho^gamma*exp(s/rho)""" - return np.power(rho, self._gamma) * np.exp(s / rho) + return np.power(rho, self.propagators.variat_ent.options.gamma) * np.exp(s / rho) class Poisson(StruphyModel): diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 68f5fa1db..bf1c4fa48 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -28,12 +28,11 @@ ) from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators -from struphy.feec.preconditioner import MassMatrixPreconditioner +from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner, MassMatrixPreconditioner from struphy.feec.projectors import L2Projector from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.variational_utilities import ( BracketOperator, - H1vecMassMatrix_density, InternalEnergyEvaluator, KineticEnergyEvaluator, ) @@ -44,6 +43,7 @@ OptsDirectSolver, OptsGenSolver, OptsMassPrecond, + OptsNonlinearSolver, OptsSaddlePointSolver, OptsSymmSolver, OptsVecSpace, @@ -54,7 +54,7 @@ from struphy.kinetic_background.maxwellians import GyroMaxwellian2D, Maxwellian3D from struphy.linear_algebra.saddle_point import SaddlePointSolver from struphy.linear_algebra.schur_solver import SchurSolver -from struphy.linear_algebra.solver import SolverParameters +from struphy.linear_algebra.solver import NonlinearSolverParameters, SolverParameters from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.ode.solvers import ODEsolverFEEC from struphy.ode.utils import ButcherTableau, OptsButcher @@ -3017,50 +3017,73 @@ class VariationalMomentumAdvection(Propagator): \hat{\mathbf{u}}_h^{n+1/2} = (\mathbf{u}^{n+1/2})^\top \vec{\boldsymbol \Lambda}^v \in (V_h^0)^3 \,, \qquad \hat{\mathbf A}^1_{\mu,h} = \nabla P_\mu((\mathbf u^{n+1/2})^\top \vec{\boldsymbol \Lambda}^v)] \in V_h^1\,, \qquad \hat{\rho}_h^{n} = (\rho^{n})^\top \vec{\boldsymbol \Lambda}^3 \in V_h^3 \,. """ - @staticmethod - def options(default=False): - dct = {} - dct["lin_solver"] = { - "tol": 1e-12, - "maxiter": 500, - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "verbose": False, - } - dct["nonlin_solver"] = { - "tol": 1e-8, - "maxiter": 100, - "type": ["Newton", "Picard"], - "info": False, - } - if default: - dct = descend_options_dict(dct, []) - return dct + class Variables: + def __init__(self): + self._u: FEECVariable = None - def __init__( - self, - u: BlockVector, - *, - mass_ops: H1vecMassMatrix_density, - lin_solver: dict = options(default=True)["lin_solver"], - nonlin_solver: dict = options(default=True)["nonlin_solver"], - ): - super().__init__(u) + @property + def u(self) -> FEECVariable: + return self._u - assert mass_ops is not None + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1vec" + self._u = new - self._lin_solver = lin_solver - self._nonlin_solver = nonlin_solver + def __init__(self): + self.variables = self.Variables() - self._info = self._nonlin_solver["info"] and (MPI.COMM_WORLD.Get_rank() == 0) + @dataclass + class Options: + # propagator options + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None - self._Mrho = mass_ops + def __post_init__(self): + # checks + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + + self._info = self._nonlin_solver.info and (MPI.COMM_WORLD.Get_rank() == 0) + + self._Mrho = self.mass_ops.WMM + self._Mrho.inv._options["pc"] = MassMatrixDiagonalPreconditioner(self._Mrho.massop) self._initialize_mass() # bunch of temporaries to avoid allocating in the loop + u = self.variables.u.spline.vector + self._tmp_un1 = u.space.zeros() self._tmp_un12 = u.space.zeros() self._tmp_diff = u.space.zeros() @@ -3076,25 +3099,25 @@ def __init__( self.inv_derivative = inverse( self._Mrho.inv @ self.derivative, "gmres", - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], - verbose=self._lin_solver["verbose"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, + verbose=self._lin_solver.verbose, recycle=True, ) def __call__(self, dt): - if self._nonlin_solver["type"] == "Newton": + if self._nonlin_solver.type == "Newton": self.__call_newton(dt) - elif self._nonlin_solver["type"] == "Picard": + elif self._nonlin_solver.type == "Picard": self.__call_picard(dt) def __call_newton(self, dt): # Initialize variable for Newton iteration - un = self.feec_vars[0] + un = self.variables.u.spline.vector mn = self._Mrho.massop.dot(un, out=self._tmp_mn) mn1 = mn.copy(out=self._tmp_mn1) un1 = un.copy(out=self._tmp_un1) - tol = self._nonlin_solver["tol"] + tol = self.options.nonlin_solver.tol err = tol + 1 # Jacobian matrix for Newton solve self._dt2_brack._scalar = dt / 2 @@ -3102,7 +3125,7 @@ def __call_newton(self, dt): print() print("Newton iteration in VariationalMomentumAdvection") - for it in range(self._nonlin_solver["maxiter"]): + for it in range(self.options.nonlin_solver.maxiter): un12 = un.copy(out=self._tmp_un12) un12 += un1 un12 *= 0.5 @@ -3136,24 +3159,24 @@ def __call_newton(self, dt): un1 -= update mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver["maxiter"] - 1 or np.isnan(err): + if it == self.options.nonlin_solver.maxiter - 1 or np.isnan(err): print( f"!!!WARNING: Maximum iteration in VariationalMomentumAdvection reached - not converged \n {err = } \n {tol**2 = }", ) - self.feec_vars_update(un1) + self.update_feec_variables(u=un1) def __call_picard(self, dt): # Initialize variable for Picard iteration - un = self.feec_vars[0] + un = self.variables.u.spline.vector mn = self._Mrho.massop.dot(un, out=self._tmp_mn) mn1 = mn.copy(out=self._tmp_mn1) un1 = un.copy(out=self._tmp_un1) - tol = self._nonlin_solver["tol"] + tol = self.options.nonlin_solver.tol err = tol + 1 # Jacobian matrix for Newton solve - for it in range(self._nonlin_solver["maxiter"]): + for it in range(self.options.nonlin_solver.maxiter): # Picard iteration if err < tol**2 or np.isnan(err): break @@ -3182,12 +3205,12 @@ def __call_picard(self, dt): # Inverse the mass matrix to get the velocity un1 = self._Mrho.inv.dot(mn1, out=self._tmp_un1) - if it == self._nonlin_solver["maxiter"] - 1 or np.isnan(err): + if it == self.options.nonlin_solver.maxiter - 1 or np.isnan(err): print( f"!!!WARNING: Maximum iteration in VariationalMomentumAdvection reached - not converged \n {err = } \n {tol**2 = }", ) - self.feec_vars_update(un1) + self.update_feec_variables(u=un1) def _initialize_mass(self): """Initialization of the mass matrix solver""" @@ -3260,48 +3283,38 @@ class VariationalDensityEvolve(Propagator): \hat{\mathbf{u}}_h^{k} = (\mathbf{u}^{k})^\top \vec{\boldsymbol \Lambda}^v \in (V_h^0)^3 \, \text{for k in} \{n, n+1/2, n+1\}, \qquad \hat{\rho}_h^{k} = (\rho^{k})^\top \vec{\boldsymbol \Lambda}^3 \in V_h^3 \, \text{for k in} \{n, n+1/2, n+1\} . """ - @staticmethod - def options(default=False): - dct = {} - dct["lin_solver"] = { - "tol": 1e-12, - "maxiter": 500, - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "verbose": False, - "recycle": True, - } - dct["nonlin_solver"] = { - "tol": 1e-8, - "maxiter": 100, - "info": False, - "linearize": False, - } - dct["physics"] = {"gamma": 5 / 3} + class Variables: + def __init__(self): + self._rho: FEECVariable = None + self._u: FEECVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def rho(self) -> FEECVariable: + return self._rho - return dct + @rho.setter + def rho(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._rho = new - def __init__( - self, - rho: StencilVector, - u: BlockVector, - *, - model: str = "barotropic", - gamma: float = options()["physics"]["gamma"], - s: StencilVector = None, - mass_ops: H1vecMassMatrix_density, - lin_solver: dict = options(default=True)["lin_solver"], - nonlin_solver: dict = options(default=True)["nonlin_solver"], - energy_evaluator: InternalEnergyEvaluator = None, - ): - super().__init__(rho, u) + @property + def u(self) -> FEECVariable: + return self._u - assert model in [ + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1vec" + self._u = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsModel = Literal[ "pressureless", "barotropic", "full", @@ -3312,27 +3325,69 @@ def __init__( "linear_q", "deltaf_q", ] - if model == "full": - assert s is not None - assert mass_ops is not None + # propagator options + model: OptsModel = "barotropic" + gamma: float = 5.0 / 3.0 + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None + s: FEECVariable = None - self._model = model - self._gamma = gamma - self._s = s - self._lin_solver = lin_solver - self._nonlin_solver = nonlin_solver - self._linearize = self._nonlin_solver["linearize"] + def __post_init__(self): + # checks + check_option(self.model, self.OptsModel) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) - self._info = self._nonlin_solver["info"] and (MPI.COMM_WORLD.Get_rank() == 0) + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() - self._Mrho = mass_ops + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + if self.options.model == "full": + assert self.options.s is not None + + self._model = self.options.model + self._gamma = self.options.gamma + self._s = self.options.s + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + self._linearize = self.options.nonlin_solver.linearize + + self._info = self.options.nonlin_solver.info and (MPI.COMM_WORLD.Get_rank() == 0) + + self._Mrho = self.mass_ops.WMM + self._Mrho.inv._options["pc"] = MassMatrixDiagonalPreconditioner(self._Mrho.massop) # Femfields for the projector self.rhof = self.derham.create_spline_function("rhof", "L2") self.rhof1 = self.derham.create_spline_function("rhof1", "L2") + rho = self.variables.rho.spline.vector + u = self.variables.u.spline.vector + # Projector - self._energy_evaluator = energy_evaluator + self._energy_evaluator = InternalEnergyEvaluator(self.derham, self._gamma) self._kinetic_evaluator = KineticEnergyEvaluator(self.derham, self.domain, self.mass_ops) self._initialize_projectors_and_mass() if self._model in ["linear", "linear_q"]: @@ -3368,6 +3423,7 @@ def __init__( if self._model in ["linear", "linear_q"]: self._update_Pirho(self.projected_equil.n3) + @profile def __call__(self, dt): self.__call_newton(dt) @@ -3379,15 +3435,15 @@ def __call_newton(self, dt): print("Newton iteration in VariationalDensityEvolve") # Initial variables - rhon = self.feec_vars[0] - un = self.feec_vars[1] + rhon = self.variables.rho.spline.vector + un = self.variables.u.spline.vector if self._model in ["linear", "linear_q"]: advection = self.divPirho.dot(un, out=self._tmp_rho_advection) advection *= dt rhon1 = rhon.copy(out=self._tmp_rhon1) rhon1 -= advection - self.feec_vars_update(rhon1, un) + self.update_feec_variables(rho=rhon1, u=un) return if self._model in ["deltaf", "deltaf_q"]: @@ -3400,7 +3456,7 @@ def __call_newton(self, dt): # Initialize variable for Newton iteration if self._model == "full": - s = self._s + s = self._s.spline.vector else: s = None @@ -3422,10 +3478,10 @@ def __call_newton(self, dt): un1 = un.copy(out=self._tmp_un1) un1 += self._tmp_un_diff mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - tol = self._nonlin_solver["tol"] + tol = self._nonlin_solver.tol err = tol + 1 - for it in range(self._nonlin_solver["maxiter"]): + for it in range(self._nonlin_solver.maxiter): # Newton iteration un12 = un.copy(out=self._tmp_un12) @@ -3501,14 +3557,14 @@ def __call_newton(self, dt): mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver["maxiter"] - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalDensityEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) self._tmp_un_diff = un1 - un self._tmp_rhon_diff = rhon1 - rhon - self.feec_vars_update(rhon1, un1) + self.update_feec_variables(rho=rhon1, u=un1) def _initialize_projectors_and_mass(self): """Initialization of all the `BasisProjectionOperator` and `CoordinateProjector` needed to compute the bracket term""" @@ -3577,17 +3633,17 @@ def _initialize_projectors_and_mass(self): self._Jacobian, "pbicgstab", pc=self._Mrho.inv, - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], - verbose=self._lin_solver["verbose"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, + verbose=self._lin_solver.verbose, recycle=True, ) # self._inv_Jacobian = inverse(self._Jacobian, # 'gmres', - # tol=self._lin_solver['tol'], - # maxiter=self._lin_solver['maxiter'], - # verbose=self._lin_solver['verbose'], + # tol=self._lin_solver.tol, + # maxiter=self._lin_solver.maxiter, + # verbose=self._lin_solver.verbose, # recycle=True) # L2-projector for V3 @@ -3640,7 +3696,7 @@ def _update_weighted_MM(self, rho): self._Mrho.update_weight(rho) def _update_linear_form_dl_drho(self, rhon, rhon1, un, un1, sn): - """Update the linearform representing integration in V3 against kynetic energy""" + """Update the linearform representing integration in V3 against kinetic energy""" self._kinetic_evaluator.get_u2_grid(un, un1, self._eval_dl_drho) @@ -3766,67 +3822,100 @@ class VariationalEntropyEvolve(Propagator): \hat{\mathbf{u}}_h^{k} = (\mathbf{u}^{k})^\top \vec{\boldsymbol \Lambda}^v \in (V_h^0)^3 \, \text{for k in} \{n, n+1/2, n+1\}, \qquad \hat{s}_h^{k} = (s^{k})^\top \vec{\boldsymbol \Lambda}^3 \in V_h^3 \, \text{for k in} \{n, n+1/2, n+1\} \qquad \hat{\rho}_h^{n} = (\rho^{n})^\top \vec{\boldsymbol \Lambda}^3 \in V_h^3 \. """ - @staticmethod - def options(default=False): - dct = {} - dct["lin_solver"] = { - "tol": 1e-12, - "maxiter": 500, - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "verbose": False, - } - dct["nonlin_solver"] = { - "tol": 1e-8, - "maxiter": 100, - "info": False, - "linearize": "False", - } - dct["physics"] = {"gamma": 5 / 3} + class Variables: + def __init__(self): + self._s: FEECVariable = None + self._u: FEECVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def s(self) -> FEECVariable: + return self._s - return dct + @s.setter + def s(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._s = new - def __init__( - self, - s: StencilVector, - u: BlockVector, - *, - model: str = "full", - gamma: float = options()["physics"]["gamma"], - rho: StencilVector, - mass_ops: H1vecMassMatrix_density, - lin_solver: dict = options(default=True)["lin_solver"], - nonlin_solver: dict = options(default=True)["nonlin_solver"], - energy_evaluator: InternalEnergyEvaluator = None, - ): - super().__init__(s, u) + @property + def u(self) -> FEECVariable: + return self._u - assert model in ["full"] - if model == "full": - assert rho is not None - assert mass_ops is not None + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1vec" + self._u = new - self._model = model - self._gamma = gamma - self._rho = rho - self._lin_solver = lin_solver - self._nonlin_solver = nonlin_solver - self._linearize = self._nonlin_solver["linearize"] + def __init__(self): + self.variables = self.Variables() - self._info = self._nonlin_solver["info"] and (MPI.COMM_WORLD.Get_rank() == 0) + @dataclass + class Options: + # specific literals + OptsModel = Literal["full"] + # propagator options + model: OptsModel = "full" + gamma: float = 5.0 / 3.0 + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None + rho: FEECVariable = None - self._Mrho = mass_ops + def __post_init__(self): + # checks + check_option(self.model, self.OptsModel) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + if self.options.model == "full": + assert self.options.rho is not None + + self._model = self.options.model + self._gamma = self.options.gamma + self._rho = self.options.rho + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + self._linearize = self.options.nonlin_solver.linearize + + self._info = self._nonlin_solver.info and (MPI.COMM_WORLD.Get_rank() == 0) + + self._Mrho = self.mass_ops.WMM + self._Mrho.inv._options["pc"] = MassMatrixDiagonalPreconditioner(self._Mrho.massop) # Projector - self._energy_evaluator = energy_evaluator + self._energy_evaluator = InternalEnergyEvaluator(self.derham, self._gamma) self._initialize_projectors_and_mass() # bunch of temporaries to avoid allocating in the loop + s = self.variables.s.spline.vector + u = self.variables.u.spline.vector + self._tmp_un1 = u.space.zeros() self._tmp_un2 = u.space.zeros() self._tmp_un12 = u.space.zeros() @@ -3854,12 +3943,12 @@ def __call_newton(self, dt): if self._info: print() print("Newton iteration in VariationalEntropyEvolve") - sn = self.feec_vars[0] - un = self.feec_vars[1] + sn = self.variables.s.spline.vector + un = self.variables.u.spline.vector sn1 = sn.copy(out=self._tmp_sn1) # Initialize variable for Newton iteration - rho = self._rho + rho = self._rho.spline.vector self._update_Pis(sn) mn = self._Mrho.massop.dot(un, out=self._tmp_mn) @@ -3868,10 +3957,10 @@ def __call_newton(self, dt): un1 = un.copy(out=self._tmp_un1) un1 += self._tmp_un_diff mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - tol = self._nonlin_solver["tol"] + tol = self._nonlin_solver.tol err = tol + 1 - for it in range(self._nonlin_solver["maxiter"]): + for it in range(self._nonlin_solver.maxiter): # Newton iteration un12 = un.copy(out=self._tmp_un12) @@ -3931,13 +4020,13 @@ def __call_newton(self, dt): # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver["maxiter"] - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalEntropyEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) self._tmp_sn_diff = sn1 - sn self._tmp_un_diff = un1 - un - self.feec_vars_update(sn1, un1) + self.update_feec_variables(s=sn1, u=un1) def _initialize_projectors_and_mass(self): """Initialization of all the `BasisProjectionOperator` and `CoordinateProjector` needed to compute the bracket term""" @@ -3990,19 +4079,19 @@ def _initialize_projectors_and_mass(self): self._inv_Jacobian = SchurSolverFull( self._Jacobian, - self._lin_solver["type"][0], + self.options.solver, pc=self._Mrho.inv, - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], - verbose=self._lin_solver["verbose"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, + verbose=self._lin_solver.verbose, recycle=True, ) # self._inv_Jacobian = inverse(self._Jacobian, # 'gmres', - # tol=self._lin_solver['tol'], - # maxiter=self._lin_solver['maxiter'], - # verbose=self._lin_solver['verbose'], + # tol=self._lin_solver.tol, + # maxiter=self._lin_solver.maxiter, + # verbose=self._lin_solver.verbose, # recycle=True) # prepare for integration of linear form @@ -4167,7 +4256,7 @@ def __init__( u: BlockVector, *, model: str = "full", - mass_ops: H1vecMassMatrix_density, + mass_ops, # H1vecMassMatrix_density, lin_solver: dict = options(default=True)["lin_solver"], nonlin_solver: dict = options(default=True)["nonlin_solver"], ): @@ -4528,7 +4617,7 @@ def __init__( *, model: str = "full", gamma: float = options()["physics"]["gamma"], - mass_ops: H1vecMassMatrix_density, + mass_ops, # H1vecMassMatrix_density, lin_solver: dict = options(default=True)["lin_solver"], nonlin_solver: dict = options(default=True)["nonlin_solver"], div_u: StencilVector | None = None, @@ -5074,7 +5163,7 @@ def __init__( *, model: str = "full", gamma: float = options()["physics"]["gamma"], - mass_ops: H1vecMassMatrix_density, + mass_ops, # H1vecMassMatrix_density, lin_solver: dict = options(default=True)["lin_solver"], nonlin_solver: dict = options(default=True)["nonlin_solver"], div_u: StencilVector | None = None, @@ -5635,7 +5724,7 @@ def __init__( mu: float = options()["physics"]["mu"], mu_a: float = options()["physics"]["mu_a"], alpha: float = options()["physics"]["alpha"], - mass_ops: H1vecMassMatrix_density, + mass_ops, # H1vecMassMatrix_density, lin_solver: dict = options(default=True)["lin_solver"], nonlin_solver: dict = options(default=True)["nonlin_solver"], energy_evaluator: InternalEnergyEvaluator = None, From 1bc3a8d23a99b3c08ab87fdd5da684e079525c34 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 9 Oct 2025 09:54:56 +0000 Subject: [PATCH 138/292] Ported three fluid models --- src/struphy/models/base.py | 4 +- src/struphy/models/fluid.py | 426 +++++----- src/struphy/models/tests/test_models.py | 3 + src/struphy/propagators/propagators_fields.py | 749 +++++++++++------- 4 files changed, 689 insertions(+), 493 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index acc198e9f..b64c7bc10 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1285,7 +1285,7 @@ def generate_default_parameter_file( file = open(path, "w") else: print("exiting ...") - return + exit() except FileNotFoundError: folder = os.path.join("/", *path.split("/")[:-1]) if not prompt: @@ -1297,7 +1297,7 @@ def generate_default_parameter_file( file = open(path, "x") else: print("exiting ...") - return + exit() file.write("from struphy.io.options import EnvironmentOptions, BaseUnits, Time\n") file.write("from struphy.geometry import domains\n") diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 26e6314ee..737e83274 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -1,6 +1,7 @@ import numpy as np from mpi4py import MPI from psydac.linalg.block import BlockVector +from psydac.linalg.stencil import StencilVector from struphy.models.base import StruphyModel from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies @@ -203,87 +204,52 @@ class LinearExtendedMHDuniform(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["em_fields"]["b_field"] = "Hcurl" - dct["fluid"]["mhd"] = { - "rho": "L2", - "u": "Hdiv", - "p": "L2", - } - return dct - - @staticmethod - def bulk_species(): - return "mhd" - - @staticmethod - def velocity_scale(): - return "alfvén" - - @staticmethod - def propagators_dct(): - return { - propagators_fields.ShearAlfvenB1: ["mhd_u", "b_field"], - propagators_fields.Hall: ["b_field"], - propagators_fields.MagnetosonicUniform: ["mhd_rho", "mhd_u", "mhd_p"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + ## species - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hcurl") + self.init_variables() - from struphy.polar.basic import PolarVector + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="Hdiv") + self.pressure = FEECVariable(space="L2") + self.init_variables() - # extract necessary parameters - alfven_solver = params["fluid"]["mhd"]["options"]["ShearAlfvenB1"]["solver"] - M1_inv = params["fluid"]["mhd"]["options"]["ShearAlfvenB1"]["solver_M1"] - hall_solver = params["em_fields"]["options"]["Hall"]["solver"] - sonic_solver = params["fluid"]["mhd"]["options"]["MagnetosonicUniform"]["solver"] + ## propagators - # project background magnetic field (1-form) and pressure (3-form) - self._b_eq = self.projected_equil.b1 - self._a_eq = self.projected_equil.a1 - self._p_eq = self.projected_equil.p3 - self._ones = self.pointer["mhd_p"].space.zeros() + class Propagators: + def __init__(self): + self.shear_alf = propagators_fields.ShearAlfvenB1() + self.hall = propagators_fields.Hall() + self.mag_sonic = propagators_fields.MagnetosonicUniform() - if isinstance(self._ones, PolarVector): - self._ones.tp[:] = 1.0 - else: - self._ones[:] = 1.0 + ## abstract methods - # compute coupling parameters - epsilon = self.equation_params["mhd"]["epsilon"] + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - if abs(epsilon - 1) < 1e-6: - epsilon = 1.0 + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() - # set keyword arguments for propagators - self._kwargs[propagators_fields.ShearAlfvenB1] = { - "solver": alfven_solver, - "solver_M1": M1_inv, - } + # 2. instantiate all propagators + self.propagators = self.Propagators() - self._kwargs[propagators_fields.Hall] = { - "solver": hall_solver, - "epsilon": epsilon, - } + # 3. assign variables to propagators + self.propagators.shear_alf.variables.u = self.mhd.velocity + self.propagators.shear_alf.variables.b = self.em_fields.b_field - self._kwargs[propagators_fields.MagnetosonicUniform] = {"solver": sonic_solver} + self.propagators.hall.variables.b = self.em_fields.b_field - # Initialize propagators used in splitting substeps - self.init_propagators() + self.propagators.mag_sonic.variables.n = self.mhd.density + self.propagators.mag_sonic.variables.u = self.mhd.velocity + self.propagators.mag_sonic.variables.p = self.mhd.pressure - # Scalar variables to be saved during simulation + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_p") self.add_scalar("en_B") @@ -293,17 +259,45 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("en_tot") self.add_scalar("helicity") - # temporary vectors for scalar quantities - self._tmp_b1 = self.derham.Vh["1"].zeros() - self._tmp_b2 = self.derham.Vh["1"].zeros() + @property + def bulk_species(self): + return self.mhd + + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + self._b_eq = self.projected_equil.b1 + self._a_eq = self.projected_equil.a1 + self._p_eq = self.projected_equil.p3 + + self._ones = self.projected_equil.p3.space.zeros() + if isinstance(self._ones, PolarVector): + self._ones.tp[:] = 1.0 + else: + self._ones[:] = 1.0 + + self._tmp_b1: BlockVector = self.derham.Vh["1"].zeros() # TODO: replace derham.Vh dict by class + self._tmp_b2: BlockVector = self.derham.Vh["1"].zeros() + + # adjust coupling parameters + epsilon = self.mhd.equation_params.epsilon + + if abs(epsilon - 1) < 1e-6: + self.mhd.equation_params.epsilon = 1.0 def update_scalar_quantities(self): # perturbed fields - en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.pointer["mhd_u"], self.pointer["mhd_u"]) - b1 = self.mass_ops.M1.dot(self.pointer["b_field"], out=self._tmp_b1) - en_B = 0.5 * self.pointer["b_field"].inner(b1) + u = self.mhd.velocity.spline.vector + p = self.mhd.pressure.spline.vector + b = self.em_fields.b_field.spline.vector + + en_U = 0.5 * self.mass_ops.M2n.dot_inner(u, u) + b1 = self.mass_ops.M1.dot(b, out=self._tmp_b1) + en_B = 0.5 * b.inner(b1) helicity = 2.0 * self._a_eq.inner(b1) - en_p_i = self.pointer["mhd_p"].inner(self._ones) / (5.0 / 3.0 - 1.0) + en_p_i = p.inner(self._ones) / (5.0 / 3.0 - 1.0) self.update_scalar("en_U", en_U) self.update_scalar("en_B", en_B) @@ -321,13 +315,30 @@ def update_scalar_quantities(self): # total magnetic field b1 = self._b_eq.copy(out=self._tmp_b1) - self._tmp_b1 += self.pointer["b_field"] + self._tmp_b1 += b b2 = self.mass_ops.M1.dot(b1, apply_bc=False, out=self._tmp_b2) en_Btot = b1.inner(b2) / 2.0 self.update_scalar("en_B_tot", en_Btot) + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "hall.Options" in line: + new_file += [ + "model.propagators.hall.options = model.propagators.hall.Options(epsilon_from=model.mhd)\n" + ] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class ColdPlasma(StruphyModel): r"""Cold plasma model. @@ -364,82 +375,74 @@ class ColdPlasma(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["em_fields"]["e_field"] = "Hcurl" - dct["em_fields"]["b_field"] = "Hdiv" - dct["fluid"]["electrons"] = {"j": "Hcurl"} - return dct + ## species - @staticmethod - def bulk_species(): - return "electrons" + class EMFields(FieldSpecies): + def __init__(self): + self.e_field = FEECVariable(space="Hcurl") + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def velocity_scale(): - return "light" + class Electrons(FluidSpecies): + def __init__(self): + self.current = FEECVariable(space="Hcurl") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_fields.Maxwell: ["e_field", "b_field"], - propagators_fields.OhmCold: ["electrons_j", "e_field"], - propagators_fields.JxBCold: ["electrons_j"], - } + ## propagators - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + class Propagators: + def __init__(self): + self.maxwell = propagators_fields.Maxwell() + self.ohm = propagators_fields.OhmCold() + self.jxb = propagators_fields.JxBCold() - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + ## abstract methods - # model parameters - self._alpha = self.equation_params["electrons"]["alpha"] - self._epsilon = self.equation_params["electrons"]["epsilon"] + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # solver parameters - params_maxwell = params["em_fields"]["options"]["Maxwell"]["solver"] - params_ohmcold = params["fluid"]["electrons"]["options"]["OhmCold"]["solver"] - params_jxbcold = params["fluid"]["electrons"]["options"]["JxBCold"]["solver"] + # 1. instantiate all species + self.em_fields = self.EMFields() + self.electrons = self.Electrons() - # set keyword arguments for propagators - self._kwargs[propagators_fields.Maxwell] = {"solver": params_maxwell} + # 2. instantiate all propagators + self.propagators = self.Propagators() - self._kwargs[propagators_fields.OhmCold] = { - "alpha": self._alpha, - "epsilon": self._epsilon, - "solver": params_ohmcold, - } + # 3. assign variables to propagators + self.propagators.maxwell.variables.e = self.em_fields.e_field + self.propagators.maxwell.variables.b = self.em_fields.b_field - self._kwargs[propagators_fields.JxBCold] = { - "epsilon": self._epsilon, - "solver": params_jxbcold, - } + self.propagators.ohm.variables.j = self.electrons.current + self.propagators.ohm.variables.e = self.em_fields.e_field - # Initialize propagators used in splitting substeps - self.init_propagators() + self.propagators.jxb.variables.j = self.electrons.current - # Scalar variables to be saved during simulation + # define scalars for update_scalar_quantities self.add_scalar("electric energy") self.add_scalar("magnetic energy") self.add_scalar("kinetic energy") self.add_scalar("total energy") + @property + def bulk_species(self): + return self.electrons + + @property + def velocity_scale(self): + return "light" + + def allocate_helpers(self): + self._alpha = self.electrons.equation_params.alpha + def update_scalar_quantities(self): - en_E = 0.5 * self.mass_ops.M1.dot_inner(self.pointer["e_field"], self.pointer["e_field"]) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) - en_J = ( - 0.5 - * self._alpha**2 - * self.mass_ops.M1ninv.dot_inner(self.pointer["electrons_j"], self.pointer["electrons_j"]) - ) + e = self.em_fields.e_field.spline.vector + b = self.em_fields.b_field.spline.vector + j = self.electrons.current.spline.vector + + en_E = 0.5 * self.mass_ops.M1.dot_inner(e, e) + en_B = 0.5 * self.mass_ops.M2.dot_inner(b, b) + en_J = 0.5 * self._alpha**2 * self.mass_ops.M1ninv.dot_inner(j, j) self.update_scalar("electric energy", en_E) self.update_scalar("magnetic energy", en_B) @@ -2330,119 +2333,102 @@ class HasegawaWakatani(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species - dct["em_fields"] = {"phi0": "H1"} - dct["fluid"]["hw"] = { - "n0": "H1", - "omega0": "H1", - } - return dct + class EMFields(FieldSpecies): + def __init__(self): + self.phi = FEECVariable(space="H1") + self.init_variables() - @staticmethod - def bulk_species(): - return "hw" + class Plasma(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="H1") + self.vorticity = FEECVariable(space="H1") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + ## propagators - # @staticmethod - # def diagnostics_dct(): - # dct = {} - # dct["projected_density"] = "L2" - # return dct + class Propagators: + def __init__(self): + self.poisson = propagators_fields.Poisson() + self.hw = propagators_fields.HasegawaWakatani() - @staticmethod - def propagators_dct(): - return { - propagators_fields.Poisson: ["phi0"], - propagators_fields.HasegawaWakatani: ["hw_n0", "hw_omega0"], - } + ## abstract methods - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + # 1. instantiate all species + self.em_fields = self.EMFields() + self.plasma = self.Plasma() - from struphy.polar.basic import PolarVector + # 2. instantiate all propagators + self.propagators = self.Propagators() - # extract necessary parameters - self._stab_eps = params["em_fields"]["options"]["Poisson"]["stabilization"]["stab_eps"] - self._stab_mat = params["em_fields"]["options"]["Poisson"]["stabilization"]["stab_mat"] - self._solver = params["em_fields"]["options"]["Poisson"]["solver"] - c_fun = params["fluid"]["hw"]["options"]["HasegawaWakatani"]["c_fun"] - kappa = params["fluid"]["hw"]["options"]["HasegawaWakatani"]["kappa"] - nu = params["fluid"]["hw"]["options"]["HasegawaWakatani"]["nu"] - algo = params["fluid"]["hw"]["options"]["HasegawaWakatani"]["algo"] - M0_solver = params["fluid"]["hw"]["options"]["HasegawaWakatani"]["M0_solver"] - - # rhs of Poisson - self._rho = self.derham.Vh["0"].zeros() - self.update_rho() + # 3. assign variables to propagators + self.propagators.poisson.variables.phi = self.em_fields.phi + self.propagators.hw.variables.n = self.plasma.density + self.propagators.hw.variables.omega = self.plasma.vorticity - # set keyword arguments for propagators - self._kwargs[propagators_fields.Poisson] = { - "stab_eps": self._stab_eps, - "stab_mat": self._stab_mat, - "rho": self.update_rho, - "solver": self._solver, - } + # define scalars for update_scalar_quantities - self._kwargs[propagators_fields.HasegawaWakatani] = { - "phi": self.em_fields["phi0"]["obj"], - "c_fun": c_fun, - "kappa": kappa, - "nu": nu, - "algo": algo, - "M0_solver": M0_solver, - } + @property + def bulk_species(self): + return self.plasma - # Initialize propagators used in splitting substeps - self.init_propagators() + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + self._rho: StencilVector = self.derham.Vh["0"].zeros() + self.update_rho() + + def update_scalar_quantities(self): + pass def update_rho(self): - self._rho = self.mass_ops.M0.dot(self.pointer["hw_omega0"], out=self._rho) + omega = self.plasma.vorticity.spline.vector + self._rho = self.mass_ops.M0.dot(omega, out=self._rho) self._rho.update_ghost_regions() return self._rho - def initialize_from_params(self): + def allocate_propagators(self): """Solve initial Poisson equation. :meta private: """ # initialize fields and particles - super().initialize_from_params() + super().allocate_propagators() - if self.rank_world == 0: + if MPI.COMM_WORLD.Get_rank() == 0: print("\nINITIAL POISSON SOLVE:") - # Instantiate Poisson solver - poisson_solver = propagators_fields.Poisson( - self.pointer["phi0"], - stab_eps=self._stab_eps, - stab_mat=self._stab_mat, - rho=self._rho, - solver=self._solver, - ) - - # Solve with dt=1. and compute electric field - if self.rank_world == 0: - print("\nSolving initial Poisson problem...") - self.update_rho() - poisson_solver(1.0) + self.propagators.poisson(1.0) - if self.rank_world == 0: + if MPI.COMM_WORLD.Get_rank() == 0: print("Done.") def update_scalar_quantities(self): pass + + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "hw.Options" in line: + new_file += [ + "model.propagators.hw.options = model.propagators.hw.Options(phi=model.em_fields.phi)\n" + ] + elif "vorticity.add_background" in line: + new_file += ["model.plasma.density.add_background(FieldsBackground())\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index 30771a2de..3a39c5944 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -24,6 +24,9 @@ fluid_models = [ "LinearMHD", "EulerSPH", + "LinearExtendedMHDuniform", + "ColdPlasma", + "HasegawaWakatani", ] # for name, obj in inspect.getmembers(fluid): # if inspect.isclass(obj) and "models.fluid" in obj.__module__: diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index bf1c4fa48..f9579553f 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -55,6 +55,7 @@ from struphy.linear_algebra.saddle_point import SaddlePointSolver from struphy.linear_algebra.schur_solver import SchurSolver from struphy.linear_algebra.solver import NonlinearSolverParameters, SolverParameters +from struphy.models.species import Species from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.ode.solvers import ODEsolverFEEC from struphy.ode.utils import ButcherTableau, OptsButcher @@ -279,39 +280,71 @@ class OhmCold(Propagator): \end{bmatrix} \,. """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) + class Variables: + def __init__(self): + self._j: FEECVariable = None + self._e: FEECVariable = None - return dct + @property + def j(self) -> FEECVariable: + return self._j - def __init__( - self, - j: BlockVector, - e: BlockVector, - *, - alpha: float = 1.0, - epsilon: float = 1.0, - solver: dict = options(default=True)["solver"], - ): - super().__init__(e, j) + @j.setter + def j(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._j = new - self._info = solver["info"] - self._alpha = alpha - self._epsilon = epsilon + @property + def e(self) -> FEECVariable: + return self._e + + @e.setter + def e(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._e = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # propagator options + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + + def __post_init__(self): + # checks + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._info = self.options.solver_params.info + + self._alpha = self.variables.j.species.equation_params.alpha + self._epsilon = self.variables.j.species.equation_params.epsilon # Define block matrix [[A B], [C I]] (without time step size dt in the diagonals) _A = self.mass_ops.M1ninv @@ -319,10 +352,10 @@ def __init__( self._B = -1 / 2 * 1 / self._epsilon * self.mass_ops.M1 # no dt # Preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(self.mass_ops.M1ninv) # Instantiate Schur solver (constant in this case) @@ -331,13 +364,14 @@ def __init__( self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, ) + j = self.variables.j.spline.vector + e = self.variables.e.spline.vector + self._tmp_j1 = j.space.zeros() self._tmp_j2 = j.space.zeros() self._tmp_e1 = e.space.zeros() @@ -345,8 +379,8 @@ def __init__( def __call__(self, dt): # current variables - en = self.feec_vars[0] - jn = self.feec_vars[1] + jn = self.variables.j.spline.vector + en = self.variables.e.spline.vector # in-place solution (no tmps created here) Ben = self._B.dot(en, out=self._tmp_e1) @@ -360,13 +394,13 @@ def __call__(self, dt): en1 += en # write new coeffs into Propagator.variables - max_de, max_dj = self.feec_vars_update(en1, jn1) + diffs = self.update_feec_variables(e=en1, j=jn1) if self._info: print("Status for OhmCold:", info["success"]) print("Iterations for OhmCold:", info["niter"]) - print("Maxdiff e1 for OhmCold:", max_de) - print("Maxdiff j1 for OhmCold:", max_dj) + print("Maxdiff e1 for OhmCold:", diffs["e"]) + print("Maxdiff j1 for OhmCold:", diffs["j"]) print() @@ -386,65 +420,90 @@ class JxBCold(Propagator): \mathbb M_{1/n_0} \left( \mathbf j^{n+1} - \mathbf j^n \right) = \frac{\Delta t}{2} \frac{1}{\varepsilon} \mathbb M_{B_0/n_0} \left( \mathbf j^{n+1} - \mathbf j^n \right)\,. """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) + class Variables: + def __init__(self): + self._j: FEECVariable = None - return dct + @property + def j(self) -> FEECVariable: + return self._j - def __init__( - self, - j: BlockVector, - *, - epsilon: float = 1.0, - solver: dict = options(default=True)["solver"], - ): - super().__init__(j) + @j.setter + def j(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._j = new - self._info = solver["info"] + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # propagator options + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + + def __post_init__(self): + # checks + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._info = self.options.solver_params.info + + epsilon = self.variables.j.species.equation_params.epsilon # mass matrix in system (M - dt/2 * A)*j^(n + 1) = (M + dt/2 * A)*j^n self._M = self.mass_ops.M1ninv self._A = -1 / epsilon * self.mass_ops.M1Bninv # no dt # Preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(self.mass_ops.M1ninv) # Instantiate linear solver self._solver = inverse( self._M, - solver["type"][0], + self.options.solver, pc=pc, - x0=self.feec_vars[0], - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], + x0=self.variables.j.spline.vector, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, ) # allocate dummy vectors to avoid temporary array allocations self._rhs_j = self._M.codomain.zeros() - self._j_new = j.space.zeros() + self._j_new = self.variables.j.spline.vector.space.zeros() + @profile def __call__(self, dt): # current variables - jn = self.feec_vars[0] + jn = self.variables.j.spline.vector # define system (M - dt/2 * A)*b^(n + 1) = (M + dt/2 * A)*b^n lhs = self._M - dt / 2.0 * self._A @@ -459,7 +518,7 @@ def __call__(self, dt): info = self._solver._info # write new coeffs into Propagator.variables - max_dj = self.feec_vars_update(jn1)[0] + max_dj = self.update_feec_variables(j=jn1) if self._info: print("Status for FluidCold:", info["success"]) @@ -695,62 +754,91 @@ class ShearAlfvenB1(Propagator): the MHD equilibirum density. """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["solver_M1"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) + class Variables: + def __init__(self): + self._u: FEECVariable = None + self._b: FEECVariable = None - return dct + @property + def u(self) -> FEECVariable: + return self._u - def __init__( - self, - u: BlockVector, - b: BlockVector, - *, - solver: dict = options(default=True)["solver"], - solver_M1: dict = options(default=True)["solver_M1"], - ): - super().__init__(u, b) + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hdiv") + self._u = new - self._info = solver["info"] + @property + def b(self) -> FEECVariable: + return self._b + + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._b = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # propagator options + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + solver_M1: OptsSymmSolver = "pcg" + precond_M1: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params_M1: SolverParameters = None + + def __post_init__(self): + # checks + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + check_option(self.solver_M1, OptsSymmSolver) + check_option(self.precond_M1, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.solver_params_M1 is None: + self.solver_params_M1 = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._info = self.options.solver_params.info # define inverse of M1 - if solver_M1["type"][1] is None: + if self.options.precond_M1 is None: pc = None else: - pc_class = getattr(preconditioner, solver_M1["type"][1]) + pc_class = getattr(preconditioner, self.options.precond_M1) pc = pc_class(self.mass_ops.M1) M1_inv = inverse( self.mass_ops.M1, - solver_M1["type"][0], + self.options.solver_M1, pc=pc, - tol=solver_M1["tol"], - maxiter=solver_M1["maxiter"], - verbose=solver_M1["verbose"], + tol=self.options.solver_params_M1.tol, + maxiter=self.options.solver_params_M1.maxiter, + verbose=self.options.solver_params_M1.verbose, ) # define block matrix [[A B], [C I]] (without time step size dt in the diagonals) @@ -760,10 +848,10 @@ def __init__( self._C = 1 / 2 * M1_inv @ self.derham.curl.T @ self.mass_ops.M2B # Preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(getattr(self.mass_ops, "M2n")) # instantiate Schur solver (constant in this case) @@ -772,24 +860,26 @@ def __init__( self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, ) # allocate dummy vectors to avoid temporary array allocations + u = self.variables.u.spline.vector + b = self.variables.b.spline.vector + self._u_tmp1 = u.space.zeros() self._u_tmp2 = u.space.zeros() self._b_tmp1 = b.space.zeros() self._byn = self._B.codomain.zeros() + @profile def __call__(self, dt): # current variables - un = self.feec_vars[0] - bn = self.feec_vars[1] + un = self.variables.u.spline.vector + bn = self.variables.b.spline.vector # solve for new u coeffs byn = self._B.dot(bn, out=self._byn) @@ -804,13 +894,13 @@ def __call__(self, dt): bn1 += bn # write new coeffs into self.feec_vars - max_du, max_db = self.feec_vars_update(un1, bn1) + max_diffs = self.update_feec_variables(u=un1, b=bn1) if self._info and MPI.COMM_WORLD.Get_rank() == 0: print("Status for ShearAlfvenB1:", info["success"]) print("Iterations for ShearAlfvenB1:", info["niter"]) - print("Maxdiff up for ShearAlfvenB1:", max_du) - print("Maxdiff b2 for ShearAlfvenB1:", max_db) + print("Maxdiff up for ShearAlfvenB1:", max_diffs["u"]) + print("Maxdiff b2 for ShearAlfvenB1:", max_diffs["b"]) print() @@ -834,38 +924,66 @@ class Hall(Propagator): The solution of the above system is based on the Pre-conditioned Biconjugate Gradient Stabilized algortihm (PBiConjugateGradientStab). """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pbicgstab", "MassMatrixPreconditioner"), - ("bicgstab", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) + class Variables: + def __init__(self): + self._b: FEECVariable = None + + @property + def b(self) -> FEECVariable: + return self._b + + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._b = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # propagator options + solver: OptsGenSolver = "pbicgstab" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + epsilon_from: Species = None + + def __post_init__(self): + # checks + check_option(self.solver, OptsGenSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options - return dct + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new - def __init__( - self, - b: BlockVector, - *, - epsilon: float = 1.0, - solver: dict = options(default=True)["solver"], - ): - super().__init__(b) + @profile + def allocate(self): + if self.options.epsilon_from is None: + epsilon = 1.0 + else: + epsilon = self.options.epsilon_from.equation_params.epsilon - self._info = solver["info"] - self._tol = solver["tol"] - self._maxiter = solver["maxiter"] - self._verbose = solver["verbose"] + self._info = self.options.solver_params.info + self._tol = self.options.solver_params.tol + self._maxiter = self.options.solver_params.maxiter + self._verbose = self.options.solver_params.verbose # mass matrix in system (M - dt/2 * A)*b^(n + 1) = (M + dt/2 * A)*b^n id_M = "M1" @@ -875,18 +993,18 @@ def __init__( self._A = 1.0 / epsilon * self.derham.curl.T @ self._M2Bn @ self.derham.curl # Preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(getattr(self.mass_ops, id_M)) # Instantiate linear solver self._solver = inverse( self._M, - solver["type"][0], + self.options.solver, pc=pc, - x0=self.feec_vars[0], + x0=self.variables.b.spline.vector, tol=self._tol, maxiter=self._maxiter, verbose=self._verbose, @@ -894,11 +1012,11 @@ def __init__( # allocate dummy vectors to avoid temporary array allocations self._rhs_b = self._M.codomain.zeros() - self._b_new = b.space.zeros() + self._b_new = self.variables.b.spline.vector.space.zeros() def __call__(self, dt): # current variables - bn = self.feec_vars[0] + bn = self.variables.b.spline.vector # define system (M - dt/2 * A)*b^(n + 1) = (M + dt/2 * A)*b^n lhs = self._M - dt / 2.0 * self._A @@ -912,12 +1030,12 @@ def __call__(self, dt): info = self._solver._info # write new coeffs into self.feec_vars - max_db = self.feec_vars_update(bn1) + max_db = self.update_feec_variables(b=bn1) if self._info and MPI.COMM_WORLD.Get_rank() == 0: print("Status for Hall:", info["success"]) print("Iterations for Hall:", info["niter"]) - print("Maxdiff b1 for Hall:", max_db) + print("Maxdiff b1 for Hall:", max_db["b"]) print() @@ -1183,36 +1301,78 @@ class MagnetosonicUniform(Propagator): Solver- and/or other parameters for this splitting step. """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pbicgstab", "MassMatrixPreconditioner"), - ("bicgstab", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) + class Variables: + def __init__(self): + self._n: FEECVariable = None + self._u: FEECVariable = None + self._p: FEECVariable = None - return dct + @property + def n(self) -> FEECVariable: + return self._n - def __init__( - self, - n: StencilVector, - u: BlockVector, - p: StencilVector, - *, - solver: dict = options(default=True)["solver"], - ): - super().__init__(n, u, p) + @n.setter + def n(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._n = new - self._info = solver["info"] + @property + def u(self) -> FEECVariable: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new + + @property + def p(self) -> FEECVariable: + return self._p + + @p.setter + def p(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._p = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + solver: OptsGenSolver = "pbicgstab" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + + def __post_init__(self): + # checks + check_option(self.solver, OptsGenSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._info = self.options.solver_params.info self._bc = self.derham.dirichlet_bc # define block matrix [[A B], [C I]] (without time step size dt in the diagonals) @@ -1228,10 +1388,10 @@ def __init__( self._QD = getattr(self.basis_ops, id_Q) @ self.derham.div # preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(getattr(self.mass_ops, id_Mn)) # instantiate Schur solver (constant in this case) @@ -1240,14 +1400,16 @@ def __init__( self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, ) # allocate dummy vectors to avoid temporary array allocations + n = self.variables.n.spline.vector + u = self.variables.u.spline.vector + p = self.variables.p.spline.vector + self._u_tmp1 = u.space.zeros() self._u_tmp2 = u.space.zeros() self._p_tmp1 = p.space.zeros() @@ -1255,11 +1417,12 @@ def __init__( self._byn1 = self._B.codomain.zeros() + @profile def __call__(self, dt): # current variables - nn = self.feec_vars[0] - un = self.feec_vars[1] - pn = self.feec_vars[2] + nn = self.variables.n.spline.vector + un = self.variables.u.spline.vector + pn = self.variables.p.spline.vector # solve for new u coeffs byn1 = self._B.dot(pn, out=self._byn1) @@ -1278,18 +1441,14 @@ def __call__(self, dt): nn1 += nn # write new coeffs into self.feec_vars - max_dn, max_du, max_dp = self.feec_vars_update( - nn1, - un1, - pn1, - ) + diffs = self.update_feec_variables(n=nn1, u=un1, p=pn1) if self._info and MPI.COMM_WORLD.Get_rank() == 0: print("Status for Magnetosonic:", info["success"]) print("Iterations for Magnetosonic:", info["niter"]) - print("Maxdiff n3 for Magnetosonic:", max_dn) - print("Maxdiff up for Magnetosonic:", max_du) - print("Maxdiff p3 for Magnetosonic:", max_dp) + print("Maxdiff n3 for Magnetosonic:", diffs["n"]) + print("Maxdiff up for Magnetosonic:", diffs["u"]) + print("Maxdiff p3 for Magnetosonic:", diffs["p"]) print() @@ -7415,59 +7574,96 @@ class HasegawaWakatani(Propagator): Solver parameters for M0 inversion. """ - @staticmethod - def options(default=False): - dct = {} - dct["c_fun"] = ["const"] - dct["kappa"] = 1.0 - dct["nu"] = 0.01 - dct["algo"] = get_args(OptsButcher) - dct["M0_solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) - return dct + class Variables: + def __init__(self): + self._n: FEECVariable = None + self._omega: FEECVariable = None - def __init__( - self, - n0: StencilVector, - omega0: StencilVector, - *, - phi: SplineFunction = None, - c_fun: str = options(default=True)["c_fun"], - kappa: float = options(default=True)["kappa"], - nu: float = options(default=True)["nu"], - algo: str = options(default=True)["algo"], - M0_solver: dict = options(default=True)["M0_solver"], - ): - super().__init__(n0, omega0) + @property + def n(self) -> FEECVariable: + return self._n + @n.setter + def n(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._n = new + + @property + def omega(self) -> FEECVariable: + return self._omega + + @omega.setter + def omega(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._omega = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsCfun = Literal["const"] + # propagator options + phi: FEECVariable = None + c_fun: OptsCfun = "const" + kappa: float = 1.0 + nu: float = 0.01 + butcher: ButcherTableau = None + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + + def __post_init__(self): + # checks + check_option(self.c_fun, self.OptsCfun) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.butcher is None: + self.butcher = ButcherTableau() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): # default phi - if phi is None: - self._phi = self.derham.create_spline_function("phi", "H1") - self._phi.vector[:] = 1.0 - self._phi.vector.update_ghost_regions() - else: - self._phi = phi + if self.options.phi is None: + self.options.phi = FEECVariable(space="H1") + self.options.phi.allocate(derham=self.derham, domain=self.domain) + + self._phi = self.options.phi.spline + self._phi.vector[:] = 1.0 + self._phi.vector.update_ghost_regions() # default c-function - if c_fun == "const": + if self.options.c_fun == "const": c_fun = lambda e1, e2, e3: 0.0 + 0.0 * e1 else: - raise NotImplementedError(f"{c_fun = } is not available.") + raise NotImplementedError(f"{self.options.c_fun = } is not available.") # expose equation parameters - self._kappa = kappa - self._nu = nu + self._kappa = self.options.kappa + self._nu = self.options.nu # get quadrature grid of V0 pts = [grid.flatten() for grid in self.derham.quad_grid_pts["0"]] @@ -7525,16 +7721,24 @@ def __init__( ) # inverse M0 mass matrix - solver = M0_solver["type"][0] - if M0_solver["type"][1] is None: + solver = self.options.solver + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, M0_solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(self.mass_ops.M0) - solver_params = deepcopy(M0_solver) # need a copy to pop, otherwise testing fails - solver_params.pop("type") - self._info = solver_params.pop("info") - M0_inv = inverse(M0, solver, pc=pc, **solver_params) + # solver_params = deepcopy(M0_solver) # need a copy to pop, otherwise testing fails + # solver_params.pop("type") + self._info = self.options.solver_params.info + M0_inv = inverse( + M0, + solver, + pc=pc, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + recycle=self.options.solver_params.recycle, + ) # basis projection operator df_12 = lambda e1, e2, e3: self.domain.jacobian_inv(e1, e2, e3)[0, 1, :, :, :] @@ -7552,6 +7756,9 @@ def __init__( # print(f"{self._BPO._dof_mat.blocks = }") # pre-allocated helper arrays + n0 = self.variables.n.spline.vector + omega0 = self.variables.omega.spline.vector + self._tmp1 = n0.space.zeros() tmp2 = n0.space.zeros() self._tmp3 = n0.space.zeros() @@ -7559,11 +7766,11 @@ def __init__( tmp5 = n0.space.zeros() # rhs-callables for explicit ode solve - terms1_n = -M0c + grad.T @ self._M1hw @ grad - nu * grad.T @ M1 @ grad + terms1_n = -M0c + grad.T @ self._M1hw @ grad - self.options.nu * grad.T @ M1 @ grad terms1_phi = M0c - terms1_phi_strong = -kappa * self._BPO @ grad + terms1_phi_strong = -self.options.kappa * self._BPO @ grad - terms2_omega = grad.T @ self._M1hw @ grad - nu * grad.T @ M1 @ grad + terms2_omega = grad.T @ self._M1hw @ grad - self.options.nu * grad.T @ M1 @ grad terms2_n = -M0c terms2_phi = M0c @@ -7591,7 +7798,7 @@ def f2(t, n, omega, out=out2): return out vector_field = {n0: f1, omega0: f2} - self._ode_solver = ODEsolverFEEC(vector_field, algo=algo) + self._ode_solver = ODEsolverFEEC(vector_field, butcher=self.options.butcher) def __call__(self, dt): # update time-dependent mass operator From 433d54bfa9afde65dc235d1245e8d5a15e4005a8 Mon Sep 17 00:00:00 2001 From: Byung Kyu Na Date: Tue, 14 Oct 2025 04:38:19 +0000 Subject: [PATCH 139/292] Porting the model LinearMHDDriftkineticCC --- doc/conf.py | 6 +- src/struphy/console/params.py | 5 +- src/struphy/fields_background/equils.py | 20 +- src/struphy/io/options.py | 5 +- src/struphy/linear_algebra/solver.py | 11 + src/struphy/models/base.py | 30 +- src/struphy/models/hybrid.py | 504 +++--- src/struphy/models/tests/test_models.py | 2 +- .../pic/accumulation/accum_kernels_gc.py | 744 ++++++--- src/struphy/pic/accumulation/filter.py | 185 +++ .../pic/accumulation/filter_kernels.py | 22 +- .../pic/accumulation/particles_to_grid.py | 137 +- src/struphy/pic/particles.py | 18 +- src/struphy/pic/pushing/pusher_kernels_gc.py | 493 ++++-- src/struphy/pic/utilities_kernels.py | 285 ++-- src/struphy/propagators/__init__.py | 1 - .../propagators/propagators_coupling.py | 1401 ++++++++++------- src/struphy/propagators/propagators_fields.py | 1040 +++++------- 18 files changed, 2755 insertions(+), 2154 deletions(-) create mode 100644 src/struphy/pic/accumulation/filter.py diff --git a/doc/conf.py b/doc/conf.py index c5a605edb..d5e8a3265 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -12,9 +12,11 @@ # import os import shutil + # import sys # sys.path.insert(0, os.path.abspath('.')) + def copy_tutorials(app): src = os.path.abspath("../tutorials") dst = os.path.abspath("source/tutorials") @@ -25,9 +27,11 @@ def copy_tutorials(app): shutil.copytree(src, dst) + def setup(app): app.connect("builder-inited", copy_tutorials) + with open("../src/struphy/console/main.py") as f: exec(f.read()) @@ -60,7 +64,7 @@ def setup(app): "sphinx_design", ] -nbsphinx_execute = 'auto' +nbsphinx_execute = "auto" napoleon_use_admonition_for_examples = True napoleon_use_admonition_for_notes = True diff --git a/src/struphy/console/params.py b/src/struphy/console/params.py index 1e5ed7d5f..e32fd4165 100644 --- a/src/struphy/console/params.py +++ b/src/struphy/console/params.py @@ -45,5 +45,6 @@ def struphy_params(model_name: str, params_path: str, yes: bool = False, check_f else: prompt = not yes - print(f"Generating default parameter file for {model_class}.") - model_class().generate_default_parameter_file(path=params_path, prompt=prompt) + model.generate_default_parameter_file(path=params_path, prompt=prompt) + # print(f"Generating default parameter file for {model_class}.") + # model_class().generate_default_parameter_file(path=params_path, prompt=prompt) diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 96febc74b..e102257c4 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -178,7 +178,8 @@ class ShearedSlab(CartesianMHDequilibrium): Ion number density at x=a (default: 1.). beta : float Plasma beta (ratio of kinematic pressure to B^2/2, default: 0.1). - + q_kind : int + Kind of safety factor profile, (0 or 1, default: 0). Note ---- In the parameter .yml, use the following in the section ``fluid_background``:: @@ -193,6 +194,7 @@ class ShearedSlab(CartesianMHDequilibrium): n2 : 0. # 2nd shape factor for ion number density profile na : 1. # number density at r=a beta : .1 # plasma beta = p*2/B^2 + q_kind : 0. # kind of safety factor profile """ def __init__( @@ -206,6 +208,7 @@ def __init__( n2: float = 0.0, na: float = 1.0, beta: float = 0.1, + q_kind: int = 0, ): # use params setter self.params = copy.deepcopy(locals()) @@ -226,10 +229,19 @@ def q_x(self, x, der=0): qout = 0 * x else: - if der == 0: - qout = self.params["q0"] + (self.params["q1"] - self.params["q0"]) * (x / self.params["a"]) ** 2 + if self.params["q_kind"] == 0: + if der == 0: + qout = self.params["q0"] + (self.params["q1"] - self.params["q0"]) * (x / self.params["a"]) ** 2 + else: + qout = 2 * (self.params["q1"] - self.params["q0"]) * x / self.params["a"] ** 2 + else: - qout = 2 * (self.params["q1"] - self.params["q0"]) * x / self.params["a"] ** 2 + if der == 0: + qout = self.params["q0"] + self.params["q1"] * np.sin(2.0 * np.pi * x / self.params["a"]) + else: + qout = ( + 2.0 * np.pi / self.params["a"] * self.params["q1"] * np.cos(2.0 * np.pi * x / self.params["a"]) + ) return qout diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index a98e2cd5c..dce8b072c 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -26,7 +26,7 @@ # solvers OptsSymmSolver = Literal["pcg", "cg"] OptsGenSolver = Literal["pbicgstab", "bicgstab", "GMRES"] -OptsMassPrecond = Literal["MassMatrixPreconditioner", None] +OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] OptsSaddlePointSolver = Literal["Uzawa", "GMRES"] OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] OptsNonlinearSolver = Literal["Picard", "Newton"] @@ -46,6 +46,9 @@ OptsSpatialLoading = Literal["uniform", "disc"] OptsMPIsort = Literal["each", "last", None] +# filters +OptsFilter = Literal["fourier_in_tor", "hybrid", "three_point", None] + # sph OptsKernel = Literal[ "trigonometric_1d", diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py index f14a98915..217326309 100644 --- a/src/struphy/linear_algebra/solver.py +++ b/src/struphy/linear_algebra/solver.py @@ -14,6 +14,17 @@ class SolverParameters: recycle: bool = True +@dataclass +class DiscreteGradientSolverParameters: + """Parameters for discrete gradient solvers.""" + + relaxation_factor: float = 0.5 + tol: float = 1e-12 + maxiter: int = 20 + verbose: bool = False + info: bool = False + + @dataclass class NonlinearSolverParameters: """Parameters for psydac solvers.""" diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index b64c7bc10..0b330a57d 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -184,9 +184,10 @@ def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): verbose=self.verbose, ) - # create weighted mass operators + # create weighted mass and basis operators if self.derham is None: self._mass_ops = None + self._basis_ops = None else: self._mass_ops = WeightedMassOperators( self.derham, @@ -195,6 +196,13 @@ def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): eq_mhd=self.equil, ) + self._basis_ops = BasisProjectionOperators( + self.derham, + self.domain, + verbose=self.verbose, + eq_mhd=self.equil, + ) + # create projected equilibrium if self.derham is None: self._projected_equil = None @@ -221,17 +229,10 @@ def allocate_propagators(self): # set propagators base class attributes (then available to all propagators) Propagator.derham = self.derham Propagator.domain = self.domain - Propagator.mass_ops = self.mass_ops - if self.derham is None: - Propagator.basis_ops = None - else: - Propagator.basis_ops = BasisProjectionOperators( - self.derham, - self.domain, - verbose=self.verbose, - eq_mhd=self.equil, - ) - Propagator.projected_equil = self.projected_equil + if self.derham is not None: + Propagator.mass_ops = self.mass_ops + Propagator.basis_ops = self.basis_ops + Propagator.projected_equil = self.projected_equil assert len(self.prop_list) > 0, "No propagators in this model, check the model class." for prop in self.prop_list: @@ -313,6 +314,11 @@ def mass_ops(self): """WeighteMassOperators object, see :ref:`mass_ops`.""" return self._mass_ops + @property + def basis_ops(self): + """Basis projection operators.""" + return self._basis_ops + @property def prop_list(self): """List of Propagator objects.""" diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index ee93662ca..26a7b2dc7 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -1,9 +1,16 @@ import numpy as np +from mpi4py import MPI +from psydac.linalg.block import BlockVector from struphy.models.base import StruphyModel +from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies +from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.pic.accumulation import accum_kernels, accum_kernels_gc +from struphy.polar.basic import PolarVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers +rank = MPI.COMM_WORLD.Get_rank() + class LinearMHDVlasovCC(StruphyModel): r""" @@ -567,13 +574,13 @@ class LinearMHDDriftkineticCC(StruphyModel): &\frac{\partial \tilde{\rho}}{\partial t}+\nabla\cdot(\rho_{0} \tilde{\mathbf{U}})=0\,, \\ \rho_{0} &\frac{\partial \tilde{\mathbf{U}}}{\partial t} - \tilde p\, \nabla - = (\nabla \times \tilde{\mathbf{B}}) \times (\mathbf{B}_0 + (\nabla \times \mathbf B_0) \times \tilde{\mathbf{B}} + = (\nabla \times \tilde{\mathbf{B}}) \times \mathbf{B} + (\nabla \times \mathbf B_0) \times \tilde{\mathbf{B}} + \frac{A_\textnormal{h}}{A_\textnormal{b}} \left[ \frac{1}{\epsilon} n_\textnormal{gc} \tilde{\mathbf{U}} - \frac{1}{\epsilon} \mathbf{J}_\textnormal{gc} - \nabla \times \mathbf{M}_\textnormal{gc} \right] \times \mathbf{B} \,, \\ &\frac{\partial \tilde p}{\partial t} + \nabla\cdot(p_0 \tilde{\mathbf{U}}) + \frac{2}{3}\,p_0\nabla\cdot \tilde{\mathbf{U}}=0\,, \\ - &\frac{\partial \tilde{\mathbf{B}}}{\partial t} - \nabla\times(\tilde{\mathbf{U}} \times \mathbf{B}_0) + &\frac{\partial \tilde{\mathbf{B}}}{\partial t} - \nabla\times(\tilde{\mathbf{U}} \times \mathbf{B}) = 0\,, \end{aligned} \right. @@ -586,7 +593,7 @@ class LinearMHDDriftkineticCC(StruphyModel): \\ & n_\textnormal{gc} = \int f_\textnormal{h} B_\parallel^* \,\textnormal dv_\parallel \textnormal d\mu \,, \\ - & \mathbf{J}_\textnormal{gc} = \int f_\textnormal{h}(v_\parallel \mathbf{B}^* - \mathbf{b}_0 \times \mathbf{E}^*) \,\textnormal dv_\parallel \textnormal d\mu \,, + & \mathbf{J}_\textnormal{gc} = \int \frac{f_\textnormal{h}}{B_\parallel^*}(v_\parallel \mathbf{B}^* - \mathbf{b}_0 \times \mathbf{E}^*) \,\textnormal dv_\parallel \textnormal d\mu \,, \\ & \mathbf{M}_\textnormal{gc} = - \int f_\textnormal{h} B_\parallel^* \mu \mathbf{b}_0 \,\textnormal dv_\parallel \textnormal d\mu \,, \end{aligned} @@ -598,9 +605,11 @@ class LinearMHDDriftkineticCC(StruphyModel): .. math:: \begin{align} - \mathbf{B}^* &= \mathbf{B} + \epsilon v_\parallel \nabla \times \mathbf{b}_0 \,,\qquad B^*_\parallel = \mathbf{b}_0 \cdot \mathbf{B}^*\,, + B^*_\parallel = \mathbf{b}_0 \cdot \mathbf{B}^*\,, \\[2mm] - \mathbf{E}^* &= - \tilde{\mathbf{U}} \times \mathbf{B} - \epsilon \mu \nabla B_\parallel \,, + \mathbf{B}^* &= \mathbf{B} + \epsilon v_\parallel \nabla \times \mathbf{b}_0 \,, + \\[2mm] + \mathbf{E}^* &= - \tilde{\mathbf{U}} \times \mathbf{B} - \epsilon \mu \nabla (\mathbf{b}_0 \cdot \mathbf{B}) \,, \end{align} with the normalization parameter @@ -617,340 +626,237 @@ class LinearMHDDriftkineticCC(StruphyModel): 4. :class:`~struphy.propagators.propagators_coupling.CurrentCoupling5DCurlb` 5. :class:`~struphy.propagators.propagators_fields.CurrentCoupling5DDensity` 6. :class:`~struphy.propagators.propagators_fields.ShearAlfvenCurrentCoupling5D` - 7. :class:`~struphy.propagators.propagators_fields.MagnetosonicCurrentCoupling5D` + 7. :class:`~struphy.propagators.propagators_fields.Magnetosonic` :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["em_fields"]["b_field"] = "Hdiv" - dct["fluid"]["mhd"] = { - "density": "L2", - "velocity": "Hdiv", - "pressure": "L2", - } - dct["kinetic"]["energetic_ions"] = "Particles5D" - return dct + ## species + class EnergeticIons(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="Particles5D") + self.init_variables() + + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() + + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.pressure = FEECVariable(space="L2") + self.velocity = FEECVariable(space="Hdiv") + self.init_variables() + + ## propagators + + class Propagators: + def __init__(self, turn_off: tuple[str, ...] = (None,)): + if not "PushGuidingCenterBxEstar" in turn_off: + self.push_bxe = propagators_markers.PushGuidingCenterBxEstar() + if not "PushGuidingCenterParallel" in turn_off: + self.push_parallel = propagators_markers.PushGuidingCenterParallel() + if not "ShearAlfvenCurrentCoupling5D" in turn_off: + self.shearalfen_cc5d = propagators_fields.ShearAlfvenCurrentCoupling5D() + if not "Magnetosonic" in turn_off: + self.magnetosonic = propagators_fields.Magnetosonic() + if not "CurrentCoupling5DDensity" in turn_off: + self.cc5d_density = propagators_fields.CurrentCoupling5DDensity() + if not "CurrentCoupling5DGradB" in turn_off: + self.cc5d_gradb = propagators_coupling.CurrentCoupling5DGradB() + if not "CurrentCoupling5DCurlb" in turn_off: + self.cc5d_curlb = propagators_coupling.CurrentCoupling5DCurlb() + + def __init__(self, turn_off: tuple[str, ...] = (None,)): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + self.energetic_ions = self.EnergeticIons() + + # 2. instantiate all propagators + self.propagators = self.Propagators(turn_off) + + # 3. assign variables to propagators + if not "ShearAlfvenCurrentCoupling5D" in turn_off: + self.propagators.shearalfen_cc5d.variables.u = self.mhd.velocity + self.propagators.shearalfen_cc5d.variables.b = self.em_fields.b_field + if not "Magnetosonic" in turn_off: + self.propagators.magnetosonic.variables.n = self.mhd.density + self.propagators.magnetosonic.variables.u = self.mhd.velocity + self.propagators.magnetosonic.variables.p = self.mhd.pressure + if not "CurrentCoupling5DDensity" in turn_off: + self.propagators.cc5d_density.variables.u = self.mhd.velocity + if not "CurrentCoupling5DGradB" in turn_off: + self.propagators.cc5d_gradb.variables.u = self.mhd.velocity + self.propagators.cc5d_gradb.variables.energetic_ions = self.energetic_ions.var + if not "CurrentCoupling5DCurlb" in turn_off: + self.propagators.cc5d_curlb.variables.u = self.mhd.velocity + self.propagators.cc5d_curlb.variables.energetic_ions = self.energetic_ions.var + if not "PushGuidingCenterBxEstar" in turn_off: + self.propagators.push_bxe.variables.ions = self.energetic_ions.var + if not "PushGuidingCenterParallel" in turn_off: + self.propagators.push_parallel.variables.ions = self.energetic_ions.var + + # define scalars for update_scalar_quantities + self.add_scalar("en_U") + self.add_scalar("en_p") + self.add_scalar("en_B") + self.add_scalar("en_fv", compute="from_particles", variable=self.energetic_ions.var) + self.add_scalar("en_fB", compute="from_particles", variable=self.energetic_ions.var) + self.add_scalar("en_tot", summands=["en_U", "en_p", "en_B", "en_fv", "en_fB"]) - @staticmethod - def bulk_species(): - return "mhd" + @property + def bulk_species(self): + return self.mhd - @staticmethod - def velocity_scale(): + @property + def velocity_scale(self): return "alfvén" - @staticmethod - def propagators_dct(): - return { - propagators_markers.PushGuidingCenterBxEstar: ["energetic_ions"], - propagators_markers.PushGuidingCenterParallel: ["energetic_ions"], - propagators_coupling.CurrentCoupling5DGradB: ["energetic_ions", "mhd_velocity"], - propagators_coupling.CurrentCoupling5DCurlb: ["energetic_ions", "mhd_velocity"], - propagators_fields.CurrentCoupling5DDensity: ["mhd_velocity"], - propagators_fields.ShearAlfvenCurrentCoupling5D: ["mhd_velocity", "b_field"], - propagators_fields.MagnetosonicCurrentCoupling5D: ["mhd_density", "mhd_velocity", "mhd_pressure"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - # add special options - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["fluid", "mhd"], - key="u_space", - option="Hdiv", - dct=dct, - ) - return dct - - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - from mpi4py.MPI import IN_PLACE, SUM - - from struphy.polar.basic import PolarVector - - # extract necessary parameters - u_space = params["fluid"]["mhd"]["options"]["u_space"] - params_alfven = params["fluid"]["mhd"]["options"]["ShearAlfvenCurrentCoupling5D"] - params_sonic = params["fluid"]["mhd"]["options"]["MagnetosonicCurrentCoupling5D"] - params_density = params["fluid"]["mhd"]["options"]["CurrentCoupling5DDensity"] - - params_bxE = params["kinetic"]["energetic_ions"]["options"]["PushGuidingCenterBxEstar"] - params_parallel = params["kinetic"]["energetic_ions"]["options"]["PushGuidingCenterParallel"] - params_cc_gradB = params["kinetic"]["energetic_ions"]["options"]["CurrentCoupling5DGradB"] - params_cc_curlb = params["kinetic"]["energetic_ions"]["options"]["CurrentCoupling5DCurlb"] - params_cc_gradB = params["kinetic"]["energetic_ions"]["options"]["CurrentCoupling5DGradB"] - - # compute coupling parameters - Ab = params["fluid"]["mhd"]["phys_params"]["A"] - Ah = params["kinetic"]["energetic_ions"]["phys_params"]["A"] - epsilon = self.equation_params["energetic_ions"]["epsilon"] - - self._coupling_params = {} - self._coupling_params["Ab"] = Ab - self._coupling_params["Ah"] = Ah - - # add control variate to mass_ops object - if self.pointer["energetic_ions"].control_variate: - self.mass_ops.weights["f0"] = self.pointer["energetic_ions"].f0 - - # Project magnetic field - self._b_eq = self.derham.P["2"]( - [ - self.equil.b2_1, - self.equil.b2_2, - self.equil.b2_3, - ] - ) - - self._absB0 = self.derham.P["0"](self.equil.absB0) - - self._unit_b1 = self.derham.P["1"]( - [ - self.equil.unit_b1_1, - self.equil.unit_b1_2, - self.equil.unit_b1_3, - ] - ) - - self._unit_b2 = self.derham.P["2"]( - [ - self.equil.unit_b2_1, - self.equil.unit_b2_2, - self.equil.unit_b2_3, - ] - ) - - self._gradB1 = self.derham.P["1"]( - [ - self.equil.gradB1_1, - self.equil.gradB1_2, - self.equil.gradB1_3, - ] - ) - - self._curl_unit_b2 = self.derham.P["2"]( - [ - self.equil.curl_unit_b2_1, - self.equil.curl_unit_b2_2, - self.equil.curl_unit_b2_3, - ] - ) - - self._p_eq = self.derham.P["3"](self.equil.p3) - self._ones = self._p_eq.space.zeros() - + def allocate_helpers(self): + self._ones = self.projected_equil.p3.space.zeros() if isinstance(self._ones, PolarVector): self._ones.tp[:] = 1.0 else: self._ones[:] = 1.0 - # set keyword arguments for propagators - self._kwargs[propagators_markers.PushGuidingCenterBxEstar] = { - "b_tilde": self.pointer["b_field"], - "algo": params_bxE["algo"], - "epsilon": epsilon, - } - - self._kwargs[propagators_markers.PushGuidingCenterParallel] = { - "b_tilde": self.pointer["b_field"], - "algo": params_parallel["algo"], - "epsilon": epsilon, - } - - if params_cc_gradB["turn_off"]: - self._kwargs[propagators_coupling.CurrentCoupling5DGradB] = None - else: - self._kwargs[propagators_coupling.CurrentCoupling5DGradB] = { - "b": self.pointer["b_field"], - "b_eq": self._b_eq, - "unit_b1": self._unit_b1, - "unit_b2": self._unit_b2, - "absB0": self._absB0, - "gradB1": self._gradB1, - "curl_unit_b2": self._curl_unit_b2, - "u_space": u_space, - "solver": params_cc_gradB["solver"], - "algo": params_cc_gradB["algo"], - "filter": params_cc_gradB["filter"], - "coupling_params": self._coupling_params, - "epsilon": epsilon, - "boundary_cut": params_cc_gradB["boundary_cut"], - } - - if params_cc_curlb["turn_off"]: - self._kwargs[propagators_coupling.CurrentCoupling5DCurlb] = None - else: - self._kwargs[propagators_coupling.CurrentCoupling5DCurlb] = { - "b": self.pointer["b_field"], - "b_eq": self._b_eq, - "unit_b1": self._unit_b1, - "absB0": self._absB0, - "gradB1": self._gradB1, - "curl_unit_b2": self._curl_unit_b2, - "u_space": u_space, - "solver": params_cc_curlb["solver"], - "filter": params_cc_curlb["filter"], - "coupling_params": self._coupling_params, - "epsilon": epsilon, - "boundary_cut": params_cc_curlb["boundary_cut"], - } - - if params_density["turn_off"]: - self._kwargs[propagators_fields.CurrentCoupling5DDensity] = None - else: - self._kwargs[propagators_fields.CurrentCoupling5DDensity] = { - "particles": self.pointer["energetic_ions"], - "b": self.pointer["b_field"], - "b_eq": self._b_eq, - "unit_b1": self._unit_b1, - "curl_unit_b2": self._curl_unit_b2, - "u_space": u_space, - "solver": params_density["solver"], - "coupling_params": self._coupling_params, - "epsilon": epsilon, - "boundary_cut": params_density["boundary_cut"], - } - - if params_alfven["turn_off"]: - self._kwargs[propagators_fields.ShearAlfvenCurrentCoupling5D] = None - else: - self._kwargs[propagators_fields.ShearAlfvenCurrentCoupling5D] = { - "particles": self.pointer["energetic_ions"], - "unit_b1": self._unit_b1, - "absB0": self._absB0, - "u_space": u_space, - "solver": params_alfven["solver"], - "filter": params_alfven["filter"], - "coupling_params": self._coupling_params, - "accumulated_magnetization": self.pointer["accumulated_magnetization"], - "boundary_cut": params_alfven["boundary_cut"], - } - - if params_sonic["turn_off"]: - self._kwargs[propagators_fields.MagnetosonicCurrentCoupling5D] = None - else: - self._kwargs[propagators_fields.MagnetosonicCurrentCoupling5D] = { - "particles": self.pointer["energetic_ions"], - "b": self.pointer["b_field"], - "unit_b1": self._unit_b1, - "absB0": self._absB0, - "u_space": u_space, - "solver": params_sonic["solver"], - "filter": params_sonic["filter"], - "coupling_params": self._coupling_params, - "boundary_cut": params_sonic["boundary_cut"], - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - # Scalar variables to be saved during simulation - self.add_scalar("en_U", compute="from_field") - self.add_scalar("en_p", compute="from_field") - self.add_scalar("en_B", compute="from_field") - self.add_scalar("en_fv", compute="from_particles", species="energetic_ions") - self.add_scalar("en_fB", compute="from_particles", species="energetic_ions") - # self.add_scalar('en_fv_lost', compute = 'from_particles', species='energetic_ions') - # self.add_scalar('en_fB_lost', compute = 'from_particles', species='energetic_ions') - # self.add_scalar('en_tot',summands = ['en_U','en_p','en_B','en_fv','en_fB','en_fv_lost','en_fB_lost']) - self.add_scalar("en_tot", summands=["en_U", "en_p", "en_B", "en_fv", "en_fB"]) - - # things needed in update_scalar_quantities - self._mpi_sum = SUM - self._mpi_in_place = IN_PLACE - - # temporaries - self._b_full1 = self._b_eq.space.zeros() - self._PBb = self._absB0.space.zeros() - self._en_fv = np.empty(1, dtype=float) self._en_fB = np.empty(1, dtype=float) - # self._en_fv_lost = np.empty(1, dtype=float) - # self._en_fB_lost = np.empty(1, dtype=float) + self._en_tot = np.empty(1, dtype=float) self._n_lost_particles = np.empty(1, dtype=float) + self._PB = getattr(self.basis_ops, "PB") + self._PBb = self._PB.codomain.zeros() + def update_scalar_quantities(self): - en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.pointer["mhd_velocity"], self.pointer["mhd_velocity"]) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) - en_p = self.pointer["mhd_pressure"].inner(self._ones) / (5 / 3 - 1) + # scaling factor + Ab = self.mhd.mass_number + Ah = self.energetic_ions.var.species.mass_number + + # perturbed fields + en_U = 0.5 * self.mass_ops.M2n.dot_inner( + self.mhd.velocity.spline.vector, + self.mhd.velocity.spline.vector, + ) + en_B = 0.5 * self.mass_ops.M2.dot_inner( + self.em_fields.b_field.spline.vector, + self.em_fields.b_field.spline.vector, + ) + en_p = self.mhd.pressure.spline.vector.inner(self._ones) / (5 / 3 - 1) self.update_scalar("en_U", en_U) - self.update_scalar("en_p", en_p) self.update_scalar("en_B", en_B) + self.update_scalar("en_p", en_p) + + # particles' energy + particles = self.energetic_ions.var.particles self._en_fv[0] = ( - self.pointer["energetic_ions"] - .markers[~self.pointer["energetic_ions"].holes, 5] - .dot( - self.pointer["energetic_ions"].markers[~self.pointer["energetic_ions"].holes, 3] ** 2, + particles.markers[~particles.holes, 5].dot( + particles.markers[~particles.holes, 3] ** 2, ) / (2.0) - * self._coupling_params["Ah"] - / self._coupling_params["Ab"] + * Ah + / Ab ) - self.update_scalar("en_fv", self._en_fv[0]) - - # self._en_fv_lost[0] = self.pointer['energetic_ions'].lost_markers[:self.pointer['energetic_ions'].n_lost_markers, 5].dot( - # self.pointer['energetic_ions'].lost_markers[:self.pointer['energetic_ions'].n_lost_markers, 3]**2) / (2.0) * self._coupling_params['Ah']/self._coupling_params['Ab'] - - # self.update_scalar('en_fv_lost', self._en_fv_lost[0]) - - # calculate particle magnetic energy - self.pointer["energetic_ions"].save_magnetic_energy( - self.pointer["b_field"], - ) + self._PBb = self._PB.dot(self.em_fields.b_field.spline.vector) + particles.save_magnetic_energy(self._PBb) self._en_fB[0] = ( - self.pointer["energetic_ions"] - .markers[~self.pointer["energetic_ions"].holes, 5] - .dot( - self.pointer["energetic_ions"].markers[~self.pointer["energetic_ions"].holes, 8], + particles.markers[~particles.holes, 5].dot( + particles.markers[~particles.holes, 8], ) - * self._coupling_params["Ah"] - / self._coupling_params["Ab"] + * Ah + / Ab ) + self.update_scalar("en_fv", self._en_fv[0]) self.update_scalar("en_fB", self._en_fB[0]) + self.update_scalar("en_tot") - # self._en_fB_lost[0] = self.pointer['energetic_ions'].lost_markers[:self.pointer['energetic_ions'].n_lost_markers, 5].dot( - # self.pointer['energetic_ions'] .lost_markers[:self.pointer['energetic_ions'].n_lost_markers, 8]) * self._coupling_params['Ah']/self._coupling_params['Ab'] + # print number of lost particles + n_lost_markers = np.array(particles.n_lost_markers) - # self.update_scalar('en_fB_lost', self._en_fB_lost[0]) + self.derham.comm.Allreduce( + MPI.IN_PLACE, + n_lost_markers, + op=MPI.SUM, + ) - self.update_scalar("en_tot") + if self.clone_config is not None: + self.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + n_lost_markers, + op=MPI.SUM, + ) - # Print number of lost ions - self._n_lost_particles[0] = self.pointer["energetic_ions"].n_lost_markers - self.derham.comm.Allreduce(self._mpi_in_place, self._n_lost_particles, op=self._mpi_sum) - if self.derham.comm.Get_rank() == 0: + if rank == 0: print( - "ratio of lost particles: ", - self._n_lost_particles[0] / self.pointer["energetic_ions"].Np * 100, - "%", + "Lost particle ratio: ", + n_lost_markers / particles.Np * 100, + "% \n", ) - @staticmethod - def diagnostics_dct(): - dct = {} - - dct["accumulated_magnetization"] = "Hdiv" - return dct - - __diagnostics__ = diagnostics_dct() + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "shearalfen_cc5d.Options" in line: + new_file += [ + """model.propagators.shearalfen_cc5d.options = model.propagators.shearalfen_cc5d.Options( + energetic_ions = model.energetic_ions.var,)\n""" + ] + + elif "magnetosonic.Options" in line: + new_file += [ + """model.propagators.magnetosonic.options = model.propagators.magnetosonic.Options( + b_field=model.em_fields.b_field,)\n""" + ] + + elif "cc5d_density.Options" in line: + new_file += [ + """model.propagators.cc5d_density.options = model.propagators.cc5d_density.Options( + energetic_ions = model.energetic_ions.var, + b_tilde = model.em_fields.b_field,)\n""" + ] + + elif "cc5d_curlb.Options" in line: + new_file += [ + """model.propagators.cc5d_curlb.options = model.propagators.cc5d_curlb.Options( + b_tilde = model.em_fields.b_field,)\n""" + ] + + elif "cc5d_gradb.Options" in line: + new_file += [ + """model.propagators.cc5d_gradb.options = model.propagators.cc5d_gradb.Options( + b_tilde = model.em_fields.b_field,)\n""" + ] + + elif "push_bxe.Options" in line: + new_file += [ + """model.propagators.push_bxe.options = model.propagators.push_bxe.Options( + b_tilde = model.em_fields.b_field,)\n""" + ] + + elif "push_parallel.Options" in line: + new_file += [ + """model.propagators.push_parallel.options = model.propagators.push_parallel.Options( + b_tilde = model.em_fields.b_field,)\n""" + ] + + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) class ColdPlasmaVlasov(StruphyModel): diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index 3a39c5944..c236f8d99 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -43,7 +43,7 @@ if rank == 0: print(f"\n{kinetic_models = }") -hybrid_models = [] +hybrid_models = ["LinearMHDDriftkineticCC"] # for name, obj in inspect.getmembers(hybrid): # if inspect.isclass(obj) and "models.hybrid" in obj.__module__: # hybrid_models += [name] diff --git a/src/struphy/pic/accumulation/accum_kernels_gc.py b/src/struphy/pic/accumulation/accum_kernels_gc.py index 628eeeab7..c5e836f81 100644 --- a/src/struphy/pic/accumulation/accum_kernels_gc.py +++ b/src/struphy/pic/accumulation/accum_kernels_gc.py @@ -8,7 +8,7 @@ These kernels are passed to :class:`struphy.pic.accumulation.particles_to_grid.Accumulator`. """ -from numpy import empty, shape, zeros +from numpy import empty, mod, shape, zeros from pyccel.decorators import stack_array import struphy.bsplines.bsplines_kernels as bsplines_kernels @@ -67,6 +67,46 @@ def gc_density_0form( # -- removed omp: #$ omp end parallel +def gc_mag_density_0form( + args_markers: "MarkerArguments", + args_derham: "DerhamArguments", + args_domain: "DomainArguments", + vec: "float[:,:,:]", + scale: "float", # model specific argument +): + r""" + Kernel for :class:`~struphy.pic.accumulation.particles_to_grid.AccumulatorVector` into V0 with the filling + + .. math:: + + B_p^\mu = \mu \frac{w_p}{N} \,. + """ + + markers = args_markers.markers + Np = args_markers.Np + + # -- removed omp: #$ omp parallel private (ip, eta1, eta2, eta3, filling) + # -- removed omp: #$ omp for reduction ( + :vec) + for ip in range(shape(markers)[0]): + # only do something if particle is a "true" particle (i.e. not a hole) + if markers[ip, 0] == -1.0: + continue + + # marker positions + eta1 = markers[ip, 0] + eta2 = markers[ip, 1] + eta3 = markers[ip, 2] + + # marker weight and magnetic moment + weight = markers[ip, 5] + mu = markers[ip, 9] + + # filling =mu*w_p/N + filling = mu * weight / Np * scale + + particle_to_mat_kernels.vec_fill_b_v0(args_derham, eta1, eta2, eta3, vec, filling) + + @stack_array("dfm", "df_inv", "df_inv_t", "g_inv", "tmp1", "tmp2", "b", "b_prod", "bstar", "norm_b1", "curl_norm_b") def cc_lin_mhd_5d_D( args_markers: "MarkerArguments", @@ -75,22 +115,19 @@ def cc_lin_mhd_5d_D( mat12: "float[:,:,:,:,:,:]", mat13: "float[:,:,:,:,:,:]", mat23: "float[:,:,:,:,:,:]", - epsilon: float, # model specific argument - b2_1: "float[:,:,:]", # model specific argument - b2_2: "float[:,:,:]", # model specific argument - b2_3: "float[:,:,:]", # model specific argument - # model specific argument + epsilon: float, + ep_scale: "float", + b2_1: "float[:,:,:]", + b2_2: "float[:,:,:]", + b2_3: "float[:,:,:]", norm_b11: "float[:,:,:]", norm_b12: "float[:,:,:]", norm_b13: "float[:,:,:]", - # model specific argument curl_norm_b1: "float[:,:,:]", curl_norm_b2: "float[:,:,:]", curl_norm_b3: "float[:,:,:]", basis_u: "int", - scale_mat: "float", - boundary_cut: float, -): # model specific argument +): r"""Accumulation kernel for the propagator :class:`~struphy.propagators.propagators_fields.CurrentCoupling5DDensity`. Accumulates :math:`\alpha`-form matrix with the filling functions (:math:`\alpha = 2`) @@ -157,9 +194,6 @@ def cc_lin_mhd_5d_D( v = markers[ip, 3] - if eta1 < boundary_cut or eta1 > 1.0 - boundary_cut: - continue - # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) @@ -186,11 +220,9 @@ def cc_lin_mhd_5d_D( # calculate Bstar and transform to H1vec b_star[:] = b + epsilon * v * curl_norm_b - b_star /= det_df # calculate b_para and b_star_para b_para = linalg_kernels.scalar_dot(norm_b1, b) - b_para /= det_df b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) @@ -202,9 +234,9 @@ def cc_lin_mhd_5d_D( if basis_u == 0: # filling functions - filling_m12 = -weight * density_const * b_prod[0, 1] * scale_mat - filling_m13 = -weight * density_const * b_prod[0, 2] * scale_mat - filling_m23 = -weight * density_const * b_prod[1, 2] * scale_mat + filling_m12 = -weight * density_const * b_prod[0, 1] * ep_scale / epsilon + filling_m13 = -weight * density_const * b_prod[0, 2] * ep_scale / epsilon + filling_m23 = -weight * density_const * b_prod[1, 2] * ep_scale / epsilon # call the appropriate matvec filler particle_to_mat_kernels.mat_fill_v0vec_asym( @@ -219,9 +251,9 @@ def cc_lin_mhd_5d_D( linalg_kernels.matrix_matrix(g_inv, b_prod, tmp1) linalg_kernels.matrix_matrix(tmp1, g_inv, tmp2) - filling_m12 = -weight * density_const * tmp2[0, 1] * scale_mat - filling_m13 = -weight * density_const * tmp2[0, 2] * scale_mat - filling_m23 = -weight * density_const * tmp2[1, 2] * scale_mat + filling_m12 = -weight * density_const * tmp2[0, 1] * ep_scale / epsilon + filling_m13 = -weight * density_const * tmp2[0, 2] * ep_scale / epsilon + filling_m23 = -weight * density_const * tmp2[1, 2] * ep_scale / epsilon # call the appropriate matvec filler particle_to_mat_kernels.mat_fill_v1_asym( @@ -230,9 +262,9 @@ def cc_lin_mhd_5d_D( elif basis_u == 2: # filling functions - filling_m12 = -weight * density_const * b_prod[0, 1] * scale_mat / det_df**2 - filling_m13 = -weight * density_const * b_prod[0, 2] * scale_mat / det_df**2 - filling_m23 = -weight * density_const * b_prod[1, 2] * scale_mat / det_df**2 + filling_m12 = -weight * density_const * b_prod[0, 1] * ep_scale / epsilon / det_df**2 + filling_m13 = -weight * density_const * b_prod[0, 2] * ep_scale / epsilon / det_df**2 + filling_m23 = -weight * density_const * b_prod[1, 2] * ep_scale / epsilon / det_df**2 # call the appropriate matvec filler particle_to_mat_kernels.mat_fill_v2_asym( @@ -248,23 +280,23 @@ def cc_lin_mhd_5d_D( @stack_array( "dfm", - "df_inv_t", "df_inv", + "df_inv_t", "g_inv", "filling_m", "filling_v", "tmp", "tmp1", - "tmp2", "tmp_m", "tmp_v", "b", + "bfull_star", "b_prod", - "b_prod_negb_star", + "b_prod_neg", "norm_b1", "curl_norm_b", ) -def cc_lin_mhd_5d_J1( +def cc_lin_mhd_5d_curlb( args_markers: "MarkerArguments", args_derham: "DerhamArguments", args_domain: "DomainArguments", @@ -277,21 +309,19 @@ def cc_lin_mhd_5d_J1( vec1: "float[:,:,:]", vec2: "float[:,:,:]", vec3: "float[:,:,:]", - epsilon: float, # model specific argument - b1: "float[:,:,:]", # model specific argument - b2: "float[:,:,:]", # model specific argument - b3: "float[:,:,:]", # model specific argument - norm_b11: "float[:,:,:]", # model specific argument - norm_b12: "float[:,:,:]", # model specific argument - norm_b13: "float[:,:,:]", # model specific argument - curl_norm_b1: "float[:,:,:]", # model specific argument - curl_norm_b2: "float[:,:,:]", # model specific argument - curl_norm_b3: "float[:,:,:]", # model specific argument - basis_u: "int", # model specific argument - scale_mat: "float", # model specific argument - scale_vec: "float", # model specific argument - boundary_cut: "float", -): # model specific argument + epsilon: float, + ep_scale: float, + b1: "float[:,:,:]", + b2: "float[:,:,:]", + b3: "float[:,:,:]", + norm_b11: "float[:,:,:]", + norm_b12: "float[:,:,:]", + norm_b13: "float[:,:,:]", + curl_norm_b1: "float[:,:,:]", + curl_norm_b2: "float[:,:,:]", + curl_norm_b3: "float[:,:,:]", + basis_u: "int", +): r"""Accumulation kernel for the propagator :class:`~struphy.propagators.propagators_coupling.CurrentCoupling5DCurlb`. Accumulates :math:`\alpha`-form matrix and vector with the filling functions (:math:`\alpha = 2`) @@ -303,21 +333,6 @@ def cc_lin_mhd_5d_J1( B_p^\mu &= w_p \left( \frac{v^2_{\parallel,p}}{g\hat B^*_\parallel} \mathbf B^2_{\times} \right)_\mu \,, where :math:`\mathbf B^2_{\times} \mathbf a := \hat{\mathbf B}^2 \times \mathbf a` for :math:`a \in \mathbb R^3`. - - Parameters - ---------- - b1, b2, b3 : array[float] - FE coefficients c_ijk of the magnetic field as a 2-form. - - norm_b11, norm_b12, norm_b13 : array[float] - FE coefficients c_ijk of the normalized magnetic field as a 1-form. - - curl_norm_b1, curl_norm_b2, curl_norm_b3 : array[float] - FE coefficients c_ijk of the curl of normalized magnetic field as a 2-form. - - Note - ---- - The above parameter list contains only the model specific input arguments. """ markers = args_markers.markers @@ -325,7 +340,7 @@ def cc_lin_mhd_5d_J1( # allocate for magnetic field evaluation b = empty(3, dtype=float) - b_star = empty(3, dtype=float) + bfull_star = empty(3, dtype=float) b_prod = zeros((3, 3), dtype=float) b_prod_neg = zeros((3, 3), dtype=float) norm_b1 = empty(3, dtype=float) @@ -338,12 +353,11 @@ def cc_lin_mhd_5d_J1( g_inv = empty((3, 3), dtype=float) # allocate for filling - filling_m = empty((3, 3), dtype=float) - filling_v = empty(3, dtype=float) + filling_m = zeros((3, 3), dtype=float) + filling_v = zeros(3, dtype=float) tmp = empty((3, 3), dtype=float) tmp1 = empty((3, 3), dtype=float) - tmp2 = empty((3, 3), dtype=float) tmp_m = empty((3, 3), dtype=float) tmp_v = empty(3, dtype=float) @@ -351,8 +365,6 @@ def cc_lin_mhd_5d_J1( # get number of markers n_markers_loc = shape(markers)[0] - # -- removed omp: #$ omp parallel firstprivate(b_prod) private(ip, boundary_cut, eta1, eta2, eta3, v, weight, span1, span2, span3, b1, b2, b3, b, b_star, b_prod_neg, norm_b1, curl_norm_b, abs_b_star_para, dfm, df_inv, df_inv_t, g_inv, det_df, tmp, tmp1, tmp2, tmp_m, tmp_v, filling_m, filling_v) - # -- removed omp: #$ omp for reduction ( + : mat11, mat12, mat13, mat22, mat23, mat33, vec1, vec2, vec3) for ip in range(n_markers_loc): # only do something if particle is a "true" particle (i.e. not a hole) if markers[ip, 0] == -1.0: @@ -367,9 +379,6 @@ def cc_lin_mhd_5d_J1( weight = markers[ip, 5] v = markers[ip, 3] - if eta1 < boundary_cut or eta1 > 1.0 - boundary_cut: - continue - # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) @@ -387,11 +396,11 @@ def cc_lin_mhd_5d_J1( # curl_norm_b; 2form eval_2form_spline_mpi(span1, span2, span3, args_derham, curl_norm_b1, curl_norm_b2, curl_norm_b3, curl_norm_b) - # b_star; 2form in H1vec - b_star[:] = (b + curl_norm_b * v * epsilon) / det_df + # b_star; 2form + bfull_star[:] = b + curl_norm_b * v * epsilon # calculate abs_b_star_para - abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) + abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, bfull_star) # calculate tensor product of two curl_norm_b linalg_kernels.outer(curl_norm_b, curl_norm_b, tmp) @@ -411,8 +420,8 @@ def cc_lin_mhd_5d_J1( linalg_kernels.matrix_matrix(tmp1, b_prod_neg, tmp_m) linalg_kernels.matrix_vector(b_prod, curl_norm_b, tmp_v) - filling_m[:, :] = weight * tmp_m * v**2 / abs_b_star_para**2 / det_df**2 * scale_mat - filling_v[:] = weight * tmp_v * v**2 / abs_b_star_para / det_df * scale_vec + filling_m[:, :] += weight * tmp_m * v**2 / abs_b_star_para**2 * ep_scale + filling_v[:] += weight * tmp_v * v**2 / abs_b_star_para * ep_scale # call the appropriate matvec filler particle_to_mat_kernels.m_v_fill_v0vec_symm( @@ -440,54 +449,13 @@ def cc_lin_mhd_5d_J1( filling_v[2], ) - elif basis_u == 1: - # needed metric coefficients - linalg_kernels.matrix_inv_with_det(dfm, det_df, df_inv) - linalg_kernels.transpose(df_inv, df_inv_t) - linalg_kernels.matrix_matrix(df_inv, df_inv_t, g_inv) - linalg_kernels.matrix_matrix(g_inv, b_prod, tmp1) - linalg_kernels.matrix_vector(tmp1, curl_norm_b, tmp_v) - - linalg_kernels.matrix_matrix(tmp1, tmp, tmp2) - linalg_kernels.matrix_matrix(tmp2, b_prod_neg, tmp1) - linalg_kernels.matrix_matrix(tmp1, g_inv, tmp_m) - - filling_m[:, :] = weight * tmp_m * v**2 / abs_b_star_para**2 / det_df**2 * scale_mat - filling_v[:] = weight * tmp_v * v**2 / abs_b_star_para / det_df * scale_vec - - # call the appropriate matvec filler - particle_to_mat_kernels.m_v_fill_v1_symm( - args_derham, - span1, - span2, - span3, - mat11, - mat12, - mat13, - mat22, - mat23, - mat33, - filling_m[0, 0], - filling_m[0, 1], - filling_m[0, 2], - filling_m[1, 1], - filling_m[1, 2], - filling_m[2, 2], - vec1, - vec2, - vec3, - filling_v[0], - filling_v[1], - filling_v[2], - ) - elif basis_u == 2: linalg_kernels.matrix_matrix(b_prod, tmp, tmp1) linalg_kernels.matrix_matrix(tmp1, b_prod_neg, tmp_m) linalg_kernels.matrix_vector(b_prod, curl_norm_b, tmp_v) - filling_m[:, :] = weight * tmp_m * v**2 / abs_b_star_para**2 / det_df**4 * scale_mat - filling_v[:] = weight * tmp_v * v**2 / abs_b_star_para / det_df**2 * scale_vec + filling_m[:, :] = weight * tmp_m * v**2 / abs_b_star_para**2 / det_df**2 * ep_scale + filling_v[:] = weight * tmp_v * v**2 / abs_b_star_para / det_df * ep_scale # call the appropriate matvec filler particle_to_mat_kernels.m_v_fill_v2_symm( @@ -526,8 +494,6 @@ def cc_lin_mhd_5d_J1( vec2 /= Np vec3 /= Np - # -- removed omp: #$ omp end parallel - @stack_array("dfm", "norm_b1", "filling_v") def cc_lin_mhd_5d_M( @@ -547,8 +513,7 @@ def cc_lin_mhd_5d_M( norm_b12: "float[:,:,:]", # model specific argument norm_b13: "float[:,:,:]", # model specific argument scale_vec: "float", # model specific argument - boundary_cut: "float", -): # model specific argument +): r"""Accumulation kernel for the propagator :class:`~struphy.propagators.propagators_fields.ShearAlfvenCurrentCoupling5D` and :class:`~struphy.propagators.propagators_fields.MagnetosonicCurrentCoupling5D`. Accumulates 2-form vector with the filling functions: @@ -600,9 +565,6 @@ def cc_lin_mhd_5d_M( weight = markers[ip, 5] mu = markers[ip, 9] - if eta1 < boundary_cut or eta1 > 1.0 - boundary_cut: - continue - # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) @@ -633,19 +595,17 @@ def cc_lin_mhd_5d_M( "df_inv", "g_inv", "filling_v", - "tmp1", - "tmp2", + "tmp", "tmp_v", "b", "b_prod", - "norm_b2_prod", + "norm_b_prod", "b_star", "curl_norm_b", "norm_b1", - "norm_b2", "grad_PB", ) -def cc_lin_mhd_5d_J2( +def cc_lin_mhd_5d_gradB( args_markers: "MarkerArguments", args_derham: "DerhamArguments", args_domain: "DomainArguments", @@ -658,27 +618,22 @@ def cc_lin_mhd_5d_J2( vec1: "float[:,:,:]", vec2: "float[:,:,:]", vec3: "float[:,:,:]", - epsilon: float, # model specific argument - b1: "float[:,:,:]", # model specific argument - b2: "float[:,:,:]", # model specific argument - b3: "float[:,:,:]", # model specific argument - norm_b11: "float[:,:,:]", # model specific argument - norm_b12: "float[:,:,:]", # model specific argument - norm_b13: "float[:,:,:]", # model specific argument - norm_b21: "float[:,:,:]", # model specific argument - norm_b22: "float[:,:,:]", # model specific argument - norm_b23: "float[:,:,:]", # model specific argument - curl_norm_b1: "float[:,:,:]", # model specific argument - curl_norm_b2: "float[:,:,:]", # model specific argument - curl_norm_b3: "float[:,:,:]", # model specific argument - grad_PB1: "float[:,:,:]", # model specific argument - grad_PB2: "float[:,:,:]", # model specific argument - grad_PB3: "float[:,:,:]", # model specific argument + epsilon: float, + ep_scale: float, + b1: "float[:,:,:]", + b2: "float[:,:,:]", + b3: "float[:,:,:]", + norm_b11: "float[:,:,:]", + norm_b12: "float[:,:,:]", + norm_b13: "float[:,:,:]", + curl_norm_b1: "float[:,:,:]", + curl_norm_b2: "float[:,:,:]", + curl_norm_b3: "float[:,:,:]", + grad_PB1: "float[:,:,:]", + grad_PB2: "float[:,:,:]", + grad_PB3: "float[:,:,:]", basis_u: "int", - scale_mat: "float", - scale_vec: "float", - boundary_cut: float, -): # model specific argument +): r"""Accumulation kernel for the propagator :class:`~struphy.propagators.propagators_coupling.CurrentCoupling5DGradB`. Accumulates math:`\alpha` -form vector with the filling functions @@ -697,9 +652,6 @@ def cc_lin_mhd_5d_J2( norm_b11, norm_b12, norm_b13 : array[float] FE coefficients c_ijk of the normalized magnetic field as a 1-form. - norm_b21, norm_b22, norm_b23 : array[float] - FE coefficients c_ijk of the normalized magnetic field as a 2-form. - curl_norm_b1, curl_norm_b2, curl_norm_b3 : array[float] FE coefficients c_ijk of the curl of normalized magnetic field as a 2-form. @@ -718,10 +670,9 @@ def cc_lin_mhd_5d_J2( b = empty(3, dtype=float) b_star = empty(3, dtype=float) b_prod = zeros((3, 3), dtype=float) - norm_b2_prod = zeros((3, 3), dtype=float) + norm_b_prod = zeros((3, 3), dtype=float) curl_norm_b = empty(3, dtype=float) norm_b1 = empty(3, dtype=float) - norm_b2 = empty(3, dtype=float) grad_PB = empty(3, dtype=float) # allocate for metric coeffs @@ -732,17 +683,13 @@ def cc_lin_mhd_5d_J2( # allocate for filling filling_v = empty(3, dtype=float) - - tmp1 = empty((3, 3), dtype=float) - tmp2 = empty((3, 3), dtype=float) + tmp = empty((3, 3), dtype=float) tmp_v = empty(3, dtype=float) # get number of markers n_markers_loc = shape(markers)[0] - # -- removed omp: #$ omp parallel firstprivate(b_prod) private(ip, boundary_cut, eta1, eta2, eta3, v, mu, weight, span1, span2, span3, b1, b2, b3, b, b_star, norm_b1, norm_b2, norm_b2_prod, curl_norm_b, grad_PB, abs_b_star_para, dfm, df_inv, df_inv_t, g_inv, det_df, tmp1, tmp2, tmp_v, filling_v) - # -- removed omp: #$ omp for reduction ( + : mat11, mat12, mat13, mat22, mat23, mat33, vec1, vec2, vec3) for ip in range(n_markers_loc): # only do something if particle is a "true" particle (i.e. not a hole) if markers[ip, 0] == -1.0: @@ -753,9 +700,173 @@ def cc_lin_mhd_5d_J2( eta2 = markers[ip, 1] eta3 = markers[ip, 2] - if eta1 < boundary_cut or eta1 > 1.0 - boundary_cut: + # marker weight and velocity + weight = markers[ip, 5] + v = markers[ip, 3] + mu = markers[ip, 9] + + # b-field evaluation + span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) + + # evaluate Jacobian, result in dfm + evaluation_kernels.df(eta1, eta2, eta3, args_domain, dfm) + + det_df = linalg_kernels.det(dfm) + + # needed metric coefficients + linalg_kernels.matrix_inv_with_det(dfm, det_df, df_inv) + linalg_kernels.transpose(df_inv, df_inv_t) + linalg_kernels.matrix_matrix(df_inv, df_inv_t, g_inv) + + # b; 2form + eval_2form_spline_mpi(span1, span2, span3, args_derham, b1, b2, b3, b) + + # norm_b1; 1form + eval_1form_spline_mpi(span1, span2, span3, args_derham, norm_b11, norm_b12, norm_b13, norm_b1) + + # curl_norm_b; 2form + eval_2form_spline_mpi(span1, span2, span3, args_derham, curl_norm_b1, curl_norm_b2, curl_norm_b3, curl_norm_b) + + # grad_PB; 1form + eval_1form_spline_mpi(span1, span2, span3, args_derham, grad_PB1, grad_PB2, grad_PB3, grad_PB) + + # b_star; 2form transformed into H1vec + b_star[:] = b + curl_norm_b * v * epsilon + + # calculate abs_b_star_para + abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) + + # operator bx() as matrix + b_prod[0, 1] = -b[2] + b_prod[0, 2] = +b[1] + b_prod[1, 0] = +b[2] + b_prod[1, 2] = -b[0] + b_prod[2, 0] = -b[1] + b_prod[2, 1] = +b[0] + + norm_b_prod[0, 1] = -norm_b1[2] + norm_b_prod[0, 2] = +norm_b1[1] + norm_b_prod[1, 0] = +norm_b1[2] + norm_b_prod[1, 2] = -norm_b1[0] + norm_b_prod[2, 0] = -norm_b1[1] + norm_b_prod[2, 1] = +norm_b1[0] + + if basis_u == 0: + linalg_kernels.matrix_matrix(b_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + + filling_v[:] = weight * tmp_v * mu / abs_b_star_para * ep_scale + + # call the appropriate matvec filler + particle_to_mat_kernels.vec_fill_v0vec( + args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + ) + + elif basis_u == 2: + linalg_kernels.matrix_matrix(b_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + + filling_v[:] = weight * tmp_v * mu / abs_b_star_para / det_df * ep_scale + + # call the appropriate matvec filler + particle_to_mat_kernels.vec_fill_v2( + args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + ) + vec1 /= Np + vec2 /= Np + vec3 /= Np + + +@stack_array( + "dfm", + "df_inv_t", + "df_inv", + "g_inv", + "filling_v", + "tmp", + "tmp_v", + "b", + "b_prod", + "beq", + "beq_prod", + "norm_b_prod", + "bfull_star", + "curl_norm_b", + "norm_b1", + "grad_PB", + "grad_PBeq", +) +def cc_lin_mhd_5d_gradB_dg_init( + args_markers: "MarkerArguments", + args_derham: "DerhamArguments", + args_domain: "DomainArguments", + vec1: "float[:,:,:]", + vec2: "float[:,:,:]", + vec3: "float[:,:,:]", + epsilon: float, + ep_scale: float, + b1: "float[:,:,:]", + b2: "float[:,:,:]", + b3: "float[:,:,:]", + beq1: "float[:,:,:]", + beq2: "float[:,:,:]", + beq3: "float[:,:,:]", + norm_b11: "float[:,:,:]", + norm_b12: "float[:,:,:]", + norm_b13: "float[:,:,:]", + curl_norm_b1: "float[:,:,:]", + curl_norm_b2: "float[:,:,:]", + curl_norm_b3: "float[:,:,:]", + grad_PB1: "float[:,:,:]", + grad_PB2: "float[:,:,:]", + grad_PB3: "float[:,:,:]", + grad_PBeq1: "float[:,:,:]", + grad_PBeq2: "float[:,:,:]", + grad_PBeq3: "float[:,:,:]", + basis_u: "int", +): + r"""TODO""" + + markers = args_markers.markers + Np = args_markers.Np + + # allocate for magnetic field evaluation + b = empty(3, dtype=float) + beq = empty(3, dtype=float) + bfull_star = empty(3, dtype=float) + b_prod = zeros((3, 3), dtype=float) + beq_prod = zeros((3, 3), dtype=float) + norm_b_prod = zeros((3, 3), dtype=float) + curl_norm_b = empty(3, dtype=float) + norm_b1 = empty(3, dtype=float) + grad_PB = empty(3, dtype=float) + grad_PBeq = empty(3, dtype=float) + + # allocate for metric coeffs + dfm = empty((3, 3), dtype=float) + df_inv = empty((3, 3), dtype=float) + df_inv_t = empty((3, 3), dtype=float) + g_inv = empty((3, 3), dtype=float) + + # allocate for filling + filling_v = empty(3, dtype=float) + tmp = empty((3, 3), dtype=float) + + tmp_v = empty(3, dtype=float) + + # get number of markers + n_markers_loc = shape(markers)[0] + + for ip in range(n_markers_loc): + # only do something if particle is a "true" particle (i.e. not a hole) + if markers[ip, 0] == -1.0: continue + # marker positions + eta1 = markers[ip, 0] + eta2 = markers[ip, 1] + eta3 = markers[ip, 2] + # marker weight and velocity weight = markers[ip, 5] v = markers[ip, 3] @@ -777,23 +888,26 @@ def cc_lin_mhd_5d_J2( # b; 2form eval_2form_spline_mpi(span1, span2, span3, args_derham, b1, b2, b3, b) + # beq; 2form + eval_2form_spline_mpi(span1, span2, span3, args_derham, beq1, beq2, beq3, beq) + # norm_b1; 1form eval_1form_spline_mpi(span1, span2, span3, args_derham, norm_b11, norm_b12, norm_b13, norm_b1) - # norm_b2; 2form - eval_2form_spline_mpi(span1, span2, span3, args_derham, norm_b21, norm_b22, norm_b23, norm_b2) - # curl_norm_b; 2form eval_2form_spline_mpi(span1, span2, span3, args_derham, curl_norm_b1, curl_norm_b2, curl_norm_b3, curl_norm_b) # grad_PB; 1form eval_1form_spline_mpi(span1, span2, span3, args_derham, grad_PB1, grad_PB2, grad_PB3, grad_PB) + # grad_PBeq; 1form + eval_1form_spline_mpi(span1, span2, span3, args_derham, grad_PBeq1, grad_PBeq2, grad_PBeq3, grad_PBeq) + # b_star; 2form transformed into H1vec - b_star[:] = (b + curl_norm_b * v * epsilon) / det_df + bfull_star[:] = b + beq + curl_norm_b * v * epsilon # calculate abs_b_star_para - abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) + abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, bfull_star) # operator bx() as matrix b_prod[0, 1] = -b[2] @@ -803,50 +917,304 @@ def cc_lin_mhd_5d_J2( b_prod[2, 0] = -b[1] b_prod[2, 1] = +b[0] - norm_b2_prod[0, 1] = -norm_b2[2] - norm_b2_prod[0, 2] = +norm_b2[1] - norm_b2_prod[1, 0] = +norm_b2[2] - norm_b2_prod[1, 2] = -norm_b2[0] - norm_b2_prod[2, 0] = -norm_b2[1] - norm_b2_prod[2, 1] = +norm_b2[0] + beq_prod[0, 1] = -beq[2] + beq_prod[0, 2] = +beq[1] + beq_prod[1, 0] = +beq[2] + beq_prod[1, 2] = -beq[0] + beq_prod[2, 0] = -beq[1] + beq_prod[2, 1] = +beq[0] + + norm_b_prod[0, 1] = -norm_b1[2] + norm_b_prod[0, 2] = +norm_b1[1] + norm_b_prod[1, 0] = +norm_b1[2] + norm_b_prod[1, 2] = -norm_b1[0] + norm_b_prod[2, 0] = -norm_b1[1] + norm_b_prod[2, 1] = +norm_b1[0] if basis_u == 0: - linalg_kernels.matrix_matrix(b_prod, g_inv, tmp1) - linalg_kernels.matrix_matrix(tmp1, norm_b2_prod, tmp2) - linalg_kernels.matrix_matrix(tmp2, g_inv, tmp1) + # beq contribution + linalg_kernels.matrix_matrix(beq_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PBeq, tmp_v) + + filling_v[:] = weight * tmp_v * mu / abs_b_star_para * ep_scale + + # b contribution + linalg_kernels.matrix_matrix(beq_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + + filling_v[:] += weight * tmp_v * mu / abs_b_star_para * ep_scale + + linalg_kernels.matrix_matrix(b_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PBeq, tmp_v) - linalg_kernels.matrix_vector(tmp1, grad_PB, tmp_v) + filling_v[:] += weight * tmp_v * mu / abs_b_star_para * ep_scale - filling_v[:] = weight * tmp_v * mu / abs_b_star_para * scale_vec + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + + filling_v[:] += weight * tmp_v * mu / abs_b_star_para * ep_scale # call the appropriate matvec filler particle_to_mat_kernels.vec_fill_v0vec( args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] ) - elif basis_u == 1: - linalg_kernels.matrix_matrix(g_inv, b_prod, tmp1) - linalg_kernels.matrix_matrix(tmp1, g_inv, tmp2) - linalg_kernels.matrix_matrix(tmp2, norm_b2_prod, tmp1) - linalg_kernels.matrix_matrix(tmp1, g_inv, tmp2) + elif basis_u == 2: + # beq contribution + linalg_kernels.matrix_matrix(beq_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PBeq, tmp_v) + + filling_v[:] = weight * tmp_v * mu / abs_b_star_para / det_df * ep_scale + + # b contribution + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + + filling_v[:] += weight * tmp_v * mu / abs_b_star_para / det_df * ep_scale + + linalg_kernels.matrix_matrix(b_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PBeq, tmp_v) + + filling_v[:] += weight * tmp_v * mu / abs_b_star_para / det_df * ep_scale + + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + + filling_v[:] += weight * tmp_v * mu / abs_b_star_para / det_df * ep_scale + + # call the appropriate matvec filler + particle_to_mat_kernels.vec_fill_v2( + args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + ) + + vec1 /= Np + vec2 /= Np + vec3 /= Np + + +@stack_array( + "dfm", + "df_inv_t", + "df_inv", + "g_inv", + "filling_v", + "tmp", + "tmp_v", + "b", + "b_prod", + "eta_diff", + "beq", + "beq_prod", + "norm_b_prod", + "bfull_star", + "curl_norm_b", + "norm_b1", + "grad_PB", + "grad_PBeq", + "eta_mid", + "eta_diff", +) +def cc_lin_mhd_5d_gradB_dg( + args_markers: "MarkerArguments", + args_derham: "DerhamArguments", + args_domain: "DomainArguments", + vec1: "float[:,:,:]", + vec2: "float[:,:,:]", + vec3: "float[:,:,:]", + epsilon: float, + ep_scale: float, + b1: "float[:,:,:]", + b2: "float[:,:,:]", + b3: "float[:,:,:]", + beq1: "float[:,:,:]", + beq2: "float[:,:,:]", + beq3: "float[:,:,:]", + norm_b11: "float[:,:,:]", + norm_b12: "float[:,:,:]", + norm_b13: "float[:,:,:]", + curl_norm_b1: "float[:,:,:]", + curl_norm_b2: "float[:,:,:]", + curl_norm_b3: "float[:,:,:]", + grad_PB1: "float[:,:,:]", + grad_PB2: "float[:,:,:]", + grad_PB3: "float[:,:,:]", + grad_PBeq1: "float[:,:,:]", + grad_PBeq2: "float[:,:,:]", + grad_PBeq3: "float[:,:,:]", + basis_u: "int", + const: "float", +): + r"""TODO""" + + markers = args_markers.markers + Np = args_markers.Np + + # allocate for magnetic field evaluation + eta_diff = empty(3, dtype=float) + eta_mid = empty(3, dtype=float) + b = empty(3, dtype=float) + beq = empty(3, dtype=float) + bfull_star = empty(3, dtype=float) + b_prod = zeros((3, 3), dtype=float) + beq_prod = zeros((3, 3), dtype=float) + norm_b_prod = zeros((3, 3), dtype=float) + curl_norm_b = empty(3, dtype=float) + norm_b1 = empty(3, dtype=float) + grad_PB = empty(3, dtype=float) + grad_PBeq = empty(3, dtype=float) + + # allocate for metric coeffs + dfm = empty((3, 3), dtype=float) + df_inv = empty((3, 3), dtype=float) + df_inv_t = empty((3, 3), dtype=float) + g_inv = empty((3, 3), dtype=float) + + # allocate for filling + filling_v = empty(3, dtype=float) + tmp = empty((3, 3), dtype=float) + + tmp_v = empty(3, dtype=float) + + # get number of markers + n_markers_loc = shape(markers)[0] + + for ip in range(n_markers_loc): + # only do something if particle is a "true" particle (i.e. not a hole) + if markers[ip, 0] == -1.0: + continue + + # marker positions, mid point + eta_mid[:] = (markers[ip, 0:3] + markers[ip, 11:14]) / 2.0 + eta_mid[:] = mod(eta_mid[:], 1.0) + + eta_diff[:] = markers[ip, 0:3] - markers[ip, 11:14] + + # marker weight and velocity + weight = markers[ip, 5] + v = markers[ip, 3] + mu = markers[ip, 9] + + # b-field evaluation + span1, span2, span3 = get_spans(eta_mid[0], eta_mid[1], eta_mid[2], args_derham) + + # evaluate Jacobian, result in dfm + evaluation_kernels.df(eta_mid[0], eta_mid[1], eta_mid[2], args_domain, dfm) + + det_df = linalg_kernels.det(dfm) - linalg_kernels.matrix_vector(tmp2, grad_PB, tmp_v) + # needed metric coefficients + linalg_kernels.matrix_inv_with_det(dfm, det_df, df_inv) + linalg_kernels.transpose(df_inv, df_inv_t) + linalg_kernels.matrix_matrix(df_inv, df_inv_t, g_inv) + + # b; 2form + eval_2form_spline_mpi(span1, span2, span3, args_derham, b1, b2, b3, b) + + # beq; 2form + eval_2form_spline_mpi(span1, span2, span3, args_derham, beq1, beq2, beq3, beq) + + # norm_b1; 1form + eval_1form_spline_mpi(span1, span2, span3, args_derham, norm_b11, norm_b12, norm_b13, norm_b1) + + # curl_norm_b; 2form + eval_2form_spline_mpi(span1, span2, span3, args_derham, curl_norm_b1, curl_norm_b2, curl_norm_b3, curl_norm_b) + + # grad_PB; 1form + eval_1form_spline_mpi(span1, span2, span3, args_derham, grad_PB1, grad_PB2, grad_PB3, grad_PB) - filling_v[:] = weight * tmp_v * mu / abs_b_star_para * scale_vec + # grad_PBeq; 1form + eval_1form_spline_mpi(span1, span2, span3, args_derham, grad_PBeq1, grad_PBeq2, grad_PBeq3, grad_PBeq) + + # b_star; 2form transformed into H1vec + bfull_star[:] = b + beq + curl_norm_b * v * epsilon + + # calculate abs_b_star_para + abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, bfull_star) + + # operator bx() as matrix + b_prod[0, 1] = -b[2] + b_prod[0, 2] = +b[1] + b_prod[1, 0] = +b[2] + b_prod[1, 2] = -b[0] + b_prod[2, 0] = -b[1] + b_prod[2, 1] = +b[0] + + beq_prod[0, 1] = -beq[2] + beq_prod[0, 2] = +beq[1] + beq_prod[1, 0] = +beq[2] + beq_prod[1, 2] = -beq[0] + beq_prod[2, 0] = -beq[1] + beq_prod[2, 1] = +beq[0] + + norm_b_prod[0, 1] = -norm_b1[2] + norm_b_prod[0, 2] = +norm_b1[1] + norm_b_prod[1, 0] = +norm_b1[2] + norm_b_prod[1, 2] = -norm_b1[0] + norm_b_prod[2, 0] = -norm_b1[1] + norm_b_prod[2, 1] = +norm_b1[0] + + if basis_u == 0: + # beq * gradPBeq contribution + linalg_kernels.matrix_matrix(beq_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PBeq, tmp_v) + + filling_v[:] = weight * tmp_v * mu / abs_b_star_para * ep_scale + + # beq * gradPB contribution + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + filling_v[:] += weight * tmp_v * mu / abs_b_star_para * ep_scale + + # beq * dg term contribution + linalg_kernels.matrix_vector(tmp, eta_diff, tmp_v) + filling_v[:] += tmp_v / abs_b_star_para * const + + # b * gradPBeq contribution + linalg_kernels.matrix_matrix(b_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PBeq, tmp_v) + filling_v[:] += weight * tmp_v * mu / abs_b_star_para * ep_scale + + # b * gradPB contribution + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + filling_v[:] += weight * tmp_v * mu / abs_b_star_para * ep_scale + + # b * dg term contribution + linalg_kernels.matrix_vector(tmp, eta_diff, tmp_v) + filling_v[:] += tmp_v / abs_b_star_para * const # call the appropriate matvec filler - particle_to_mat_kernels.vec_fill_v1( + particle_to_mat_kernels.vec_fill_v0vec( args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] ) elif basis_u == 2: - linalg_kernels.matrix_matrix(b_prod, g_inv, tmp1) - linalg_kernels.matrix_matrix(tmp1, norm_b2_prod, tmp2) - linalg_kernels.matrix_matrix(tmp2, g_inv, tmp1) + # beq * gradPBeq contribution + linalg_kernels.matrix_matrix(beq_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PBeq, tmp_v) + + filling_v[:] = weight * tmp_v * mu / abs_b_star_para / det_df * ep_scale + + # beq * gradPB contribution + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + + filling_v[:] += weight * tmp_v * mu / abs_b_star_para / det_df * ep_scale + + # beq * dg term contribution + linalg_kernels.matrix_vector(tmp, eta_diff, tmp_v) + + filling_v[:] += tmp_v / abs_b_star_para / det_df * const - linalg_kernels.matrix_vector(tmp1, grad_PB, tmp_v) + # b * gradPBeq contribtuion + linalg_kernels.matrix_matrix(b_prod, norm_b_prod, tmp) + linalg_kernels.matrix_vector(tmp, grad_PBeq, tmp_v) - filling_v[:] = weight * tmp_v * mu / abs_b_star_para / det_df * scale_vec + filling_v[:] += weight * tmp_v * mu / abs_b_star_para / det_df * ep_scale + + # b * gradPB contribution + linalg_kernels.matrix_vector(tmp, grad_PB, tmp_v) + + filling_v[:] += weight * tmp_v * mu / abs_b_star_para / det_df * ep_scale + + # b * dg term contribution + linalg_kernels.matrix_vector(tmp, eta_diff, tmp_v) + + filling_v[:] += tmp_v / abs_b_star_para / det_df * const # call the appropriate matvec filler particle_to_mat_kernels.vec_fill_v2( @@ -856,5 +1224,3 @@ def cc_lin_mhd_5d_J2( vec1 /= Np vec2 /= Np vec3 /= Np - - # -- removed omp: #$ omp end parallel diff --git a/src/struphy/pic/accumulation/filter.py b/src/struphy/pic/accumulation/filter.py new file mode 100644 index 000000000..b80da8c7a --- /dev/null +++ b/src/struphy/pic/accumulation/filter.py @@ -0,0 +1,185 @@ +from dataclasses import dataclass + +import numpy as np +from scipy.fft import irfft, rfft + +from struphy.feec.psydac_derham import Derham +from struphy.io.options import OptsFilter +from struphy.pic.accumulation.filter_kernels import apply_three_point_filter_3d + + +@dataclass +class FilterParameters: + """Parameters for the AccumFilter class""" + + use_filter: OptsFilter | None = None + modes: tuple[int, ...] = (1,) + repeat: int = 1 + alpha: float = 0.5 + + +class AccumFilter: + """ + Callable filter that applies one of: + - 'fourier_in_tor' + - 'three_point' + - 'hybrid' (three_point, then fourier_in_tor) + """ + + def __init__(self, params: FilterParameters, derham: Derham, space_id: str): + self._params = params if params is not None else FilterParameters() + self._derham = derham + self._space_id = space_id + + self._form = derham.space_to_form[space_id] + self._form_int = 0 if self._form == "v" else int(self._form) + + @property + def params(self) -> FilterParameters: + return self._params + + @property + def derham(self): + """Discrete Derham complex on the logical unit cube.""" + return self._derham + + @property + def space_id(self): + """Space identifier for the matrix/vector (H1, Hcurl, Hdiv, L2 or H1vec) to be accumulated into.""" + return self._space_id + + @property + def form(self): + """p-form("0", "1", "2", "3") to be accumulated into.""" + return self._form + + @property + def form_int(self): + """Integer notation of p-form("0", "1", "2", "3") to be accumulated into.""" + return self._form_int + + def __call__(self, vec): + """ + Apply the chosen filter to `vec` in-place and return it. + + Parameters + ---------- + vec : BlockVector + Accumulated vector object. + """ + use = self.params.use_filter + if use is None: + return vec # nothing to do + + if use == "fourier_in_tor": + self._apply_toroidal_fourier_filter(vec, self._params.modes) + + elif use == "three_point": + self._apply_three_point(vec, repeat=self._params.repeat, alpha=self._params.alpha) + + elif use == "hybrid": + self._apply_three_point(vec, repeat=self._params.repeat, alpha=self._params.alpha) + self._apply_toroidal_fourier_filter(vec, self._params.modes) + + else: + raise NotImplementedError("The type of filter must be 'fourier_in_tor', 'three_point', or 'hybrid'.") + + return vec + + def _yield_dir_components(self, vec): + """ + Yields (axis, comp_vec, starts, ends) for each directions. + - For scalar accumulations ('H1','L2'): yields (0, vec, starts, ends). + - Otherwise: yields (axis, vec[axis], starts, ends) for axis=0,1,2. + """ + if self.space_id in ("H1", "L2"): + starts = self.derham.Vh[self.form].starts + ends = self.derham.Vh[self.form].ends + + yield 0, vec, starts, ends + + else: + for axis in range(3): + starts = self.derham.Vh[self.form][axis].starts + ends = self.derham.Vh[self.form][axis].ends + + yield axis, vec[axis], starts, ends + + def _apply_three_point(self, vec, repeat: int, alpha: float): + """ + Applying three point smoothing filter to the spline coefficients of the accumulated vector (``._data`` of the StencilVector): + + Parameters + ---------- + vec : BlockVector + + repeat : int + Number of repeatition. + + alpha : float + Alpha factor of the smoothing filter. + + """ + + for _ in range(repeat): + for axis, comp, starts, ends in self._yield_dir_components(vec): + apply_three_point_filter_3d( + comp._data, + axis, + self.form_int, + np.array(self.derham.Nel), + np.array(self.derham.spl_kind), + np.array(self.derham.p), + np.array(starts), + np.array(ends), + alpha=alpha, + ) + + vec.update_ghost_regions() + + def _apply_toroidal_fourier_filter(self, vec, modes: tuple[int, ...]): + """ + Applying fourier filter to the spline coefficients of the accumulated vector (toroidal direction). + + Parameters + ---------- + vec : BlockVector + + modes : tuple[int, ...] + Mode numbers which are not filtered out. + """ + + tor_Nel = self.derham.Nel[2] + modes = np.asarray(modes, dtype=int) + + assert tor_Nel >= 2 * int(np.max(modes)), "Nel[2] must be at least 2*max(modes)" + assert self.derham.domain_decomposition.nprocs[2] == 1, "No domain decomposition along toroidal direction" + + pn = np.asarray(self.derham.p, dtype=int) + ir = np.empty(3, dtype=int) + + # rfft output length + if (tor_Nel % 2) == 0: + vec_temp = np.zeros(int(tor_Nel / 2) + 1, dtype=complex) + else: + vec_temp = np.zeros(int((tor_Nel - 1) / 2) + 1, dtype=complex) + + for axis, comp, starts, ends in self._yield_dir_components(vec): + for i in range(3): + ir[i] = int(ends[i] + 1 - starts[i]) + + # filter along toroidal index (k direction) + for i in range(ir[0]): + ii = pn[0] + i + for j in range(ir[1]): + jj = pn[1] + j + + # forward FFT along toroidal line + line = rfft(comp._data[ii, jj, pn[2] : pn[2] + ir[2]]) + vec_temp[:] = 0 + vec_temp[modes] = line[modes] # keep selected modes only + + # inverse FFT back to real space, write in-place + comp._data[ii, jj, pn[2] : pn[2] + ir[2]] = irfft(vec_temp, n=tor_Nel) + + vec.update_ghost_regions() diff --git a/src/struphy/pic/accumulation/filter_kernels.py b/src/struphy/pic/accumulation/filter_kernels.py index e24c7ad5d..a6c498ca8 100644 --- a/src/struphy/pic/accumulation/filter_kernels.py +++ b/src/struphy/pic/accumulation/filter_kernels.py @@ -5,8 +5,10 @@ @stack_array("vec_copy", "mask1d", "mask", "top", "i_bottom", "i_top", "fi", "ir") -def apply_three_point_filter( +def apply_three_point_filter_3d( vec: "float[:,:,:]", + dir: "int", + form: "int", Nel: "int[:]", spl_kind: "bool[:]", pn: "int[:]", @@ -47,6 +49,7 @@ def apply_three_point_filter( i_top = zeros(3, dtype=int) fi = empty(3, dtype=int) ir = empty(3, dtype=int) + isDspline = zeros(3, dtype=int) # copy vectors vec_copy[:, :, :] = vec[:, :, :] @@ -62,22 +65,33 @@ def apply_three_point_filter( mask[i, j, k] *= mask1d[i] * mask1d[j] * mask1d[k] # consider left and right boundary + if form == 1: + isDspline[dir] = 1 + elif form == 2: + isDspline[:] = 1 + isDspline[dir] = 0 + elif form == 3: + isDspline[:] = 1 + for i in range(3): if spl_kind[i]: top[i] = Nel[i] - 1 else: - top[i] = Nel[i] + pn[i] - 1 + if isDspline[i] == 1: + top[i] = Nel[i] + pn[i] - 2 + else: + top[i] = Nel[i] + pn[i] - 1 for i in range(3): if starts[i] == 0: if spl_kind[i]: - i_bottom[i] = -1 + i_bottom[i] = 0 else: i_bottom[i] = +1 if ends[i] == top[i]: if spl_kind[i]: - i_top[i] = +1 + i_top[i] = 0 else: i_top[i] = -1 diff --git a/src/struphy/pic/accumulation/particles_to_grid.py b/src/struphy/pic/accumulation/particles_to_grid.py index 23345df23..6c9c0cc67 100644 --- a/src/struphy/pic/accumulation/particles_to_grid.py +++ b/src/struphy/pic/accumulation/particles_to_grid.py @@ -7,10 +7,10 @@ import struphy.pic.accumulation.accum_kernels as accums import struphy.pic.accumulation.accum_kernels_gc as accums_gc -import struphy.pic.accumulation.filter_kernels as filters from struphy.feec.mass import WeightedMassOperators from struphy.feec.psydac_derham import Derham from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments +from struphy.pic.accumulation.filter import AccumFilter, FilterParameters from struphy.pic.base import Particles from struphy.profiling.profiling import ProfileManager @@ -83,12 +83,7 @@ def __init__( *, add_vector: bool = False, symmetry: str = None, - filter_params: dict = { - "use_filter": None, - "modes": None, - "repeat": None, - "alpha": None, - }, + filter_params: FilterParameters = None, ): self._particles = particles self._space_id = space_id @@ -98,8 +93,6 @@ def __init__( self._symmetry = symmetry - self._filter_params = filter_params - self._form = self.derham.space_to_form[space_id] # initialize matrices (instances of WeightedMassOperator) @@ -176,6 +169,9 @@ def __init__( for bl in vec.blocks: self._args_data += (bl._data,) + # initialize filter + self._accfilter = AccumFilter(filter_params, self._derham, self._space_id) + def __call__(self, *optional_args, **args_control): """ Performs the accumulation into the matrix/vector by calling the chosen accumulation kernel and additional analytical contributions (control variate, optional). @@ -214,52 +210,13 @@ def __call__(self, *optional_args, **args_control): ) # apply filter - if self.filter_params["use_filter"] is not None: + if self.accfilter.params.use_filter is not None: for vec in self._vectors: vec.exchange_assembly_data() vec.update_ghost_regions() - if self.filter_params["use_filter"] == "fourier_in_tor": - self.apply_toroidal_fourier_filter(vec, self.filter_params["modes"]) - - elif self.filter_params["use_filter"] == "three_point": - for _ in range(self.filter_params["repeat"]): - for i in range(3): - filters.apply_three_point_filter( - vec[i]._data, - np.array(self.derham.Nel), - np.array(self.derham.spl_kind), - np.array(self.derham.p), - np.array(self.derham.Vh[self.form][i].starts), - np.array(self.derham.Vh[self.form][i].ends), - alpha=self.filter_params["alpha"], - ) - - vec.update_ghost_regions() - - elif self.filter_params["use_filter"] == "hybrid": - self.apply_toroidal_fourier_filter(vec, self.filter_params["modes"]) - - for _ in range(self.filter_params["repeat"]): - for i in range(2): - filters.apply_three_point_filter( - vec[i]._data, - np.array(self.derham.Nel), - np.array(self.derham.spl_kind), - np.array(self.derham.p), - np.array(self.derham.Vh[self.form][i].starts), - np.array(self.derham.Vh[self.form][i].ends), - alpha=self.filter_params["alpha"], - ) - - vec.update_ghost_regions() - - else: - raise NotImplemented( - "The type of filter must be fourier or three_point.", - ) - - vec_finished = True + self.accfilter(vec) + vec_finished = True if self.particles.clone_config is None: num_clones = 1 @@ -393,14 +350,9 @@ def vectors(self): return out @property - def filter_params(self): - """Dict of three components for the accumulation filter parameters: use_filter(string), repeat(int) and alpha(float).""" - return self._filter_params - - @property - def filter_params(self): - """Dict of three components for the accumulation filter parameters: use_filter(string), repeat(int) and alpha(float).""" - return self._filter_params + def accfilter(self): + """Callable filters""" + return self._accfilter def init_control_variate(self, mass_ops): """Set up the use of noise reduction by control variate.""" @@ -410,55 +362,6 @@ def init_control_variate(self, mass_ops): # L2 projector for dofs self._get_L2dofs = L2Projector(self.space_id, mass_ops).get_dofs - def apply_toroidal_fourier_filter(self, vec, modes): - """ - Applying fourier filter to the spline coefficients of the accumulated vector (toroidal direction). - - Parameters - ---------- - vec : BlockVector - - modes : list - Mode numbers which are not filtered out. - """ - - from scipy.fft import irfft, rfft - - tor_Nel = self.derham.Nel[2] - - # Nel along the toroidal direction must be equal or bigger than 2*maximum mode - assert tor_Nel >= 2 * max(modes) - - pn = self.derham.p - ir = np.empty(3, dtype=int) - - if (tor_Nel % 2) == 0: - vec_temp = np.zeros(int(tor_Nel / 2) + 1, dtype=complex) - else: - vec_temp = np.zeros(int((tor_Nel - 1) / 2) + 1, dtype=complex) - - # no domain decomposition along the toroidal direction - assert self.derham.domain_decomposition.nprocs[2] == 1 - - for axis in range(3): - starts = self.derham.Vh[ſelf.form][axis].starts - ends = self.derham.Vh[self.form][axis].ends - - # index range - for i in range(3): - ir[i] = ends[i] + 1 - starts[i] - - # filtering - for i in range(ir[0]): - for j in range(ir[1]): - vec_temp[:] = 0 - vec_temp[modes] = rfft( - vec[axis]._data[pn[0] + i, pn[1] + j, pn[2] : pn[2] + ir[2]], - )[modes] - vec[axis]._data[pn[0] + i, pn[1] + j, pn[2] : pn[2] + ir[2]] = irfft(vec_temp, n=tor_Nel) - - vec.update_ghost_regions() - def show_accumulated_spline_field(self, mass_ops: WeightedMassOperators, eta_direction=0, component=0): r"""1D plot of the spline field corresponding to the accumulated vector. The latter can be viewed as the rhs of an L2-projection: @@ -530,6 +433,7 @@ def __init__( kernel, mass_ops: WeightedMassOperators, args_domain: DomainArguments, + filter_params: FilterParameters = None, ): self._particles = particles self._space_id = space_id @@ -582,6 +486,9 @@ def __init__( for bl in vec.blocks: self._args_data += (bl._data,) + # initialize filter + self._accfilter = AccumFilter(filter_params, self._derham, self._space_id) + def __call__(self, *optional_args, **args_control): """ Performs the accumulation into the vector by calling the chosen accumulation kernel @@ -618,6 +525,15 @@ def __call__(self, *optional_args, **args_control): *optional_args, ) + # apply filter + if self.accfilter.params.use_filter is not None: + for vec in self._vectors: + vec.exchange_assembly_data() + vec.update_ghost_regions() + + self.accfilter(vec) + vec_finished = True + if self.particles.clone_config is None: num_clones = 1 else: @@ -687,6 +603,11 @@ def vectors(self): return out + @property + def accfilter(self): + """Callable filters""" + return self._accfilter + def init_control_variate(self, mass_ops): """Set up the use of noise reduction by control variate.""" diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 3489e701a..93f5c701e 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -350,6 +350,7 @@ def __init__( self._unit_b1_h = self.projected_equil.unit_b1 self._derham = self.projected_equil.derham + self._tmp0 = self.derham.Vh["0"].zeros() self._tmp2 = self.derham.Vh["2"].zeros() @property @@ -559,7 +560,7 @@ def save_constants_of_motion(self): self.absB0_h._data, ) - def save_magnetic_energy(self, b2): + def save_magnetic_energy(self, PBb): r""" Calculate magnetic field energy at each particles' position and assign it into markers[:,self.first_diagnostics_idx]. @@ -570,22 +571,17 @@ def save_magnetic_energy(self, b2): Finite element coefficients of the time-dependent magnetic field. """ - E2T = self.derham.extraction_ops["2"].transpose() - b2t = E2T.dot(b2, out=self._tmp2) - b2t.update_ghost_regions() + E0T = self.derham.extraction_ops["0"].transpose() + PBbt = E0T.dot(PBb, out=self._tmp0) + PBbt.update_ghost_regions() - utilities_kernels.eval_magnetic_energy( + utilities_kernels.eval_magnetic_energy_PBb( self.markers, self.derham.args_derham, self.domain.args_domain, self.first_diagnostics_idx, self.absB0_h._data, - self.unit_b1_h[0]._data, - self.unit_b1_h[1]._data, - self.unit_b1_h[2]._data, - b2t[0]._data, - b2t[1]._data, - b2t[2]._data, + PBbt._data, ) def save_magnetic_background_energy(self): diff --git a/src/struphy/pic/pushing/pusher_kernels_gc.py b/src/struphy/pic/pushing/pusher_kernels_gc.py index 0b6c9b3c7..5dfee707b 100644 --- a/src/struphy/pic/pushing/pusher_kernels_gc.py +++ b/src/struphy/pic/pushing/pusher_kernels_gc.py @@ -1896,7 +1896,7 @@ def push_gc_cc_J1_H1vec( ) # b_star; in H1vec - b_star[:] = (b + curl_norm_b * v * epsilon) / det_df + b_star[:] = b + curl_norm_b * v * epsilon # calculate abs_b_star_para abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) @@ -1905,7 +1905,7 @@ def push_gc_cc_J1_H1vec( linalg_kernels.cross(b, u, e) # curl_norm_b dot electric field - temp = linalg_kernels.scalar_dot(e, curl_norm_b) / det_df + temp = linalg_kernels.scalar_dot(e, curl_norm_b) markers[ip, 3] += temp / abs_b_star_para * v * dt @@ -2077,7 +2077,6 @@ def push_gc_cc_J1_Hdiv( u1: "float[:,:,:]", u2: "float[:,:,:]", u3: "float[:,:,:]", - boundary_cut: float, ): r"""Velocity update step for the `CurrentCoupling5DCurlb `_ @@ -2105,8 +2104,6 @@ def push_gc_cc_J1_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers - # -- removed omp: #$ omp parallel private(ip, boundary_cut, eta1, eta2, eta3, v, det_df, dfm, span1, span2, span3, b, u, e, curl_norm_b, norm_b1, b_star, temp, abs_b_star_para) - # -- removed omp: #$ omp for for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) if markers[ip, 0] == -1.0: @@ -2117,9 +2114,6 @@ def push_gc_cc_J1_Hdiv( eta3 = markers[ip, 2] v = markers[ip, 3] - if eta1 < boundary_cut or eta1 > 1.0 - boundary_cut: - continue - # evaluate Jacobian, result in dfm evaluation_kernels.df( eta1, @@ -2183,10 +2177,10 @@ def push_gc_cc_J1_Hdiv( curl_norm_b, ) - # b_star; 2form in H1vec - b_star[:] = (b + curl_norm_b * v * epsilon) / det_df + # b_star; 2form + b_star[:] = b + curl_norm_b * v * epsilon - # calculate abs_b_star_para + # calculate 3form abs_b_star_para abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) # transform u into H1vec @@ -2196,12 +2190,10 @@ def push_gc_cc_J1_Hdiv( linalg_kernels.cross(b, u, e) # curl_norm_b dot electric field - temp = linalg_kernels.scalar_dot(e, curl_norm_b) / det_df + temp = linalg_kernels.scalar_dot(e, curl_norm_b) markers[ip, 3] += temp / abs_b_star_para * v * dt - # -- removed omp: #$ omp end parallel - @stack_array( "dfm", @@ -2212,13 +2204,11 @@ def push_gc_cc_J1_Hdiv( "u", "bb", "b_star", - "norm_b1", - "norm_b2", + "norm_b", "curl_norm_b", - "tmp1", - "tmp2", + "tmp", "b_prod", - "norm_b2_prod", + "norm_b_prod", ) def push_gc_cc_J2_stage_H1vec( dt: float, @@ -2233,9 +2223,6 @@ def push_gc_cc_J2_stage_H1vec( norm_b11: "float[:,:,:]", norm_b12: "float[:,:,:]", norm_b13: "float[:,:,:]", - norm_b21: "float[:,:,:]", - norm_b22: "float[:,:,:]", - norm_b23: "float[:,:,:]", curl_norm_b1: "float[:,:,:]", curl_norm_b2: "float[:,:,:]", curl_norm_b3: "float[:,:,:]", @@ -2264,16 +2251,14 @@ def push_gc_cc_J2_stage_H1vec( g_inv = empty((3, 3), dtype=float) # containers for fields - tmp1 = empty((3, 3), dtype=float) - tmp2 = empty((3, 3), dtype=float) + tmp = empty((3, 3), dtype=float) b_prod = zeros((3, 3), dtype=float) - norm_b2_prod = empty((3, 3), dtype=float) + norm_b_prod = empty((3, 3), dtype=float) e = empty(3, dtype=float) u = empty(3, dtype=float) bb = empty(3, dtype=float) b_star = empty(3, dtype=float) norm_b1 = empty(3, dtype=float) - norm_b2 = empty(3, dtype=float) curl_norm_b = empty(3, dtype=float) # get marker arguments @@ -2354,18 +2339,6 @@ def push_gc_cc_J2_stage_H1vec( norm_b1, ) - # norm_b; 2form - eval_2form_spline_mpi( - span1, - span2, - span3, - args_derham, - norm_b21, - norm_b22, - norm_b23, - norm_b2, - ) - # curl_norm_b; 2form eval_2form_spline_mpi( span1, @@ -2386,24 +2359,21 @@ def push_gc_cc_J2_stage_H1vec( b_prod[2, 0] = -bb[1] b_prod[2, 1] = +bb[0] - norm_b2_prod[0, 1] = -norm_b2[2] - norm_b2_prod[0, 2] = +norm_b2[1] - norm_b2_prod[1, 0] = +norm_b2[2] - norm_b2_prod[1, 2] = -norm_b2[0] - norm_b2_prod[2, 0] = -norm_b2[1] - norm_b2_prod[2, 1] = +norm_b2[0] + norm_b_prod[0, 1] = -norm_b1[2] + norm_b_prod[0, 2] = +norm_b1[1] + norm_b_prod[1, 0] = +norm_b1[2] + norm_b_prod[1, 2] = -norm_b1[0] + norm_b_prod[2, 0] = -norm_b1[1] + norm_b_prod[2, 1] = +norm_b1[0] # b_star; 2form in H1vec - b_star[:] = (bb + curl_norm_b * v * epsilon) / det_df + b_star[:] = bb + curl_norm_b * v * epsilon - # calculate abs_b_star_para + # calculate 3form abs_b_star_para abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) - linalg_kernels.matrix_matrix(g_inv, norm_b2_prod, tmp1) - linalg_kernels.matrix_matrix(tmp1, g_inv, tmp2) - linalg_kernels.matrix_matrix(tmp2, b_prod, tmp1) - - linalg_kernels.matrix_vector(tmp1, u, e) + linalg_kernels.matrix_matrix(norm_b_prod, b_prod, tmp) + linalg_kernels.matrix_vector(tmp, u, e) e /= abs_b_star_para @@ -2428,12 +2398,10 @@ def push_gc_cc_J2_stage_H1vec( "bb", "b_star", "norm_b1", - "norm_b2", "curl_norm_b", - "tmp1", - "tmp2", + "tmp", "b_prod", - "norm_b2_prod", + "norm_b_prod", ) def push_gc_cc_J2_stage_Hdiv( dt: float, @@ -2448,9 +2416,6 @@ def push_gc_cc_J2_stage_Hdiv( norm_b11: "float[:,:,:]", norm_b12: "float[:,:,:]", norm_b13: "float[:,:,:]", - norm_b21: "float[:,:,:]", - norm_b22: "float[:,:,:]", - norm_b23: "float[:,:,:]", curl_norm_b1: "float[:,:,:]", curl_norm_b2: "float[:,:,:]", curl_norm_b3: "float[:,:,:]", @@ -2460,7 +2425,6 @@ def push_gc_cc_J2_stage_Hdiv( a: "float[:]", b: "float[:]", c: "float[:]", - boundary_cut: float, ): r"""Single stage of a s-stage explicit pushing step for the `CurrentCoupling5DGradB `_ @@ -2480,16 +2444,14 @@ def push_gc_cc_J2_stage_Hdiv( g_inv = empty((3, 3), dtype=float) # containers for fields - tmp1 = zeros((3, 3), dtype=float) - tmp2 = zeros((3, 3), dtype=float) + tmp = zeros((3, 3), dtype=float) b_prod = zeros((3, 3), dtype=float) - norm_b2_prod = zeros((3, 3), dtype=float) + norm_b_prod = zeros((3, 3), dtype=float) e = empty(3, dtype=float) u = empty(3, dtype=float) bb = empty(3, dtype=float) b_star = empty(3, dtype=float) norm_b1 = empty(3, dtype=float) - norm_b2 = empty(3, dtype=float) curl_norm_b = empty(3, dtype=float) # get marker arguments @@ -2507,8 +2469,6 @@ def push_gc_cc_J2_stage_Hdiv( else: last = 0.0 - # -- removed omp: #$ omp parallel firstprivate(b_prod, norm_b2_prod) private(ip, boundary_cut, eta1, eta2, eta3, v, det_df, dfm, df_inv, df_inv_t, g_inv, span1, span2, span3, bb, u, e, curl_norm_b, norm_b1, norm_b2, b_star, tmp1, tmp2, abs_b_star_para) - # -- removed omp: #$ omp for for ip in range(n_markers): # check if marker is a hole if markers[ip, first_init_idx] == -1.0: @@ -2519,9 +2479,6 @@ def push_gc_cc_J2_stage_Hdiv( eta3 = markers[ip, 2] v = markers[ip, 3] - if eta1 < boundary_cut or eta2 > 1.0 - boundary_cut: - continue - # evaluate Jacobian, result in dfm evaluation_kernels.df( eta1, @@ -2576,18 +2533,6 @@ def push_gc_cc_J2_stage_Hdiv( norm_b1, ) - # norm_b; 2form - eval_2form_spline_mpi( - span1, - span2, - span3, - args_derham, - norm_b21, - norm_b22, - norm_b23, - norm_b2, - ) - # curl_norm_b; 2form eval_2form_spline_mpi( span1, @@ -2608,24 +2553,21 @@ def push_gc_cc_J2_stage_Hdiv( b_prod[2, 0] = -bb[1] b_prod[2, 1] = +bb[0] - norm_b2_prod[0, 1] = -norm_b2[2] - norm_b2_prod[0, 2] = +norm_b2[1] - norm_b2_prod[1, 0] = +norm_b2[2] - norm_b2_prod[1, 2] = -norm_b2[0] - norm_b2_prod[2, 0] = -norm_b2[1] - norm_b2_prod[2, 1] = +norm_b2[0] + norm_b_prod[0, 1] = -norm_b1[2] + norm_b_prod[0, 2] = +norm_b1[1] + norm_b_prod[1, 0] = +norm_b1[2] + norm_b_prod[1, 2] = -norm_b1[0] + norm_b_prod[2, 0] = -norm_b1[1] + norm_b_prod[2, 1] = +norm_b1[0] - # b_star; 2form in H1vec - b_star[:] = (bb + curl_norm_b * v * epsilon) / det_df + # b_star; 2form + b_star[:] = bb + curl_norm_b * v * epsilon # calculate abs_b_star_para abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) - linalg_kernels.matrix_matrix(g_inv, norm_b2_prod, tmp1) - linalg_kernels.matrix_matrix(tmp1, g_inv, tmp2) - linalg_kernels.matrix_matrix(tmp2, b_prod, tmp1) - - linalg_kernels.matrix_vector(tmp1, u, e) + linalg_kernels.matrix_matrix(norm_b_prod, b_prod, tmp) + linalg_kernels.matrix_vector(tmp, u, e) e /= abs_b_star_para e /= det_df @@ -2640,4 +2582,367 @@ def push_gc_cc_J2_stage_Hdiv( + last * markers[ip, first_free_idx : first_free_idx + 3] ) - # -- removed omp: #$ omp end parallel + +@stack_array( + "dfm", + "df_inv", + "df_inv_t", + "g_inv", + "e", + "u", + "bb", + "b_star", + "norm_b1", + "curl_norm_b", + "tmp1", + "b_prod", + "norm_b_prod", +) +def push_gc_cc_J2_dg_init_Hdiv( + dt: float, + args_markers: "MarkerArguments", + args_domain: "DomainArguments", + args_derham: "DerhamArguments", + epsilon: float, + b1: "float[:,:,:]", + b2: "float[:,:,:]", + b3: "float[:,:,:]", + norm_b11: "float[:,:,:]", + norm_b12: "float[:,:,:]", + norm_b13: "float[:,:,:]", + curl_norm_b1: "float[:,:,:]", + curl_norm_b2: "float[:,:,:]", + curl_norm_b3: "float[:,:,:]", + u1: "float[:,:,:]", + u2: "float[:,:,:]", + u3: "float[:,:,:]", +): + r"""TODO""" + + # allocate metric coeffs + dfm = empty((3, 3), dtype=float) + df_inv = empty((3, 3), dtype=float) + df_inv_t = empty((3, 3), dtype=float) + g_inv = empty((3, 3), dtype=float) + + # containers for fields + tmp1 = zeros((3, 3), dtype=float) + b_prod = zeros((3, 3), dtype=float) + norm_b_prod = zeros((3, 3), dtype=float) + e = empty(3, dtype=float) + u = empty(3, dtype=float) + bb = empty(3, dtype=float) + b_star = empty(3, dtype=float) + norm_b1 = empty(3, dtype=float) + curl_norm_b = empty(3, dtype=float) + + # get marker arguments + markers = args_markers.markers + n_markers = args_markers.n_markers + mu_idx = args_markers.mu_idx + first_init_idx = args_markers.first_init_idx + first_free_idx = args_markers.first_free_idx + + for ip in range(n_markers): + # check if marker is a hole + if markers[ip, first_init_idx] == -1.0: + continue + + eta1 = markers[ip, 0] + eta2 = markers[ip, 1] + eta3 = markers[ip, 2] + v = markers[ip, 3] + + # evaluate Jacobian, result in dfm + evaluation_kernels.df( + eta1, + eta2, + eta3, + args_domain, + dfm, + ) + + # metric coeffs + det_df = linalg_kernels.det(dfm) + linalg_kernels.matrix_inv_with_det(dfm, det_df, df_inv) + linalg_kernels.transpose(df_inv, df_inv_t) + linalg_kernels.matrix_matrix(df_inv, df_inv_t, g_inv) + + # spline evaluation + span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) + + # b; 2form + eval_2form_spline_mpi( + span1, + span2, + span3, + args_derham, + b1, + b2, + b3, + bb, + ) + + # u; 2form + eval_2form_spline_mpi( + span1, + span2, + span3, + args_derham, + u1, + u2, + u3, + u, + ) + + # norm_b1; 1form + eval_1form_spline_mpi( + span1, + span2, + span3, + args_derham, + norm_b11, + norm_b12, + norm_b13, + norm_b1, + ) + + # curl_norm_b; 2form + eval_2form_spline_mpi( + span1, + span2, + span3, + args_derham, + curl_norm_b1, + curl_norm_b2, + curl_norm_b3, + curl_norm_b, + ) + + # operator bx() as matrix + b_prod[0, 1] = -bb[2] + b_prod[0, 2] = +bb[1] + b_prod[1, 0] = +bb[2] + b_prod[1, 2] = -bb[0] + b_prod[2, 0] = -bb[1] + b_prod[2, 1] = +bb[0] + + norm_b_prod[0, 1] = -norm_b1[2] + norm_b_prod[0, 2] = +norm_b1[1] + norm_b_prod[1, 0] = +norm_b1[2] + norm_b_prod[1, 2] = -norm_b1[0] + norm_b_prod[2, 0] = -norm_b1[1] + norm_b_prod[2, 1] = +norm_b1[0] + + # b_star; 2form + b_star[:] = bb + curl_norm_b * v * epsilon + + # calculate 3form abs_b_star_para + abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) + + linalg_kernels.matrix_matrix(norm_b_prod, b_prod, tmp1) + linalg_kernels.matrix_vector(tmp1, u, e) + + e /= abs_b_star_para + e /= det_df + + markers[ip, 0:3] -= dt * e + + +@stack_array( + "dfm", + "df_inv", + "df_inv_t", + "g_inv", + "e", + "u", + "ud", + "bb", + "b_star", + "norm_b1", + "curl_norm_b", + "tmp1", + "tmp2", + "b_prod", + "norm_b_prod", + "eta_old", + "eta_mid", +) +def push_gc_cc_J2_dg_Hdiv( + dt: float, + args_markers: "MarkerArguments", + args_domain: "DomainArguments", + args_derham: "DerhamArguments", + epsilon: float, + b1: "float[:,:,:]", + b2: "float[:,:,:]", + b3: "float[:,:,:]", + norm_b11: "float[:,:,:]", + norm_b12: "float[:,:,:]", + norm_b13: "float[:,:,:]", + curl_norm_b1: "float[:,:,:]", + curl_norm_b2: "float[:,:,:]", + curl_norm_b3: "float[:,:,:]", + u1: "float[:,:,:]", + u2: "float[:,:,:]", + u3: "float[:,:,:]", + ud1: "float[:,:,:]", + ud2: "float[:,:,:]", + ud3: "float[:,:,:]", + const: float, + alpha: float, +): + r"""TODO""" + + # allocate metric coeffs + dfm = empty((3, 3), dtype=float) + df_inv = empty((3, 3), dtype=float) + df_inv_t = empty((3, 3), dtype=float) + g_inv = empty((3, 3), dtype=float) + + # containers for fields + tmp1 = zeros((3, 3), dtype=float) + tmp2 = zeros(3, dtype=float) + b_prod = zeros((3, 3), dtype=float) + norm_b_prod = zeros((3, 3), dtype=float) + e = empty(3, dtype=float) + u = empty(3, dtype=float) + ud = empty(3, dtype=float) + bb = empty(3, dtype=float) + b_star = empty(3, dtype=float) + norm_b1 = empty(3, dtype=float) + curl_norm_b = empty(3, dtype=float) + eta_old = empty(3, dtype=float) + eta_mid = empty(3, dtype=float) + + # get marker arguments + markers = args_markers.markers + n_markers = args_markers.n_markers + mu_idx = args_markers.mu_idx + first_init_idx = args_markers.first_init_idx + first_free_idx = args_markers.first_free_idx + + for ip in range(n_markers): + # check if marker is a hole + if markers[ip, 0] == -1.0: + continue + + # marker positions, mid point + eta_old[:] = markers[ip, 0:3] + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = mod(eta_mid[:], 1.0) + + v = markers[ip, 3] + + # evaluate Jacobian, result in dfm + evaluation_kernels.df( + eta_mid[0], + eta_mid[1], + eta_mid[2], + args_domain, + dfm, + ) + + # metric coeffs + det_df = linalg_kernels.det(dfm) + linalg_kernels.matrix_inv_with_det(dfm, det_df, df_inv) + linalg_kernels.transpose(df_inv, df_inv_t) + linalg_kernels.matrix_matrix(df_inv, df_inv_t, g_inv) + + # spline evaluation + span1, span2, span3 = get_spans(eta_mid[0], eta_mid[1], eta_mid[2], args_derham) + + # b; 2form + eval_2form_spline_mpi( + span1, + span2, + span3, + args_derham, + b1, + b2, + b3, + bb, + ) + + # u; 2form + eval_2form_spline_mpi( + span1, + span2, + span3, + args_derham, + u1, + u2, + u3, + u, + ) + + # ud; 2form + eval_2form_spline_mpi( + span1, + span2, + span3, + args_derham, + ud1, + ud2, + ud3, + ud, + ) + + # norm_b1; 1form + eval_1form_spline_mpi( + span1, + span2, + span3, + args_derham, + norm_b11, + norm_b12, + norm_b13, + norm_b1, + ) + + # curl_norm_b; 2form + eval_2form_spline_mpi( + span1, + span2, + span3, + args_derham, + curl_norm_b1, + curl_norm_b2, + curl_norm_b3, + curl_norm_b, + ) + + # operator bx() as matrix + b_prod[0, 1] = -bb[2] + b_prod[0, 2] = +bb[1] + b_prod[1, 0] = +bb[2] + b_prod[1, 2] = -bb[0] + b_prod[2, 0] = -bb[1] + b_prod[2, 1] = +bb[0] + + norm_b_prod[0, 1] = -norm_b1[2] + norm_b_prod[0, 2] = +norm_b1[1] + norm_b_prod[1, 0] = +norm_b1[2] + norm_b_prod[1, 2] = -norm_b1[0] + norm_b_prod[2, 0] = -norm_b1[1] + norm_b_prod[2, 1] = +norm_b1[0] + + # b_star; 2form + b_star[:] = bb + curl_norm_b * v * epsilon + + # calculate 3form abs_b_star_para + abs_b_star_para = linalg_kernels.scalar_dot(norm_b1, b_star) + + linalg_kernels.matrix_matrix(norm_b_prod, b_prod, tmp1) + linalg_kernels.matrix_vector(tmp1, u, e) + linalg_kernels.matrix_vector(tmp1, ud, tmp2) + tmp2 *= const + + e += tmp2 + + e /= abs_b_star_para + e /= det_df + + markers[ip, 0:3] = markers[ip, first_init_idx : first_init_idx + 3] - dt * e + markers[ip, 0:3] *= alpha + markers[ip, 0:3] += eta_old * (1.0 - alpha) diff --git a/src/struphy/pic/utilities_kernels.py b/src/struphy/pic/utilities_kernels.py index d0f3c4e92..cb25cc05f 100644 --- a/src/struphy/pic/utilities_kernels.py +++ b/src/struphy/pic/utilities_kernels.py @@ -1,4 +1,4 @@ -from numpy import abs, empty, log, pi, shape, sign, sqrt, zeros +from numpy import abs, empty, log, mod, pi, shape, sign, sqrt, zeros from pyccel.decorators import stack_array import struphy.bsplines.bsplines_kernels as bsplines_kernels @@ -14,7 +14,7 @@ eval_vectorfield_spline_mpi, get_spans, ) -from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments +from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments, MarkerArguments def eval_magnetic_moment_5d( @@ -331,6 +331,71 @@ def eval_magnetic_energy( # -- removed omp: #$ omp end parallel +@stack_array("dfm", "eta") +def eval_magnetic_energy_PBb( + markers: "float[:,:]", + args_derham: "DerhamArguments", + args_domain: "DomainArguments", + first_diagnostics_idx: int, + abs_B0: "float[:,:,:]", + PBb: "float[:,:,:]", +): + r""" + Evaluate :math:`mu_p |B(\boldsymbol \eta_p)_\parallel|` for each marker. + The result is stored at markers[:, first_diagnostics_idx]. + """ + eta = empty(3, dtype=float) + + dfm = empty((3, 3), dtype=float) + + # get number of markers + n_markers = shape(markers)[0] + + for ip in range(n_markers): + # only do something if particle is a "true" particle (i.e. not a hole) + if markers[ip, 0] == -1.0: + continue + + eta[:] = mod(markers[ip, 0:3], 1.0) + + weight = markers[ip, 7] + dweight = markers[ip, 5] + + mu = markers[ip, first_diagnostics_idx + 1] + + # spline evaluation + span1, span2, span3 = get_spans(eta[0], eta[1], eta[2], args_derham) + + # evaluate Jacobian, result in dfm + evaluation_kernels.df( + eta[0], + eta[1], + eta[2], + args_domain, + dfm, + ) + + # abs_B0; 0form + abs_B = eval_0form_spline_mpi( + span1, + span2, + span3, + args_derham, + abs_B0, + ) + + # PBb; 0form + PB_b = eval_0form_spline_mpi( + span1, + span2, + span3, + args_derham, + PBb, + ) + + markers[ip, first_diagnostics_idx] = mu * (abs_B + PB_b) + + @stack_array("v", "dfm", "b2", "norm_b_cart", "temp", "v_perp", "Larmor_r") def eval_guiding_center_from_6d( markers: "float[:,:]", @@ -441,189 +506,101 @@ def eval_guiding_center_from_6d( markers[ip, first_diagnostics_idx + 2] = z - Larmor_r[2] -@stack_array("grad_PB", "tmp") -def accum_gradI_const( - markers: "float[:,:]", - Np: "int", +@stack_array("dfm", "df_t", "g", "g_inv", "gradB, grad_PB_b", "tmp", "eta_mid", "eta_diff") +def eval_gradB_ediff( + args_markers: "MarkerArguments", + args_domain: "DomainArguments", args_derham: "DerhamArguments", - grad_PB1: "float[:,:,:]", - grad_PB2: "float[:,:,:]", - grad_PB3: "float[:,:,:]", - scale: "float", + gradB1: "float[:,:,:]", + gradB2: "float[:,:,:]", + gradB3: "float[:,:,:]", + grad_PB_b1: "float[:,:,:]", + grad_PB_b2: "float[:,:,:]", + grad_PB_b3: "float[:,:,:]", + idx: int, ): r"""TODO""" + + # allocate metric coeffs + dfm = empty((3, 3), dtype=float) + df_t = empty((3, 3), dtype=float) + g = empty((3, 3), dtype=float) + g_inv = empty((3, 3), dtype=float) + # allocate for magnetic field evaluation - grad_PB = empty(3, dtype=float) + gradB = empty(3, dtype=float) + grad_PB_b = empty(3, dtype=float) tmp = empty(3, dtype=float) + eta_mid = empty(3, dtype=float) + eta_diff = empty(3, dtype=float) - # allocate for filling - res = zeros(1, dtype=float) - - # get number of markers - n_markers_loc = shape(markers)[0] + # get marker arguments + markers = args_markers.markers + n_markers = args_markers.n_markers + mu_idx = args_markers.mu_idx + first_init_idx = args_markers.first_init_idx + first_free_idx = args_markers.first_free_idx - for ip in range(n_markers_loc): + for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) if markers[ip, 0] == -1.0: continue - # marker positions - eta1 = markers[ip, 0] # mid - eta2 = markers[ip, 1] # mid - eta3 = markers[ip, 2] # mid + # marker positions, mid point + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = mod(eta_mid[:], 1.0) + + eta_diff = markers[ip, 0:3] - markers[ip, first_init_idx : first_init_idx + 3] # marker weight and velocity weight = markers[ip, 5] - mu = markers[ip, 9] + mu = markers[ip, mu_idx] # b-field evaluation - span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) + span1, span2, span3 = get_spans(eta_mid[0], eta_mid[1], eta_mid[2], args_derham) + # print(span1, span2, span3) - # grad_PB; 1form + # evaluate Jacobian, result in dfm + evaluation_kernels.df( + eta_mid[0], + eta_mid[1], + eta_mid[2], + args_domain, + dfm, + ) + + linalg_kernels.transpose(dfm, df_t) + linalg_kernels.matrix_matrix(df_t, dfm, g) + linalg_kernels.matrix_inv(g, g_inv) + + # gradB; 1form eval_1form_spline_mpi( span1, span2, span3, args_derham, - grad_PB1, - grad_PB2, - grad_PB3, - grad_PB, + gradB1, + gradB2, + gradB3, + gradB, ) - tmp[:] = markers[ip, 15:18] - res += linalg_kernels.scalar_dot(tmp, grad_PB) * weight * mu * scale - - return res / Np - - -def accum_en_fB( - markers: "float[:,:]", - Np: "int", - args_derham: "DerhamArguments", - PB: "float[:,:,:]", -): - r"""TODO""" - - # allocate for filling - res = zeros(1, dtype=float) - - # get number of markers - n_markers_loc = shape(markers)[0] - - for ip in range(n_markers_loc): - # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: - continue - - # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - - # marker weight and velocity - mu = markers[ip, 9] - weight = markers[ip, 5] - - # b-field evaluation - span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) - - B0 = eval_0form_spline_mpi( + # grad_PB_b; 1form + eval_1form_spline_mpi( span1, span2, span3, args_derham, - PB, + grad_PB_b1, + grad_PB_b2, + grad_PB_b3, + grad_PB_b, ) - res += abs(B0) * mu * weight - - return res / Np - - -@stack_array("e", "e_diff") -def check_eta_diff(markers: "float[:,:]"): - r"""TODO""" - # marker position e - e = empty(3, dtype=float) - e_diff = empty(3, dtype=float) - - # get number of markers - n_markers_loc = shape(markers)[0] - - for ip in range(n_markers_loc): - # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: - continue - - e[:] = markers[ip, 0:3] - e_diff[:] = e[:] - markers[ip, 9:12] - - for axis in range(3): - if e_diff[axis] > 0.5: - e_diff[axis] -= 1.0 - elif e_diff[axis] < -0.5: - e_diff[axis] += 1.0 - - markers[ip, 15:18] = e_diff[:] - - -@stack_array("e", "e_diff") -def check_eta_diff2(markers: "float[:,:]"): - r"""TODO""" - # marker position e - e = empty(3, dtype=float) - e_diff = empty(3, dtype=float) - - # get number of markers - n_markers_loc = shape(markers)[0] - - for ip in range(n_markers_loc): - # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: - continue - - e[:] = markers[ip, 0:3] - e_diff[:] = e[:] - markers[ip, 12:15] - - for axis in range(3): - if e_diff[axis] > 0.5: - e_diff[axis] -= 1.0 - elif e_diff[axis] < -0.5: - e_diff[axis] += 1.0 - - markers[ip, 15:18] = e_diff[:] - - -@stack_array("e", "e_diff", "e_mid") -def check_eta_mid(markers: "float[:,:]"): - r"""TODO""" - # marker position e - e = empty(3, dtype=float) - e_diff = empty(3, dtype=float) - e_mid = empty(3, dtype=float) - - # get number of markers - n_markers_loc = shape(markers)[0] - - for ip in range(n_markers_loc): - # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: - continue - - e[:] = markers[ip, 0:3] - markers[ip, 12:15] = e[:] - - e_diff[:] = e[:] - markers[ip, 9:12] - e_mid[:] = (e[:] + markers[ip, 9:12]) / 2.0 - - for axis in range(3): - if e_diff[axis] > 0.5: - e_mid[axis] += 0.5 - elif e_diff[axis] < -0.5: - e_mid[axis] += 0.5 + tmp = gradB + grad_PB_b - markers[ip, 0:3] = e_mid[:] + markers[ip, idx] = linalg_kernels.scalar_dot(eta_diff, tmp) + markers[ip, idx] *= mu @stack_array("dfm", "dfinv", "dfinv_t", "v", "a_form", "dfta_form") diff --git a/src/struphy/propagators/__init__.py b/src/struphy/propagators/__init__.py index d8f1bd3bd..04418745c 100644 --- a/src/struphy/propagators/__init__.py +++ b/src/struphy/propagators/__init__.py @@ -67,7 +67,6 @@ # "FaradayExtended", # "CurrentCoupling6DDensity", # "ShearAlfvenCurrentCoupling5D", -# "MagnetosonicCurrentCoupling5D", # "CurrentCoupling5DDensity", # "ImplicitDiffusion", # "Poisson", diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index cd629ffe4..5d2391534 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -7,6 +7,7 @@ from line_profiler import profile from mpi4py import MPI from psydac.linalg.block import BlockVector +from psydac.linalg.solvers import inverse from psydac.linalg.stencil import StencilVector from struphy.feec import preconditioner @@ -16,10 +17,13 @@ from struphy.kinetic_background.base import Maxwellian from struphy.kinetic_background.maxwellians import Maxwellian3D from struphy.linear_algebra.schur_solver import SchurSolver -from struphy.linear_algebra.solver import SolverParameters +from struphy.linear_algebra.solver import DiscreteGradientSolverParameters, SolverParameters from struphy.models.variables import FEECVariable, PICVariable +from struphy.ode.utils import ButcherTableau +from struphy.pic import utilities_kernels from struphy.pic.accumulation import accum_kernels, accum_kernels_gc -from struphy.pic.accumulation.particles_to_grid import Accumulator +from struphy.pic.accumulation.filter import FilterParameters +from struphy.pic.accumulation.particles_to_grid import Accumulator, AccumulatorVector from struphy.pic.particles import Particles5D, Particles6D from struphy.pic.pushing import pusher_kernels, pusher_kernels_gc from struphy.pic.pushing.pusher import Pusher @@ -1130,288 +1134,198 @@ class CurrentCoupling5DCurlb(Propagator): For the detail explanation of the notations, see `2022_DriftKineticCurrentCoupling `_. """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["filter"] = { - "use_filter": None, - "modes": (1), - "repeat": 1, - "alpha": 0.5, - } - dct["boundary_cut"] = { - "e1": 0.0, - "e2": 0.0, - "e3": 0.0, - } - dct["turn_off"] = False + class Variables: + def __init__(self): + self._u: FEECVariable = None + self._energetic_ions: PICVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def u(self) -> FEECVariable: + return self._u - return dct + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new - def __init__( - self, - particles: Particles5D, - u: BlockVector, - *, - b: BlockVector, - b_eq: BlockVector, - unit_b1: BlockVector, - absB0: StencilVector, - gradB1: BlockVector, - curl_unit_b2: BlockVector, - u_space: str, - solver: dict = options(default=True)["solver"], - filter: dict = options(default=True)["filter"], - coupling_params: dict, - epsilon: float = 1.0, - boundary_cut: dict = options(default=True)["boundary_cut"], - ): - super().__init__(particles, u) + @property + def energetic_ions(self) -> PICVariable: + return self._energetic_ions - assert u_space in {"Hcurl", "Hdiv", "H1vec"} + @energetic_ions.setter + def energetic_ions(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles5D" + self._energetic_ions = new - if u_space == "H1vec": - self._space_key_int = 0 - else: - self._space_key_int = int( - self.derham.space_to_form[u_space], - ) + def __init__(self): + self.variables = self.Variables() - self._epsilon = epsilon - self._b = b - self._b_eq = b_eq - self._unit_b1 = unit_b1 - self._absB0 = absB0 - self._gradB1 = gradB1 - self._curl_norm_b = curl_unit_b2 + @dataclass + class Options: + # propagator options + b_tilde: FEECVariable = None + ep_scale: float = 1.0 + u_space: OptsVecSpace = "Hdiv" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + filter_params: FilterParameters = None - self._info = solver["info"] - self._rank = self.derham.comm.Get_rank() + def __post_init__(self): + # checks + check_option(self.u_space, OptsVecSpace) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + assert isinstance(self.b_tilde, FEECVariable) + assert isinstance(self.ep_scale, float) - self._coupling_mat = coupling_params["Ah"] / coupling_params["Ab"] - self._coupling_vec = coupling_params["Ah"] / coupling_params["Ab"] - self._scale_push = 1 + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() - self._boundary_cut_e1 = boundary_cut["e1"] + if self.filter_params is None: + self.filter_params = FilterParameters() - u_id = self.derham.space_to_form[u_space] - self._E0T = self.derham.extraction_ops["0"].transpose() - self._EuT = self.derham.extraction_ops[u_id].transpose() - self._E2T = self.derham.extraction_ops["2"].transpose() - self._E1T = self.derham.extraction_ops["1"].transpose() + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options - self._unit_b1 = self._E1T.dot(self._unit_b1) - self._curl_norm_b = self._E2T.dot(self._curl_norm_b) - self._curl_norm_b.update_ghost_regions() - self._absB0 = self._E0T.dot(self._absB0) + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new - # define system [[A B], [C I]] [u_new, v_new] = [[A -B], [-C I]] [u_old, v_old] (without time step size dt) - _A = getattr(self.mass_ops, "M" + u_id + "n") + @profile + def allocate(self): + if self.options.u_space == "H1vec": + self._u_form_int = 0 + else: + self._u_form_int = int(self.derham.space_to_form[self.options.u_space]) - # preconditioner - if solver["type"][1] is None: + # call operatros + id_M = "M" + self.derham.space_to_form[self.options.u_space] + "n" + _A = getattr(self.mass_ops, id_M) + + # Preconditioner + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) - pc = pc_class(_A) + pc_class = getattr(preconditioner, self.options.precond) + pc = pc_class(getattr(self.mass_ops, id_M)) + + # magnetic equilibrium field + unit_b1 = self.projected_equil.unit_b1 + curl_unit_b1 = self.projected_equil.curl_unit_b1 + self._b2 = self.projected_equil.b2 + + # magnetic field + self._b_tilde = self.options.b_tilde.spline.vector + + # scaling factor + epsilon = self.variables.energetic_ions.species.equation_params.epsilon # temporary vectors to avoid memory allocation - self._b_full1 = self._b_eq.space.zeros() - self._b_full2 = self._E2T.codomain.zeros() - self._u_new = u.space.zeros() - self._u_avg1 = u.space.zeros() - self._u_avg2 = self._EuT.codomain.zeros() + self._b_full = self._b2.space.zeros() + self._u_new = self.variables.u.spline.vector.space.zeros() + self._u_avg = self.variables.u.spline.vector.space.zeros() - # Call the accumulation and Pusher class + # define Accumulator and arguments self._ACC = Accumulator( - particles, - u_space, - accum_kernels_gc.cc_lin_mhd_5d_J1, + self.variables.energetic_ions.particles, + self.options.u_space, + accum_kernels_gc.cc_lin_mhd_5d_curlb, self.mass_ops, self.domain.args_domain, add_vector=True, symmetry="symm", - filter_params=filter, + filter_params=self.options.filter_params, ) - if u_space == "Hcurl": - kernel = pusher_kernels_gc.push_gc_cc_J1_Hcurl - elif u_space == "Hdiv": - kernel = pusher_kernels_gc.push_gc_cc_J1_Hdiv - elif u_space == "H1vec": - kernel = pusher_kernels_gc.push_gc_cc_J1_H1vec + self._args_accum_kernel = ( + epsilon, + self.options.ep_scale, + self._b_full[0]._data, + self._b_full[1]._data, + self._b_full[2]._data, + unit_b1[0]._data, + unit_b1[1]._data, + unit_b1[2]._data, + curl_unit_b1[0]._data, + curl_unit_b1[1]._data, + curl_unit_b1[2]._data, + self._u_form_int, + ) + + # define Pusher + if self.options.u_space == "Hcurl": + pusher_kernel = pusher_kernels_gc.push_gc_cc_J1_Hcurl + elif self.options.u_space == "Hdiv": + pusher_kernel = pusher_kernels_gc.push_gc_cc_J1_Hdiv + elif self.options.u_space == "H1vec": + pusher_kernel = pusher_kernels_gc.push_gc_cc_J1_H1vec else: raise ValueError( - f'{u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', + f'{self.options.u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', ) - # instantiate Pusher - args_kernel = ( + args_pusher_kernel = ( self.derham.args_derham, - self._epsilon, - self._b_full2[0]._data, - self._b_full2[1]._data, - self._b_full2[2]._data, - self._unit_b1[0]._data, - self._unit_b1[1]._data, - self._unit_b1[2]._data, - self._curl_norm_b[0]._data, - self._curl_norm_b[1]._data, - self._curl_norm_b[2]._data, - self._u_avg2[0]._data, - self._u_avg2[1]._data, - self._u_avg2[2]._data, - 0.0, + epsilon, + self._b_full[0]._data, + self._b_full[1]._data, + self._b_full[2]._data, + unit_b1[0]._data, + unit_b1[1]._data, + unit_b1[2]._data, + curl_unit_b1[0]._data, + curl_unit_b1[1]._data, + curl_unit_b1[2]._data, + self._u_avg[0]._data, + self._u_avg[1]._data, + self._u_avg[2]._data, ) self._pusher = Pusher( - particles, - kernel, - args_kernel, + self.variables.energetic_ions.particles, + pusher_kernel, + args_pusher_kernel, self.domain.args_domain, alpha_in_kernel=1.0, ) - # define BC and B dot V of the Schur block matrix [[A, B], [C, I]] _BC = -1 / 4 * self._ACC.operators[0] - # call SchurSolver class self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, ) def __call__(self, dt): - un = self.feec_vars[0] + # current FE coeffs + un = self.variables.u.spline.vector # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) - b_full = self._b_eq.copy(out=self._b_full1) - - if self._b is not None: - self._b_full1 += self._b - - # extract coefficients to tensor product space (in-place) - Eb_full = self._E2T.dot(b_full, out=self._b_full2) - - # update ghost regions because of non-local access in accumulation kernel! - Eb_full.update_ghost_regions() + b_full = self._b2.copy(out=self._b_full) - # perform accumulation (either with or without control variate) - # if self.particles[0].control_variate: - - # # evaluate magnetic field at quadrature points (in-place) - # WeightedMassOperator.eval_quad(self.derham.Vh_fem['2'], self._b_full2, - # out=[self._b_at_quad[0], self._b_at_quad[1], self._b_at_quad[2]]) - - # # evaluate B_parallel - # self._B_para_at_quad = np.sum( - # p * q for p, q in zip(self._unit_b1_at_quad, self._b_at_quad)) - # self._B_para_at_quad += self._unit_b1_dot_curl_norm_b_at_quad - - # # assemble (B x)(curl norm_b)(curl norm_b)(B x) / B_star_para² / det_df³ * (f0.u_para² + f0.vth_para²) * f0.n - # self._mat11[:, :, :] = (self._b_at_quad[1]*self._curl_norm_b_at_quad[2] - - # self._b_at_quad[2]*self._curl_norm_b_at_quad[1])**2 * \ - # self._control_const * self._coupling_mat / \ - # self._det_df_at_quad**3 / self._B_para_at_quad**2 - # self._mat12[:, :, :] = (self._b_at_quad[1]*self._curl_norm_b_at_quad[2] - - # self._b_at_quad[2]*self._curl_norm_b_at_quad[1]) * \ - # (self._b_at_quad[2]*self._curl_norm_b_at_quad[0] - - # self._b_at_quad[0]*self._curl_norm_b_at_quad[2]) * \ - # self._control_const * self._coupling_mat / \ - # self._det_df_at_quad**3 / self._B_para_at_quad**2 - # self._mat13[:, :, :] = (self._b_at_quad[1]*self._curl_norm_b_at_quad[2] - - # self._b_at_quad[2]*self._curl_norm_b_at_quad[1]) * \ - # (self._b_at_quad[0]*self._curl_norm_b_at_quad[1] - - # self._b_at_quad[1]*self._curl_norm_b_at_quad[0]) * \ - # self._control_const * self._coupling_mat / \ - # self._det_df_at_quad**3 / self._B_para_at_quad**2 - # self._mat22[:, :, :] = (self._b_at_quad[2]*self._curl_norm_b_at_quad[0] - - # self._b_at_quad[0]*self._curl_norm_b_at_quad[2])**2 * \ - # self._control_const * self._coupling_mat / \ - # self._det_df_at_quad**3 / self._B_para_at_quad**2 - # self._mat23[:, :, :] = (self._b_at_quad[2]*self._curl_norm_b_at_quad[0] - - # self._b_at_quad[0]*self._curl_norm_b_at_quad[2]) * \ - # (self._b_at_quad[0]*self._curl_norm_b_at_quad[1] - - # self._b_at_quad[1]*self._curl_norm_b_at_quad[0]) * \ - # self._control_const * self._coupling_mat / \ - # self._det_df_at_quad**3 / self._B_para_at_quad**2 - # self._mat33[:, :, :] = (self._b_at_quad[0]*self._curl_norm_b_at_quad[1] - - # self._b_at_quad[1]*self._curl_norm_b_at_quad[0])**2 * \ - # self._control_const * self._coupling_mat / \ - # self._det_df_at_quad**3 / self._B_para_at_quad**2 - - # self._mat21[:, :, :] = -self._mat12 - # self._mat31[:, :, :] = -self._mat13 - # self._mat32[:, :, :] = -self._mat23 - - # # assemble (B x)(curl norm_b) / B_star_para / det_df * (f0.u_para² + f0.vth_para²) * f0.n - # self._vec1[:, :, :] = (self._b_at_quad[1]*self._curl_norm_b_at_quad[2] - - # self._b_at_quad[2]*self._curl_norm_b_at_quad[1]) * \ - # self._control_const * self._coupling_vec / \ - # self._det_df_at_quad / self._B_para_at_quad - # self._vec2[:, :, :] = (self._b_at_quad[2]*self._curl_norm_b_at_quad[0] - - # self._b_at_quad[0]*self._curl_norm_b_at_quad[2]) * \ - # self._control_const * self._coupling_vec / \ - # self._det_df_at_quad / self._B_para_at_quad - # self._vec3[:, :, :] = (self._b_at_quad[0]*self._curl_norm_b_at_quad[1] - - # self._b_at_quad[1]*self._curl_norm_b_at_quad[0]) * \ - # self._control_const * self._coupling_vec / \ - # self._det_df_at_quad / self._B_para_at_quad - - # self._ACC.accumulate(self.particles[0], self._epsilon, - # Eb_full[0]._data, Eb_full[1]._data, Eb_full[2]._data, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._curl_norm_b[0]._data, self._curl_norm_b[1]._data, self._curl_norm_b[2]._data, - # self._space_key_int, self._coupling_mat, self._coupling_vec, 0.1, - # control_mat=[[None, self._mat12, self._mat13], - # [self._mat21, None, self._mat23], - # [self._mat31, self._mat32, None]], - # control_vec=[self._vec1, self._vec2, self._vec3]) - # else: - # self._ACC.accumulate(self.particles[0], self._epsilon, - # Eb_full[0]._data, Eb_full[1]._data, Eb_full[2]._data, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._curl_norm_b[0]._data, self._curl_norm_b[1]._data, self._curl_norm_b[2]._data, - # self._space_key_int, self._coupling_mat, self._coupling_vec, 0.1) + b_full += self._b_tilde + b_full.update_ghost_regions() self._ACC( - self._epsilon, - Eb_full[0]._data, - Eb_full[1]._data, - Eb_full[2]._data, - self._unit_b1[0]._data, - self._unit_b1[1]._data, - self._unit_b1[2]._data, - self._curl_norm_b[0]._data, - self._curl_norm_b[1]._data, - self._curl_norm_b[2]._data, - self._space_key_int, - self._coupling_mat, - self._coupling_vec, - self._boundary_cut_e1, + *self._args_accum_kernel, ) - # update u coefficients + # solve un1, info = self._schur_solver( un, -self._ACC.vectors[0] / 2, @@ -1420,27 +1334,25 @@ def __call__(self, dt): ) # call pusher kernel with average field (u_new + u_old)/2 and update ghost regions because of non-local access in kernel - _u = un.copy(out=self._u_avg1) + _u = un.copy(out=self._u_avg) _u += un1 _u *= 0.5 - _Eu = self._EuT.dot(_u, out=self._u_avg2) - - _Eu.update_ghost_regions() + _u.update_ghost_regions() - self._pusher(self._scale_push * dt) + self._pusher(dt) - # write new coeffs into Propagator.variables - (max_du,) = self.feec_vars_update(un1) + # update u coefficients + diffs = self.update_feec_variables(u=un1) # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() + if self.variables.energetic_ions.species.weights_params.control_variate: + self.variables.energetic_ions.particles.update_weights() - if self._info and self._rank == 0: + if self.options.solver_params.info and MPI.COMM_WORLD.Get_rank() == 0: print("Status for CurrentCoupling5DCurlb:", info["success"]) print("Iterations for CurrentCoupling5DCurlb:", info["niter"]) - print("Maxdiff up for CurrentCoupling5DCurlb:", max_du) + print("Maxdiff up for CurrentCoupling5DCurlb:", diffs["u"]) print() @@ -1480,435 +1392,722 @@ class CurrentCoupling5DGradB(Propagator): For the detail explanation of the notations, see `2022_DriftKineticCurrentCoupling `_. """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["algo"] = ["rk4", "forward_euler", "heun2", "rk2", "heun3"] - dct["filter"] = { - "use_filter": None, - "modes": (1), - "repeat": 1, - "alpha": 0.5, - } - dct["boundary_cut"] = { - "e1": 0.0, - "e2": 0.0, - "e3": 0.0, - } - dct["turn_off"] = False + class Variables: + def __init__(self): + self._u: FEECVariable = None + self._energetic_ions: PICVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def u(self) -> FEECVariable: + return self._u - return dct + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new - def __init__( - self, - particles: Particles5D, - u: BlockVector, - *, - b: BlockVector, - b_eq: BlockVector, - unit_b1: BlockVector, - unit_b2: BlockVector, - absB0: StencilVector, - gradB1: BlockVector, - curl_unit_b2: BlockVector, - u_space: str, - solver: dict = options(default=True)["solver"], - algo: dict = options(default=True)["algo"], - filter: dict = options(default=True)["filter"], - coupling_params: dict, - epsilon: float = 1.0, - boundary_cut: dict = options(default=True)["boundary_cut"], - ): - from psydac.linalg.solvers import inverse + @property + def energetic_ions(self) -> PICVariable: + return self._energetic_ions - from struphy.ode.utils import ButcherTableau + @energetic_ions.setter + def energetic_ions(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles5D" + self._energetic_ions = new - super().__init__(particles, u) + def __init__(self): + self.variables = self.Variables() - assert u_space in {"Hcurl", "Hdiv", "H1vec"} + @dataclass + class Options: + # specific literals + OptsAlgo = Literal[ + "discrete_gradient", + "explicit", + ] + # propagator options + b_tilde: FEECVariable = None + ep_scale: float = 1.0 + algo: OptsAlgo = "explicit" + butcher: ButcherTableau = None + u_space: OptsVecSpace = "Hdiv" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + filter_params: FilterParameters = None + dg_solver_params: DiscreteGradientSolverParameters = None - if u_space == "H1vec": - self._space_key_int = 0 - else: - self._space_key_int = int( - self.derham.space_to_form[u_space], - ) + def __post_init__(self): + # checks + check_option(self.algo, self.OptsAlgo) + check_option(self.u_space, OptsVecSpace) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + assert isinstance(self.b_tilde, FEECVariable) + assert isinstance(self.ep_scale, float) - self._epsilon = epsilon - self._b = b - self._b_eq = b_eq - self._unit_b1 = unit_b1 - self._unit_b2 = unit_b2 - self._absB0 = absB0 - self._gradB1 = gradB1 - self._curl_norm_b = curl_unit_b2 + # defaults + if self.algo == "explicit" and self.butcher is None: + self.butcher = ButcherTableau() - self._info = solver["info"] - self._rank = self.derham.comm.Get_rank() + if self.algo == "discrete_gradient" and self.dg_solver_params is None: + self.dg_solver_params = DiscreteGradientSolverParameters() - self._coupling_mat = coupling_params["Ah"] / coupling_params["Ab"] - self._coupling_vec = coupling_params["Ah"] / coupling_params["Ab"] - self._scale_push = 1 + if self.solver_params is None: + self.solver_params = SolverParameters() - self._boundary_cut_e1 = boundary_cut["e1"] + if self.filter_params is None: + self.filter_params = FilterParameters() - u_id = self.derham.space_to_form[u_space] - self._E0T = self.derham.extraction_ops["0"].transpose() - self._EuT = self.derham.extraction_ops[u_id].transpose() - self._E1T = self.derham.extraction_ops["1"].transpose() - self._E2T = self.derham.extraction_ops["2"].transpose() + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options - self._PB = getattr(self.basis_ops, "PB") + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new - self._unit_b1 = self._E1T.dot(self._unit_b1) - self._unit_b2 = self._E2T.dot(self._unit_b2) - self._curl_norm_b = self._E2T.dot(self._curl_norm_b) - self._absB0 = self._E0T.dot(self._absB0) + @profile + def allocate(self): + if self.options.u_space == "H1vec": + self._u_form_int = 0 + else: + self._u_form_int = int(self.derham.space_to_form[self.options.u_space]) - _A = getattr(self.mass_ops, "M" + u_id + "n") + # call operatros + id_M = "M" + self.derham.space_to_form[self.options.u_space] + "n" + self._A = getattr(self.mass_ops, id_M) + self._PB = getattr(self.basis_ops, "PB") - # preconditioner - if solver["type"][1] is None: + # Preconditioner + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) - pc = pc_class(_A) + pc_class = getattr(preconditioner, self.options.precond) + pc = pc_class(getattr(self.mass_ops, id_M)) - self._solver = inverse( - _A, - solver["type"][0], + # linear solver + self._A_inv = inverse( + self._A, + self.options.solver, pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, ) + # magnetic equilibrium field + unit_b1 = self.projected_equil.unit_b1 + curl_unit_b1 = self.projected_equil.curl_unit_b1 + self._b2 = self.projected_equil.b2 + gradB1 = self.projected_equil.gradB1 + absB0 = self.projected_equil.absB0 + + # magnetic field + self._b_tilde = self.options.b_tilde.spline.vector + + # scaling factor + epsilon = self.variables.energetic_ions.species.equation_params.epsilon + + if self.options.algo == "explicit": + # temporary vectors to avoid memory allocation + self._b_full = self._b2.space.zeros() + self._u_new = self.variables.u.spline.vector.space.zeros() + self._u_temp = self.variables.u.spline.vector.space.zeros() + self._ku = self.variables.u.spline.vector.space.zeros() + self._PB_b = self._PB.codomain.zeros() + self._grad_PB_b = self.derham.grad.codomain.zeros() + + # define Accumulator and arguments + self._ACC = Accumulator( + self.variables.energetic_ions.particles, + self.options.u_space, + accum_kernels_gc.cc_lin_mhd_5d_gradB, + self.mass_ops, + self.domain.args_domain, + add_vector=True, + symmetry="symm", + filter_params=self.options.filter_params, + ) - # Call the accumulation and Pusher class - self._ACC = Accumulator( - particles, - u_space, - accum_kernels_gc.cc_lin_mhd_5d_J2, - self.mass_ops, - self.domain.args_domain, - add_vector=True, - symmetry="symm", - filter_params=filter, - ) + self._args_accum_kernel = ( + epsilon, + self.options.ep_scale, + self._b_full[0]._data, + self._b_full[1]._data, + self._b_full[2]._data, + unit_b1[0]._data, + unit_b1[1]._data, + unit_b1[2]._data, + curl_unit_b1[0]._data, + curl_unit_b1[1]._data, + curl_unit_b1[2]._data, + self._grad_PB_b[0]._data, + self._grad_PB_b[1]._data, + self._grad_PB_b[2]._data, + self._u_form_int, + ) - # if self.particles[0].control_variate: + # define Pusher + if self.options.u_space == "Hdiv": + self._pusher_kernel = pusher_kernels_gc.push_gc_cc_J2_stage_Hdiv + elif self.options.u_space == "H1vec": + self._pusher_kernel = pusher_kernels_gc.push_gc_cc_J2_stage_H1vec + else: + raise ValueError( + f'{self.options.u_space = } not valid, choose from "Hdiv" or "H1vec.', + ) - # # control variate method is only valid with Maxwellian distributions - # assert isinstance(self.particles[0].f0, Maxwellian) - # assert params['u_space'] == 'Hdiv' + # temp fix due to refactoring of ButcherTableau: + butcher = self.options.butcher + import numpy as np + + butcher._a = np.diag(butcher.a, k=-1) + butcher._a = np.array(list(butcher.a) + [0.0]) + + self._args_pusher_kernel = ( + self.domain.args_domain, + self.derham.args_derham, + epsilon, + self._b_full[0]._data, + self._b_full[1]._data, + self._b_full[2]._data, + unit_b1[0]._data, + unit_b1[1]._data, + unit_b1[2]._data, + curl_unit_b1[0]._data, + curl_unit_b1[1]._data, + curl_unit_b1[2]._data, + self._u_temp[0]._data, + self._u_temp[1]._data, + self._u_temp[2]._data, + self.options.butcher.a, + self.options.butcher.b, + self.options.butcher.c, + ) + + else: + # temporary vectors to avoid memory allocation + self._b_full = self._b2.space.zeros() + self._PB_b = self._PB.codomain.zeros() + self._grad_PB_b = self.derham.grad.codomain.zeros() + self._u_old = self.variables.u.spline.vector.space.zeros() + self._u_new = self.variables.u.spline.vector.space.zeros() + self._u_diff = self.variables.u.spline.vector.space.zeros() + self._u_mid = self.variables.u.spline.vector.space.zeros() + self._M2n_dot_u = self.variables.u.spline.vector.space.zeros() + self._ku = self.variables.u.spline.vector.space.zeros() + self._u_temp = self.variables.u.spline.vector.space.zeros() + + # Call the accumulation and Pusher class + accum_kernel_init = accum_kernels_gc.cc_lin_mhd_5d_gradB_dg_init + accum_kernel = accum_kernels_gc.cc_lin_mhd_5d_gradB_dg + self._accum_kernel_en_fB_mid = utilities_kernels.eval_gradB_ediff + + self._args_accum_kernel = ( + epsilon, + self.options.ep_scale, + self._b_tilde[0]._data, + self._b_tilde[1]._data, + self._b_tilde[2]._data, + self._b2[0]._data, + self._b2[1]._data, + self._b2[2]._data, + unit_b1[0]._data, + unit_b1[1]._data, + unit_b1[2]._data, + curl_unit_b1[0]._data, + curl_unit_b1[1]._data, + curl_unit_b1[2]._data, + self._grad_PB_b[0]._data, + self._grad_PB_b[1]._data, + self._grad_PB_b[2]._data, + gradB1[0]._data, + gradB1[1]._data, + gradB1[2]._data, + self._u_form_int, + ) + + self._args_accum_kernel_en_fB_mid = ( + self.domain.args_domain, + self.derham.args_derham, + gradB1[0]._data, + gradB1[1]._data, + gradB1[2]._data, + self._grad_PB_b[0]._data, + self._grad_PB_b[1]._data, + self._grad_PB_b[2]._data, + ) + + self._ACC_init = AccumulatorVector( + self.variables.energetic_ions.particles, + self.options.u_space, + accum_kernel_init, + self.mass_ops, + self.domain.args_domain, + filter_params=self.options.filter_params, + ) - # self._ACC.init_control_variate(self.mass_ops) + self._ACC = AccumulatorVector( + self.variables.energetic_ions.particles, + self.options.u_space, + accum_kernel, + self.mass_ops, + self.domain.args_domain, + filter_params=self.options.filter_params, + ) - # # evaluate and save n0 at quadrature points - # quad_pts = [quad_grid[nquad].points.flatten() - # for quad_grid, nquad in zip(self.derham.get_quad_grids(self.derham.Vh_fem['0']), self.derham.nquads)] + self._args_pusher_kernel_init = ( + self.domain.args_domain, + self.derham.args_derham, + epsilon, + self._b_full[0]._data, + self._b_full[1]._data, + self._b_full[2]._data, + unit_b1[0]._data, + unit_b1[1]._data, + unit_b1[2]._data, + curl_unit_b1[0]._data, + curl_unit_b1[1]._data, + curl_unit_b1[2]._data, + self.variables.u.spline.vector[0]._data, + self.variables.u.spline.vector[1]._data, + self.variables.u.spline.vector[2]._data, + ) - # self._n0_at_quad = self.domain.push( - # self.particles[0].f0.n, *quad_pts, kind='0', squeeze_out=False) + self._args_pusher_kernel = ( + self.domain.args_domain, + self.derham.args_derham, + epsilon, + self._b_full[0]._data, + self._b_full[1]._data, + self._b_full[2]._data, + unit_b1[0]._data, + unit_b1[1]._data, + unit_b1[2]._data, + curl_unit_b1[0]._data, + curl_unit_b1[1]._data, + curl_unit_b1[2]._data, + self._u_mid[0]._data, + self._u_mid[1]._data, + self._u_mid[2]._data, + self._u_temp[0]._data, + self._u_temp[1]._data, + self._u_temp[2]._data, + ) - # # evaluate unit_b1 (1form) dot epsilon * u0_parallel * curl_norm_b/|det(DF)| at quadrature points - # quad_pts_array = self.domain.prepare_eval_pts(*quad_pts)[:3] + self._pusher_kernel_init = pusher_kernels_gc.push_gc_cc_J2_dg_init_Hdiv + self._pusher_kernel = pusher_kernels_gc.push_gc_cc_J2_dg_Hdiv - # u0_parallel_at_quad = self.particles[0].f0.u( - # *quad_pts_array)[0] + def __call__(self, dt): + # current FE coeffs + un = self.variables.u.spline.vector - # vth_perp = self.particles[0].f0.vth(*quad_pts_array)[1] + # particle markers and idx + particles = self.variables.energetic_ions.particles + holes = particles.holes + args_markers = particles.args_markers + markers = args_markers.markers + first_init_idx = args_markers.first_init_idx + first_free_idx = args_markers.first_free_idx - # absB0_at_quad = WeightedMassOperator.eval_quad( - # self.derham.Vh_fem['0'], self._absB0) + # clear buffer + markers[:, first_init_idx:-2] = 0.0 - # self._det_df_at_quad = self.domain.jacobian_det( - # *quad_pts, squeeze_out=False) + # save old marker positions + markers[:, first_init_idx : first_init_idx + 3] = markers[:, :3] - # self._unit_b1_at_quad = WeightedMassOperator.eval_quad( - # self.derham.Vh_fem['1'], self._unit_b1) + # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) + b_full = self._b2.copy(out=self._b_full) - # curl_norm_b_at_quad = WeightedMassOperator.eval_quad( - # self.derham.Vh_fem['2'], self._curl_norm_b) + b_full += self._b_tilde + b_full.update_ghost_regions() - # self._unit_b1_dot_curl_norm_b_at_quad = np.sum( - # p * q for p, q in zip(self._unit_b1_at_quad, curl_norm_b_at_quad)) + if self.options.algo == "explicit": + PB_b = self._PB.dot(b_full, out=self._PB_b) + grad_PB_b = self.derham.grad.dot(PB_b, out=self._grad_PB_b) + grad_PB_b.update_ghost_regions() - # self._unit_b1_dot_curl_norm_b_at_quad /= self._det_df_at_quad - # self._unit_b1_dot_curl_norm_b_at_quad *= self._epsilon - # self._unit_b1_dot_curl_norm_b_at_quad *= u0_parallel_at_quad + # save old u + u_new = un.copy(out=self._u_new) - # # precalculate constant 2 * f0.vth_perp² / B0 * f0.n for control MAT and VEC - # self._control_const = vth_perp**2 / absB0_at_quad * self._n0_at_quad + for stage in range(self.options.butcher.n_stages): + # accumulate + self._ACC( + *self._args_accum_kernel, + ) - # # assemble the matrix (G_inv)(unit_b1 x)(G_inv) - # G_inv_at_quad = self.domain.metric_inv( - # *quad_pts, squeeze_out=False) + # push particles + self._pusher_kernel( + dt, + stage, + args_markers, + *self._args_pusher_kernel, + ) - # self._G_inv_bx_G_inv_at_quad = [[np.zeros_like(self._n0_at_quad), np.zeros_like(self._n0_at_quad), np.zeros_like(self._n0_at_quad)], - # [np.zeros_like(self._n0_at_quad), np.zeros_like( - # self._n0_at_quad), np.zeros_like(self._n0_at_quad)], - # [np.zeros_like(self._n0_at_quad), np.zeros_like(self._n0_at_quad), np.zeros_like(self._n0_at_quad)]] + if particles.mpi_comm is not None: + particles.mpi_sort_markers() + else: + particles.apply_kinetic_bc() - # for j in range(3): - # temp = (-self._unit_b1_at_quad[2]*G_inv_at_quad[1, j] + self._unit_b1_at_quad[1]*G_inv_at_quad[2, j], - # self._unit_b1_at_quad[2]*G_inv_at_quad[0, j] - - # self._unit_b1_at_quad[0]*G_inv_at_quad[2, j], - # -self._unit_b1_at_quad[1]*G_inv_at_quad[0, j] + self._unit_b1_at_quad[0]*G_inv_at_quad[1, j]) + # solve linear system for updating u coefficients + ku = self._A_inv.dot(self._ACC.vectors[0], out=self._ku) + info = self._A_inv._info - # for i in range(3): - # self._G_inv_bx_G_inv_at_quad[i][j] = np.sum( - # p * q for p, q in zip(G_inv_at_quad[i], temp[:])) + # calculate u^{n+1}_k + u_temp = un.copy(out=self._u_temp) + u_temp += ku * dt * self.options.butcher.a[stage] - # # memory allocation of magnetic field at quadrature points - # self._b_at_quad = [np.zeros_like(self._n0_at_quad), - # np.zeros_like(self._n0_at_quad), - # np.zeros_like(self._n0_at_quad)] + u_temp.update_ghost_regions() - # # memory allocation of parallel magnetic field at quadrature points - # self._B_para_at_quad = np.zeros_like(self._n0_at_quad) + # calculate u^{n+1} + u_new += ku * dt * self.options.butcher.b[stage] - # # memory allocation of gradient of parallel magnetic field at quadrature points - # self._grad_PBb_at_quad = (np.zeros_like(self._n0_at_quad), - # np.zeros_like(self._n0_at_quad), - # np.zeros_like(self._n0_at_quad)) - # # memory allocation for temporary matrix - # self._temp = [[np.zeros_like(self._n0_at_quad), np.zeros_like(self._n0_at_quad), np.zeros_like(self._n0_at_quad)], - # [np.zeros_like(self._n0_at_quad), np.zeros_like( - # self._n0_at_quad), np.zeros_like(self._n0_at_quad)], - # [np.zeros_like(self._n0_at_quad), np.zeros_like(self._n0_at_quad), np.zeros_like(self._n0_at_quad)]] + if self.options.solver_params.info and MPI.COMM_WORLD.Get_rank() == 0: + print("Stage: ", stage) + print("Status for CurrentCoupling5DGradB:", info["success"]) + print("Iterations for CurrentCoupling5DGradB:", info["niter"]) + print() - # # memory allocation for control VEC - # self._vec1 = np.zeros_like(self._n0_at_quad) - # self._vec2 = np.zeros_like(self._n0_at_quad) - # self._vec3 = np.zeros_like(self._n0_at_quad) + # update u coefficients + diffs = self.update_feec_variables(u=u_new) - # choose algorithm - self._butcher = ButcherTableau(algo) - # temp fix due to refactoring of ButcherTableau: - self._butcher._a = np.diag(self._butcher.a, k=-1) - self._butcher._a = np.array(list(self._butcher.a) + [0.0]) + # clear the buffer + markers[:, first_init_idx:-2] = 0.0 + + # update_weights + if self.variables.energetic_ions.species.weights_params.control_variate: + particles.update_weights() + + if self.options.solver_params.info and MPI.COMM_WORLD.Get_rank() == 0: + print("Maxdiff up for CurrentCoupling5DGradB:", diffs["u"]) + print() - # instantiate Pusher - if u_space == "Hdiv": - kernel = pusher_kernels_gc.push_gc_cc_J2_stage_Hdiv - elif u_space == "H1vec": - kernel = pusher_kernels_gc.push_gc_cc_J2_stage_H1vec else: - raise ValueError( - f'{u_space = } not valid, choose from "Hdiv" or "H1vec.', + # total number of markers + n_mks_tot = particles.Np + + # relaxation factor + alpha = self.options.dg_solver_params.relaxation_factor + + # eval parallel tilde b and its gradient + PB_b = self._PB.dot(self._b_tilde, out=self._PB_b) + PB_b.update_ghost_regions() + grad_PB_b = self.derham.grad.dot(PB_b, out=self._grad_PB_b) + grad_PB_b.update_ghost_regions() + + # save old u + u_old = un.copy(out=self._u_old) + u_new = un.copy(out=self._u_new) + + # save en_U_old + self._A.dot(un, out=self._M2n_dot_u) + en_U_old = un.inner(self._M2n_dot_u) / 2.0 + + # save en_fB_old + particles.save_magnetic_energy(PB_b) + en_fB_old = np.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_old /= n_mks_tot + + buffer_array = np.array([en_fB_old]) + + if particles.mpi_comm is not None: + particles.mpi_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + if particles.clone_config is not None: + particles.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + en_fB_old = buffer_array[0] + en_tot_old = en_U_old + en_fB_old + + # initial guess + self._ACC_init(*self._args_accum_kernel) + + ku = self._A_inv.dot(self._ACC_init.vectors[0], out=self._ku) + u_new += ku * dt + + u_new.update_ghost_regions() + + # save en_U_new + self._A.dot(u_new, out=self._M2n_dot_u) + en_U_new = u_new.inner(self._M2n_dot_u) / 2.0 + + # push eta + self._pusher_kernel_init( + dt, + args_markers, + *self._args_pusher_kernel_init, ) - args_kernel = (self.derham.args_derham,) + if particles.mpi_comm is not None: + particles.mpi_sort_markers(apply_bc=False) - self._pusher = Pusher( - particles, - kernel, - args_kernel, - self.domain.args_domain, - alpha_in_kernel=1.0, - ) + # save en_fB_new + particles.save_magnetic_energy(PB_b) + en_fB_new = np.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_new /= n_mks_tot - # temporary vectors to avoid memory allocation - self._b_full1 = self._b_eq.space.zeros() - self._b_full2 = self._E2T.codomain.zeros() - self._u_new = u.space.zeros() - self._Eu_new = self._EuT.codomain.zeros() - self._u_temp1 = u.space.zeros() - self._u_temp2 = u.space.zeros() - self._Eu_temp = self._EuT.codomain.zeros() - self._tmp1 = self._E0T.codomain.zeros() - self._tmp2 = self._gradB1.space.zeros() - self._tmp3 = self._E1T.codomain.zeros() + buffer_array = np.array([en_fB_new]) - def __call__(self, dt): - un = self.feec_vars[0] + if particles.mpi_comm is not None: + particles.mpi_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) - # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) - b_full = self._b_eq.copy(out=self._b_full1) + if particles.clone_config is not None: + particles.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) - if self._b is not None: - self._b_full1 += self._b + en_fB_new = buffer_array[0] - PBb = self._PB.dot(self._b, out=self._tmp1) - grad_PBb = self.derham.grad.dot(PBb, out=self._tmp2) - grad_PBb += self._gradB1 + # fixed-point iterations + iter_num = 0 - Eb_full = self._E2T.dot(b_full, out=self._b_full2) - Eb_full.update_ghost_regions() + while True: + iter_num += 1 - Egrad_PBb = self._E1T.dot(grad_PBb, out=self._tmp3) - Egrad_PBb.update_ghost_regions() + if self.options.dg_solver_params.verbose and MPI.COMM_WORLD.Get_rank() == 0: + print("# of iteration: ", iter_num) - # perform accumulation (either with or without control variate) - # if self.particles[0].control_variate: + # calculate discrete gradient + # save u^{n+1, k} + u_old = u_new.copy(out=self._u_old) - # # evaluate magnetic field at quadrature points (in-place) - # WeightedMassOperator.eval_quad(self.derham.Vh_fem['2'], self._b_full2, - # out=[self._b_at_quad[0], self._b_at_quad[1], self._b_at_quad[2]]) - - # # evaluate B_parallel - # self._B_para_at_quad = np.sum( - # p * q for p, q in zip(self._unit_b1_at_quad, self._b_at_quad)) - # self._B_para_at_quad += self._unit_b1_dot_curl_norm_b_at_quad - - # # evaluate grad B_parallel - # WeightedMassOperator.eval_quad(self.derham.Vh_fem['1'], self._tmp3, - # out=[self._grad_PBb_at_quad[0], self._grad_PBb_at_quad[1], self._grad_PBb_at_quad[2]]) - - # # assemble temp = (B x)(G_inv)(unit_b1 x)(G_inv) - # for i in range(3): - # self._temp[0][i] = -self._b_at_quad[2]*self._G_inv_bx_G_inv_at_quad[1][i] + \ - # self._b_at_quad[1]*self._G_inv_bx_G_inv_at_quad[2][i] - # self._temp[1][i] = +self._b_at_quad[2]*self._G_inv_bx_G_inv_at_quad[0][i] - \ - # self._b_at_quad[0]*self._G_inv_bx_G_inv_at_quad[2][i] - # self._temp[2][i] = -self._b_at_quad[1]*self._G_inv_bx_G_inv_at_quad[0][i] + \ - # self._b_at_quad[0]*self._G_inv_bx_G_inv_at_quad[1][i] - - # # assemble (temp)(grad B_parallel) / B_star_para * 2 * f0.vth_perp² / B0 * f0.n - # self._vec1[:, :, :] = np.sum(p * q for p, q in zip(self._temp[0][:], self._grad_PBb_at_quad)) * \ - # self._control_const * self._coupling_vec / self._B_para_at_quad - # self._vec2[:, :, :] = np.sum(p * q for p, q in zip(self._temp[1][:], self._grad_PBb_at_quad)) * \ - # self._control_const * self._coupling_vec / self._B_para_at_quad - # self._vec3[:, :, :] = np.sum(p * q for p, q in zip(self._temp[2][:], self._grad_PBb_at_quad)) * \ - # self._control_const * self._coupling_vec / self._B_para_at_quad - - # save old u - _u_new = un.copy(out=self._u_new) - _u_temp = un.copy(out=self._u_temp1) + u_diff = u_old.copy(out=self._u_diff) + u_diff -= un + u_diff.update_ghost_regions() - # save old marker positions - self.particles[0].markers[ - ~self.particles[0].holes, - 11:14, - ] = self.particles[0].markers[~self.particles[0].holes, 0:3] - - for stage in range(self._butcher.n_stages): - # accumulate RHS - # if self.particles[0].control_variate: - # self._ACC.accumulate(self.particles[0], self._epsilon, - # Eb_full[0]._data, Eb_full[1]._data, Eb_full[2]._data, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._unit_b2[0]._data, self._unit_b2[1]._data, self._unit_b2[2]._data, - # self._curl_norm_b[0]._data, self._curl_norm_b[1]._data, self._curl_norm_b[2]._data, - # Egrad_PBb[0]._data, Egrad_PBb[1]._data, Egrad_PBb[2]._data, - # self._space_key_int, self._coupling_mat, self._coupling_vec, 0., - # control_vec=[self._vec1, self._vec2, self._vec3]) - # else: - # self._ACC.accumulate(self.particles[0], self._epsilon, - # Eb_full[0]._data, Eb_full[1]._data, Eb_full[2]._data, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._unit_b2[0]._data, self._unit_b2[1]._data, self._unit_b2[2]._data, - # self._curl_norm_b[0]._data, self._curl_norm_b[1]._data, self._curl_norm_b[2]._data, - # Egrad_PBb[0]._data, Egrad_PBb[1]._data, Egrad_PBb[2]._data, - # self._space_key_int, self._coupling_mat, self._coupling_vec, 0.) - - self._ACC( - self._epsilon, - Eb_full[0]._data, - Eb_full[1]._data, - Eb_full[2]._data, - self._unit_b1[0]._data, - self._unit_b1[1]._data, - self._unit_b1[2]._data, - self._unit_b2[0]._data, - self._unit_b2[1]._data, - self._unit_b2[2]._data, - self._curl_norm_b[0]._data, - self._curl_norm_b[1]._data, - self._curl_norm_b[2]._data, - Egrad_PBb[0]._data, - Egrad_PBb[1]._data, - Egrad_PBb[2]._data, - self._space_key_int, - self._coupling_mat, - self._coupling_vec, - self._boundary_cut_e1, - ) + u_mid = u_old.copy(out=self._u_mid) + u_mid += un + u_mid /= 2.0 + u_mid.update_ghost_regions() - # push particles - Eu = self._EuT.dot(_u_temp, out=self._Eu_temp) - Eu.update_ghost_regions() + # save H^{n+1, k} + markers[~holes, first_free_idx : first_free_idx + 3] = markers[~holes, 0:3] - self._pusher.kernel( - dt, - stage, - self.particles[0].args_markers, - self.domain.args_domain, - self.derham.args_derham, - self._epsilon, - Eb_full[0]._data, - Eb_full[1]._data, - Eb_full[2]._data, - self._unit_b1[0]._data, - self._unit_b1[1]._data, - self._unit_b1[2]._data, - self._unit_b2[0]._data, - self._unit_b2[1]._data, - self._unit_b2[2]._data, - self._curl_norm_b[0]._data, - self._curl_norm_b[1]._data, - self._curl_norm_b[2]._data, - Eu[0]._data, - Eu[1]._data, - Eu[2]._data, - self._butcher.a, - self._butcher.b, - self._butcher.c, - self._boundary_cut_e1, - ) + # calculate denominator ||z^{n+1, k} - z^n||^2 + sum_u_diff_loc = np.sum((u_diff.toarray() ** 2)) - self.particles[0].mpi_sort_markers() + sum_H_diff_loc = np.sum( + (markers[~holes, :3] - markers[~holes, first_init_idx : first_init_idx + 3]) ** 2 + ) + + buffer_array = np.array([sum_u_diff_loc]) + + if particles.mpi_comm is not None: + particles.mpi_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) - # solve linear system for updated u coefficients - _ku = self._solver.dot(self._ACC.vectors[0], out=self._u_temp2) + denominator = buffer_array[0] - # calculate u^{n+1}_k - _u_temp = un.copy(out=self._u_temp1) - _u_temp += _ku * dt * self._butcher.a[stage] + buffer_array = np.array([sum_H_diff_loc]) - # calculate u^{n+1} - _u_new += _ku * dt * self._butcher.b[stage] + if particles.mpi_comm is not None: + particles.mpi_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + if particles.clone_config is not None: + particles.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) - if self._info and self._rank == 0: - print("Stage:", stage) - print( - "Status for CurrentCoupling5DGradB:", - self._solver._info["success"], + denominator += buffer_array[0] + + # sorting markers at mid-point + if particles.mpi_comm is not None: + particles.mpi_sort_markers(apply_bc=False, alpha=0.5) + + self._accum_kernel_en_fB_mid( + args_markers, + *self._args_accum_kernel_en_fB_mid, + first_free_idx + 3, ) - print( - "Iterations for CurrentCoupling5DGradB:", - self._solver._info["niter"], + en_fB_mid = np.sum(markers[~holes, first_free_idx + 3].dot(markers[~holes, 5])) * self.options.ep_scale + + en_fB_mid /= n_mks_tot + + buffer_array = np.array([en_fB_mid]) + + if particles.mpi_comm is not None: + particles.mpi_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + if particles.clone_config is not None: + particles.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + en_fB_mid = buffer_array[0] + + if denominator == 0.0: + const = 0.0 + else: + const = (en_fB_new - en_fB_old - en_fB_mid) / denominator + + # update u^{n+1, k} + self._ACC(*self._args_accum_kernel, const) + + ku = self._A_inv.dot(self._ACC.vectors[0], out=self._ku) + + u_new = un.copy(out=self._u_new) + u_new += ku * dt + u_new *= alpha + u_new += u_old * (1.0 - alpha) + + u_new.update_ghost_regions() + + # update en_U_new + self._A.dot(u_new, out=self._M2n_dot_u) + en_U_new = u_new.inner(self._M2n_dot_u) / 2.0 + + # update H^{n+1, k} + self._pusher_kernel( + dt, + args_markers, + *self._args_pusher_kernel, + const, + alpha, ) - # clear the buffer - if stage == self._butcher.n_stages - 1: - self.particles[0].markers[ - ~self.particles[0].holes, - 11:-1, - ] = 0.0 + sum_H_diff_loc = np.sum( + np.abs(markers[~holes, 0:3] - markers[~holes, first_free_idx : first_free_idx + 3]) + ) - # write new coeffs into Propagator.variables - (max_du,) = self.feec_vars_update(_u_new) + if particles.mpi_comm is not None: + particles.mpi_sort_markers(apply_bc=False) - # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() + # update en_fB_new + particles.save_magnetic_energy(PB_b) + en_fB_new = np.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_new /= n_mks_tot - if self._info and self._rank == 0: - print("Maxdiff up for CurrentCoupling5DGradB:", max_du) - print() + buffer_array = np.array([en_fB_new]) + + if particles.mpi_comm is not None: + particles.mpi_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + if particles.clone_config is not None: + particles.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + en_fB_new = buffer_array[0] + + # calculate total energy difference + e_diff = np.abs(en_U_new + en_fB_new - en_tot_old) + + # calculate ||z^{n+1, k} - z^{n+1, k-1|| + sum_u_diff_loc = np.sum(np.abs(u_new.toarray() - u_old.toarray())) + + buffer_array = np.array([sum_u_diff_loc]) + + if particles.mpi_comm is not None: + particles.mpi_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + diff = buffer_array[0] + + buffer_array = np.array([sum_H_diff_loc]) + + if particles.mpi_comm is not None: + particles.mpi_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + if particles.clone_config is not None: + particles.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + buffer_array, + op=MPI.SUM, + ) + + diff += buffer_array[0] + + # check convergence + if diff < self.options.dg_solver_params.tol: + if self.options.dg_solver_params.verbose and MPI.COMM_WORLD.Get_rank() == 0: + print("converged diff: ", diff) + print("converged e_diff: ", e_diff) + + if particles.mpi_comm is not None: + particles.mpi_comm.Barrier() + break + + else: + if self.options.dg_solver_params.verbose and MPI.COMM_WORLD.Get_rank() == 0: + print("not converged diff: ", diff) + print("not converged e_diff: ", e_diff) + + if iter_num == self.options.dg_solver_params.maxiter: + if self.options.dg_solver_params.info and MPI.COMM_WORLD.Get_rank() == 0: + print( + f"{iter_num = }, maxiter={self.options.dg_solver_params.maxiter} reached! diff: {diff}, e_diff: {e_diff}", + ) + if particles.mpi_comm is not None: + particles.mpi_comm.Barrier() + break + + # sorting markers + if particles.mpi_comm is not None: + particles.mpi_sort_markers() + else: + particles.apply_kinetic_bc() + + # update u coefficients + diffs = self.update_feec_variables(u=u_new) + + # clear the buffer + markers[:, first_init_idx:-2] = 0.0 + + # update_weights + if self.variables.energetic_ions.species.weights_params.control_variate: + particles.update_weights() + + if self.options.dg_solver_params.info and MPI.COMM_WORLD.Get_rank() == 0: + print("Maxdiff up for CurrentCoupling5DGradB:", diffs["u"]) + print() diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index f9579553f..c7515c10c 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -60,6 +60,7 @@ from struphy.ode.solvers import ODEsolverFEEC from struphy.ode.utils import ButcherTableau, OptsButcher from struphy.pic.accumulation import accum_kernels, accum_kernels_gc +from struphy.pic.accumulation.filter import FilterParameters from struphy.pic.accumulation.particles_to_grid import Accumulator, AccumulatorVector from struphy.pic.base import Particles from struphy.pic.particles import Particles5D, Particles6D @@ -1907,9 +1908,9 @@ class ShearAlfvenCurrentCoupling5D(Propagator): \left\{ \begin{aligned} - \int \rho_0 &\frac{\partial \tilde{\mathbf U}}{\partial t} \cdot \mathbf V \, \textnormal{d} \mathbf{x} = \int \left(\tilde{\mathbf B} - \frac{A_\textnormal{h}}{A_b} \iint f^\text{vol} \mu \mathbf{b}_0\textnormal{d} \mu \textnormal{d} v_\parallel \right) \cdot \nabla \times (\mathbf B_0 \times \mathbf V) \, \textnormal{d} \mathbf{x} \quad \forall \, \mathbf V \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\}\,, \,, + \int \rho_0 &\frac{\partial \tilde{\mathbf U}}{\partial t} \cdot \mathbf V \, \textnormal{d} \mathbf{x} = \int \left(\tilde{\mathbf B} - \frac{A_\textnormal{h}}{A_b} \iint f^\text{vol} \mu \mathbf{b}_0\textnormal{d} \mu \textnormal{d} v_\parallel \right) \cdot \nabla \times (\tilde{\mathbf B} \times \mathbf V) \, \textnormal{d} \mathbf{x} \quad \forall \, \mathbf V \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\}\,, \,, \\ - &\frac{\partial \tilde{\mathbf B}}{\partial t} = - \nabla \times (\mathbf B_0 \times \tilde{\mathbf U}) \,. + &\frac{\partial \tilde{\mathbf B}}{\partial t} = - \nabla \times (\tilde{\mathbf B} \times \tilde{\mathbf U}) \,. \end{aligned} \right. @@ -1922,499 +1923,242 @@ class ShearAlfvenCurrentCoupling5D(Propagator): \end{bmatrix} = \frac{\Delta t}{2} \,. \begin{bmatrix} - 0 & (\mathbb M^{\alpha,n})^{-1} \mathcal {T^\alpha}^\top \mathbb C^\top \\ - \mathbb C \mathcal {T^\alpha} (\mathbb M^{\alpha,n})^{-1} & 0 + 0 & (\mathbb M^{2,n})^{-1} \mathcal {T^2}^\top \mathbb C^\top \\ - \mathbb C \mathcal {T^2} (\mathbb M^{2,n})^{-1} & 0 \end{bmatrix} \begin{bmatrix} - {\mathbb M^{\alpha,n}}(\mathbf u^{n+1} + \mathbf u^n) \\ \mathbb M_2(\mathbf b^{n+1} + \mathbf b^n) + \sum_k^{N_p} \omega_k \mu_k \hat{\mathbf b}¹_0 (\boldsymbol \eta_k) \cdot \left(\frac{1}{\sqrt{g(\boldsymbol \eta_k)}} \vec \Lambda² (\boldsymbol \eta_k) \right) + {\mathbb M^{2,n}}(\mathbf u^{n+1} + \mathbf u^n) \\ \mathbb M_2(\mathbf b^{n+1} + \mathbf b^n) + \sum_k^{N_p} \omega_k \mu_k \hat{\mathbf b}¹_0 (\boldsymbol \eta_k) \cdot \left(\frac{1}{\sqrt{g(\boldsymbol \eta_k)}} \vec \Lambda² (\boldsymbol \eta_k) \right) \end{bmatrix} \,, where - :math:`\mathcal{T}^\alpha` is a :class:`~struphy.feec.basis_projection_ops.BasisProjectionOperators` and - :math:`\mathbb M^{\alpha,n}` is a :class:`~struphy.feec.mass.WeightedMassOperators` being weighted with :math:`\rho_\text{eq}`, the MHD equilibirum density. - :math:`\alpha \in \{1, 2, v\}` denotes the :math:`\alpha`-form space where the operators correspond to. - Moreover, :math:`\sum_k^{N_p} \omega_k \mu_k \hat{\mathbf b}¹_0 (\boldsymbol \eta_k) \cdot \left(\frac{1}{\sqrt{g(\boldsymbol \eta_k)}} \vec \Lambda² (\boldsymbol \eta_k)\right)` is accumulated by the kernel :class:`~struphy.pic.accumulation.accum_kernels_gc.cc_lin_mhd_5d_M`. + :math:`\mathcal{T}^2 = \hat \Pi \left[\frac{\tilde{\mathbf B}^2}{\sqrt{g} \times \vec \Lambda^2\right]` and + :math:`\mathbb M^{2,n}` is a :class:`~struphy.feec.mass.WeightedMassOperators` being weighted with :math:`\rho_\text{eq}`, the MHD equilibirum density. """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["filter"] = { - "use_filter": None, - "modes": (1), - "repeat": 1, - "alpha": 0.5, - } - dct["boundary_cut"] = { - "e1": 0.0, - "e2": 0.0, - "e3": 0.0, - } - dct["turn_off"] = False - - if default: - dct = descend_options_dict(dct, []) - - return dct - - def __init__( - self, - u: BlockVector, - b: BlockVector, - *, - particles: Particles5D, - absB0: StencilVector, - unit_b1: BlockVector, - u_space: str, - solver: dict = options(default=True)["solver"], - filter: dict = options(default=True)["filter"], - coupling_params: dict, - accumulated_magnetization: BlockVector, - boundary_cut: dict = options(default=True)["boundary_cut"], - ): - super().__init__(u, b) - - self._particles = particles - self._unit_b1 = unit_b1 - self._absB0 = absB0 - - self._info = solver["info"] - - self._scale_vec = coupling_params["Ah"] / coupling_params["Ab"] - - self._E1T = self.derham.extraction_ops["1"].transpose() - self._unit_b1 = self._E1T.dot(self._unit_b1) + class Variables: + def __init__(self): + self._u: FEECVariable = None + self._b: FEECVariable = None - self._accumulated_magnetization = accumulated_magnetization + @property + def u(self) -> FEECVariable: + return self._u - self._boundary_cut_e1 = boundary_cut["e1"] + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new - self._ACC = Accumulator( - particles, - u_space, - accum_kernels_gc.cc_lin_mhd_5d_M, - self.mass_ops, - self.domain.args_domain, - add_vector=True, - symmetry="symm", - filter_params=filter, - ) + @property + def b(self) -> FEECVariable: + return self._b - # if self._particles.control_variate: + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._b = new - # # control variate method is only valid with Maxwellian distributions with "zero perp mean velocity". - # assert isinstance(self._particles.f0, Maxwellian) + def __init__(self): + self.variables = self.Variables() - # self._ACC.init_control_variate(self.mass_ops) + @dataclass + class Options: + # specific literals + OptsAlgo = Literal["implicit", "explicit"] + # propagator options + energetic_ions: PICVariable = None + ep_scale: float = 1.0 + u_space: OptsVecSpace = "Hdiv" + algo: OptsAlgo = "implicit" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixDiagonalPreconditioner" + solver_params: SolverParameters = None + filter_params: FilterParameters = None + butcher: ButcherTableau = None + nonlinear: bool = True - # # evaluate and save f0.n at quadrature points - # quad_pts = [quad_grid[nquad].points.flatten() - # for quad_grid, nquad in zip(self.derham.get_quad_grids(self.derham.Vh_fem['0']), self.derham.nquads)] + def __post_init__(self): + # checks + check_option(self.u_space, OptsVecSpace) + check_option(self.algo, self.OptsAlgo) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + assert isinstance(self.energetic_ions, PICVariable) + assert self.energetic_ions.space == "Particles5D" + assert isinstance(self.ep_scale, float) + assert isinstance(self.nonlinear, bool) - # n0_at_quad = self.domain.push( - # self._particles.f0.n, *quad_pts, kind='0', squeeze_out=False) + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() - # # evaluate M0 = unit_b1 (1form) / absB0 (0form) * 2 * vth_perp² at quadrature points - # quad_pts_array = self.domain.prepare_eval_pts(*quad_pts)[:3] + if self.filter_params is None: + self.filter_params = FilterParameters() - # vth_perp = self.particles.f0.vth(*quad_pts_array)[1] + if self.algo == "explicit" and self.butcher is None: + self.butcher = ButcherTableau() - # absB0_at_quad = WeightedMassOperator.eval_quad(self.derham.Vh_fem['0'], self._absB0) + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options - # unit_b1_at_quad = WeightedMassOperator.eval_quad(self.derham.Vh_fem['1'], self._unit_b1) + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new - # self._M0_at_quad = unit_b1_at_quad / absB0_at_quad * vth_perp**2 * n0_at_quad * self._scale_vec + @profile + def allocate(self): + self._u_form = self.derham.space_to_form[self.options.u_space] - # define block matrix [[A B], [C I]] (without time step size dt in the diagonals) - id_M = "M" + self.derham.space_to_form[u_space] + "n" - id_T = "T" + self.derham.space_to_form[u_space] + # call operatros + id_M = "M" + self._u_form + "n" + id_T = "T" + self._u_form _A = getattr(self.mass_ops, id_M) _T = getattr(self.basis_ops, id_T) + M2 = self.mass_ops.M2 + curl = self.derham.curl + PB = getattr(self.basis_ops, "PB") - self._B = -1 / 2 * _T.T @ self.derham.curl.T @ self.mass_ops.M2 - self._C = 1 / 2 * self.derham.curl @ _T - self._B2 = -1 / 2 * _T.T @ self.derham.curl.T + # define Accumulator and arguments + self._ACC = AccumulatorVector( + self.options.energetic_ions.particles, + "H1", + accum_kernels_gc.gc_mag_density_0form, + self.mass_ops, + self.domain.args_domain, + filter_params=self.options.filter_params, + ) # Preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(getattr(self.mass_ops, id_M)) - # Instantiate Schur solver (constant in this case) - _BC = self._B @ self._C - - self._schur_solver = SchurSolver( - _A, - _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], - ) - - # allocate dummy vectors to avoid temporary array allocations - self._u_tmp1 = u.space.zeros() - self._u_tmp2 = u.space.zeros() - self._b_tmp1 = b.space.zeros() - - self._byn = self._B.codomain.zeros() - self._tmp_acc = self._B2.codomain.zeros() - - def __call__(self, dt): - # current variables - un = self.feec_vars[0] - bn = self.feec_vars[1] - - # perform accumulation (either with or without control variate) - # if self._particles.control_variate: - - # self._ACC.accumulate(self._particles, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._scale_vec, 0., - # control_vec=[self._M0_at_quad[0], self._M0_at_quad[1], self._M0_at_quad[2]]) - # else: - # self._ACC.accumulate(self._particles, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._scale_vec, 0.) - - self._ACC( - self._unit_b1[0]._data, - self._unit_b1[1]._data, - self._unit_b1[2]._data, - self._scale_vec, - self._boundary_cut_e1, - ) - - self._ACC.vectors[0].copy(out=self._accumulated_magnetization) - - # solve for new u coeffs (no tmps created here) - byn = self._B.dot(bn, out=self._byn) - b2acc = self._B2.dot(self._ACC.vectors[0], out=self._tmp_acc) - byn += b2acc - - # b2acc.copy(out=self._accumulated_magnetization) - - un1, info = self._schur_solver(un, byn, dt, out=self._u_tmp1) - - # new b coeffs (no tmps created here) - _u = un.copy(out=self._u_tmp2) - _u += un1 - bn1 = self._C.dot(_u, out=self._b_tmp1) - bn1 *= -dt - bn1 += bn - - # write new coeffs into self.feec_vars - max_du, max_db = self.feec_vars_update(un1, bn1) - - if self._info and MPI.COMM_WORLD.Get_rank() == 0: - print("Status for ShearAlfven:", info["success"]) - print("Iterations for ShearAlfven:", info["niter"]) - print("Maxdiff up for ShearAlfven:", max_du) - print("Maxdiff b2 for ShearAlfven:", max_db) - print() - - -class MagnetosonicCurrentCoupling5D(Propagator): - r""" - :ref:`FEEC ` discretization of the following equations: - find :math:`\tilde \rho \in L^2, \tilde{\mathbf U} \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\}, \tilde p \in L^2` such that - - .. math:: - - \left\{ - \begin{aligned} - &\frac{\partial \tilde{\rho}}{\partial t} = - \nabla \cdot (\rho_0 \tilde{\mathbf U}) \,, - \\ - \int \rho_0 &\frac{\partial \tilde{\mathbf U}}{\partial t} \cdot \mathbf V \, \textnormal{d} \mathbf{x} = \int (\nabla \times \mathbf B_0) \times \tilde{\mathbf B} \cdot \mathbf V \, \textnormal{d} \mathbf x + \frac{A_\textnormal{h}}{A_b}\iint f^\text{vol} \mu \mathbf b_0 \cdot \nabla \times (\tilde{\mathbf B} \times \mathbf V) \, \textnormal{d} \mathbf x \textnormal{d} v_\parallel \textnormal{d} \mu + \int \tilde p \nabla \cdot \mathbf V \, \textnormal{d} \mathbf x \qquad \forall \, \mathbf V \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\}\,, - \\ - &\frac{\partial \tilde p}{\partial t} = - \nabla \cdot (p_0 \tilde{\mathbf U}) - (\gamma - 1) p_0 \nabla \cdot \tilde{\mathbf U} \,. - \end{aligned} - \right. - - :ref:`time_discret`: Crank-Nicolson (implicit mid-point). System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`: - - .. math:: - - \boldsymbol{\rho}^{n+1} - \boldsymbol{\rho}^n = - \frac{\Delta t}{2} \mathbb D \mathcal Q^\alpha (\mathbf u^{n+1} + \mathbf u^n) \,, - - .. math:: - - \begin{bmatrix} - \mathbf u^{n+1} - \mathbf u^n \\ \mathbf p^{n+1} - \mathbf p^n - \end{bmatrix} - = \frac{\Delta t}{2} - \begin{bmatrix} - 0 & (\mathbb M^{\alpha,n})^{-1} {\mathcal U^\alpha}^\top \mathbb D^\top \mathbb M_3 \\ - \mathbb D \mathcal S^\alpha - (\gamma - 1) \mathcal K^\alpha \mathbb D \mathcal U^\alpha & 0 - \end{bmatrix} - \begin{bmatrix} - (\mathbf u^{n+1} + \mathbf u^n) \\ (\mathbf p^{n+1} + \mathbf p^n) - \end{bmatrix} + - \begin{bmatrix} - \Delta t (\mathbb M^{\alpha,n})^{-1}\left[\mathbb M^{\alpha,J} \mathbf b^n + \frac{A_\textnormal{h}}{A_b}{\mathcal{T}^B}^\top \mathbb{C}^\top \sum_k^{N_p} \omega_k \mu_k \hat{\mathbf b}¹_0 (\boldsymbol \eta_k) \cdot \left(\frac{1}{\sqrt{g(\boldsymbol \eta_k)}} \vec \Lambda² (\boldsymbol \eta_k) \right)\right] \\ 0 - \end{bmatrix} \,, - - where - :math:`\mathcal U^\alpha`, :math:`\mathcal S^\alpha`, :math:`\mathcal K^\alpha` and :math:`\mathcal Q^\alpha` are :class:`~struphy.feec.basis_projection_ops.BasisProjectionOperators` and - :math:`\mathbb M^{\alpha,n}` and :math:`\mathbb M^{\alpha,J}` are :class:`~struphy.feec.mass.WeightedMassOperators` being weighted with :math:`\rho_0` the MHD equilibrium density. - :math:`\alpha \in \{1, 2, v\}` denotes the :math:`\alpha`-form space where the operators correspond to. - Moreover, :math:`\sum_k^{N_p} \omega_k \mu_k \hat{\mathbf b}¹_0 (\boldsymbol \eta_k) \cdot \left(\frac{1}{\sqrt{g(\boldsymbol \eta_k)}} \vec \Lambda² (\boldsymbol \eta_k)\right)` is accumulated by the kernel :class:`~struphy.pic.accumulation.accum_kernels_gc.cc_lin_mhd_5d_M` and - the time-varying projection operator :math:`\mathcal{T}^B` is defined as - - .. math:: - - \mathcal{T}^B_{(\mu,ijk),(\nu,mno)} := \hat \Pi¹_{(\mu,ijk)} \left[ \epsilon_{\mu \alpha \nu} \frac{\tilde{B}^2_\alpha}{\sqrt{g}} \Lambda²_{\nu,mno} \right] \,. - """ - - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pbicgstab", "MassMatrixPreconditioner"), - ("bicgstab", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["filter"] = { - "use_filter": None, - "modes": (0, 1), - "repeat": 3, - "alpha": 0.5, - } - dct["boundary_cut"] = { - "e1": 0.0, - "e2": 0.0, - "e3": 0.0, - } - dct["turn_off"] = False - - if default: - dct = descend_options_dict(dct, []) + if self.options.nonlinear: + # initialize operator TB + self._initialize_projection_operator_TB() - return dct + _T = _T + self._TB + _TT = _T.T + self._TBT - def __init__( - self, - n: StencilVector, - u: BlockVector, - p: StencilVector, - *, - particles: Particles5D, - b: BlockVector, - absB0: StencilVector, - unit_b1: BlockVector, - u_space: str, - solver: dict = options(default=True)["solver"], - filter: dict = options(default=True)["filter"], - coupling_params: dict, - boundary_cut: dict = options(default=True)["boundary_cut"], - ): - super().__init__(n, u, p) - - self._particles = particles - self._b = b - self._unit_b1 = unit_b1 - self._absB0 = absB0 - - self._info = solver["info"] - - self._scale_vec = coupling_params["Ah"] / coupling_params["Ab"] - - self._E1T = self.derham.extraction_ops["1"].transpose() - self._unit_b1 = self._E1T.dot(self._unit_b1) - - self._u_id = self.derham.space_to_form[u_space] - if self._u_id == "v": - self._space_key_int = 0 else: - self._space_key_int = int(self._u_id) + _TT = _T.T - self._boundary_cut_e1 = boundary_cut["e1"] - - self._ACC = Accumulator( - particles, - u_space, - accum_kernels_gc.cc_lin_mhd_5d_M, - self.mass_ops, - self.domain.args_domain, - add_vector=True, - symmetry="symm", - filter_params=filter, - ) - - # if self._particles.control_variate: - - # # control variate method is only valid with Maxwellian distributions with "zero perp mean velocity". - # assert isinstance(self._particles.f0, Maxwellian) - - # self._ACC.init_control_variate(self.mass_ops) - - # # evaluate and save f0.n at quadrature points - # quad_pts = [quad_grid[nquad].points.flatten() - # for quad_grid, nquad in zip(self.derham.get_quad_grids(self.derham.Vh_fem['0']), self.derham.nquads)] - - # n0_at_quad = self.domain.push( - # self._particles.f0.n, *quad_pts, kind='0', squeeze_out=False) - - # # evaluate M0 = unit_b1 (1form) / absB0 (0form) * 2 * vth_perp² at quadrature points - # quad_pts_array = self.domain.prepare_eval_pts(*quad_pts)[:3] - - # vth_perp = self.particles.f0.vth(*quad_pts_array)[1] - - # absB0_at_quad = WeightedMassOperator.eval_quad(self.derham.Vh_fem['0'], self._absB0) - - # unit_b1_at_quad = WeightedMassOperator.eval_quad(self.derham.Vh_fem['1'], self._unit_b1) - - # self._M0_at_quad = unit_b1_at_quad / absB0_at_quad * vth_perp**2 * n0_at_quad * self._scale_vec - - # define block matrix [[A B], [C I]] (without time step size dt in the diagonals) - id_Mn = "M" + self._u_id + "n" - id_MJ = "M" + self._u_id + "J" - - if self._u_id == "1": - id_S, id_U, id_K, id_Q = "S1", "U1", "K3", "Q1" - elif self._u_id == "2": - id_S, id_U, id_K, id_Q = "S2", None, "K3", "Q2" - elif self._u_id == "v": - id_S, id_U, id_K, id_Q = "Sv", "Uv", "K3", "Qv" + if self.options.algo == "implicit": + self._info = self.options.solver_params.info - self._E2T = self.derham.extraction_ops["2"].transpose() + # define block matrix [[A B], [C I]] (without time step size dt in the diagonals) + self._B = -1 / 2 * _TT @ curl.T @ M2 + self._B2 = -1 / 2 * _TT @ curl.T @ PB.T - _A = getattr(self.mass_ops, id_Mn) - _S = getattr(self.basis_ops, id_S) - _K = getattr(self.basis_ops, id_K) + self._C = 1 / 2 * curl @ _T - # initialize projection operator TB - self._initialize_projection_operator_TB() + # Instantiate Schur solver (constant in this case) + _BC = self._B @ self._C - if id_U is None: - _U, _UT = IdentityOperator(u.space), IdentityOperator(u.space) - else: - _U = getattr(self.basis_ops, id_U) - _UT = _U.T + self._schur_solver = SchurSolver( + _A, + _BC, + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, + ) - self._B = -1 / 2.0 * _UT @ self.derham.div.T @ self.mass_ops.M3 - self._C = 1 / 2.0 * (self.derham.div @ _S + 2 / 3.0 * _K @ self.derham.div @ _U) + # allocate dummy vectors to avoid temporary array allocations + self._u_tmp1 = self.variables.u.spline.vector.space.zeros() + self._u_tmp2 = self.variables.u.spline.vector.space.zeros() + self._b_tmp1 = self.variables.b.spline.vector.space.zeros() - self._MJ = getattr(self.mass_ops, id_MJ) - self._DQ = self.derham.div @ getattr(self.basis_ops, id_Q) - - self._TC = self._TB.T @ self.derham.curl.T + self._byn = self._B.codomain.zeros() + self._tmp_acc = self._B2.codomain.zeros() - # preconditioner - if solver["type"][1] is None: - pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) - pc = pc_class(getattr(self.mass_ops, id_Mn)) - - # instantiate Schur solver (constant in this case) - _BC = self._B @ self._C + self._info = False - self._schur_solver = SchurSolver( - _A, - _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], - ) + # define vector field + A_inv = inverse( + _A, + self.options.solver, + pc=pc, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + ) + _f1 = A_inv @ _TT @ curl.T @ M2 + _f1_acc = A_inv @ _TT @ curl.T @ PB.T + _f2 = curl @ _T - # allocate dummy vectors to avoid temporary array allocations - self._u_tmp1 = u.space.zeros() - self._u_tmp2 = u.space.zeros() - self._p_tmp1 = p.space.zeros() - self._n_tmp1 = n.space.zeros() - self._byn1 = self._B.codomain.zeros() - self._byn2 = self._B.codomain.zeros() - self._tmp_acc = self._TC.codomain.zeros() + # allocate output of vector field + out_acc = self.variables.u.spline.vector.space.zeros() + out1 = self.variables.u.spline.vector.space.zeros() + out2 = self.variables.b.spline.vector.space.zeros() - def __call__(self, dt): - # current variables - nn = self.feec_vars[0] - un = self.feec_vars[1] - pn = self.feec_vars[2] + def f1(t, y1, y2, out: BlockVector = out1): + _f1.dot(y2, out=out) + _f1_acc.dot(self._ACC.vectors[0], out=out_acc) + out += out_acc + out.update_ghost_regions() + return out - # perform accumulation (either with or without control variate) - # if self._particles.control_variate: + def f2(t, y1, y2, out: BlockVector = out2): + _f2.dot(y1, out=out) + out *= -1.0 + out.update_ghost_regions() + return out - # self._ACC.accumulate(self._particles, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._scale_vec, 0., - # control_vec=[self._M0_at_quad[0], self._M0_at_quad[1], self._M0_at_quad[2]]) - # else: - # self._ACC.accumulate(self._particles, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._scale_vec, 0.) + vector_field = {self.variables.u.spline.vector: f1, self.variables.b.spline.vector: f2} + self._ode_solver = ODEsolverFEEC(vector_field, butcher=self.options.butcher) - self._ACC( - self._unit_b1[0]._data, - self._unit_b1[1]._data, - self._unit_b1[2]._data, - self._scale_vec, - self._boundary_cut_e1, - ) + def __call__(self, dt): + # update time-dependent operator TB + if self.options.nonlinear: + self._update_weights_TB() - # update time-dependent operator - self._b.update_ghost_regions() - self._update_weights_TB() + # current FE coeffs + un = self.variables.u.spline.vector + bn = self.variables.b.spline.vector - # solve for new u coeffs (no tmps created here) - byn1 = self._B.dot(pn, out=self._byn1) - byn2 = self._MJ.dot(self._b, out=self._byn2) - b2acc = self._TC.dot(self._ACC.vectors[0], out=self._tmp_acc) - byn2 += b2acc - byn2 *= 1 / 2 - byn1 -= byn2 + # accumulate + self._ACC(self.options.ep_scale) - un1, info = self._schur_solver(un, byn1, dt, out=self._u_tmp1) + if self.options.algo == "implicit": + # solve for new u coeffs (no tmps created here) + byn = self._B.dot(bn, out=self._byn) + b2acc = self._B2.dot(self._ACC.vectors[0], out=self._tmp_acc) + byn += b2acc - # new p, n, b coeffs (no tmps created here) - _u = un.copy(out=self._u_tmp2) - _u += un1 - pn1 = self._C.dot(_u, out=self._p_tmp1) - pn1 *= -dt - pn1 += pn + un1, info = self._schur_solver(un, byn, dt, out=self._u_tmp1) - nn1 = self._DQ.dot(_u, out=self._n_tmp1) - nn1 *= -dt / 2 - nn1 += nn + # new b coeffs (no tmps created here) + _u = un.copy(out=self._u_tmp2) + _u += un1 + bn1 = self._C.dot(_u, out=self._b_tmp1) + bn1 *= -dt + bn1 += bn - # write new coeffs into self.feec_vars - max_dn, max_du, max_dp = self.feec_vars_update( - nn1, - un1, - pn1, - ) + diffs = self.update_feec_variables(u=un1, b=bn1) + + else: + self._ode_solver(0.0, dt) if self._info and MPI.COMM_WORLD.Get_rank() == 0: - print("Status for Magnetosonic:", info["success"]) - print("Iterations for Magnetosonic:", info["niter"]) - print("Maxdiff n3 for Magnetosonic:", max_dn) - print("Maxdiff up for Magnetosonic:", max_du) - print("Maxdiff p3 for Magnetosonic:", max_dp) - print() + if self.options.algo == "implicit": + print("Status for ShearAlfvenCurrentCoupling5D:", info["success"]) + print("Iterations for ShearAlfvenCurrentCoupling5D:", info["niter"]) + print("Maxdiff up for ShearAlfvenCurrentCoupling5D:", diffs["u"]) + print("Maxdiff b2 for ShearAlfvenCurrentCoupling5D:", diffs["b"]) + print() def _initialize_projection_operator_TB(self): r"""Initialize BasisProjectionOperator TB with the time-varying weight. @@ -2427,27 +2171,80 @@ def _initialize_projection_operator_TB(self): # Call the projector and the space P1 = self.derham.P["1"] - Vh = self.derham.Vh_fem[self._u_id] + Vh = self.derham.Vh_fem[self._u_form] # Femfield for the field evaluation self._bf = self.derham.create_spline_function("bf", "Hdiv") - # define temp callable - def tmp(x, y, z): - return 0 * x - # Initialize BasisProjectionOperator - if self.derham._with_local_projectors: - self._TB = BasisProjectionOperatorLocal(P1, Vh, [[tmp, tmp, tmp]]) + if self.derham._with_local_projectors == True: + self._TB = BasisProjectionOperatorLocal( + P1, + Vh, + [ + [None, None, None], + [None, None, None], + [None, None, None], + ], + transposed=False, + use_cache=True, + polar_shift=True, + V_extraction_op=self.derham.extraction_ops[self._u_form], + V_boundary_op=self.derham.boundary_ops[self._u_form], + P_boundary_op=self.derham.boundary_ops["1"], + ) + self._TBT = BasisProjectionOperatorLocal( + P1, + Vh, + [ + [None, None, None], + [None, None, None], + [None, None, None], + ], + transposed=True, + use_cache=True, + polar_shift=True, + V_extraction_op=self.derham.extraction_ops[self._u_form], + V_boundary_op=self.derham.boundary_ops[self._u_form], + P_boundary_op=self.derham.boundary_ops["1"], + ) else: - self._TB = BasisProjectionOperator(P1, Vh, [[tmp, tmp, tmp]]) + self._TB = BasisProjectionOperator( + P1, + Vh, + [ + [None, None, None], + [None, None, None], + [None, None, None], + ], + transposed=False, + use_cache=True, + polar_shift=True, + V_extraction_op=self.derham.extraction_ops[self._u_form], + V_boundary_op=self.derham.boundary_ops[self._u_form], + P_boundary_op=self.derham.boundary_ops["1"], + ) + self._TBT = BasisProjectionOperator( + P1, + Vh, + [ + [None, None, None], + [None, None, None], + [None, None, None], + ], + transposed=True, + use_cache=True, + polar_shift=True, + V_extraction_op=self.derham.extraction_ops[self._u_form], + V_boundary_op=self.derham.boundary_ops[self._u_form], + P_boundary_op=self.derham.boundary_ops["1"], + ) def _update_weights_TB(self): """Updats time-dependent weights of the BasisProjectionOperator TB""" # Update Femfield - self._bf.vector = self._b - self._bf.vector.update_ghost_regions() + self.variables.b.spline.vector.copy(out=self._bf.vector) # define callable weights def bf1(x, y, z): @@ -2465,7 +2262,7 @@ def bf3(x, y, z): fun = [] - if self._u_id == "v": + if self._u_form == "v": for m in range(3): fun += [[]] for n in range(3): @@ -2473,7 +2270,7 @@ def bf3(x, y, z): lambda e1, e2, e3, m=m, n=n: rot_B(e1, e2, e3)[:, :, :, m, n], ] - elif self._u_id == "1": + elif self._u_form == "1": for m in range(3): fun += [[]] for n in range(3): @@ -2499,8 +2296,9 @@ def bf3(x, y, z): / abs(self.domain.jacobian_det(e1, e2, e3, squeeze_out=False)), ] - # Initialize BasisProjectionOperator + # update BasisProjectionOperator self._TB.update_weights(fun) + self._TBT.update_weights(fun) class CurrentCoupling5DDensity(Propagator): @@ -2520,268 +2318,166 @@ class CurrentCoupling5DDensity(Propagator): For the detail explanation of the notations, see `2022_DriftKineticCurrentCoupling `_. """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pbicgstab", "MassMatrixPreconditioner"), - ("bicgstab", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["filter"] = { - "use_filter": None, - "modes": (1), - "repeat": 1, - "alpha": 0.5, - } - dct["boundary_cut"] = { - "e1": 0.0, - "e2": 0.0, - "e3": 0.0, - } - dct["turn_off"] = False + class Variables: + def __init__(self): + self._u: FEECVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def u(self) -> FEECVariable: + return self._u - return dct + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new - def __init__( - self, - u: BlockVector, - *, - particles: Particles5D, - b: BlockVector, - b_eq: BlockVector, - unit_b1: BlockVector, - curl_unit_b2: BlockVector, - u_space: str, - solver: dict = options(default=True)["solver"], - coupling_params: dict, - epsilon: float = 1.0, - filter: dict = options(default=True)["filter"], - boundary_cut: dict = options(default=True)["boundary_cut"], - ): - super().__init__(u) + def __init__(self): + self.variables = self.Variables() - # assert parameters and expose some quantities to self - assert isinstance(particles, (Particles5D)) + @dataclass + class Options: + # propagator options + energetic_ions: PICVariable = None + b_tilde: FEECVariable = None + ep_scale: float = 1.0 + u_space: OptsVecSpace = "Hdiv" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + filter_params: FilterParameters = None + + def __post_init__(self): + # checks + check_option(self.u_space, OptsVecSpace) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + assert isinstance(self.energetic_ions, PICVariable) + assert self.energetic_ions.space == "Particles5D" + assert isinstance(self.b_tilde, FEECVariable) + assert isinstance(self.ep_scale, float) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() - assert u_space in {"Hcurl", "Hdiv", "H1vec"} + if self.filter_params is None: + self.filter_params = FilterParameters() - if u_space == "H1vec": - self._space_key_int = 0 + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + if self.options.u_space == "H1vec": + self._u_form_int = 0 else: - self._space_key_int = int( - self.derham.space_to_form[u_space], - ) + self._u_form_int = int(self.derham.space_to_form[self.options.u_space]) - self._epsilon = epsilon - self._particles = particles - self._b = b - self._b_eq = b_eq - self._unit_b1 = unit_b1 - self._curl_norm_b = curl_unit_b2 + # call operatros + id_M = "M" + self.derham.space_to_form[self.options.u_space] + "n" + self._A = getattr(self.mass_ops, id_M) - self._info = solver["info"] + # magnetic equilibrium field + unit_b1 = self.projected_equil.unit_b1 + curl_unit_b1 = self.projected_equil.curl_unit_b1 + self._b2 = self.projected_equil.b2 - self._scale_mat = coupling_params["Ah"] / coupling_params["Ab"] / self._epsilon + # scaling factor + epsilon = self.options.energetic_ions.species.equation_params.epsilon - self._boundary_cut_e1 = boundary_cut["e1"] + # temporary vectors to avoid memory allocation + self._b_full = self._b2.space.zeros() + self._rhs_v = self.variables.u.spline.vector.space.zeros() + self._u_new = self.variables.u.spline.vector.space.zeros() - self._accumulator = Accumulator( - particles, - u_space, + # define Accumulator and arguments + self._ACC = Accumulator( + self.options.energetic_ions.particles, + self.options.u_space, accum_kernels_gc.cc_lin_mhd_5d_D, self.mass_ops, self.domain.args_domain, add_vector=False, symmetry="asym", - filter_params=filter, + filter_params=self.options.filter_params, ) - # if self._particles.control_variate: - - # # control variate method is only valid with Maxwellian distributions - # assert isinstance(self._particles.f0, Maxwellian) - # assert params['u_space'] == 'Hdiv' - - # # evaluate and save f0.n / |det(DF)| at quadrature points - # quad_pts = [quad_grid[nquad].points.flatten() - # for quad_grid, nquad in zip(self.derham.get_quad_grids(self.derham.Vh_fem['0']), self.derham.nquads)] - - # self._n0_at_quad = self.domain.push( - # self._particles.f0.n, *quad_pts, kind='3', squeeze_out=False) - - # # prepare field evaluation - # quad_pts_array = self.domain.prepare_eval_pts(*quad_pts)[:3] - - # u0_parallel = self._particles.f0.u(*quad_pts_array)[0] - - # det_df_at_quad = self.domain.jacobian_det(*quad_pts, squeeze_out=False) - - # # evaluate unit_b1 / |det(DF)| at quadrature points - # self._unit_b1_at_quad = WeightedMassOperator.eval_quad(self.derham.Vh_fem['1'], self._unit_b1) - # self._unit_b1_at_quad /= det_df_at_quad - - # # evaluate unit_b1 (1form) dot epsilon * f0.u * curl_norm_b (2form) / |det(DF)| at quadrature points - # curl_norm_b_at_quad = WeightedMassOperator.eval_quad(self.derham.Vh_fem['2'], self._curl_norm_b) - - # self._unit_b1_dot_curl_norm_b_at_quad = np.sum(p * q for p, q in zip(self._unit_b1_at_quad, curl_norm_b_at_quad)) - - # self._unit_b1_dot_curl_norm_b_at_quad /= det_df_at_quad - # self._unit_b1_dot_curl_norm_b_at_quad *= self._epsilon - # self._unit_b1_dot_curl_norm_b_at_quad *= u0_parallel - - # # memory allocation for magnetic field at quadrature points - # self._b_quad1 = np.zeros_like(self._n0_at_quad) - # self._b_quad2 = np.zeros_like(self._n0_at_quad) - # self._b_quad3 = np.zeros_like(self._n0_at_quad) - - # # memory allocation for parallel magnetic field at quadrature points - # self._B_para = np.zeros_like(self._n0_at_quad) - - # # memory allocation for control_const at quadrature points - # self._control_const = np.zeros_like(self._n0_at_quad) - - # # memory allocation for self._b_quad x self._nh0_at_quad * self._coupling_const - # self._mat12 = np.zeros_like(self._n0_at_quad) - # self._mat13 = np.zeros_like(self._n0_at_quad) - # self._mat23 = np.zeros_like(self._n0_at_quad) - - # self._mat21 = np.zeros_like(self._n0_at_quad) - # self._mat31 = np.zeros_like(self._n0_at_quad) - # self._mat32 = np.zeros_like(self._n0_at_quad) - - u_id = self.derham.space_to_form[u_space] - self._M = getattr(self.mass_ops, "M" + u_id + "n") - - self._E0T = self.derham.extraction_ops["0"].transpose() - self._EuT = self.derham.extraction_ops[u_id].transpose() - self._E1T = self.derham.extraction_ops["1"].transpose() - self._E2T = self.derham.extraction_ops["2"].transpose() - - self._PB = getattr(self.basis_ops, "PB") - self._unit_b1 = self._E1T.dot(self._unit_b1) + self._args_accum_kernel = ( + epsilon, + self.options.ep_scale, + self._b_full[0]._data, + self._b_full[1]._data, + self._b_full[2]._data, + unit_b1[0]._data, + unit_b1[1]._data, + unit_b1[2]._data, + curl_unit_b1[0]._data, + curl_unit_b1[1]._data, + curl_unit_b1[2]._data, + self._u_form_int, + ) - # preconditioner - if solver["type"][1] is None: - self._pc = None + # Preconditioner + if self.options.precond is None: + pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) - self._pc = pc_class(self._M) + pc_class = getattr(preconditioner, self.options.precond) + pc = pc_class(getattr(self.mass_ops, id_M)) # linear solver - self._solver = inverse( - self._M, - solver["type"][0], - pc=self._pc, - x0=self.feec_vars[0], - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], + self._A_inv = inverse( + self._A, + self.options.solver, + pc=pc, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, ) - # temporary vectors to avoid memory allocation - self._b_full1 = self._b_eq.space.zeros() - self._b_full2 = self._E2T.codomain.zeros() - self._rhs_v = u.space.zeros() - self._u_new = u.space.zeros() - def __call__(self, dt): - # pointer to old coefficients - un = self.feec_vars[0] + # current FE coeffs + un = self.variables.u.spline.vector # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) - b_full = self._b_eq.copy(out=self._b_full1) - - if self._b is not None: - b_full += self._b - - Eb_full = self._E2T.dot(b_full, out=self._b_full2) - Eb_full.update_ghost_regions() - - # perform accumulation (either with or without control variate) - # if self._particles.control_variate: - - # # evaluate magnetic field at quadrature points (in-place) - # WeightedMassOperator.eval_quad(self.derham.Vh_fem['2'], self._b_full2, - # out=[self._b_quad1, self._b_quad2, self._b_quad3]) - - # # evaluate B_parallel - # self._B_para = np.sum(p * q for p, q in zip(self._unit_b1_at_quad, [self._b_quad1, self._b_quad2, self._b_quad3])) - - # # evaluate coupling_const 1 - B_parallel / B^star_parallel - # self._control_const = 1 - (self._B_para / (self._B_para + self._unit_b1_dot_curl_norm_b_at_quad)) - - # # assemble (B x) - # self._mat12[:, :, :] = self._scale_mat * \ - # self._b_quad3 * self._n0_at_quad * self._control_const - # self._mat13[:, :, :] = -self._scale_mat * \ - # self._b_quad2 * self._n0_at_quad * self._control_const - # self._mat23[:, :, :] = self._scale_mat * \ - # self._b_quad1 * self._n0_at_quad * self._control_const - - # self._mat21[:, :, :] = -self._mat12 - # self._mat31[:, :, :] = -self._mat13 - # self._mat32[:, :, :] = -self._mat23 + b_full = self._b2.copy(out=self._b_full) - # self._accumulator.accumulate(self._particles, self._epsilon, - # Eb_full[0]._data, Eb_full[1]._data, Eb_full[2]._data, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._curl_norm_b[0]._data, self._curl_norm_b[1]._data, self._curl_norm_b[2]._data, - # self._space_key_int, self._scale_mat, 0.1, - # control_mat=[[None, self._mat12, self._mat13], - # [self._mat21, None, self._mat23], - # [self._mat31, self._mat32, None]]) - # else: - # self._accumulator.accumulate(self._particles, self._epsilon, - # Eb_full[0]._data, Eb_full[1]._data, Eb_full[2]._data, - # self._unit_b1[0]._data, self._unit_b1[1]._data, self._unit_b1[2]._data, - # self._curl_norm_b[0]._data, self._curl_norm_b[1]._data, self._curl_norm_b[2]._data, - # self._space_key_int, self._scale_mat, 0.) + b_full += self.options.b_tilde.spline.vector + b_full.update_ghost_regions() - self._accumulator( - self._epsilon, - Eb_full[0]._data, - Eb_full[1]._data, - Eb_full[2]._data, - self._unit_b1[0]._data, - self._unit_b1[1]._data, - self._unit_b1[2]._data, - self._curl_norm_b[0]._data, - self._curl_norm_b[1]._data, - self._curl_norm_b[2]._data, - self._space_key_int, - self._scale_mat, - self._boundary_cut_e1, + self._ACC( + *self._args_accum_kernel, ) # define system (M - dt/2 * A)*u^(n + 1) = (M + dt/2 * A)*u^n - lhs = self._M - dt / 2 * self._accumulator.operators[0] - rhs = self._M + dt / 2 * self._accumulator.operators[0] + lhs = self._A - dt / 2 * self._ACC.operators[0] + rhs = self._A + dt / 2 * self._ACC.operators[0] # solve linear system for updated u coefficients (in-place) rhs = rhs.dot(un, out=self._rhs_v) - self._solver.linop = lhs + self._A_inv.linop = lhs - un1 = self._solver.solve(rhs, out=self._u_new) - info = self._solver._info + _u = self._A_inv.solve(rhs, out=self._u_new) + info = self._A_inv._info - # write new coeffs into Propagator.variables - max_du = self.feec_vars_update(un1) + diffs = self.update_feec_variables(u=_u) - if self._info and MPI.COMM_WORLD.Get_rank() == 0: + if self.options.solver_params.info and MPI.COMM_WORLD.Get_rank() == 0: print("Status for CurrentCoupling5DDensity:", info["success"]) print("Iterations for CurrentCoupling5DDensity:", info["niter"]) - print("Maxdiff up for CurrentCoupling5DDensity:", max_du) + print("Maxdiff up for CurrentCoupling5DDensity:", diffs["u"]) print() From 3c6a2686d946511e2522e530d4fe141585aca3ea Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Tue, 14 Oct 2025 05:16:35 +0000 Subject: [PATCH 140/292] Added `--oversubscribe` to the test mpirun commands --- src/struphy/console/test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 938571d75..4e01a33ec 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -52,6 +52,7 @@ def struphy_test( # now run parallel unit tests cmd = [ "mpirun", + "--oversubscribe", "-n", str(mpi), "pytest", @@ -70,6 +71,7 @@ def struphy_test( elif group in {"models", "fluid", "kinetic", "hybrid", "toy"}: cmd = [ "mpirun", + "--oversubscribe", "-n", str(mpi), "pytest", @@ -91,6 +93,7 @@ def struphy_test( elif "verification" in group: cmd = [ "mpirun", + "--oversubscribe", "-n", str(mpi), "pytest", @@ -110,6 +113,7 @@ def struphy_test( else: cmd = [ "mpirun", + "--oversubscribe", "-n", str(mpi), "pytest", From f41eb71ef286970b83e239e04065a9dea862e3ef Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 09:22:04 +0000 Subject: [PATCH 141/292] Resolve "Porting the rest of fluid models" --- src/struphy/console/params.py | 2 + src/struphy/models/base.py | 142 +- src/struphy/models/fluid.py | 2241 +++++++++-------- src/struphy/models/species.py | 4 - src/struphy/models/tests/test_models.py | 14 +- src/struphy/propagators/propagators_fields.py | 885 ++++--- 6 files changed, 1678 insertions(+), 1610 deletions(-) diff --git a/src/struphy/console/params.py b/src/struphy/console/params.py index e32fd4165..e5164a10c 100644 --- a/src/struphy/console/params.py +++ b/src/struphy/console/params.py @@ -29,6 +29,8 @@ def struphy_params(model_name: str, params_path: str, yes: bool = False, check_f except AttributeError: pass + print(f"{model_name = }") + # print units if check_file: print(f"Checking {check_file} with model {model_class}") diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 0b330a57d..78ee02bc3 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -274,11 +274,6 @@ def clone_config(self, new): assert isinstance(new, CloneConfig) or new is None self._clone_config = new - @property - def diagnostics(self): - """Dictionary of diagnostics.""" - return self._diagnostics - @property def domain(self): """Domain object, see :ref:`avail_mappings`.""" @@ -644,6 +639,18 @@ def allocate_variables(self, verbose: bool = False): verbose=verbose, ) + # allocate memory for FE coeffs of fluid variables + if self.diagnostic_species: + for species, spec in self.diagnostic_species.items(): + assert isinstance(spec, DiagnosticSpecies) + for k, v in spec.variables.items(): + assert isinstance(v, FEECVariable) + v.allocate( + derham=self.derham, + domain=self.domain, + equil=self.equil, + ) + # TODO: allocate memory for FE coeffs of diagnostics # if self.params.diagnostic_fields is not None: # for key, val in self.diagnostics.items(): @@ -1490,131 +1497,6 @@ def generate_default_parameter_file( # Private methods : ################### - def _init_variable_dicts(self): - """ - Initialize em-fields, fluid and kinetic dictionaries for information on the model variables. - """ - - # electromagnetic fields, fluid and/or kinetic species - self._em_fields = {} - self._fluid = {} - self._kinetic = {} - self._diagnostics = {} - - if self.rank_world == 0 and self.verbose: - print("\nMODEL SPECIES:") - - # create dictionaries for each em-field/species and fill in space/class name and parameters - for var_name, space in self.species()["em_fields"].items(): - assert space in {"H1", "Hcurl", "Hdiv", "L2", "H1vec"} - assert self.params.em_fields is not None, '"em_fields" is missing in parameter file.' - - if self.rank_world == 0 and self.verbose: - print("em_field:".ljust(25), f'"{var_name}" ({space})') - - self._em_fields[var_name] = {} - - # space - self._em_fields[var_name]["space"] = space - - # initial conditions - if "background" in self.params.em_fields: - # background= self.params.em_fields["background"].get(var_name) - self._em_fields[var_name]["background"] = self.params.em_fields["background"].get(var_name) - # else: - # background = None - - if "perturbation" in self.params.em_fields: - # perturbation = self.params.em_fields["perturbation"].get(var_name) - self._em_fields[var_name]["perturbation"] = self.params.em_fields["perturbation"].get(var_name) - # else: - # perturbation = None - - # which components to save - if "save_data" in self.params.em_fields: - # save_data = self.params.em_fields["save_data"]["comps"][var_name] - self._em_fields[var_name]["save_data"] = self.params.em_fields["save_data"]["comps"][var_name] - else: - self._em_fields[var_name]["save_data"] = True - # save_data = True - - # self._em_fields[var_name] = Variable(name=var_name, - # space=space, - # background=background, - # perturbation=perturbation, - # save_data=save_data,) - - # overall parameters - # print(f'{self._em_fields = }') - self._em_fields["params"] = self.params.em_fields - - for var_name, space in self.species()["fluid"].items(): - assert isinstance(space, dict) - assert "fluid" in self.params, 'Top-level key "fluid" is missing in parameter file.' - assert var_name in self.params["fluid"], f"Fluid species {var_name} is missing in parameter file." - - if self.rank_world == 0 and self.verbose: - print("fluid:".ljust(25), f'"{var_name}" ({space})') - - self._fluid[var_name] = {} - for sub_var_name, sub_space in space.items(): - self._fluid[var_name][sub_var_name] = {} - - # space - self._fluid[var_name][sub_var_name]["space"] = sub_space - - # initial conditions - if "background" in self.params["fluid"][var_name]: - self._fluid[var_name][sub_var_name]["background"] = self.params["fluid"][var_name][ - "background" - ].get(sub_var_name) - if "perturbation" in self.params["fluid"][var_name]: - self._fluid[var_name][sub_var_name]["perturbation"] = self.params["fluid"][var_name][ - "perturbation" - ].get(sub_var_name) - - # which components to save - if "save_data" in self.params["fluid"][var_name]: - self._fluid[var_name][sub_var_name]["save_data"] = self.params["fluid"][var_name]["save_data"][ - "comps" - ][sub_var_name] - - else: - self._fluid[var_name][sub_var_name]["save_data"] = True - - # overall parameters - self._fluid[var_name]["params"] = self.params["fluid"][var_name] - - for var_name, space in self.species()["kinetic"].items(): - assert "Particles" in space - assert "kinetic" in self.params, 'Top-level key "kinetic" is missing in parameter file.' - assert var_name in self.params["kinetic"], f"Kinetic species {var_name} is missing in parameter file." - - if self.rank_world == 0 and self.verbose: - print("kinetic:".ljust(25), f'"{var_name}" ({space})') - - self._kinetic[var_name] = {} - self._kinetic[var_name]["space"] = space - self._kinetic[var_name]["params"] = self.params["kinetic"][var_name] - - if self.diagnostics_dct() is not None: - for var_name, space in self.diagnostics_dct().items(): - assert space in {"H1", "Hcurl", "Hdiv", "L2", "H1vec"} - - if self.rank_world == 0 and self.verbose: - print("diagnostics:".ljust(25), f'"{var_name}" ({space})') - - self._diagnostics[var_name] = {} - self._diagnostics[var_name]["space"] = space - self._diagnostics["params"] = self.params["diagnostics"][var_name] - - # which components to save - if "save_data" in self.params["diagnostics"][var_name]: - self._diagnostics[var_name]["save_data"] = self.params["diagnostics"][var_name]["save_data"] - - else: - self._diagnostics[var_name]["save_data"] = True - def compute_plasma_params(self, verbose=True): """ Compute and print volume averaged plasma parameters for each species of the model. diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 737e83274..0c953ff97 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -3,8 +3,10 @@ from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector +from struphy.feec.projectors import L2Projector +from struphy.feec.variational_utilities import H1vecMassMatrix_density, InternalEnergyEvaluator from struphy.models.base import StruphyModel -from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies +from struphy.models.species import DiagnosticSpecies, FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.polar.basic import PolarVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers @@ -450,7 +452,7 @@ def update_scalar_quantities(self): self.update_scalar("total energy", en_E + en_B + en_J) -class ViscoresistiveMHD(StruphyModel): +class ViscoResistiveMHD(StruphyModel): r"""Full (non-linear) visco-resistive MHD equations discretized with a variational method. :ref:`normalization`: @@ -486,139 +488,73 @@ class ViscoresistiveMHD(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["em_fields"]["b2"] = "Hdiv" - dct["fluid"]["mhd"] = {"rho3": "L2", "s3": "L2", "uv": "H1vec"} - return dct + ## species - @staticmethod - def bulk_species(): - return "mhd" + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.entropy = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["mhd_rho3", "mhd_uv"], - propagators_fields.VariationalMomentumAdvection: ["mhd_uv"], - propagators_fields.VariationalEntropyEvolve: ["mhd_s3", "mhd_uv"], - propagators_fields.VariationalMagFieldEvolve: ["b2", "mhd_uv"], - propagators_fields.VariationalViscosity: ["mhd_s3", "mhd_uv"], - propagators_fields.VariationalResistivity: ["mhd_s3", "b2"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import H1vecMassMatrix_density, InternalEnergyEvaluator - from struphy.polar.basic import PolarVector - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["lin_solver"] - nonlin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["nonlin_solver"] - lin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - lin_solver_entropy = params["fluid"]["mhd"]["options"]["VariationalEntropyEvolve"]["lin_solver"] - nonlin_solver_entropy = params["fluid"]["mhd"]["options"]["VariationalEntropyEvolve"]["nonlin_solver"] - lin_solver_magfield = params["em_fields"]["options"]["VariationalMagFieldEvolve"]["lin_solver"] - nonlin_solver_magfield = params["em_fields"]["options"]["VariationalMagFieldEvolve"]["nonlin_solver"] - lin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["lin_solver"] - nonlin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["nonlin_solver"] - lin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["lin_solver"] - nonlin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["nonlin_solver"] - if "linearize_current" in params["fluid"]["mhd"]["options"]["VariationalResistivity"].keys(): - self._linearize_current = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["linearize_current"] - else: - self._linearize_current = False - self._gamma = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - self._mu = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu"] - self._mu_a = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu_a"] - self._alpha = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["alpha"] - self._eta = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta"] - self._eta_a = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta_a"] - model = "full" - - self._energy_evaluator = InternalEnergyEvaluator(self.derham, self._gamma) - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": model, - "s": self.pointer["mhd_s3"], - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - "energy_evaluator": self._energy_evaluator, - } - - self._kwargs[propagators_fields.VariationalMomentumAdvection] = { - "mass_ops": self.WMM, - "lin_solver": lin_solver_momentum, - "nonlin_solver": nonlin_solver_momentum, - } - - self._kwargs[propagators_fields.VariationalEntropyEvolve] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_entropy, - "nonlin_solver": nonlin_solver_entropy, - "energy_evaluator": self._energy_evaluator, - } - - self._kwargs[propagators_fields.VariationalMagFieldEvolve] = { - "model": model, - "mass_ops": self.WMM, - "lin_solver": lin_solver_magfield, - "nonlin_solver": nonlin_solver_magfield, - } - - self._kwargs[propagators_fields.VariationalViscosity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "mu": self._mu, - "mu_a": self._mu_a, - "alpha": self._alpha, - "mass_ops": self.WMM, - "lin_solver": lin_solver_viscosity, - "nonlin_solver": nonlin_solver_viscosity, - "energy_evaluator": self._energy_evaluator, - } - - self._kwargs[propagators_fields.VariationalResistivity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "eta": self._eta, - "eta_a": self._eta_a, - "lin_solver": lin_solver_resistivity, - "nonlin_solver": nonlin_solver_resistivity, - "linearize_current": self._linearize_current, - "energy_evaluator": self._energy_evaluator, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation + ## propagators + + class Propagators: + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_mom = propagators_fields.VariationalMomentumAdvection() + self.variat_ent = propagators_fields.VariationalEntropyEvolve() + self.variat_mag = propagators_fields.VariationalMagFieldEvolve() + if with_viscosity: + self.variat_viscous = propagators_fields.VariationalViscosity() + if with_resistivity: + self.variat_resist = propagators_fields.VariationalResistivity() + + ## abstract methods + + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + + # 2. instantiate all propagators + self.propagators = self.Propagators( + with_viscosity=with_viscosity, + with_resistivity=with_resistivity, + ) + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.mhd.density + self.propagators.variat_dens.variables.u = self.mhd.velocity + self.propagators.variat_mom.variables.u = self.mhd.velocity + self.propagators.variat_ent.variables.s = self.mhd.entropy + self.propagators.variat_ent.variables.u = self.mhd.velocity + self.propagators.variat_mag.variables.u = self.mhd.velocity + self.propagators.variat_mag.variables.b = self.em_fields.b_field + if with_viscosity: + self.propagators.variat_viscous.variables.s = self.mhd.entropy + self.propagators.variat_viscous.variables.u = self.mhd.velocity + if with_resistivity: + self.propagators.variat_resist.variables.s = self.mhd.entropy + self.propagators.variat_resist.variables.b = self.em_fields.b_field + + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_thermo") self.add_scalar("en_mag") @@ -627,16 +563,24 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("entr_tot") self.add_scalar("tot_div_B") - # temporary vectors for scalar quantities - self._tmp_div_B = self.derham.Vh_pol["3"].zeros() - tmp_dof = self.derham.Vh_pol["3"].zeros() - projV3 = L2Projector("L2", self.mass_ops) + @property + def bulk_species(self): + return self.mhd + + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + projV3 = L2Projector("L2", self._mass_ops) def f(e1, e2, e3): return 1 f = np.vectorize(f) - self._integrator = projV3(f, dofs=tmp_dof) + self._integrator = projV3(f) + + self._energy_evaluator = InternalEnergyEvaluator(self.derham, self.propagators.variat_ent.options.gamma) self._ones = self.derham.Vh_pol["3"].zeros() if isinstance(self._ones, PolarVector): @@ -644,12 +588,18 @@ def f(e1, e2, e3): else: self._ones[:] = 1.0 + self._tmp_div_B = self.derham.Vh_pol["3"].zeros() + def update_scalar_quantities(self): - # Update mass matrix - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["mhd_uv"], self.pointer["mhd_uv"]) + rho = self.mhd.density.spline.vector + u = self.mhd.velocity.spline.vector + s = self.mhd.entropy.spline.vector + b = self.em_fields.b_field.spline.vector + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) - en_mag = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b2"], self.pointer["b2"]) + en_mag = 0.5 * self.mass_ops.M2.dot_inner(b, b) self.update_scalar("en_mag", en_mag) en_thermo = self.update_thermo_energy() @@ -657,12 +607,12 @@ def update_scalar_quantities(self): en_tot = en_U + en_thermo + en_mag self.update_scalar("en_tot", en_tot) - dens_tot = self._ones.inner(self.pointer["mhd_rho3"]) + dens_tot = self._ones.inner(rho) self.update_scalar("dens_tot", dens_tot) - entr_tot = self._ones.inner(self.pointer["mhd_s3"]) + entr_tot = self._ones.inner(s) self.update_scalar("entr_tot", entr_tot) - div_B = self.derham.div.dot(self.pointer["b2"], out=self._tmp_div_B) + div_B = self.derham.div.dot(b, out=self._tmp_div_B) L2_div_B = self._mass_ops.M3.dot_inner(div_B, div_B) self.update_scalar("tot_div_B", L2_div_B) @@ -671,9 +621,12 @@ def update_thermo_energy(self): :meta private: """ - en_prop = self._propagators[0] - self._energy_evaluator.sf.vector = self.pointer["mhd_s3"] - self._energy_evaluator.rhof.vector = self.pointer["mhd_rho3"] + rho = self.mhd.density.spline.vector + s = self.mhd.entropy.spline.vector + en_prop = self.propagators.variat_dens + + self._energy_evaluator.sf.vector = s + self._energy_evaluator.rhof.vector = rho sf_values = self._energy_evaluator.sf.eval_tp_fixed_loc( self._energy_evaluator.integration_grid_spans, self._energy_evaluator.integration_grid_bd, @@ -691,6 +644,44 @@ def update_thermo_energy(self): self.update_scalar("en_thermo", en_thermo) return en_thermo + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full',\n" + ] + new_file += [ + " s=model.mhd.entropy)\n" + ] + elif "variat_ent.Options" in line: + new_file += [ + "model.propagators.variat_ent.options = model.propagators.variat_ent.Options(model='full',\n" + ] + new_file += [ + " rho=model.mhd.density)\n" + ] + elif "variat_viscous.Options" in line: + new_file += [ + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(rho=model.mhd.density)\n" + ] + elif "variat_resist.Options" in line: + new_file += [ + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(rho=model.mhd.density)\n" + ] + elif "entropy.add_background" in line: + new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] + new_file += [line] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class ViscousFluid(StruphyModel): r"""Full (non-linear) viscous Navier-Stokes equations discretized with a variational method. @@ -724,122 +715,72 @@ class ViscousFluid(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["fluid"]["fluid"] = {"rho3": "L2", "s3": "L2", "uv": "H1vec"} - return dct + ## species + + class Fluid(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.entropy = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def bulk_species(): - return "fluid" + ## propagators - @staticmethod - def velocity_scale(): - return "alfvén" + class Propagators: + def __init__(self, with_viscosity: bool = True): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_mom = propagators_fields.VariationalMomentumAdvection() + self.variat_ent = propagators_fields.VariationalEntropyEvolve() + if with_viscosity: + self.variat_viscous = propagators_fields.VariationalViscosity() + + ## abstract methods + + def __init__(self, with_viscosity: bool = True): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.fluid = self.Fluid() - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["fluid_rho3", "fluid_uv"], - propagators_fields.VariationalMomentumAdvection: ["fluid_uv"], - propagators_fields.VariationalEntropyEvolve: ["fluid_s3", "fluid_uv"], - propagators_fields.VariationalViscosity: ["fluid_s3", "fluid_uv"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.projectors import L2Projector - from struphy.polar.basic import PolarVector - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - from struphy.feec.variational_utilities import H1vecMassMatrix_density, InternalEnergyEvaluator - - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_momentum = params["fluid"]["fluid"]["options"]["VariationalMomentumAdvection"]["lin_solver"] - nonlin_solver_momentum = params["fluid"]["fluid"]["options"]["VariationalMomentumAdvection"]["nonlin_solver"] - lin_solver_density = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - lin_solver_entropy = params["fluid"]["fluid"]["options"]["VariationalEntropyEvolve"]["lin_solver"] - nonlin_solver_entropy = params["fluid"]["fluid"]["options"]["VariationalEntropyEvolve"]["nonlin_solver"] - lin_solver_viscosity = params["fluid"]["fluid"]["options"]["VariationalViscosity"]["lin_solver"] - nonlin_solver_viscosity = params["fluid"]["fluid"]["options"]["VariationalViscosity"]["nonlin_solver"] - - self._gamma = params["fluid"]["fluid"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - self._mu = params["fluid"]["fluid"]["options"]["VariationalViscosity"]["physics"]["mu"] - self._mu_a = params["fluid"]["fluid"]["options"]["VariationalViscosity"]["physics"]["mu_a"] - model = "full" - - self._energy_evaluator = InternalEnergyEvaluator(self.derham, self._gamma) - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": model, - "s": self.pointer["fluid_s3"], - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - "energy_evaluator": self._energy_evaluator, - } - - self._kwargs[propagators_fields.VariationalMomentumAdvection] = { - "mass_ops": self.WMM, - "lin_solver": lin_solver_momentum, - "nonlin_solver": nonlin_solver_momentum, - } - - self._kwargs[propagators_fields.VariationalEntropyEvolve] = { - "model": model, - "rho": self.pointer["fluid_rho3"], - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_entropy, - "nonlin_solver": nonlin_solver_entropy, - "energy_evaluator": self._energy_evaluator, - } - - self._kwargs[propagators_fields.VariationalViscosity] = { - "model": model, - "gamma": self._gamma, - "rho": self.pointer["fluid_rho3"], - "mu": self._mu, - "mu_a": self._mu_a, - "mass_ops": self.WMM, - "lin_solver": lin_solver_viscosity, - "nonlin_solver": nonlin_solver_viscosity, - "energy_evaluator": self._energy_evaluator, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation + # 2. instantiate all propagators + self.propagators = self.Propagators(with_viscosity=with_viscosity) + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.fluid.density + self.propagators.variat_dens.variables.u = self.fluid.velocity + self.propagators.variat_mom.variables.u = self.fluid.velocity + self.propagators.variat_ent.variables.s = self.fluid.entropy + self.propagators.variat_ent.variables.u = self.fluid.velocity + if with_viscosity: + self.propagators.variat_viscous.variables.s = self.fluid.entropy + self.propagators.variat_viscous.variables.u = self.fluid.velocity + + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_thermo") self.add_scalar("en_tot") self.add_scalar("dens_tot") self.add_scalar("entr_tot") - # temporary vectors for scalar quantities - tmp_dof = self.derham.Vh_pol["3"].zeros() - projV3 = L2Projector("L2", self.mass_ops) + @property + def bulk_species(self): + return self.fluid + + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + projV3 = L2Projector("L2", self._mass_ops) def f(e1, e2, e3): return 1 f = np.vectorize(f) - self._integrator = projV3(f, dofs=tmp_dof) + self._integrator = projV3(f) + + self._energy_evaluator = InternalEnergyEvaluator(self.derham, self.propagators.variat_ent.options.gamma) self._ones = self.derham.Vh_pol["3"].zeros() if isinstance(self._ones, PolarVector): @@ -848,8 +789,11 @@ def f(e1, e2, e3): self._ones[:] = 1.0 def update_scalar_quantities(self): - # Update mass matrix - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["fluid_uv"], self.pointer["fluid_uv"]) + rho = self.fluid.density.spline.vector + u = self.fluid.velocity.spline.vector + s = self.fluid.entropy.spline.vector + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) en_thermo = self.update_thermo_energy() @@ -857,9 +801,9 @@ def update_scalar_quantities(self): en_tot = en_U + en_thermo self.update_scalar("en_tot", en_tot) - dens_tot = self._ones.inner(self.pointer["fluid_rho3"]) + dens_tot = self._ones.inner(rho) self.update_scalar("dens_tot", dens_tot) - entr_tot = self._ones.inner(self.pointer["fluid_s3"]) + entr_tot = self._ones.inner(s) self.update_scalar("entr_tot", entr_tot) def update_thermo_energy(self): @@ -867,9 +811,12 @@ def update_thermo_energy(self): :meta private: """ - en_prop = self._propagators[0] - self._energy_evaluator.sf.vector = self.pointer["fluid_s3"] - self._energy_evaluator.rhof.vector = self.pointer["fluid_rho3"] + rho = self.fluid.density.spline.vector + s = self.fluid.entropy.spline.vector + en_prop = self.propagators.variat_dens + + self._energy_evaluator.sf.vector = s + self._energy_evaluator.rhof.vector = rho sf_values = self._energy_evaluator.sf.eval_tp_fixed_loc( self._energy_evaluator.integration_grid_spans, self._energy_evaluator.integration_grid_bd, @@ -887,8 +834,42 @@ def update_thermo_energy(self): self.update_scalar("en_thermo", en_thermo) return en_thermo + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full',\n" + ] + new_file += [ + " s=model.fluid.entropy)\n" + ] + elif "variat_ent.Options" in line: + new_file += [ + "model.propagators.variat_ent.options = model.propagators.variat_ent.Options(model='full',\n" + ] + new_file += [ + " rho=model.fluid.density)\n" + ] + elif "variat_viscous.Options" in line: + new_file += [ + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(rho=model.fluid.density)\n" + ] + elif "entropy.add_background" in line: + new_file += ["model.fluid.density.add_background(FieldsBackground())\n"] + new_file += [line] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + -class ViscoresistiveMHD_with_p(StruphyModel): +class ViscoResistiveMHD_with_p(StruphyModel): r"""Full (non-linear) visco-resistive MHD equations, with the pressure variable discretized with a variational method. :ref:`normalization`: @@ -922,121 +903,78 @@ class ViscoresistiveMHD_with_p(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["em_fields"]["b2"] = "Hdiv" - dct["fluid"]["mhd"] = {"rho3": "L2", "p3": "L2", "uv": "H1vec"} - return dct + ## species - @staticmethod - def bulk_species(): - return "mhd" + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.pressure = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["mhd_rho3", "mhd_uv"], - propagators_fields.VariationalMomentumAdvection: ["mhd_uv"], - propagators_fields.VariationalPBEvolve: ["mhd_p3", "b2", "mhd_uv"], - propagators_fields.VariationalViscosity: ["mhd_p3", "mhd_uv"], - propagators_fields.VariationalResistivity: ["mhd_p3", "b2"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import H1vecMassMatrix_density - from struphy.polar.basic import PolarVector - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["lin_solver"] - nonlin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["nonlin_solver"] - lin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - lin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalPBEvolve"]["lin_solver"] - nonlin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalPBEvolve"]["nonlin_solver"] - lin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["lin_solver"] - nonlin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["nonlin_solver"] - lin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["lin_solver"] - nonlin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["nonlin_solver"] - if "linearize_current" in params["fluid"]["mhd"]["options"]["VariationalResistivity"].keys(): - self._linearize_current = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["linearize_current"] - else: - self._linearize_current = False - self._gamma = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - self._mu = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu"] - self._mu_a = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu_a"] - self._alpha = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["alpha"] - self._eta = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta"] - self._eta_a = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta_a"] - model = "full_p" - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": model, - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - } - - self._kwargs[propagators_fields.VariationalMomentumAdvection] = { - "mass_ops": self.WMM, - "lin_solver": lin_solver_momentum, - "nonlin_solver": nonlin_solver_momentum, - } - - self._kwargs[propagators_fields.VariationalPBEvolve] = { - "model": model, - "mass_ops": self.WMM, - "lin_solver": lin_solver_magfield, - "nonlin_solver": nonlin_solver_magfield, - "gamma": self._gamma, - } - - self._kwargs[propagators_fields.VariationalViscosity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "mu": self._mu, - "mu_a": self._mu_a, - "alpha": self._alpha, - "mass_ops": self.WMM, - "lin_solver": lin_solver_viscosity, - "nonlin_solver": nonlin_solver_viscosity, - } - - self._kwargs[propagators_fields.VariationalResistivity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "eta": self._eta, - "eta_a": self._eta_a, - "lin_solver": lin_solver_resistivity, - "nonlin_solver": nonlin_solver_resistivity, - "linearize_current": self._linearize_current, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation + class Diagnostics(DiagnosticSpecies): + def __init__(self): + self.div_u = FEECVariable(space="L2") + self.u2 = FEECVariable(space="Hdiv") + self.init_variables() + + ## propagators + + class Propagators: + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_mom = propagators_fields.VariationalMomentumAdvection() + self.variat_pb = propagators_fields.VariationalPBEvolve() + if with_viscosity: + self.variat_viscous = propagators_fields.VariationalViscosity() + if with_resistivity: + self.variat_resist = propagators_fields.VariationalResistivity() + + ## abstract methods + + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + self.diagnostics = self.Diagnostics() + + # 2. instantiate all propagators + self.propagators = self.Propagators( + with_viscosity=with_viscosity, + with_resistivity=with_resistivity, + ) + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.mhd.density + self.propagators.variat_dens.variables.u = self.mhd.velocity + self.propagators.variat_mom.variables.u = self.mhd.velocity + self.propagators.variat_pb.variables.u = self.mhd.velocity + self.propagators.variat_pb.variables.p = self.mhd.pressure + self.propagators.variat_pb.variables.b = self.em_fields.b_field + if with_viscosity: + self.propagators.variat_viscous.variables.s = self.mhd.pressure + self.propagators.variat_viscous.variables.u = self.mhd.velocity + if with_resistivity: + self.propagators.variat_resist.variables.s = self.mhd.pressure + self.propagators.variat_resist.variables.b = self.em_fields.b_field + + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_thermo") self.add_scalar("en_mag") @@ -1044,12 +982,22 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("dens_tot") self.add_scalar("tot_div_B") - # temporary vectors for scalar quantities - self._tmp_div_B = self.derham.Vh_pol["3"].zeros() - tmp_dof = self.derham.Vh_pol["3"].zeros() - projV3 = L2Projector("L2", self.mass_ops) + @property + def bulk_species(self): + return self.mhd + + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + projV3 = L2Projector("L2", self._mass_ops) - self._integrator = projV3(self.domain.jacobian_det, dofs=tmp_dof) + def f(e1, e2, e3): + return 1 + + f = np.vectorize(f) + self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() if isinstance(self._ones, PolarVector): @@ -1057,39 +1005,68 @@ def __init__(self, params, comm, clone_config=None): else: self._ones[:] = 1.0 + self._tmp_div_B = self.derham.Vh_pol["3"].zeros() + def update_scalar_quantities(self): - # Update mass matrix - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["mhd_uv"], self.pointer["mhd_uv"]) + rho = self.mhd.density.spline.vector + u = self.mhd.velocity.spline.vector + p = self.mhd.pressure.spline.vector + b = self.em_fields.b_field.spline.vector + + gamma = self.propagators.variat_pb.options.gamma + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) - en_mag = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b2"], self.pointer["b2"]) + en_mag = 0.5 * self.mass_ops.M2.dot_inner(b, b) self.update_scalar("en_mag", en_mag) - en_thermo = self.mass_ops.M3.dot_inner(self.pointer["mhd_p3"], self._integrator) / (self._gamma - 1.0) + en_thermo = self.mass_ops.M3.dot_inner(p, self._integrator) / (gamma - 1.0) self.update_scalar("en_thermo", en_thermo) en_tot = en_U + en_thermo + en_mag self.update_scalar("en_tot", en_tot) - dens_tot = self._ones.inner(self.pointer["mhd_rho3"]) + dens_tot = self._ones.inner(rho) self.update_scalar("dens_tot", dens_tot) - div_B = self.derham.div.dot(self.pointer["b2"], out=self._tmp_div_B) + div_B = self.derham.div.dot(b, out=self._tmp_div_B) L2_div_B = self._mass_ops.M3.dot_inner(div_B, div_B) self.update_scalar("tot_div_B", L2_div_B) - @staticmethod - def diagnostics_dct(): - dct = {} - - dct["div_u"] = "L2" - dct["u2"] = "Hdiv" - return dct + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_pb.Options" in line: + new_file += [ + "model.propagators.variat_pb.options = model.propagators.variat_pb.Options(div_u=model.diagnostics.div_u,\n" + ] + new_file += [ + " u2=model.diagnostics.u2)\n" + ] + elif "variat_viscous.Options" in line: + new_file += [ + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(rho=model.mhd.density)\n" + ] + elif "variat_resist.Options" in line: + new_file += [ + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(rho=model.mhd.density)\n" + ] + elif "pressure.add_background" in line: + new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] + new_file += [line] + else: + new_file += [line] - __diagnostics__ = diagnostics_dct() + with open(params_path, "w") as f: + for line in new_file: + f.write(line) -class ViscoresistiveLinearMHD(StruphyModel): +class ViscoResistiveLinearMHD(StruphyModel): r"""Linear visco-resistive MHD equations discretized with a variational method. :ref:`normalization`: @@ -1122,136 +1099,104 @@ class ViscoresistiveLinearMHD(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["em_fields"]["b2"] = "Hdiv" - dct["fluid"]["mhd"] = {"rho3": "L2", "p3": "L2", "uv": "H1vec"} - return dct + ## species + + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() + + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.pressure = FEECVariable(space="L2") + self.init_variables() + + class Diagnostics(DiagnosticSpecies): + def __init__(self): + self.div_u = FEECVariable(space="L2") + self.u2 = FEECVariable(space="Hdiv") + self.pt3 = FEECVariable(space="L2") + self.bt2 = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def bulk_species(): - return "mhd" + ## propagators - @staticmethod - def velocity_scale(): - return "alfvén" + class Propagators: + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_pb = propagators_fields.VariationalPBEvolve() + if with_viscosity: + self.variat_viscous = propagators_fields.VariationalViscosity() + if with_resistivity: + self.variat_resist = propagators_fields.VariationalResistivity() - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["mhd_rho3", "mhd_uv"], - propagators_fields.VariationalPBEvolve: ["mhd_p3", "b2", "mhd_uv"], - propagators_fields.VariationalViscosity: ["mhd_p3", "mhd_uv"], - propagators_fields.VariationalResistivity: ["mhd_p3", "b2"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import H1vecMassMatrix_density - from struphy.polar.basic import PolarVector - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - lin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalPBEvolve"]["lin_solver"] - nonlin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalPBEvolve"]["nonlin_solver"] - lin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["lin_solver"] - nonlin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["nonlin_solver"] - lin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["lin_solver"] - nonlin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["nonlin_solver"] - if "linearize_current" in params["fluid"]["mhd"]["options"]["VariationalResistivity"].keys(): - self._linearize_current = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["linearize_current"] - else: - self._linearize_current = False - self._gamma = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - self._mu = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu"] - self._mu_a = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu_a"] - self._alpha = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["alpha"] - self._eta = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta"] - self._eta_a = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta_a"] - model = "linear" - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": model, - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - } - - self._kwargs[propagators_fields.VariationalPBEvolve] = { - "model": model, - "mass_ops": self.WMM, - "lin_solver": lin_solver_magfield, - "nonlin_solver": nonlin_solver_magfield, - "gamma": self._gamma, - "div_u": self.pointer["div_u"], - "u2": self.pointer["u2"], - "bt2": self.pointer["bt2"], - "pt3": self.pointer["pt3"], - } - - self._kwargs[propagators_fields.VariationalViscosity] = { - "model": "linear_p", - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "mu": self._mu, - "mu_a": self._mu_a, - "alpha": self._alpha, - "mass_ops": self.WMM, - "lin_solver": lin_solver_viscosity, - "nonlin_solver": nonlin_solver_viscosity, - } - - self._kwargs[propagators_fields.VariationalResistivity] = { - "model": "linear_p", - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "eta": self._eta, - "eta_a": self._eta_a, - "lin_solver": lin_solver_resistivity, - "nonlin_solver": nonlin_solver_resistivity, - "linearize_current": self._linearize_current, - "pt3": self.pointer["pt3"], - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation + ## abstract methods + + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + self.diagnostics = self.Diagnostics() + + # 2. instantiate all propagators + self.propagators = self.Propagators( + with_viscosity=with_viscosity, + with_resistivity=with_resistivity, + ) + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.mhd.density + self.propagators.variat_dens.variables.u = self.mhd.velocity + self.propagators.variat_pb.variables.u = self.mhd.velocity + self.propagators.variat_pb.variables.p = self.mhd.pressure + self.propagators.variat_pb.variables.b = self.em_fields.b_field + if with_viscosity: + self.propagators.variat_viscous.variables.s = self.mhd.pressure + self.propagators.variat_viscous.variables.u = self.mhd.velocity + if with_resistivity: + self.propagators.variat_resist.variables.s = self.mhd.pressure + self.propagators.variat_resist.variables.b = self.em_fields.b_field + + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_thermo") self.add_scalar("en_mag_1") self.add_scalar("en_mag_2") self.add_scalar("en_tot") - # self.add_scalar("dens_tot") - # self.add_scalar("tot_div_B") - self.add_scalar("en_tot_l1") self.add_scalar("en_thermo_l1") self.add_scalar("en_mag_l1") - # temporary vectors for scalar quantities - self._tmp_div_B = self.derham.Vh_pol["3"].zeros() - tmp_dof = self.derham.Vh_pol["3"].zeros() - projV3 = L2Projector("L2", self.mass_ops) + @property + def bulk_species(self): + return self.mhd + + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + projV3 = L2Projector("L2", self._mass_ops) + + def f(e1, e2, e3): + return 1 - self._integrator = projV3(self.domain.jacobian_det, dofs=tmp_dof) + f = np.vectorize(f) + self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() if isinstance(self._ones, PolarVector): @@ -1259,52 +1204,104 @@ def __init__(self, params, comm, clone_config=None): else: self._ones[:] = 1.0 - def update_scalar_quantities(self): - # Update mass matrix - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["mhd_uv"], self.pointer["mhd_uv"]) - self.update_scalar("en_U", en_U) + self._tmp_div_B = self.derham.Vh_pol["3"].zeros() - en_mag1 = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b2"], self.pointer["b2"]) + def update_scalar_quantities(self): + rho = self.mhd.density.spline.vector + u = self.mhd.velocity.spline.vector + p = self.mhd.pressure.spline.vector + b = self.em_fields.b_field.spline.vector + bt2 = self.propagators.variat_pb.options.bt2.spline.vector + pt3 = self.propagators.variat_pb.options.pt3.spline.vector + + gamma = self.propagators.variat_pb.options.gamma + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) + self.update_scalar("en_U", en_U) + + en_mag1 = 0.5 * self.mass_ops.M2.dot_inner(b, b) self.update_scalar("en_mag_1", en_mag1) - en_mag2 = self.mass_ops.M2.dot_inner(self.pointer["bt2"], self.projected_equil.b2) + en_mag2 = self.mass_ops.M2.dot_inner(bt2, self.projected_equil.b2) self.update_scalar("en_mag_2", en_mag2) - en_thermo = self.mass_ops.M3.dot_inner(self.pointer["pt3"], self._integrator) / (self._gamma - 1.0) + en_thermo = self.mass_ops.M3.dot_inner(pt3, self._integrator) / (gamma - 1.0) self.update_scalar("en_thermo", en_thermo) en_tot = en_U + en_thermo + en_mag1 + en_mag2 self.update_scalar("en_tot", en_tot) - # dens_tot = self._ones.inner(self.pointer["mhd_rho3"]) + # dens_tot = self._ones.inner(rho) # self.update_scalar("dens_tot", dens_tot) - # div_B = self.derham.div.dot(self.pointer["b2"], out=self._tmp_div_B) + # div_B = self.derham.div.dot(b, out=self._tmp_div_B) # L2_div_B = self._mass_ops.M3.dot_inner(div_B, div_B) # self.update_scalar("tot_div_B", L2_div_B) - en_thermo_l1 = self.mass_ops.M3.dot_inner(self.pointer["mhd_p3"], self._integrator) / (self._gamma - 1.0) + en_thermo_l1 = self.mass_ops.M3.dot_inner(p, self._integrator) / (gamma - 1.0) self.update_scalar("en_thermo_l1", en_thermo_l1) - en_mag_l1 = self.mass_ops.M2.dot_inner(self.pointer["b2"], self.projected_equil.b2) + en_mag_l1 = self.mass_ops.M2.dot_inner(b, self.projected_equil.b2) self.update_scalar("en_mag_l1", en_mag_l1) en_tot_l1 = en_thermo_l1 + en_mag_l1 self.update_scalar("en_tot_l1", en_tot_l1) - @staticmethod - def diagnostics_dct(): - dct = {} - dct["bt2"] = "Hdiv" - dct["pt3"] = "L2" - dct["div_u"] = "L2" - dct["u2"] = "Hdiv" - return dct + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='linear')\n" + ] + elif "variat_pb.Options" in line: + new_file += [ + "model.propagators.variat_pb.options = model.propagators.variat_pb.Options(model='linear',\n" + ] + new_file += [ + " div_u=model.diagnostics.div_u,\n" + ] + new_file += [ + " u2=model.diagnostics.u2,\n" + ] + new_file += [ + " pt3=model.diagnostics.pt3,\n" + ] + new_file += [ + " bt2=model.diagnostics.bt2)\n" + ] + elif "variat_viscous.Options" in line: + new_file += [ + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='linear_p',\n" + ] + new_file += [ + " rho=model.mhd.density)\n" + ] + elif "variat_resist.Options" in line: + new_file += [ + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='linear_p',\n" + ] + new_file += [ + " rho=model.mhd.density,\n" + ] + new_file += [ + " pt3=model.diagnostics.pt3)\n" + ] + elif "pressure.add_background" in line: + new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] + new_file += [line] + else: + new_file += [line] - __diagnostics__ = diagnostics_dct() + with open(params_path, "w") as f: + for line in new_file: + f.write(line) -class ViscoresistiveDeltafMHD(StruphyModel): +class ViscoResistiveDeltafMHD(StruphyModel): r""":math:`\delta f` visco-resistive MHD equations discretized with a variational method. :ref:`normalization`: @@ -1338,141 +1335,106 @@ class ViscoresistiveDeltafMHD(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["em_fields"]["b2"] = "Hdiv" - dct["fluid"]["mhd"] = {"rho3": "L2", "p3": "L2", "uv": "H1vec"} - return dct + ## species - @staticmethod - def bulk_species(): - return "mhd" + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.pressure = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["mhd_rho3", "mhd_uv"], - propagators_fields.VariationalMomentumAdvection: ["mhd_uv"], - propagators_fields.VariationalPBEvolve: ["mhd_p3", "b2", "mhd_uv"], - propagators_fields.VariationalViscosity: ["mhd_p3", "mhd_uv"], - propagators_fields.VariationalResistivity: ["mhd_p3", "b2"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import H1vecMassMatrix_density - from struphy.polar.basic import PolarVector - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["lin_solver"] - nonlin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["nonlin_solver"] - lin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - lin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalPBEvolve"]["lin_solver"] - nonlin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalPBEvolve"]["nonlin_solver"] - lin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["lin_solver"] - nonlin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["nonlin_solver"] - lin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["lin_solver"] - nonlin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["nonlin_solver"] - if "linearize_current" in params["fluid"]["mhd"]["options"]["VariationalResistivity"].keys(): - self._linearize_current = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["linearize_current"] - else: - self._linearize_current = False - self._gamma = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - self._mu = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu"] - self._mu_a = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu_a"] - self._alpha = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["alpha"] - self._eta = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta"] - self._eta_a = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta_a"] - model = "deltaf" - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": model, - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - } - - self._kwargs[propagators_fields.VariationalMomentumAdvection] = { - "mass_ops": self.WMM, - "lin_solver": lin_solver_momentum, - "nonlin_solver": nonlin_solver_momentum, - } - - self._kwargs[propagators_fields.VariationalPBEvolve] = { - "model": model, - "mass_ops": self.WMM, - "lin_solver": lin_solver_magfield, - "nonlin_solver": nonlin_solver_magfield, - "gamma": self._gamma, - "bt2": self.pointer["bt2"], - "pt3": self.pointer["pt3"], - } - - self._kwargs[propagators_fields.VariationalViscosity] = { - "model": "full_p", - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "mu": self._mu, - "mu_a": self._mu_a, - "alpha": self._alpha, - "mass_ops": self.WMM, - "lin_solver": lin_solver_viscosity, - "nonlin_solver": nonlin_solver_viscosity, - } - - self._kwargs[propagators_fields.VariationalResistivity] = { - "model": "delta_p", - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "eta": self._eta, - "eta_a": self._eta_a, - "lin_solver": lin_solver_resistivity, - "nonlin_solver": nonlin_solver_resistivity, - "linearize_current": self._linearize_current, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation + class Diagnostics(DiagnosticSpecies): + def __init__(self): + self.div_u = FEECVariable(space="L2") + self.u2 = FEECVariable(space="Hdiv") + self.pt3 = FEECVariable(space="L2") + self.bt2 = FEECVariable(space="Hdiv") + self.init_variables() + + ## propagators + + class Propagators: + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_mom = propagators_fields.VariationalMomentumAdvection() + self.variat_pb = propagators_fields.VariationalPBEvolve() + if with_viscosity: + self.variat_viscous = propagators_fields.VariationalViscosity() + if with_resistivity: + self.variat_resist = propagators_fields.VariationalResistivity() + + ## abstract methods + + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + self.diagnostics = self.Diagnostics() + + # 2. instantiate all propagators + self.propagators = self.Propagators( + with_viscosity=with_viscosity, + with_resistivity=with_resistivity, + ) + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.mhd.density + self.propagators.variat_dens.variables.u = self.mhd.velocity + self.propagators.variat_mom.variables.u = self.mhd.velocity + self.propagators.variat_pb.variables.u = self.mhd.velocity + self.propagators.variat_pb.variables.p = self.mhd.pressure + self.propagators.variat_pb.variables.b = self.em_fields.b_field + if with_viscosity: + self.propagators.variat_viscous.variables.s = self.mhd.pressure + self.propagators.variat_viscous.variables.u = self.mhd.velocity + if with_resistivity: + self.propagators.variat_resist.variables.s = self.mhd.pressure + self.propagators.variat_resist.variables.b = self.em_fields.b_field + + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_thermo") self.add_scalar("en_mag_1") self.add_scalar("en_mag_2") self.add_scalar("en_tot") - # self.add_scalar("dens_tot") - # self.add_scalar("tot_div_B") - self.add_scalar("en_tot_l1") self.add_scalar("en_thermo_l1") self.add_scalar("en_mag_l1") - # temporary vectors for scalar quantities - tmp_dof = self.derham.Vh_pol["3"].zeros() - projV3 = L2Projector("L2", self.mass_ops) + @property + def bulk_species(self): + return self.mhd + + @property + def velocity_scale(self): + return "alfvén" - self._integrator = projV3(self.domain.jacobian_det, dofs=tmp_dof) + def allocate_helpers(self): + projV3 = L2Projector("L2", self._mass_ops) + + def f(e1, e2, e3): + return 1 + + f = np.vectorize(f) + self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() if isinstance(self._ones, PolarVector): @@ -1480,52 +1442,95 @@ def __init__(self, params, comm, clone_config=None): else: self._ones[:] = 1.0 + self._tmp_div_B = self.derham.Vh_pol["3"].zeros() + def update_scalar_quantities(self): - # Update mass matrix - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["mhd_uv"], self.pointer["mhd_uv"]) + rho = self.mhd.density.spline.vector + u = self.mhd.velocity.spline.vector + p = self.mhd.pressure.spline.vector + b = self.em_fields.b_field.spline.vector + bt2 = self.propagators.variat_pb.options.bt2.spline.vector + pt3 = self.propagators.variat_pb.options.pt3.spline.vector + + gamma = self.propagators.variat_pb.options.gamma + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) - en_mag1 = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b2"], self.pointer["b2"]) + en_mag1 = 0.5 * self.mass_ops.M2.dot_inner(b, b) self.update_scalar("en_mag_1", en_mag1) - en_mag2 = self.mass_ops.M2.dot_inner(self.pointer["bt2"], self.projected_equil.b2) + en_mag2 = self.mass_ops.M2.dot_inner(bt2, self.projected_equil.b2) self.update_scalar("en_mag_2", en_mag2) - en_thermo = self.mass_ops.M3.dot_inner(self.pointer["pt3"], self._integrator) / (self._gamma - 1.0) + en_thermo = self.mass_ops.M3.dot_inner(pt3, self._integrator) / (gamma - 1.0) self.update_scalar("en_thermo", en_thermo) en_tot = en_U + en_thermo + en_mag1 + en_mag2 self.update_scalar("en_tot", en_tot) - # dens_tot = self._ones.inner(self.pointer["mhd_rho3"]) + # dens_tot = self._ones.inner(rho) # self.update_scalar("dens_tot", dens_tot) - # div_B = self.derham.div.dot(self.pointer["b2"], out=self._tmp_div_B) + # div_B = self.derham.div.dot(b, out=self._tmp_div_B) # L2_div_B = self._mass_ops.M3.dot_inner(div_B, div_B) # self.update_scalar("tot_div_B", L2_div_B) - en_thermo_l1 = self.mass_ops.M3.dot_inner(self.pointer["mhd_p3"], self._integrator) / (self._gamma - 1.0) + en_thermo_l1 = self.mass_ops.M3.dot_inner(p, self._integrator) / (gamma - 1.0) self.update_scalar("en_thermo_l1", en_thermo_l1) - en_mag_l1 = self.mass_ops.M2.dot_inner(self.pointer["b2"], self.projected_equil.b2) + en_mag_l1 = self.mass_ops.M2.dot_inner(b, self.projected_equil.b2) self.update_scalar("en_mag_l1", en_mag_l1) en_tot_l1 = en_thermo_l1 + en_mag_l1 self.update_scalar("en_tot_l1", en_tot_l1) - @staticmethod - def diagnostics_dct(): - dct = {} - dct["bt2"] = "Hdiv" - dct["pt3"] = "L2" - dct["div_u"] = "L2" - dct["u2"] = "Hdiv" - return dct + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='deltaf')\n" + ] + elif "variat_pb.Options" in line: + new_file += [ + "model.propagators.variat_pb.options = model.propagators.variat_pb.Options(model='deltaf',\n" + ] + new_file += [ + " pt3=model.diagnostics.pt3,\n" + ] + new_file += [ + " bt2=model.diagnostics.bt2)\n" + ] + elif "variat_viscous.Options" in line: + new_file += [ + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='full_p',\n" + ] + new_file += [ + " rho=model.mhd.density)\n" + ] + elif "variat_resist.Options" in line: + new_file += [ + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='full_p',\n" + ] + new_file += [ + " rho=model.mhd.density)\n" + ] + elif "pressure.add_background" in line: + new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] + new_file += [line] + else: + new_file += [line] - __diagnostics__ = diagnostics_dct() + with open(params_path, "w") as f: + for line in new_file: + f.write(line) -class ViscoresistiveMHD_with_q(StruphyModel): +class ViscoResistiveMHD_with_q(StruphyModel): r"""Full (non-linear) visco-resistive MHD equations, with the q variable (square root of the pressure) discretized with a variational method. :ref:`normalization`: @@ -1561,121 +1566,78 @@ class ViscoresistiveMHD_with_q(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["em_fields"]["b2"] = "Hdiv" - dct["fluid"]["mhd"] = {"rho3": "L2", "q3": "L2", "uv": "H1vec"} - return dct + ## species - @staticmethod - def bulk_species(): - return "mhd" + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.sqrt_p = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["mhd_rho3", "mhd_uv"], - propagators_fields.VariationalMomentumAdvection: ["mhd_uv"], - propagators_fields.VariationalQBEvolve: ["mhd_q3", "b2", "mhd_uv"], - propagators_fields.VariationalViscosity: ["mhd_q3", "mhd_uv"], - propagators_fields.VariationalResistivity: ["mhd_q3", "b2"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import H1vecMassMatrix_density - from struphy.polar.basic import PolarVector - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["lin_solver"] - nonlin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["nonlin_solver"] - lin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - lin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalQBEvolve"]["lin_solver"] - nonlin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalQBEvolve"]["nonlin_solver"] - lin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["lin_solver"] - nonlin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["nonlin_solver"] - lin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["lin_solver"] - nonlin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["nonlin_solver"] - if "linearize_current" in params["fluid"]["mhd"]["options"]["VariationalResistivity"].keys(): - self._linearize_current = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["linearize_current"] - else: - self._linearize_current = False - self._gamma = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - self._mu = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu"] - self._mu_a = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu_a"] - self._alpha = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["alpha"] - self._eta = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta"] - self._eta_a = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta_a"] - model = "full_q" - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": model, - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - } - - self._kwargs[propagators_fields.VariationalMomentumAdvection] = { - "mass_ops": self.WMM, - "lin_solver": lin_solver_momentum, - "nonlin_solver": nonlin_solver_momentum, - } - - self._kwargs[propagators_fields.VariationalQBEvolve] = { - "model": model, - "mass_ops": self.WMM, - "lin_solver": lin_solver_magfield, - "nonlin_solver": nonlin_solver_magfield, - "gamma": self._gamma, - } - - self._kwargs[propagators_fields.VariationalViscosity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "mu": self._mu, - "mu_a": self._mu_a, - "alpha": self._alpha, - "mass_ops": self.WMM, - "lin_solver": lin_solver_viscosity, - "nonlin_solver": nonlin_solver_viscosity, - } - - self._kwargs[propagators_fields.VariationalResistivity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "eta": self._eta, - "eta_a": self._eta_a, - "lin_solver": lin_solver_resistivity, - "nonlin_solver": nonlin_solver_resistivity, - "linearize_current": self._linearize_current, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation + class Diagnostics(DiagnosticSpecies): + def __init__(self): + self.div_u = FEECVariable(space="L2") + self.u2 = FEECVariable(space="Hdiv") + self.init_variables() + + ## propagators + + class Propagators: + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_mom = propagators_fields.VariationalMomentumAdvection() + self.variat_qb = propagators_fields.VariationalQBEvolve() + if with_viscosity: + self.variat_viscous = propagators_fields.VariationalViscosity() + if with_resistivity: + self.variat_resist = propagators_fields.VariationalResistivity() + + ## abstract methods + + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + self.diagnostics = self.Diagnostics() + + # 2. instantiate all propagators + self.propagators = self.Propagators( + with_viscosity=with_viscosity, + with_resistivity=with_resistivity, + ) + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.mhd.density + self.propagators.variat_dens.variables.u = self.mhd.velocity + self.propagators.variat_mom.variables.u = self.mhd.velocity + self.propagators.variat_qb.variables.u = self.mhd.velocity + self.propagators.variat_qb.variables.q = self.mhd.sqrt_p + self.propagators.variat_qb.variables.b = self.em_fields.b_field + if with_viscosity: + self.propagators.variat_viscous.variables.s = self.mhd.sqrt_p + self.propagators.variat_viscous.variables.u = self.mhd.velocity + if with_resistivity: + self.propagators.variat_resist.variables.s = self.mhd.sqrt_p + self.propagators.variat_resist.variables.b = self.em_fields.b_field + + # define scalars for update_scalar_quantities self.add_scalar("en_U") self.add_scalar("en_thermo") self.add_scalar("en_mag") @@ -1683,12 +1645,22 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("dens_tot") self.add_scalar("tot_div_B") - # temporary vectors for scalar quantities - self._tmp_div_B = self.derham.Vh_pol["3"].zeros() - tmp_dof = self.derham.Vh_pol["3"].zeros() - projV3 = L2Projector("L2", self.mass_ops) + @property + def bulk_species(self): + return self.mhd - self._integrator = projV3(self.domain.jacobian_det, dofs=tmp_dof) + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + projV3 = L2Projector("L2", self._mass_ops) + + def f(e1, e2, e3): + return 1 + + f = np.vectorize(f) + self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() if isinstance(self._ones, PolarVector): @@ -1696,39 +1668,75 @@ def __init__(self, params, comm, clone_config=None): else: self._ones[:] = 1.0 + self._tmp_div_B = self.derham.Vh_pol["3"].zeros() + def update_scalar_quantities(self): - # Update mass matrix - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["mhd_uv"], self.pointer["mhd_uv"]) + rho = self.mhd.density.spline.vector + u = self.mhd.velocity.spline.vector + q = self.mhd.sqrt_p.spline.vector + b = self.em_fields.b_field.spline.vector + + gamma = self.propagators.variat_qb.options.gamma + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) - en_mag = 0.5 * self._mass_ops.M2.dot_inner(self.pointer["b2"], self.pointer["b2"]) + en_mag = 0.5 * self.mass_ops.M2.dot_inner(b, b) self.update_scalar("en_mag", en_mag) - en_thermo = 1 / (self._gamma - 1) * self._mass_ops.M3.dot_inner(self.pointer["mhd_q3"], self.pointer["mhd_q3"]) + en_thermo = 1.0 / (gamma - 1.0) * self._mass_ops.M3.dot_inner(q, q) self.update_scalar("en_thermo", en_thermo) en_tot = en_U + en_thermo + en_mag self.update_scalar("en_tot", en_tot) - dens_tot = self._ones.inner(self.pointer["mhd_rho3"]) + dens_tot = self._ones.inner(rho) self.update_scalar("dens_tot", dens_tot) - div_B = self.derham.div.dot(self.pointer["b2"], out=self._tmp_div_B) + div_B = self.derham.div.dot(b, out=self._tmp_div_B) L2_div_B = self._mass_ops.M3.dot_inner(div_B, div_B) self.update_scalar("tot_div_B", L2_div_B) - @staticmethod - def diagnostics_dct(): - dct = {} - - dct["div_u"] = "L2" - dct["u2"] = "Hdiv" - return dct + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full_q')\n" + ] + elif "variat_qb.Options" in line: + new_file += [ + "model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='full_q')\n" + ] + elif "variat_viscous.Options" in line: + new_file += [ + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='full_q',\n" + ] + new_file += [ + " rho=model.mhd.density)\n" + ] + elif "variat_resist.Options" in line: + new_file += [ + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='full_q',\n" + ] + new_file += [ + " rho=model.mhd.density)\n" + ] + elif "sqrt_p.add_background" in line: + new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] + new_file += [line] + else: + new_file += [line] - __diagnostics__ = diagnostics_dct() + with open(params_path, "w") as f: + for line in new_file: + f.write(line) -class ViscoresistiveLinearMHD_with_q(StruphyModel): +class ViscoResistiveLinearMHD_with_q(StruphyModel): r"""Linear visco-resistive MHD equations, with the q variable (square root of the pressure), discretized with a variational method. :ref:`normalization`: @@ -1761,138 +1769,101 @@ class ViscoresistiveLinearMHD_with_q(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["em_fields"]["b2"] = "Hdiv" - dct["fluid"]["mhd"] = {"rho3": "L2", "q3": "L2", "uv": "H1vec"} - return dct + ## species - @staticmethod - def bulk_species(): - return "mhd" + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.sqrt_p = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["mhd_rho3", "mhd_uv"], - propagators_fields.VariationalQBEvolve: ["mhd_q3", "b2", "mhd_uv"], - propagators_fields.VariationalViscosity: ["mhd_q3", "mhd_uv"], - propagators_fields.VariationalResistivity: ["mhd_q3", "b2"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import H1vecMassMatrix_density - from struphy.polar.basic import PolarVector - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - lin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalQBEvolve"]["lin_solver"] - nonlin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalQBEvolve"]["nonlin_solver"] - lin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["lin_solver"] - nonlin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["nonlin_solver"] - lin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["lin_solver"] - nonlin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["nonlin_solver"] - if "linearize_current" in params["fluid"]["mhd"]["options"]["VariationalResistivity"].keys(): - self._linearize_current = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["linearize_current"] - else: - self._linearize_current = False - self._gamma = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - self._mu = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu"] - self._mu_a = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu_a"] - self._alpha = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["alpha"] - self._eta = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta"] - self._eta_a = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta_a"] - model = "linear_q" - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": model, - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - } - - self._kwargs[propagators_fields.VariationalQBEvolve] = { - "model": model, - "mass_ops": self.WMM, - "lin_solver": lin_solver_magfield, - "nonlin_solver": nonlin_solver_magfield, - "gamma": self._gamma, - "div_u": self.pointer["div_u"], - "u2": self.pointer["u2"], - "bt2": self.pointer["bt2"], - "qt3": self.pointer["qt3"], - } - - self._kwargs[propagators_fields.VariationalViscosity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "mu": self._mu, - "mu_a": self._mu_a, - "alpha": self._alpha, - "mass_ops": self.WMM, - "lin_solver": lin_solver_viscosity, - "nonlin_solver": nonlin_solver_viscosity, - "pt3": self.pointer["qt3"], - } - - self._kwargs[propagators_fields.VariationalResistivity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "eta": self._eta, - "eta_a": self._eta_a, - "lin_solver": lin_solver_resistivity, - "nonlin_solver": nonlin_solver_resistivity, - "linearize_current": self._linearize_current, - "pt3": self.pointer["qt3"], - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation + class Diagnostics(DiagnosticSpecies): + def __init__(self): + self.div_u = FEECVariable(space="L2") + self.u2 = FEECVariable(space="Hdiv") + self.qt3 = FEECVariable(space="L2") + self.bt2 = FEECVariable(space="Hdiv") + self.init_variables() + + ## propagators + + class Propagators: + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_qb = propagators_fields.VariationalQBEvolve() + if with_viscosity: + self.variat_viscous = propagators_fields.VariationalViscosity() + if with_resistivity: + self.variat_resist = propagators_fields.VariationalResistivity() + + ## abstract methods + + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + self.diagnostics = self.Diagnostics() + + # 2. instantiate all propagators + self.propagators = self.Propagators( + with_viscosity=with_viscosity, + with_resistivity=with_resistivity, + ) + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.mhd.density + self.propagators.variat_dens.variables.u = self.mhd.velocity + self.propagators.variat_qb.variables.u = self.mhd.velocity + self.propagators.variat_qb.variables.q = self.mhd.sqrt_p + self.propagators.variat_qb.variables.b = self.em_fields.b_field + if with_viscosity: + self.propagators.variat_viscous.variables.s = self.mhd.sqrt_p + self.propagators.variat_viscous.variables.u = self.mhd.velocity + if with_resistivity: + self.propagators.variat_resist.variables.s = self.mhd.sqrt_p + self.propagators.variat_resist.variables.b = self.em_fields.b_field + + # define scalars for update_scalar_quantities self.add_scalar("en_U") - # self.add_scalar("en_thermo_1") - # self.add_scalar("en_thermo_2") - # self.add_scalar("en_mag_1") - # self.add_scalar("en_mag_2") + self.add_scalar("en_mag_1") + self.add_scalar("en_mag_2") + self.add_scalar("en_thermo_1") + self.add_scalar("en_thermo_2") self.add_scalar("en_tot") - # self.add_scalar("dens_tot") - # self.add_scalar("tot_div_B") + @property + def bulk_species(self): + return self.mhd - # self.add_scalar("en_tot_l1") - # self.add_scalar("en_thermo_l1") - # self.add_scalar("en_mag_l1") + @property + def velocity_scale(self): + return "alfvén" - # temporary vectors for scalar quantities - self._tmp_div_B = self.derham.Vh_pol["3"].zeros() - tmp_dof = self.derham.Vh_pol["3"].zeros() - projV3 = L2Projector("L2", self.mass_ops) + def allocate_helpers(self): + projV3 = L2Projector("L2", self._mass_ops) - self._integrator = projV3(self.domain.jacobian_det, dofs=tmp_dof) + def f(e1, e2, e3): + return 1 + + f = np.vectorize(f) + self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() if isinstance(self._ones, PolarVector): @@ -1900,56 +1871,94 @@ def __init__(self, params, comm, clone_config=None): else: self._ones[:] = 1.0 + self._tmp_div_B = self.derham.Vh_pol["3"].zeros() + def update_scalar_quantities(self): - # Update mass matrix - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["mhd_uv"], self.pointer["mhd_uv"]) + rho = self.mhd.density.spline.vector + u = self.mhd.velocity.spline.vector + q = self.mhd.sqrt_p.spline.vector + b = self.em_fields.b_field.spline.vector + bt2 = self.propagators.variat_qb.options.bt2.spline.vector + qt3 = self.propagators.variat_qb.options.qt3.spline.vector + + gamma = self.propagators.variat_qb.options.gamma + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) - en_mag1 = self._mass_ops.M2.dot_inner(self.pointer["b2"], self.pointer["b2"]) - # self.update_scalar("en_mag_1", en_mag1) + en_mag1 = 0.5 * self.mass_ops.M2.dot_inner(b, b) + self.update_scalar("en_mag_1", en_mag1) - en_mag2 = self._mass_ops.M2.dot_inner(self.pointer["bt2"], self.projected_equil.b2) - # self.update_scalar("en_mag_2", en_mag2) + en_mag2 = self.mass_ops.M2.dot_inner(bt2, self.projected_equil.b2) + self.update_scalar("en_mag_2", en_mag2) - en_th_1 = 1 / (self._gamma - 1) * self._mass_ops.M3.dot_inner(self.pointer["mhd_q3"], self.pointer["mhd_q3"]) - # self.update_scalar("en_thermo_1", en_th_1) + en_th_1 = 1.0 / (gamma - 1.0) * self.mass_ops.M3.dot_inner(q, q) + self.update_scalar("en_thermo_1", en_th_1) - en_th_2 = 2 / (self._gamma - 1) * self._mass_ops.M3.dot_inner(self.pointer["qt3"], self.projected_equil.q3) - # self.update_scalar("en_thermo_2", en_th_2) + en_th_2 = 2.0 / (gamma - 1.0) * self.mass_ops.M3.dot_inner(qt3, self.projected_equil.q3) + self.update_scalar("en_thermo_2", en_th_2) en_tot = en_U + en_th_1 + en_th_2 + en_mag1 + en_mag2 self.update_scalar("en_tot", en_tot) - # dens_tot = self._ones.dot(self.pointer["mhd_rho3"]) - # self.update_scalar("dens_tot", dens_tot) - - # div_B = self.derham.div.dot(self.pointer["b2"], out=self._tmp_div_B) - # L2_div_B = self._mass_ops.M3.dot_inner(div_B, div_B) - # self.update_scalar("tot_div_B", L2_div_B) - - # en_thermo_l1 = self._integrator.dot(self.mass_ops.M3.dot(self.pointer["mhd_p3"])) / (self._gamma - 1.0) - # self.update_scalar("en_thermo_l1", en_thermo_l1) - - # wb2 = self._mass_ops.M2.dot(self.pointer["b2"], out=self._tmp_wb2) - # en_mag_l1 = wb2.dot(self.projected_equil.b2) - # self.update_scalar("en_mag_l1", en_mag_l1) - - # en_tot_l1 = en_thermo_l1 + en_mag_l1 - # self.update_scalar("en_tot_l1", en_tot_l1) - - @staticmethod - def diagnostics_dct(): - dct = {} - dct["bt2"] = "Hdiv" - dct["qt3"] = "L2" - dct["div_u"] = "L2" - dct["u2"] = "Hdiv" - return dct + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='linear_q')\n" + ] + elif "variat_qb.Options" in line: + new_file += [ + "model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='linear_q',\n" + ] + new_file += [ + " div_u=model.diagnostics.div_u,\n" + ] + new_file += [ + " u2=model.diagnostics.u2,\n" + ] + new_file += [ + " qt3=model.diagnostics.qt3,\n" + ] + new_file += [ + " bt2=model.diagnostics.bt2)\n" + ] + elif "variat_viscous.Options" in line: + new_file += [ + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='linear_q',\n" + ] + new_file += [ + " rho=model.mhd.density,\n" + ] + new_file += [ + " pt3=model.diagnostics.qt3)\n" + ] + elif "variat_resist.Options" in line: + new_file += [ + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='linear_q',\n" + ] + new_file += [ + " rho=model.mhd.density,\n" + ] + new_file += [ + " pt3=model.diagnostics.qt3)\n" + ] + elif "sqrt_p.add_background" in line: + new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] + new_file += [line] + else: + new_file += [line] - __diagnostics__ = diagnostics_dct() + with open(params_path, "w") as f: + for line in new_file: + f.write(line) -class ViscoresistiveDeltafMHD_with_q(StruphyModel): +class ViscoResistiveDeltafMHD_with_q(StruphyModel): r"""Linear visco-resistive MHD equations discretized with a variational method. :ref:`normalization`: @@ -1983,147 +1992,103 @@ class ViscoresistiveDeltafMHD_with_q(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - dct["em_fields"]["b2"] = "Hdiv" - dct["fluid"]["mhd"] = {"rho3": "L2", "q3": "L2", "uv": "H1vec"} - return dct + ## species - @staticmethod - def bulk_species(): - return "mhd" + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="H1vec") + self.sqrt_p = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_fields.VariationalDensityEvolve: ["mhd_rho3", "mhd_uv"], - propagators_fields.VariationalMomentumAdvection: ["mhd_uv"], - propagators_fields.VariationalQBEvolve: ["mhd_q3", "b2", "mhd_uv"], - propagators_fields.VariationalViscosity: ["mhd_q3", "mhd_uv"], - propagators_fields.VariationalResistivity: ["mhd_q3", "b2"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - def __init__(self, params, comm, clone_config=None): - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import H1vecMassMatrix_density - from struphy.polar.basic import PolarVector - - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) - - self.WMM = H1vecMassMatrix_density(self.derham, self.mass_ops, self.domain) - - # Initialize propagators/integrators used in splitting substeps - lin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["lin_solver"] - nonlin_solver_density = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["nonlin_solver"] - lin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["lin_solver"] - nonlin_solver_momentum = params["fluid"]["mhd"]["options"]["VariationalMomentumAdvection"]["nonlin_solver"] - lin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalQBEvolve"]["lin_solver"] - nonlin_solver_magfield = params["fluid"]["mhd"]["options"]["VariationalQBEvolve"]["nonlin_solver"] - lin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["lin_solver"] - nonlin_solver_viscosity = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["nonlin_solver"] - lin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["lin_solver"] - nonlin_solver_resistivity = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["nonlin_solver"] - if "linearize_current" in params["fluid"]["mhd"]["options"]["VariationalResistivity"].keys(): - self._linearize_current = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["linearize_current"] - else: - self._linearize_current = False - self._gamma = params["fluid"]["mhd"]["options"]["VariationalDensityEvolve"]["physics"]["gamma"] - self._mu = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu"] - self._mu_a = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["mu_a"] - self._alpha = params["fluid"]["mhd"]["options"]["VariationalViscosity"]["physics"]["alpha"] - self._eta = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta"] - self._eta_a = params["fluid"]["mhd"]["options"]["VariationalResistivity"]["physics"]["eta_a"] - model = "deltaf_q" - - # set keyword arguments for propagators - self._kwargs[propagators_fields.VariationalDensityEvolve] = { - "model": model, - "gamma": self._gamma, - "mass_ops": self.WMM, - "lin_solver": lin_solver_density, - "nonlin_solver": nonlin_solver_density, - } - - self._kwargs[propagators_fields.VariationalMomentumAdvection] = { - "mass_ops": self.WMM, - "lin_solver": lin_solver_momentum, - "nonlin_solver": nonlin_solver_momentum, - } - - self._kwargs[propagators_fields.VariationalQBEvolve] = { - "model": model, - "mass_ops": self.WMM, - "lin_solver": lin_solver_magfield, - "nonlin_solver": nonlin_solver_magfield, - "gamma": self._gamma, - "div_u": self.pointer["div_u"], - "u2": self.pointer["u2"], - "bt2": self.pointer["bt2"], - "qt3": self.pointer["qt3"], - } - - self._kwargs[propagators_fields.VariationalViscosity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "mu": self._mu, - "mu_a": self._mu_a, - "alpha": self._alpha, - "mass_ops": self.WMM, - "lin_solver": lin_solver_viscosity, - "nonlin_solver": nonlin_solver_viscosity, - "pt3": self.pointer["qt3"], - } - - self._kwargs[propagators_fields.VariationalResistivity] = { - "model": model, - "rho": self.pointer["mhd_rho3"], - "gamma": self._gamma, - "eta": self._eta, - "eta_a": self._eta_a, - "lin_solver": lin_solver_resistivity, - "nonlin_solver": nonlin_solver_resistivity, - "linearize_current": self._linearize_current, - "pt3": self.pointer["qt3"], - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation + class Diagnostics(DiagnosticSpecies): + def __init__(self): + self.div_u = FEECVariable(space="L2") + self.u2 = FEECVariable(space="Hdiv") + self.qt3 = FEECVariable(space="L2") + self.bt2 = FEECVariable(space="Hdiv") + self.init_variables() + + ## propagators + + class Propagators: + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + self.variat_dens = propagators_fields.VariationalDensityEvolve() + self.variat_mom = propagators_fields.VariationalMomentumAdvection() + self.variat_qb = propagators_fields.VariationalQBEvolve() + if with_viscosity: + self.variat_viscous = propagators_fields.VariationalViscosity() + if with_resistivity: + self.variat_resist = propagators_fields.VariationalResistivity() + + ## abstract methods + + def __init__( + self, + with_viscosity: bool = True, + with_resistivity: bool = True, + ): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + self.diagnostics = self.Diagnostics() + + # 2. instantiate all propagators + self.propagators = self.Propagators( + with_viscosity=with_viscosity, + with_resistivity=with_resistivity, + ) + + # 3. assign variables to propagators + self.propagators.variat_dens.variables.rho = self.mhd.density + self.propagators.variat_dens.variables.u = self.mhd.velocity + self.propagators.variat_mom.variables.u = self.mhd.velocity + self.propagators.variat_qb.variables.u = self.mhd.velocity + self.propagators.variat_qb.variables.q = self.mhd.sqrt_p + self.propagators.variat_qb.variables.b = self.em_fields.b_field + if with_viscosity: + self.propagators.variat_viscous.variables.s = self.mhd.sqrt_p + self.propagators.variat_viscous.variables.u = self.mhd.velocity + if with_resistivity: + self.propagators.variat_resist.variables.s = self.mhd.sqrt_p + self.propagators.variat_resist.variables.b = self.em_fields.b_field + + # define scalars for update_scalar_quantities self.add_scalar("en_U") - self.add_scalar("en_thermo_1") - self.add_scalar("en_thermo_2") self.add_scalar("en_mag_1") self.add_scalar("en_mag_2") + self.add_scalar("en_thermo_1") + self.add_scalar("en_thermo_2") self.add_scalar("en_tot") - # self.add_scalar("dens_tot") - # self.add_scalar("tot_div_B") + @property + def bulk_species(self): + return self.mhd + + @property + def velocity_scale(self): + return "alfvén" - # self.add_scalar("en_tot_l1") - # self.add_scalar("en_thermo_l1") - # self.add_scalar("en_mag_l1") + def allocate_helpers(self): + projV3 = L2Projector("L2", self._mass_ops) - # temporary vectors for scalar quantities - self._tmp_div_B = self.derham.Vh_pol["3"].zeros() - tmp_dof = self.derham.Vh_pol["3"].zeros() - projV3 = L2Projector("L2", self.mass_ops) + def f(e1, e2, e3): + return 1 - self._integrator = projV3(self.domain.jacobian_det, dofs=tmp_dof) + f = np.vectorize(f) + self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() if isinstance(self._ones, PolarVector): @@ -2131,53 +2096,91 @@ def __init__(self, params, comm, clone_config=None): else: self._ones[:] = 1.0 + self._tmp_div_B = self.derham.Vh_pol["3"].zeros() + def update_scalar_quantities(self): - # Update mass matrix - en_U = 0.5 * self.WMM.massop.dot_inner(self.pointer["mhd_uv"], self.pointer["mhd_uv"]) + rho = self.mhd.density.spline.vector + u = self.mhd.velocity.spline.vector + q = self.mhd.sqrt_p.spline.vector + b = self.em_fields.b_field.spline.vector + bt2 = self.propagators.variat_qb.options.bt2.spline.vector + qt3 = self.propagators.variat_qb.options.qt3.spline.vector + + gamma = self.propagators.variat_qb.options.gamma + + en_U = 0.5 * self.mass_ops.WMM.massop.dot_inner(u, u) self.update_scalar("en_U", en_U) - en_mag1 = 0.5 * self._mass_ops.M2.dot_inner(self.pointer["b2"], self.pointer["b2"]) + en_mag1 = 0.5 * self.mass_ops.M2.dot_inner(b, b) self.update_scalar("en_mag_1", en_mag1) - en_mag2 = 0.5 * self._mass_ops.M2.dot_inner(self.pointer["bt2"], self.projected_equil.b2) + en_mag2 = self.mass_ops.M2.dot_inner(bt2, self.projected_equil.b2) self.update_scalar("en_mag_2", en_mag2) - en_th_1 = 1 / (self._gamma - 1) * self._mass_ops.M3.dot_inner(self.pointer["mhd_q3"], self.pointer["mhd_q3"]) + en_th_1 = 1.0 / (gamma - 1.0) * self.mass_ops.M3.dot_inner(q, q) self.update_scalar("en_thermo_1", en_th_1) - en_th_2 = 2 / (self._gamma - 1) * self._mass_ops.M3.dot_inner(self.pointer["qt3"], self.projected_equil.q3) + en_th_2 = 2.0 / (gamma - 1.0) * self.mass_ops.M3.dot_inner(qt3, self.projected_equil.q3) self.update_scalar("en_thermo_2", en_th_2) en_tot = en_U + en_th_1 + en_th_2 + en_mag1 + en_mag2 self.update_scalar("en_tot", en_tot) - # dens_tot = self._ones.dot(self.pointer["mhd_rho3"]) - # self.update_scalar("dens_tot", dens_tot) - - # div_B = self.derham.div.dot(self.pointer["b2"], out=self._tmp_div_B) - # L2_div_B = self._mass_ops.M3.dot_inner(div_B, div_B) - # self.update_scalar("tot_div_B", L2_div_B) - - # en_thermo_l1 = self._integrator.dot(self.mass_ops.M3.dot(self.pointer["mhd_p3"])) / (self._gamma - 1.0) - # self.update_scalar("en_thermo_l1", en_thermo_l1) - - # wb2 = self._mass_ops.M2.dot(self.pointer["b2"], out=self._tmp_wb2) - # en_mag_l1 = wb2.dot(self.projected_equil.b2) - # self.update_scalar("en_mag_l1", en_mag_l1) - - # en_tot_l1 = en_thermo_l1 + en_mag_l1 - # self.update_scalar("en_tot_l1", en_tot_l1) - - @staticmethod - def diagnostics_dct(): - dct = {} - dct["bt2"] = "Hdiv" - dct["qt3"] = "L2" - dct["div_u"] = "L2" - dct["u2"] = "Hdiv" - return dct + # default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "variat_dens.Options" in line: + new_file += [ + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='deltaf_q')\n" + ] + elif "variat_qb.Options" in line: + new_file += [ + "model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='deltaf_q',\n" + ] + new_file += [ + " div_u=model.diagnostics.div_u,\n" + ] + new_file += [ + " u2=model.diagnostics.u2,\n" + ] + new_file += [ + " qt3=model.diagnostics.qt3,\n" + ] + new_file += [ + " bt2=model.diagnostics.bt2)\n" + ] + elif "variat_viscous.Options" in line: + new_file += [ + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='deltaf_q',\n" + ] + new_file += [ + " rho=model.mhd.density,\n" + ] + new_file += [ + " pt3=model.diagnostics.qt3)\n" + ] + elif "variat_resist.Options" in line: + new_file += [ + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='deltaf_q',\n" + ] + new_file += [ + " rho=model.mhd.density,\n" + ] + new_file += [ + " pt3=model.diagnostics.qt3)\n" + ] + elif "sqrt_p.add_background" in line: + new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] + new_file += [line] + else: + new_file += [line] - __diagnostics__ = diagnostics_dct() + with open(params_path, "w") as f: + for line in new_file: + f.write(line) class EulerSPH(StruphyModel): diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index d70de2c0e..1ef696788 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -145,8 +145,6 @@ class FieldSpecies(Species): class FluidSpecies(Species): """Single fluid species in 3d configuration space.""" - pass - class ParticleSpecies(Species): """Single kinetic species in 3d + vdim phase space.""" @@ -216,5 +214,3 @@ def set_save_data( class DiagnosticSpecies(Species): """Diagnostic species (fields) without mass and charge.""" - - pass diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index c236f8d99..0f46d947b 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -21,16 +21,10 @@ if rank == 0: print(f"\n{toy_models = }") -fluid_models = [ - "LinearMHD", - "EulerSPH", - "LinearExtendedMHDuniform", - "ColdPlasma", - "HasegawaWakatani", -] -# for name, obj in inspect.getmembers(fluid): -# if inspect.isclass(obj) and "models.fluid" in obj.__module__: -# fluid_models += [name] +fluid_models = [] +for name, obj in inspect.getmembers(fluid): + if inspect.isclass(obj) and "models.fluid" in obj.__module__: + fluid_models += [name] if rank == 0: print(f"\n{fluid_models = }") diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index c7515c10c..43f68c262 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -33,8 +33,10 @@ from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.variational_utilities import ( BracketOperator, + Hdiv0_transport_operator, InternalEnergyEvaluator, KineticEnergyEvaluator, + Pressure_transport_operator, ) from struphy.fields_background.equils import set_defaults from struphy.geometry.utilities import TransformedPformComponent @@ -53,7 +55,7 @@ from struphy.kinetic_background.base import Maxwellian from struphy.kinetic_background.maxwellians import GyroMaxwellian2D, Maxwellian3D from struphy.linear_algebra.saddle_point import SaddlePointSolver -from struphy.linear_algebra.schur_solver import SchurSolver +from struphy.linear_algebra.schur_solver import SchurSolver, SchurSolverFull from struphy.linear_algebra.solver import NonlinearSolverParameters, SolverParameters from struphy.models.species import Species from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable @@ -4080,58 +4082,91 @@ class VariationalMagFieldEvolve(Propagator): """ - @staticmethod - def options(default=False): - dct = {} - dct["lin_solver"] = { - "tol": 1e-12, - "maxiter": 500, - "non_linear_maxiter": 100, - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "verbose": False, - } - dct["nonlin_solver"] = { - "tol": 1e-8, - "maxiter": 100, - "info": False, - "linearize": False, - } + class Variables: + def __init__(self): + self._u: FEECVariable = None + self._b: FEECVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def u(self) -> FEECVariable: + return self._u - return dct + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1vec" + self._u = new - def __init__( - self, - b: BlockVector, - u: BlockVector, - *, - model: str = "full", - mass_ops, # H1vecMassMatrix_density, - lin_solver: dict = options(default=True)["lin_solver"], - nonlin_solver: dict = options(default=True)["nonlin_solver"], - ): - super().__init__(b, u) + @property + def b(self) -> FEECVariable: + return self._b - assert model in ["full", "full_p", "linear"] - self._model = model - self._mass_ops = mass_ops - self._lin_solver = lin_solver - self._nonlin_solver = nonlin_solver - self._linearize = self._nonlin_solver["linearize"] + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._b = new - self._info = self._nonlin_solver["info"] and (MPI.COMM_WORLD.Get_rank() == 0) + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + OptsModel = Literal["full", "full_p", "linear"] + # propagator options + model: OptsModel = "full" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None - self._Mrho = mass_ops + def __post_init__(self): + # checks + check_option(self.model, self.OptsModel) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters(type="Newton") + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._model = self.options.model + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + self._linearize = self._nonlin_solver.linearize + + self._info = self._nonlin_solver.info and (MPI.COMM_WORLD.Get_rank() == 0) + + self._Mrho = self.mass_ops.WMM + self._Mrho.inv._options["pc"] = MassMatrixDiagonalPreconditioner(self._Mrho.massop) # Projector self._initialize_projectors_and_mass() # bunch of temporaries to avoid allocating in the loop + u = self.variables.u.spline.vector + b = self.variables.b.spline.vector + self._tmp_un1 = u.space.zeros() self._tmp_un2 = u.space.zeros() self._tmp_un12 = u.space.zeros() @@ -4162,8 +4197,8 @@ def __call_newton(self, dt): print() print("Newton iteration in VariationalMagFieldEvolve") # Compute implicit approximation of s^{n+1} - bn = self.feec_vars[0] - un = self.feec_vars[1] + un = self.variables.u.spline.vector + bn = self.variables.b.spline.vector bn1 = bn.copy(out=self._tmp_bn1) # Initialize variable for Newton iteration @@ -4176,10 +4211,10 @@ def __call_newton(self, dt): un1 = un.copy(out=self._tmp_un1) un1 += self._tmp_un_diff mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - tol = self._nonlin_solver["tol"] + tol = self._nonlin_solver.tol err = tol + 1 - for it in range(self._nonlin_solver["maxiter"]): + for it in range(self._nonlin_solver.maxiter): # Newton iteration # half time step approximation bn12 = bn.copy(out=self._tmp_bn12) @@ -4262,20 +4297,18 @@ def __call_newton(self, dt): # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver["maxiter"] - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalMagFieldEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) self._tmp_un_diff = un1 - un self._tmp_bn_diff = bn1 - bn - self.feec_vars_update(bn1, un1) + self.update_feec_variables(b=bn1, u=un1) def _initialize_projectors_and_mass(self): """Initialization of all the `BasisProjectionOperator` and needed to compute the bracket term""" - from struphy.feec.variational_utilities import Hdiv0_transport_operator - self.curlPib = Hdiv0_transport_operator(self.derham) self.curlPibT = self.curlPib.T @@ -4324,15 +4357,13 @@ def _initialize_projectors_and_mass(self): self._Jacobian[1, 0] = self._dt2_curlPib self._Jacobian[1, 1] = self._I2 - from struphy.linear_algebra.schur_solver import SchurSolverFull - self._inv_Jacobian = SchurSolverFull( self._Jacobian, - self._lin_solver["type"][0], + self.options.solver, pc=self._Mrho.inv, - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], - verbose=self._lin_solver["verbose"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, + verbose=self._lin_solver.verbose, recycle=True, ) @@ -4350,8 +4381,6 @@ def _update_Pib(self, b): self.curlPibT.update_coeffs(b) def _create_Pib0(self): - from struphy.feec.variational_utilities import Hdiv0_transport_operator - self.curlPib0 = Hdiv0_transport_operator(self.derham) self.curlPibT0 = self.curlPib0.T @@ -4437,72 +4466,130 @@ class VariationalPBEvolve(Propagator): and :math:`\mathcal{U}^v` is :class:`~struphy.feec.basis_projection_ops.BasisProjectionOperators`. """ - @staticmethod - def options(default=False): - dct = {} - dct["lin_solver"] = { - "tol": 1e-12, - "maxiter": 500, - "non_linear_maxiter": 100, - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "verbose": False, - } - dct["nonlin_solver"] = { - "tol": 1e-8, - "maxiter": 100, - "type": ["Picard"], - "info": False, - "linearize": False, - } - dct["physics"] = {"gamma": 5 / 3} + class Variables: + def __init__(self): + self._p: FEECVariable = None + self._u: FEECVariable = None + self._b: FEECVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def p(self) -> FEECVariable: + return self._p - return dct + @p.setter + def p(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._p = new - def __init__( - self, - p: StencilVector, - b: BlockVector, - u: BlockVector, - *, - model: str = "full", - gamma: float = options()["physics"]["gamma"], - mass_ops, # H1vecMassMatrix_density, - lin_solver: dict = options(default=True)["lin_solver"], - nonlin_solver: dict = options(default=True)["nonlin_solver"], - div_u: StencilVector | None = None, - u2: BlockVector | None = None, - pt3: StencilVector | None = None, - bt2: BlockVector | None = None, - ): - super().__init__(p, b, u) + @property + def u(self) -> FEECVariable: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1vec" + self._u = new + + @property + def b(self) -> FEECVariable: + return self._b + + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._b = new - assert model in ["full_p", "linear", "deltaf"] - self._model = model - self._mass_ops = mass_ops - self._lin_solver = lin_solver - self._nonlin_solver = nonlin_solver - self._linearize = self._nonlin_solver["linearize"] - self._gamma = gamma + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsModel = Literal["full_p", "linear", "deltaf"] + # propagator options + model: OptsModel = "full_p" + gamma: float = 5.0 / 3.0 + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None + div_u: FEECVariable = None + u2: FEECVariable = None + pt3: FEECVariable = None + bt2: FEECVariable = None + + def __post_init__(self): + # checks + check_option(self.model, self.OptsModel) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() - self._divu = div_u - self._u2 = u2 - self._pt3 = pt3 - self._bt2 = bt2 + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters() - self._info = self._nonlin_solver["info"] and (MPI.COMM_WORLD.Get_rank() == 0) + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options - self._Mrho = mass_ops + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._model = self.options.model + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + self._linearize = self.options.nonlin_solver.linearize + self._gamma = self.options.gamma + + if self.options.div_u is None: + self._divu = None + else: + self._divu = self.options.div_u.spline.vector + + if self.options.u2 is None: + self._u2 = None + else: + self._u2 = self.options.u2.spline.vector + + if self.options.pt3 is None: + self._pt3 = None + else: + self._pt3 = self.options.pt3.spline.vector + + if self.options.bt2 is None: + self._bt2 = None + else: + self._bt2 = self.options.bt2.spline.vector + + self._info = self._nonlin_solver.info and (MPI.COMM_WORLD.Get_rank() == 0) + + self._Mrho = self.mass_ops.WMM + self._Mrho.inv._options["pc"] = MassMatrixDiagonalPreconditioner(self._Mrho.massop) # Projector self._initialize_projectors_and_mass() # bunch of temporaries to avoid allocating in the loop + u = self.variables.u.spline.vector + p = self.variables.p.spline.vector + b = self.variables.b.spline.vector + self._tmp_un1 = u.space.zeros() self._tmp_un2 = u.space.zeros() self._tmp_un12 = u.space.zeros() @@ -4536,7 +4623,7 @@ def __init__( self._extracted_b2 = self.derham.extraction_ops["2"].dot(self.projected_equil.b2) def __call__(self, dt): - if self._nonlin_solver["type"] == "Picard": + if self._nonlin_solver.type == "Picard": self.__call_picard(dt) else: raise ValueError("Only Picard solver is implemented for VariationalPBEvolve") @@ -4548,9 +4635,9 @@ def __call_picard(self, dt): print() print("Newton iteration in VariationalPBEvolve") - pn = self.feec_vars[0] - bn = self.feec_vars[1] - un = self.feec_vars[2] + un = self.variables.u.spline.vector + pn = self.variables.p.spline.vector + bn = self.variables.b.spline.vector self._update_Pib(bn) self._update_Projp(pn) @@ -4563,10 +4650,10 @@ def __call_picard(self, dt): un1 = un.copy(out=self._tmp_un1) un1 += self._tmp_un_diff mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - tol = self._nonlin_solver["tol"] + tol = self._nonlin_solver.tol err = tol + 1 - for it in range(self._nonlin_solver["maxiter"]): + for it in range(self._nonlin_solver.maxiter): # Picard iteration # half time step approximation @@ -4713,7 +4800,7 @@ def __call_picard(self, dt): # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver["maxiter"] - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalPBEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) @@ -4721,7 +4808,7 @@ def __call_picard(self, dt): self._tmp_un_diff = un1 - un self._tmp_bn_diff = bn1 - bn self._tmp_pn_diff = pn1 - pn - self.feec_vars_update(pn1, bn1, un1) + self.update_feec_variables(p=pn1, b=bn1, u=un1) self._transop_p.div.dot(un12, out=self._divu) self._transop_p._Uv.dot(un1, out=self._u2) @@ -4747,9 +4834,6 @@ def __call_picard(self, dt): def _initialize_projectors_and_mass(self): """Initialization of all the `BasisProjectionOperator` and needed to compute the bracket term""" - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import Hdiv0_transport_operator, Pressure_transport_operator - self.curlPib = Hdiv0_transport_operator(self.derham) self.curlPibT = self.curlPib.T self._transop_p = Pressure_transport_operator(self.derham, self.domain, self.basis_ops.Uv, self._gamma) @@ -4837,11 +4921,11 @@ def _initialize_projectors_and_mass(self): self._inv_Jacobian = SchurSolverFull( self._Jacobian, - self._lin_solver["type"][0], + self.options.solver, pc=self._Mrho.inv, - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], - verbose=self._lin_solver["verbose"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, + verbose=self._lin_solver.verbose, recycle=True, ) @@ -4860,8 +4944,6 @@ def _update_Pib(self, b): self.curlPibT.update_coeffs(b) def _create_Pib0(self): - from struphy.feec.variational_utilities import Hdiv0_transport_operator - self.curlPib0 = Hdiv0_transport_operator(self.derham) self.curlPibT0 = self.curlPib.T self.curlPib0.update_coeffs(self.projected_equil.b2) @@ -4874,7 +4956,6 @@ def _update_Projp(self, p): def _create_transop0(self): """Update the weights of the `BasisProjectionOperator`""" - from struphy.feec.variational_utilities import Pressure_transport_operator self._transop_p0 = Pressure_transport_operator(self.derham, self.domain, self.basis_ops.Uv, self._gamma) self._transop_p0T = self._transop_p0.T @@ -4983,72 +5064,130 @@ class VariationalQBEvolve(Propagator): and :math:`\mathcal{U}^v` is :class:`~struphy.feec.basis_projection_ops.BasisProjectionOperators`. """ - @staticmethod - def options(default=False): - dct = {} - dct["lin_solver"] = { - "tol": 1e-12, - "maxiter": 500, - "non_linear_maxiter": 100, - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "verbose": False, - } - dct["nonlin_solver"] = { - "tol": 1e-8, - "maxiter": 100, - "type": ["Picard"], - "info": False, - "linearize": False, - } - dct["physics"] = {"gamma": 5 / 3} + class Variables: + def __init__(self): + self._q: FEECVariable = None + self._u: FEECVariable = None + self._b: FEECVariable = None + + @property + def q(self) -> FEECVariable: + return self._q + + @q.setter + def q(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._q = new + + @property + def u(self) -> FEECVariable: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1vec" + self._u = new + + @property + def b(self) -> FEECVariable: + return self._b + + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._b = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # specific literals + OptsModel = Literal["full_q", "linear_q", "deltaf_q"] + # propagator options + model: OptsModel = "full_q" + gamma: float = 5.0 / 3.0 + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None + div_u: FEECVariable = None + u2: FEECVariable = None + qt3: FEECVariable = None + bt2: FEECVariable = None + + def __post_init__(self): + # checks + check_option(self.model, self.OptsModel) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new - if default: - dct = descend_options_dict(dct, []) + @profile + def allocate(self): + self._model = self.options.model + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + self._linearize = self.options.nonlin_solver.linearize + self._gamma = self.options.gamma - return dct + if self.options.div_u is None: + self._divu = None + else: + self._divu = self.options.div_u.spline.vector - def __init__( - self, - q: StencilVector, - b: BlockVector, - u: BlockVector, - *, - model: str = "full", - gamma: float = options()["physics"]["gamma"], - mass_ops, # H1vecMassMatrix_density, - lin_solver: dict = options(default=True)["lin_solver"], - nonlin_solver: dict = options(default=True)["nonlin_solver"], - div_u: StencilVector | None = None, - u2: BlockVector | None = None, - qt3: StencilVector | None = None, - bt2: BlockVector | None = None, - ): - super().__init__(q, b, u) + if self.options.u2 is None: + self._u2 = None + else: + self._u2 = self.options.u2.spline.vector - assert model in ["full_q", "linear_q", "deltaf_q"] - self._model = model - self._mass_ops = mass_ops - self._lin_solver = lin_solver - self._nonlin_solver = nonlin_solver - self._linearize = self._nonlin_solver["linearize"] - self._gamma = gamma + if self.options.qt3 is None: + self._qt3 = None + else: + self._qt3 = self.options.qt3.spline.vector - self._divu = div_u - self._u2 = u2 - self._qt3 = qt3 - self._bt2 = bt2 + if self.options.bt2 is None: + self._bt2 = None + else: + self._bt2 = self.options.bt2.spline.vector - self._info = self._nonlin_solver["info"] and (self.rank == 0) + self._info = self._nonlin_solver.info and (self.rank == 0) - self._Mrho = mass_ops + self._Mrho = self.mass_ops.WMM + self._Mrho.inv._options["pc"] = MassMatrixDiagonalPreconditioner(self._Mrho.massop) # Projector self._initialize_projectors_and_mass() # bunch of temporaries to avoid allocating in the loop + u = self.variables.u.spline.vector + q = self.variables.q.spline.vector + b = self.variables.b.spline.vector + self._tmp_un1 = u.space.zeros() self._tmp_un12 = u.space.zeros() self._tmp_bn1 = b.space.zeros() @@ -5080,7 +5219,7 @@ def __init__( self._extracted_q3 = self.derham.extraction_ops["3"].dot(self.projected_equil.q3) def __call__(self, dt): - if self._nonlin_solver["type"] == "Picard": + if self._nonlin_solver.type == "Picard": self.__call_picard(dt) else: raise ValueError("Only Picard solver is implemented for VariationalQBEvolve") @@ -5092,9 +5231,9 @@ def __call_picard(self, dt): print() print("Newton iteration in VariationalQBEvolve") - qn = self.feec_vars[0] - bn = self.feec_vars[1] - un = self.feec_vars[2] + un = self.variables.u.spline.vector + qn = self.variables.q.spline.vector + bn = self.variables.b.spline.vector self._update_Pib(bn) self._update_Projq(qn) @@ -5107,10 +5246,10 @@ def __call_picard(self, dt): un1 = un.copy(out=self._tmp_un1) un1 += self._tmp_un_diff mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - tol = self._nonlin_solver["tol"] + tol = self._nonlin_solver.tol err = tol + 1 - for it in range(self._nonlin_solver["maxiter"]): + for it in range(self._nonlin_solver.maxiter): # Picard iteration # half time step approximation @@ -5252,7 +5391,7 @@ def __call_picard(self, dt): # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver["maxiter"] - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalPBEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) @@ -5260,7 +5399,7 @@ def __call_picard(self, dt): self._tmp_un_diff = un1 - un self._tmp_bn_diff = bn1 - bn self._tmp_qn_diff = qn1 - qn - self.feec_vars_update(qn1, bn1, un1) + self.update_feec_variables(q=qn1, b=bn1, u=un1) self._transop_q.div.dot(un12, out=self._divu) self._transop_q._Uv.dot(un1, out=self._u2) @@ -5286,9 +5425,6 @@ def __call_picard(self, dt): def _initialize_projectors_and_mass(self): """Initialization of all the `BasisProjectionOperator` and needed to compute the bracket term""" - from struphy.feec.projectors import L2Projector - from struphy.feec.variational_utilities import Hdiv0_transport_operator, Pressure_transport_operator - self.curlPib = Hdiv0_transport_operator(self.derham) self.curlPibT = self.curlPib.T self._transop_q = Pressure_transport_operator(self.derham, self.domain, self.basis_ops.Uv, self._gamma / 2.0) @@ -5395,11 +5531,11 @@ def _initialize_projectors_and_mass(self): self._inv_Jacobian = SchurSolverFull3( self._Jacobian, - self._lin_solver["type"][0], + self.options.solver, pc=self._Mrho.inv, - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], - verbose=self._lin_solver["verbose"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, + verbose=self._lin_solver.verbose, recycle=True, ) @@ -5418,8 +5554,6 @@ def _update_Pib(self, b): self.curlPibT.update_coeffs(b) def _create_Pib0(self): - from struphy.feec.variational_utilities import Hdiv0_transport_operator - self.curlPib0 = Hdiv0_transport_operator(self.derham) self.curlPibT0 = self.curlPib.T self.curlPib0.update_coeffs(self.projected_equil.b2) @@ -5432,7 +5566,6 @@ def _update_Projq(self, q): def _create_transop0(self): """Update the weights of the `BasisProjectionOperator`""" - from struphy.feec.variational_utilities import Pressure_transport_operator self._transop_q0 = Pressure_transport_operator(self.derham, self.domain, self.basis_ops.Uv, self._gamma / 2.0) self._transop_q0T = self._transop_q0.T @@ -5537,72 +5670,95 @@ class VariationalViscosity(Propagator): """ - @staticmethod - def options(default=False): - dct = {} - dct["lin_solver"] = { - "tol": 1e-12, - "maxiter": 500, - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "verbose": False, - } - dct["nonlin_solver"] = { - "tol": 1e-8, - "maxiter": 100, - "type": ["Newton"], - "info": False, - "fast": False, - } - dct["physics"] = { - "gamma": 1.66666666667, - "mu": 0.0, - "mu_a": 0.0, - "alpha": 0.0, - } + class Variables: + def __init__(self): + self._s: FEECVariable = None + self._u: FEECVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def s(self) -> FEECVariable: + return self._s - return dct + @s.setter + def s(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._s = new - def __init__( - self, - s: StencilVector, - u: BlockVector, - *, - model: str = "barotropic", - gamma: float = options()["physics"]["gamma"], - rho: StencilVector, - mu: float = options()["physics"]["mu"], - mu_a: float = options()["physics"]["mu_a"], - alpha: float = options()["physics"]["alpha"], - mass_ops, # H1vecMassMatrix_density, - lin_solver: dict = options(default=True)["lin_solver"], - nonlin_solver: dict = options(default=True)["nonlin_solver"], - energy_evaluator: InternalEnergyEvaluator = None, - pt3: StencilVector | None = None, - ): - super().__init__(s, u) + @property + def u(self) -> FEECVariable: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1vec" + self._u = new - assert model in ["full", "full_p", "full_q", "linear_p", "linear_q", "deltaf_q"] + def __init__(self): + self.variables = self.Variables() - self._model = model - self._gamma = gamma - self._lin_solver = lin_solver - self._nonlin_solver = nonlin_solver - self._mu_a = mu_a - self._alpha = alpha - self._mu = mu - self._rho = rho - self._pt3 = pt3 - self._energy_evaluator = energy_evaluator + @dataclass + class Options: + # specific literals + OptsModel = Literal["full", "full_p", "full_q", "linear_p", "linear_q", "deltaf_q"] + # propagator options + model: OptsModel = "full" + gamma: float = 5.0 / 3.0 + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixDiagonalPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None + rho: FEECVariable = None + pt3: FEECVariable = None + mu: float = 0.0 + mu_a: float = 0.0 + alpha: float = 0.0 + + def __post_init__(self): + # checks + check_option(self.model, self.OptsModel) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters(type="Newton") + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._model = self.options.model + self._gamma = self.options.gamma + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + self._mu_a = self.options.mu_a + self._alpha = self.options.alpha + self._mu = self.options.mu + self._rho = self.options.rho + self._pt3 = self.options.pt3 - self._info = self._nonlin_solver["info"] and (MPI.COMM_WORLD.Get_rank() == 0) + self._info = self._nonlin_solver.info and (MPI.COMM_WORLD.Get_rank() == 0) - self._Mrho = mass_ops + self._Mrho = self.mass_ops.WMM + self._Mrho.inv._options["pc"] = MassMatrixDiagonalPreconditioner(self._Mrho.massop) # Femfields for the projector self.sf = self.derham.create_spline_function("sf", "L2") @@ -5617,9 +5773,13 @@ def __init__( self.gu122f = self.derham.create_spline_function("gu122", "Hcurl") # Projector + self._energy_evaluator = InternalEnergyEvaluator(self.derham, self._gamma) self._initialize_projectors_and_mass() # bunch of temporaries to avoid allocating in the loop + u = self.variables.u.spline.vector + s = self.variables.s.spline.vector + self._tmp_un1 = u.space.zeros() self._tmp_un12 = u.space.zeros() self._tmp_sn1 = s.space.zeros() @@ -5636,7 +5796,7 @@ def __init__( self.tot_rhs = s.space.zeros() def __call__(self, dt): - if self._nonlin_solver["type"] == "Newton": + if self._nonlin_solver.type == "Newton": self.__call_newton(dt) else: raise ValueError( @@ -5646,10 +5806,11 @@ def __call__(self, dt): def __call_newton(self, dt): """Solve the non linear system for updating the variables using Newton iteration method""" # Compute dissipation implicitely - sn = self.feec_vars[0] - un = self.feec_vars[1] + sn = self.variables.s.spline.vector + un = self.variables.u.spline.vector + if self._mu < 1.0e-15 and self._mu_a < 1.0e-15 and self._alpha < 1.0e-15: - self.feec_vars_update(sn, un) + self.update_feec_variables(s=sn, u=un) return if self._info: @@ -5667,7 +5828,7 @@ def __call_newton(self, dt): print("information on the linear solver : ", self.inv_lop._info) if self._model == "linear_p" or (self._model == "linear_q" and self._nonlin_solver["fast"]): - self.feec_vars_update(sn, un1) + self.update_feec_variables(s=sn, u=un1) return # Energy balance term @@ -5676,7 +5837,7 @@ def __call_newton(self, dt): # 2) Initial energy and linear form rho = self._rho if self._model in ["deltaf_q", "linear_q"]: - self.sf.vector = self._pt3 + self.sf.vector = self._pt3.spline.vector else: self.sf.vector = sn @@ -5732,7 +5893,7 @@ def __call_newton(self, dt): for it in range(self._nonlin_solver["maxiter"]): if self._model in ["deltaf_q", "linear_q"]: - self.sf1.vector = self._pt3 + self.sf1.vector = self._pt3.spline.vector else: self.sf1.vector = sn1 @@ -5827,13 +5988,11 @@ def __call_newton(self, dt): f"!!!Warning: Maximum iteration in VariationalViscosity reached - not converged:\n {err = } \n {tol**2 = }", ) - self.feec_vars_update(sn1, un1) + self.update_feec_variables(s=sn1, u=un1) def _initialize_projectors_and_mass(self): """Initialization of all the `BasisProjectionOperator` and needed to compute the bracket term""" - from struphy.feec.projectors import L2Projector - Xv = getattr(self.basis_ops, "Xv") Pcoord0 = CoordinateProjector( 0, @@ -5868,12 +6027,12 @@ def _initialize_projectors_and_mass(self): self.M_de_ds = self.mass_ops.create_weighted_mass("L2", "L2") - if self._lin_solver["type"][1] is None: + if self.options.precond is None: self.pc_jac = None else: pc_class = getattr( preconditioner, - self._lin_solver["type"][1], + self.options.precond, ) self.pc_jac = pc_class(self.M_de_ds) @@ -5881,8 +6040,8 @@ def _initialize_projectors_and_mass(self): self.M_de_ds, "pcg", pc=self.pc_jac, - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, verbose=False, recycle=True, ) @@ -5923,8 +6082,8 @@ def _initialize_projectors_and_mass(self): self.l_op, "pcg", pc=self._Mrho.inv, - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, verbose=False, recycle=True, ) @@ -6269,63 +6428,92 @@ class VariationalResistivity(Propagator): """ - @staticmethod - def options(default=False): - dct = {} - dct["lin_solver"] = { - "tol": 1e-12, - "maxiter": 500, - "type": [ - ("pcg", "MassMatrixDiagonalPreconditioner"), - ("cg", None), - ], - "verbose": False, - } - dct["nonlin_solver"] = {"tol": 1e-8, "maxiter": 100, "type": ["Newton"], "info": False, "fast": False} - dct["physics"] = { - "eta": 0.0, - "eta_a": 0.0, - "gamma": 5 / 3, - } - dct["linearize_current"] = False + class Variables: + def __init__(self): + self._s: FEECVariable = None + self._b: FEECVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def s(self) -> FEECVariable: + return self._s - return dct + @s.setter + def s(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._s = new - def __init__( - self, - s: StencilVector, - b: BlockVector, - *, - model: str = "full", - gamma: float = options()["physics"]["gamma"], - rho: StencilVector, - eta: float = options()["physics"]["eta"], - eta_a: float = options()["physics"]["eta_a"], - lin_solver: dict = options(default=True)["lin_solver"], - nonlin_solver: dict = options(default=True)["nonlin_solver"], - linearize_current: dict = options(default=True)["linearize_current"], - energy_evaluator: InternalEnergyEvaluator = None, - pt3: StencilVector | None = None, - ): - super().__init__(s, b) + @property + def b(self) -> FEECVariable: + return self._b + + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._b = new - assert model in ["full", "full_p", "full_q", "linear_p", "delta_p", "linear_q", "deltaf_q"] + def __init__(self): + self.variables = self.Variables() - self._energy_evaluator = energy_evaluator - self._model = model - self._gamma = gamma - self._eta = eta - self._eta_a = eta_a - self._lin_solver = lin_solver - self._nonlin_solver = nonlin_solver - self._rho = rho - self._linearize_current = linearize_current - self._pt3 = pt3 + @dataclass + class Options: + # specific literals + OptsModel = Literal["full", "full_p", "full_q", "linear_p", "linear_q", "deltaf_q"] + # propagator options + model: OptsModel = "full" + gamma: float = 5.0 / 3.0 + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixDiagonalPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None + linearize_current: bool = False + rho: FEECVariable = None + pt3: FEECVariable = None + eta: float = 0.0 + eta_a: float = 0.0 + + def __post_init__(self): + # checks + check_option(self.model, self.OptsModel) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters(type="Newton") + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._model = self.options.model + self._gamma = self.options.gamma + self._eta = self.options.eta + self._eta_a = self.options.eta_a + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + self._linearize_current = self.options.linearize_current + self._rho = self.options.rho + self._pt3 = self.options.pt3 - self._info = self._nonlin_solver["info"] and (MPI.COMM_WORLD.Get_rank() == 0) + self._info = self._nonlin_solver.info and (MPI.COMM_WORLD.Get_rank() == 0) # Femfields for the projector self.rhof = self.derham.create_spline_function("rhof", "L2") @@ -6337,9 +6525,13 @@ def __init__( self.cbf12 = self.derham.create_spline_function("cBf", "Hcurl") # Projector + self._energy_evaluator = InternalEnergyEvaluator(self.derham, self._gamma) self._initialize_projectors_and_mass() # bunch of temporaries to avoid allocating in the loop + s = self.variables.s.spline.vector + b = self.variables.b.spline.vector + self._tmp_bn1 = b.space.zeros() self._tmp_bn12 = b.space.zeros() self._tmp_sn1 = s.space.zeros() @@ -6356,7 +6548,7 @@ def __init__( ) def __call__(self, dt): - if self._nonlin_solver["type"] == "Newton": + if self._nonlin_solver.type == "Newton": self.__call_newton(dt) else: raise ValueError( @@ -6366,10 +6558,11 @@ def __call__(self, dt): def __call_newton(self, dt): """Solve the non linear system for updating the variables using Newton iteration method""" # Compute dissipation implicitely - sn = self.feec_vars[0] - bn = self.feec_vars[1] + sn = self.variables.s.spline.vector + bn = self.variables.b.spline.vector + if self._eta < 1.0e-15 and self._eta_a < 1.0e-15: - self.feec_vars_update(sn, bn) + self.update_feec_variables(s=sn, b=bn) return if self._info: @@ -6396,17 +6589,17 @@ def __call_newton(self, dt): print("information on the linear solver : ", self.inv_lop._info) if self._model == "linear_p" or (self._model == "linear_q" and self._nonlin_solver["fast"]): - self.feec_vars_update(sn, bn1) + self.update_feec_variables(s=sn, b=bn1) return # Energy balance term # 1) Pointwize energy change energy_change = self._get_energy_change(bn, bn1, total_resistivity) # 2) Initial energy and linear form - rho = self._rho + rho = self._rho.spline.vector self.rhof.vector = rho if self._model in ["deltaf_q", "linear_q"]: - self.sf.vector = self._pt3 + self.sf.vector = self._pt3.spline.vector else: self.sf.vector = sn @@ -6466,7 +6659,7 @@ def __call_newton(self, dt): for it in range(self._nonlin_solver["maxiter"]): if self._model in ["deltaf_q", "linear_q"]: - self.sf1.vector = self._pt3 + self.sf1.vector = self._pt3.spline.vector else: self.sf1.vector = sn1 @@ -6559,7 +6752,7 @@ def __call_newton(self, dt): f"!!!Warning: Maximum iteration in VariationalResistivity reached - not converged:\n {err = } \n {tol**2 = }", ) - self.feec_vars_update(sn1, bn1) + self.update_feec_variables(s=sn1, b=bn1) # if self._pt3 is not None: # bn12 = bn.copy(out=self._tmp_bn12) @@ -6650,8 +6843,6 @@ def __call_newton(self, dt): def _initialize_projectors_and_mass(self): """Initialization of all the `BasisProjectionOperator` and needed to compute the bracket term""" - from struphy.feec.projectors import L2Projector - pc_M1 = preconditioner.MassMatrixDiagonalPreconditioner( self.mass_ops.M1, ) @@ -6682,12 +6873,12 @@ def _initialize_projectors_and_mass(self): D = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] self.M1_cb = self.mass_ops.create_weighted_mass("Hcurl", "Hcurl", weights=[D, "sqrt_g"]) - if self._lin_solver["type"][1] is None: + if self.options.precond is None: self.pc = None else: pc_class = getattr( preconditioner, - self._lin_solver["type"][1], + self.options.precond, ) self.pc_jac = pc_class(self.M_de_ds) @@ -6695,8 +6886,8 @@ def _initialize_projectors_and_mass(self): self.M_de_ds, "pcg", pc=self.pc_jac, - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, verbose=False, recycle=True, ) @@ -6712,12 +6903,12 @@ def _initialize_projectors_and_mass(self): self.r_op = M2 # - self._scaled_stiffness self.l_op = M2 + self._scaled_stiffness + self.phy_cb_stiffness - if self._lin_solver["type"][1] is None: + if self.options.precond is None: self.pc = None else: pc_class = getattr( preconditioner, - self._lin_solver["type"][1], + self.options.precond, ) self.pc = pc_class(M2) @@ -6725,8 +6916,8 @@ def _initialize_projectors_and_mass(self): self.l_op, "pcg", pc=self.pc, - tol=self._lin_solver["tol"], - maxiter=self._lin_solver["maxiter"], + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, verbose=False, recycle=True, ) From 6e9ce14d78b2fc963938302c6a06bb60d891c400 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 14:00:16 +0200 Subject: [PATCH 142/292] re-add model_tests_mpi --- .gitlab-ci.yml | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 3d6a17d39..236d015b5 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -413,14 +413,19 @@ stages: - struphy compile --status - struphy test LinearMHD - struphy test toy - - struphy test models --mpi 1 - - struphy test models --mpi 2 - - struphy test models --mpi 3 - - struphy test models --mpi 4 - - struphy test verification --mpi 4 + - struphy test models + - struphy test verification # - struphy test models --fast --verification --mpi 4 # - struphy test models --fast --verification --mpi 4 --nclones 2 # - struphy test DriftKineticElectrostaticAdiabatic --mpi 2 --nclones 2 + model_tests_mpi: + - struphy compile --status + - struphy test models + - struphy test models --mpi 2 + - struphy test models --verification --mpi 1 + - struphy test models --verification --mpi 4 + - struphy test models --verification --mpi 4 --nclones 2 + - struphy test DriftKineticElectrostaticAdiabatic --mpi 2 --nclones 2 quickstart_tests: - struphy -p - struphy -h From 353d6922b42f512336b3e7970b2a6a18697b5f13 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 14:18:37 +0200 Subject: [PATCH 143/292] remove test_tutorials.py again --- src/struphy/tutorials/tests/test_tutorials.py | 172 ------------------ 1 file changed, 172 deletions(-) delete mode 100644 src/struphy/tutorials/tests/test_tutorials.py diff --git a/src/struphy/tutorials/tests/test_tutorials.py b/src/struphy/tutorials/tests/test_tutorials.py deleted file mode 100644 index c8de4d5c8..000000000 --- a/src/struphy/tutorials/tests/test_tutorials.py +++ /dev/null @@ -1,172 +0,0 @@ -import os - -import pytest -import yaml -from psydac.ddm.mpi import mpi as MPI - -import struphy -from struphy.main import main -from struphy.post_processing import pproc_struphy - -comm = MPI.COMM_WORLD -rank = comm.Get_rank() - -libpath = struphy.__path__[0] -i_path = os.path.join(libpath, "io", "inp") -o_path = os.path.join(libpath, "io", "out") - - -def test_tutorial_02(): - main( - "LinearMHDVlasovCC", - os.path.join(i_path, "tutorials", "params_02.yml"), - os.path.join(o_path, "tutorial_02"), - supress_out=True, - ) - - -def test_tutorial_03(): - main( - "LinearMHD", - os.path.join(i_path, "tutorials", "params_03.yml"), - os.path.join(o_path, "tutorial_03"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_03"), physical=True) - - -def test_tutorial_04(fast): - main( - "Maxwell", - os.path.join(i_path, "tutorials", "params_04a.yml"), - os.path.join(o_path, "tutorial_04a"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_04a")) - - main( - "LinearMHD", - os.path.join(i_path, "tutorials", "params_04b.yml"), - os.path.join(o_path, "tutorial_04b"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_04b")) - - if not fast: - main( - "VariationalMHD", - os.path.join(i_path, "tutorials", "params_04c.yml"), - os.path.join(o_path, "tutorial_04c"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_04c")) - - -def test_tutorial_05(): - main( - "Vlasov", - os.path.join(i_path, "tutorials", "params_05a.yml"), - os.path.join(o_path, "tutorial_05a"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05a")) - - main( - "Vlasov", - os.path.join(i_path, "tutorials", "params_05b.yml"), - os.path.join(o_path, "tutorial_05b"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05b")) - - main( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_05c.yml"), - os.path.join(o_path, "tutorial_05c"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05c")) - - main( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_05d.yml"), - os.path.join(o_path, "tutorial_05d"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05d")) - - main( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_05e.yml"), - os.path.join(o_path, "tutorial_05e"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05e")) - - main( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_05f.yml"), - os.path.join(o_path, "tutorial_05f"), - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_05f")) - - -def test_tutorial_12(): - main( - "Vlasov", - os.path.join(i_path, "tutorials", "params_12a.yml"), - os.path.join(o_path, "tutorial_12a"), - save_step=100, - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_12a")) - - main( - "GuidingCenter", - os.path.join(i_path, "tutorials", "params_12b.yml"), - os.path.join(o_path, "tutorial_12b"), - save_step=10, - supress_out=True, - ) - - comm.Barrier() - if rank == 0: - pproc_struphy.main(os.path.join(o_path, "tutorial_12b")) - - -if __name__ == "__main__": - test_tutorial_04(True) From ac598e38d7eb978e7246f19416186fe76dc4c2ef Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 14:21:49 +0200 Subject: [PATCH 144/292] add new mpi and xp to species.py --- src/struphy/models/species.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 1ef696788..92710cbe7 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,15 +1,10 @@ import warnings from abc import ABCMeta, abstractmethod -from copy import deepcopy -from dataclasses import dataclass -from typing import Callable -import numpy as np -from mpi4py import MPI +from psydac.ddm.mpi import mpi as MPI +from struphy.utils.arrays import xp as np -from struphy.fields_background.base import FluidEquilibrium from struphy.io.options import Units -from struphy.kinetic_background.base import KineticBackground from struphy.models.variables import Variable from struphy.physics.physics import ConstantsOfNature from struphy.pic.utilities import ( From 7be337dcb3d7c7914883737240196fb05d004de8 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 14:27:26 +0200 Subject: [PATCH 145/292] add mpi and xp to some new files which do not exist in devel --- src/struphy/io/options.py | 5 +- src/struphy/main.py | 2 - src/struphy/models/tests/test_models.py | 2 +- .../models/tests/test_verif_EulerSPH.py | 6 +- .../models/tests/test_verif_LinearMHD.py | 6 +- .../models/tests/test_verif_Maxwell.py | 5 +- .../models/tests/test_verif_Poisson.py | 7 +- .../test_verif_VlasovAmpereOneSpecies.py | 6 +- src/struphy/models/variables.py | 4 +- src/struphy/propagators/base.py | 3 +- src/struphy/utils/mpi.py | 102 ------------------ 11 files changed, 21 insertions(+), 127 deletions(-) delete mode 100644 src/struphy/utils/mpi.py diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index dce8b072c..d3b3d4a18 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -2,7 +2,8 @@ from dataclasses import dataclass from typing import Literal, get_args -import numpy as np +from psydac.ddm.mpi import mpi as MPI +from struphy.utils.arrays import xp as np from struphy.physics.physics import ConstantsOfNature @@ -177,8 +178,6 @@ def j(self): 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.""" - from mpi4py import MPI - con = ConstantsOfNature() # velocity (m/s) diff --git a/src/struphy/main.py b/src/struphy/main.py index c1c049029..731a9e8b9 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -10,8 +10,6 @@ import h5py from line_profiler import profile -from mpi4py import MPI -from pyevtk.hl import gridToVTK from psydac.ddm.mpi import MockMPI from psydac.ddm.mpi import mpi as MPI diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index 0f46d947b..dec9aa532 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -3,7 +3,7 @@ from types import ModuleType import pytest -from mpi4py import MPI +from psydac.ddm.mpi import mpi as MPI from struphy import main from struphy.io.options import EnvironmentOptions diff --git a/src/struphy/models/tests/test_verif_EulerSPH.py b/src/struphy/models/tests/test_verif_EulerSPH.py index fa83945e6..c9660d58c 100644 --- a/src/struphy/models/tests/test_verif_EulerSPH.py +++ b/src/struphy/models/tests/test_verif_EulerSPH.py @@ -1,10 +1,10 @@ import os - -import numpy as np import pytest from matplotlib import pyplot as plt from matplotlib.ticker import FormatStrFormatter -from mpi4py import MPI + +from psydac.ddm.mpi import mpi as MPI +from struphy.utils.arrays import xp as np from struphy import main from struphy.fields_background import equils diff --git a/src/struphy/models/tests/test_verif_LinearMHD.py b/src/struphy/models/tests/test_verif_LinearMHD.py index 8ff0d28b8..9396d2fe4 100644 --- a/src/struphy/models/tests/test_verif_LinearMHD.py +++ b/src/struphy/models/tests/test_verif_LinearMHD.py @@ -1,8 +1,8 @@ import os - -import numpy as np import pytest -from mpi4py import MPI + +from psydac.ddm.mpi import mpi as MPI +from struphy.utils.arrays import xp as np from struphy import main from struphy.diagnostics.diagn_tools import power_spectrum_2d diff --git a/src/struphy/models/tests/test_verif_Maxwell.py b/src/struphy/models/tests/test_verif_Maxwell.py index 9068eccb8..94cf2352a 100644 --- a/src/struphy/models/tests/test_verif_Maxwell.py +++ b/src/struphy/models/tests/test_verif_Maxwell.py @@ -1,11 +1,12 @@ import os -import numpy as np import pytest from matplotlib import pyplot as plt -from mpi4py import MPI from scipy.special import jv, yn +from psydac.ddm.mpi import mpi as MPI +from struphy.utils.arrays import xp as np + from struphy import main from struphy.diagnostics.diagn_tools import power_spectrum_2d from struphy.fields_background import equils diff --git a/src/struphy/models/tests/test_verif_Poisson.py b/src/struphy/models/tests/test_verif_Poisson.py index 2abda9dec..29f2354b1 100644 --- a/src/struphy/models/tests/test_verif_Poisson.py +++ b/src/struphy/models/tests/test_verif_Poisson.py @@ -1,9 +1,8 @@ import os - -import numpy as np -import pytest from matplotlib import pyplot as plt -from mpi4py import MPI + +from psydac.ddm.mpi import mpi as MPI +from struphy.utils.arrays import xp as np from struphy import main from struphy.fields_background import equils diff --git a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py index d2420294e..f1a67b182 100644 --- a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py @@ -1,11 +1,11 @@ import os import h5py -import numpy as np -import pytest from matplotlib import pyplot as plt from matplotlib.ticker import FormatStrFormatter -from mpi4py import MPI + +from psydac.ddm.mpi import mpi as MPI +from struphy.utils.arrays import xp as np from struphy import main from struphy.fields_background import equils diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 101e0d497..c1c9874ce 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -4,8 +4,8 @@ from abc import ABCMeta, abstractmethod from typing import TYPE_CHECKING -import numpy as np -from mpi4py import MPI +from psydac.ddm.mpi import mpi as MPI +from struphy.utils.arrays import xp as np from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.fields_background.base import FluidEquilibrium diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 10d8eba6e..f5e2de8a1 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -5,7 +5,6 @@ from typing import Literal from struphy.utils.arrays import xp as np -from mpi4py import MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -72,7 +71,7 @@ def options(self) -> Options: @abstractmethod def options(self, new): assert isinstance(new, self.Options) - if MPI.COMM_WORLD.Get_rank() == 0: + if True: print(f"\nNew options for propagator '{self.__class__.__name__}':") for k, v in new.__dict__.items(): print(f" {k}: {v}") diff --git a/src/struphy/utils/mpi.py b/src/struphy/utils/mpi.py deleted file mode 100644 index bb4eb10d4..000000000 --- a/src/struphy/utils/mpi.py +++ /dev/null @@ -1,102 +0,0 @@ -from dataclasses import dataclass -from time import time -from typing import TYPE_CHECKING - - -# Might not be needed -class MPICommWrapper: - def __init__(self, use_mpi=True): - self.use_mpi = use_mpi - if use_mpi: - from mpi4py import MPI - - self.comm = MPI.COMM_WORLD - else: - self.comm = MockComm() - - def __getattr__(self, name): - return getattr(self.comm, name) - - -class MockComm: - def __getattr__(self, name): - # Return a function that does nothing and returns None - def dummy(*args, **kwargs): - return None - - return dummy - - # Override some functions - def Get_rank(self): - return 0 - - def Get_size(self): - return 1 - - def Barrier(self): - return - - -class MPIwrapper: - def __init__(self, use_mpi: bool = False): - self.use_mpi = use_mpi - if use_mpi: - from mpi4py import MPI - - self._MPI = MPI - print("MPI is enabled") - else: - self._MPI = MockMPI() - print("MPI is NOT enabled") - - @property - def MPI(self): - return self._MPI - - -class MockMPI: - def __getattr__(self, name): - # Return a function that does nothing and returns None - def dummy(*args, **kwargs): - return None - - return dummy - - # Override some functions - @property - def COMM_WORLD(self): - return MockComm() - - # def comm_Get_rank(self): - # return 0 - - # def comm_Get_size(self): - # return 1 - - -try: - from mpi4py import MPI - - _comm = MPI.COMM_WORLD - rank = _comm.Get_rank() - size = _comm.Get_size() - mpi_enabled = size > 1 -except ImportError: - # mpi4py not installed - mpi_enabled = False -except Exception: - # mpi4py installed but not running under mpirun - mpi_enabled = False - -# TODO: add environment variable for mpi use -mpi_wrapper = MPIwrapper(use_mpi=mpi_enabled) - -# TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime. -if TYPE_CHECKING: - from mpi4py import MPI - - mpi = MPI -else: - mpi = mpi_wrapper.MPI - -print(f"{mpi = }") From 9d92d1fa1ba935a5a32406106cf72917a66a5c82 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 14:38:01 +0200 Subject: [PATCH 146/292] add new mpi to new models files --- src/struphy/models/fluid.py | 4 ++++ src/struphy/models/hybrid.py | 5 +++++ src/struphy/models/kinetic.py | 5 +++++ src/struphy/models/toy.py | 5 ++++- 4 files changed, 18 insertions(+), 1 deletion(-) diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 7b5ffca54..2c8d04b6d 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -1,5 +1,6 @@ from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector +from psydac.ddm.mpi import mpi as MPI from struphy.feec.projectors import L2Projector from struphy.feec.variational_utilities import H1vecMassMatrix_density, InternalEnergyEvaluator @@ -11,6 +12,9 @@ from struphy.utils.arrays import xp as np +rank = MPI.COMM_WORLD.Get_rank() + + class LinearMHD(StruphyModel): r"""Linear ideal MHD with zero-flow equilibrium (:math:`\mathbf U_0 = 0`). diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index a85f83462..0e69f39c7 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -1,3 +1,5 @@ +from psydac.ddm.mpi import mpi as MPI + from struphy.models.base import StruphyModel from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable @@ -8,6 +10,9 @@ from struphy.utils.pyccel import Pyccelkernel +rank = MPI.COMM_WORLD.Get_rank() + + class LinearMHDVlasovCC(StruphyModel): r""" Hybrid linear MHD + energetic ions (6D Vlasov) with **current coupling scheme**. diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 8017480d1..a426ab7ab 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -1,3 +1,5 @@ +from psydac.ddm.mpi import mpi as MPI + from struphy.kinetic_background.base import KineticBackground from struphy.models.base import StruphyModel from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies @@ -9,6 +11,9 @@ from struphy.utils.pyccel import Pyccelkernel +rank = MPI.COMM_WORLD.Get_rank() + + class VlasovAmpereOneSpecies(StruphyModel): r"""Vlasov-Ampère equations for one species. diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 5addd0dd8..a232e09a6 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,5 +1,5 @@ -from psydac.ddm.mpi import mpi as MPI +from psydac.ddm.mpi import mpi as MPI from struphy.feec.projectors import L2Projector from struphy.feec.variational_utilities import InternalEnergyEvaluator @@ -10,6 +10,9 @@ from struphy.utils.arrays import xp as np +rank = MPI.COMM_WORLD.Get_rank() + + class Maxwell(StruphyModel): r"""Maxwell's equations in vacuum. From 0e764a78979e01e4301db997baca15db49646ed6 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 14:53:11 +0200 Subject: [PATCH 147/292] add missing Pyccelkernel --- src/struphy/propagators/propagators_markers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 11a31522b..bdbe67ae7 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -1724,9 +1724,9 @@ def allocate(self): # pusher kernel if self.options.thermodynamics == "isothermal": - kernel = pusher_kernels.push_v_sph_pressure + kernel = Pyccelkernel(pusher_kernels.push_v_sph_pressure) elif self.options.thermodynamics == "polytropic": - kernel = pusher_kernels.push_v_sph_pressure_ideal_gas + kernel = Pyccelkernel(pusher_kernels.push_v_sph_pressure_ideal_gas) gravity = np.array(self.options.gravity, dtype=float) @@ -1861,7 +1861,7 @@ def __init__( args_init, ) - kernel = pusher_kernels.push_v_viscosity + kernel = Pyccelkernel(pusher_kernels.push_v_viscosity) args_kernel = ( boxes, From 0a85779594fa1d3637ae00a9b155a10f90909edc Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 15:02:01 +0200 Subject: [PATCH 148/292] check for MockMPI in StruphyModel --- src/struphy/models/base.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index dbdfb58dc..fb4705464 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -8,6 +8,7 @@ import yaml from line_profiler import profile from psydac.ddm.mpi import mpi as MPI +from psydac.ddm.mpi import MockMPI from psydac.linalg.stencil import StencilVector import struphy @@ -526,8 +527,8 @@ def update_scalar(self, name, value=None): value_array = np.array([value], dtype=np.float64) # Perform MPI operations based on the compute flags - if "sum_world" in compute_operations and self.comm_world is not None: - self.comm_world.Allreduce( + if "sum_world" in compute_operations and not isinstance(MPI, MockMPI): + MPI.COMM_WORLD.Allreduce( MPI.IN_PLACE, value_array, op=MPI.SUM, From 9d08cd5d28281c7bca6b43db7a529af3620d42b1 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 15:08:15 +0200 Subject: [PATCH 149/292] remove deprecated --verification flag from tests --- .gitlab-ci.yml | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 236d015b5..3f04888d2 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -415,16 +415,13 @@ stages: - struphy test toy - struphy test models - struphy test verification - # - struphy test models --fast --verification --mpi 4 - # - struphy test models --fast --verification --mpi 4 --nclones 2 - # - struphy test DriftKineticElectrostaticAdiabatic --mpi 2 --nclones 2 model_tests_mpi: - struphy compile --status - struphy test models - struphy test models --mpi 2 - - struphy test models --verification --mpi 1 - - struphy test models --verification --mpi 4 - - struphy test models --verification --mpi 4 --nclones 2 + - struphy test verification --mpi 1 + - struphy test verification --mpi 4 + - struphy test verification --mpi 4 --nclones 2 - struphy test DriftKineticElectrostaticAdiabatic --mpi 2 --nclones 2 quickstart_tests: - struphy -p From 74c5e0679f3fc1abfcf528e55781bcca6e6a21c6 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 15:09:10 +0200 Subject: [PATCH 150/292] formatting --- src/struphy/initial/perturbations.py | 1 - src/struphy/io/options.py | 2 +- src/struphy/io/setup.py | 4 ++-- src/struphy/main.py | 5 ++--- src/struphy/models/base.py | 8 ++++---- src/struphy/models/fluid.py | 3 +-- src/struphy/models/hybrid.py | 1 - src/struphy/models/kinetic.py | 1 - src/struphy/models/species.py | 2 +- src/struphy/models/tests/test_verif_EulerSPH.py | 4 ++-- src/struphy/models/tests/test_verif_LinearMHD.py | 4 ++-- src/struphy/models/tests/test_verif_Maxwell.py | 5 ++--- src/struphy/models/tests/test_verif_Poisson.py | 4 ++-- .../models/tests/test_verif_VlasovAmpereOneSpecies.py | 3 +-- src/struphy/models/toy.py | 2 -- src/struphy/models/variables.py | 2 +- src/struphy/pic/base.py | 2 ++ src/struphy/pic/particles.py | 1 - src/struphy/post_processing/pproc_struphy.py | 2 +- src/struphy/propagators/base.py | 2 +- src/struphy/propagators/propagators_markers.py | 2 +- 21 files changed, 26 insertions(+), 34 deletions(-) diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 3becdd85e..1be8d4191 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -11,7 +11,6 @@ from struphy.utils.arrays import xp as np - @dataclass class Noise(Perturbation): """White noise for FEEC coefficients. diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index d3b3d4a18..3a9cd8e61 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -3,9 +3,9 @@ from typing import Literal, get_args from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np from struphy.physics.physics import ConstantsOfNature +from struphy.utils.arrays import xp as np ## Literal options diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index a6c0d730a..375a76b4a 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -5,11 +5,11 @@ import sys from types import ModuleType +from psydac.ddm.mpi import mpi as MPI + from struphy.geometry.base import Domain from struphy.io.options import DerhamOptions from struphy.topology.grids import TensorProductGrid -from psydac.ddm.mpi import mpi as MPI - from struphy.utils.arrays import xp as np diff --git a/src/struphy/main.py b/src/struphy/main.py index 731a9e8b9..b8d990cb6 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -10,12 +10,10 @@ import h5py from line_profiler import profile - from psydac.ddm.mpi import MockMPI from psydac.ddm.mpi import mpi as MPI from pyevtk.hl import gridToVTK -from struphy.utils.arrays import xp as np from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB from struphy.fields_background.equils import HomogenSlab from struphy.geometry import domains @@ -40,6 +38,7 @@ from struphy.profiling.profiling import ProfileManager from struphy.topology import grids from struphy.topology.grids import TensorProductGrid +from struphy.utils.arrays import xp as np from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml @@ -104,7 +103,7 @@ def run( save_step = env.save_step sort_step = env.sort_step num_clones = env.num_clones - use_mpi = not comm is None, + use_mpi = (not comm is None,) meta = {} meta["platform"] = sysconfig.get_platform() diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index fb4705464..1e5ebff3d 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -7,8 +7,8 @@ import yaml from line_profiler import profile -from psydac.ddm.mpi import mpi as MPI from psydac.ddm.mpi import MockMPI +from psydac.ddm.mpi import mpi as MPI from psydac.linalg.stencil import StencilVector import struphy @@ -962,9 +962,9 @@ def print_scalar_quantities(self): # if obj.coords == "vpara_mu": # obj.save_magnetic_moment() - # obj.draw_markers(sort=True, verbose=self.verbose) - # if self.comm_world is not None: - # obj.mpi_sort_markers(do_test=True) + # obj.draw_markers(sort=True, verbose=self.verbose) + # if self.comm_world is not None: + # obj.mpi_sort_markers(do_test=True) # obj.initialize_weights( # reject_weights=obj.weights_params["reject_weights"], diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 2c8d04b6d..2ddb72972 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -1,6 +1,6 @@ +from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector -from psydac.ddm.mpi import mpi as MPI from struphy.feec.projectors import L2Projector from struphy.feec.variational_utilities import H1vecMassMatrix_density, InternalEnergyEvaluator @@ -11,7 +11,6 @@ from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers from struphy.utils.arrays import xp as np - rank = MPI.COMM_WORLD.Get_rank() diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index 0e69f39c7..70968d643 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -9,7 +9,6 @@ from struphy.utils.arrays import xp as np from struphy.utils.pyccel import Pyccelkernel - rank = MPI.COMM_WORLD.Get_rank() diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index a426ab7ab..792bc14cb 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -10,7 +10,6 @@ from struphy.utils.arrays import xp as np from struphy.utils.pyccel import Pyccelkernel - rank = MPI.COMM_WORLD.Get_rank() diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 92710cbe7..3c33d1b71 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -2,7 +2,6 @@ from abc import ABCMeta, abstractmethod from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np from struphy.io.options import Units from struphy.models.variables import Variable @@ -14,6 +13,7 @@ LoadingParameters, WeightsParameters, ) +from struphy.utils.arrays import xp as np class Species(metaclass=ABCMeta): diff --git a/src/struphy/models/tests/test_verif_EulerSPH.py b/src/struphy/models/tests/test_verif_EulerSPH.py index c9660d58c..d16e76569 100644 --- a/src/struphy/models/tests/test_verif_EulerSPH.py +++ b/src/struphy/models/tests/test_verif_EulerSPH.py @@ -1,10 +1,9 @@ import os + import pytest from matplotlib import pyplot as plt from matplotlib.ticker import FormatStrFormatter - from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np from struphy import main from struphy.fields_background import equils @@ -20,6 +19,7 @@ WeightsParameters, ) from struphy.topology import grids +from struphy.utils.arrays import xp as np test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") diff --git a/src/struphy/models/tests/test_verif_LinearMHD.py b/src/struphy/models/tests/test_verif_LinearMHD.py index 9396d2fe4..faa59fc1a 100644 --- a/src/struphy/models/tests/test_verif_LinearMHD.py +++ b/src/struphy/models/tests/test_verif_LinearMHD.py @@ -1,8 +1,7 @@ import os -import pytest +import pytest from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np from struphy import main from struphy.diagnostics.diagn_tools import power_spectrum_2d @@ -12,6 +11,7 @@ from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time from struphy.kinetic_background import maxwellians from struphy.topology import grids +from struphy.utils.arrays import xp as np test_folder = os.path.join(os.getcwd(), "verification_tests") diff --git a/src/struphy/models/tests/test_verif_Maxwell.py b/src/struphy/models/tests/test_verif_Maxwell.py index 94cf2352a..0cb4e778e 100644 --- a/src/struphy/models/tests/test_verif_Maxwell.py +++ b/src/struphy/models/tests/test_verif_Maxwell.py @@ -2,10 +2,8 @@ import pytest from matplotlib import pyplot as plt -from scipy.special import jv, yn - from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np +from scipy.special import jv, yn from struphy import main from struphy.diagnostics.diagn_tools import power_spectrum_2d @@ -16,6 +14,7 @@ from struphy.kinetic_background import maxwellians from struphy.models.toy import Maxwell from struphy.topology import grids +from struphy.utils.arrays import xp as np test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") diff --git a/src/struphy/models/tests/test_verif_Poisson.py b/src/struphy/models/tests/test_verif_Poisson.py index 29f2354b1..0264d4b6d 100644 --- a/src/struphy/models/tests/test_verif_Poisson.py +++ b/src/struphy/models/tests/test_verif_Poisson.py @@ -1,8 +1,7 @@ import os -from matplotlib import pyplot as plt +from matplotlib import pyplot as plt from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np from struphy import main from struphy.fields_background import equils @@ -19,6 +18,7 @@ WeightsParameters, ) from struphy.topology import grids +from struphy.utils.arrays import xp as np test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") diff --git a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py index f1a67b182..75a6f1ddc 100644 --- a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py @@ -3,9 +3,7 @@ import h5py from matplotlib import pyplot as plt from matplotlib.ticker import FormatStrFormatter - from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np from struphy import main from struphy.fields_background import equils @@ -21,6 +19,7 @@ WeightsParameters, ) from struphy.topology import grids +from struphy.utils.arrays import xp as np test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index a232e09a6..fdeac6df8 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,4 +1,3 @@ - from psydac.ddm.mpi import mpi as MPI from struphy.feec.projectors import L2Projector @@ -9,7 +8,6 @@ from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers from struphy.utils.arrays import xp as np - rank = MPI.COMM_WORLD.Get_rank() diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index c1c9874ce..b39fcc38b 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -5,7 +5,6 @@ from typing import TYPE_CHECKING from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.fields_background.base import FluidEquilibrium @@ -22,6 +21,7 @@ from struphy.pic import particles from struphy.pic.base import Particles from struphy.pic.particles import ParticlesSPH +from struphy.utils.arrays import xp as np from struphy.utils.clone_config import CloneConfig if TYPE_CHECKING: diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index f1031f113..e9ab1ba7b 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -9,9 +9,11 @@ try: from mpi4py.MPI import Intracomm except ModuleNotFoundError: + class Intracomm: x = None + from line_profiler import profile from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 62553c3cd..39de5e4e6 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -2,7 +2,6 @@ from struphy.fields_background import equils from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB -from struphy.fields_background.base import FluidEquilibriumWithB from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB from struphy.geometry.base import Domain from struphy.geometry.utilities import TransformedPformComponent diff --git a/src/struphy/post_processing/pproc_struphy.py b/src/struphy/post_processing/pproc_struphy.py index 132042f51..e9fe1b9b4 100644 --- a/src/struphy/post_processing/pproc_struphy.py +++ b/src/struphy/post_processing/pproc_struphy.py @@ -5,10 +5,10 @@ import h5py import yaml -from struphy.utils.arrays import xp as np import struphy.post_processing.orbits.orbits_tools as orbits_pproc import struphy.post_processing.post_processing_tools as pproc from struphy.io.setup import import_parameters_py +from struphy.utils.arrays import xp as np def main( diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index f5e2de8a1..5f92f7373 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -4,7 +4,6 @@ from dataclasses import dataclass from typing import Literal -from struphy.utils.arrays import xp as np from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -15,6 +14,7 @@ from struphy.geometry.base import Domain from struphy.io.options import check_option from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable +from struphy.utils.arrays import xp as np class Propagator(metaclass=ABCMeta): diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index bdbe67ae7..02cccb306 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -5,8 +5,8 @@ from typing import Callable, Literal, get_args from line_profiler import profile -from psydac.ddm.mpi import mpi as MPI from numpy import array, polynomial, random +from psydac.ddm.mpi import mpi as MPI from psydac.linalg.basic import LinearOperator from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector From 77fd7e415b0cf791920bb6bf3f7dcb18ccffe026 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 15:19:18 +0200 Subject: [PATCH 151/292] test an already ported model for the single test --- .gitlab-ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 3f04888d2..3385d72d3 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -422,7 +422,7 @@ stages: - struphy test verification --mpi 1 - struphy test verification --mpi 4 - struphy test verification --mpi 4 --nclones 2 - - struphy test DriftKineticElectrostaticAdiabatic --mpi 2 --nclones 2 + - struphy test VlasovAmpereOneSpecies --mpi 2 --nclones 2 quickstart_tests: - struphy -p - struphy -h From 111775d1d386f46dcf79cc100f0d997e6f7fa467 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 15:20:55 +0200 Subject: [PATCH 152/292] use Barrier() in main.py --- src/struphy/main.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/struphy/main.py b/src/struphy/main.py index b8d990cb6..e40059537 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -185,7 +185,7 @@ def run( clone_config.print_particle_config() model.clone_config = clone_config - comm.Barrier() + Barrier() ## configure model instance @@ -226,7 +226,7 @@ def run( if rank < 32: if rank == 0: print("") - comm.Barrier() + Barrier() print(f"Rank {rank}: executing main.run() for model {model_name} ...") if size > 32 and rank == 32: From 533223c29a7d250265f766f5cbc83d84ab1d08a8 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 16:55:32 +0200 Subject: [PATCH 153/292] fix verification test --- src/struphy/console/test.py | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 6ac2ed0e3..4ab2cb5dd 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -95,17 +95,26 @@ def struphy_test( subp_run(cmd) elif "verification" in group: - cmd = [ - "mpirun", - "--oversubscribe", - "-n", - str(mpi), - "pytest", - "-k", - "_verif_", - "-s", - "--with-mpi", - ] + if mpi > 1: + cmd = [ + "mpirun", + "--oversubscribe", + "-n", + str(mpi), + "pytest", + "-k", + "_verif_", + "-s", + "--with-mpi", + ] + else: + cmd = [ + "pytest", + "-k", + "_verif_", + "-s", + ] + if vrbose: cmd += ["--vrbose"] if nclones > 1: From 06e86f4985e1d51e236b94d293a4917464ae40b2 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 17:05:10 +0200 Subject: [PATCH 154/292] fix gyrokinetic_poisson --- src/struphy/propagators/tests/test_gyrokinetic_poisson.py | 1 + src/struphy/propagators/tests/test_poisson.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index fb2f444d1..a2dc81519 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -11,6 +11,7 @@ from struphy.models.variables import FEECVariable from struphy.propagators.base import Propagator from struphy.utils.arrays import xp as np +from struphy.propagators.propagators_fields import ImplicitDiffusion comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 65eb3b720..0b1e14de6 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -23,6 +23,7 @@ from struphy.propagators.base import Propagator from struphy.propagators.propagators_fields import ImplicitDiffusion, Poisson from struphy.utils.arrays import xp as np +from struphy.utils.pyccel import Pyccelkernel comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -339,7 +340,7 @@ def test_poisson_accum_1d(mapping, do_plot=False): particles.initialize_weights() # particle to grid coupling - kernel = charge_density_0form + kernel = Pyccelkernel(charge_density_0form) accum = AccumulatorVector(particles, "H1", kernel, mass_ops, domain.args_domain) # accum() # if do_plot: From 613cd583f1d226f675dd4811c597a13e6dd934d0 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 21 Oct 2025 17:06:00 +0200 Subject: [PATCH 155/292] formatting --- src/struphy/console/test.py | 2 +- src/struphy/propagators/tests/test_gyrokinetic_poisson.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 4ab2cb5dd..9e82ee146 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -114,7 +114,7 @@ def struphy_test( "_verif_", "-s", ] - + if vrbose: cmd += ["--vrbose"] if nclones > 1: diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index a2dc81519..67e84decd 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -10,8 +10,8 @@ from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable from struphy.propagators.base import Propagator -from struphy.utils.arrays import xp as np from struphy.propagators.propagators_fields import ImplicitDiffusion +from struphy.utils.arrays import xp as np comm = MPI.COMM_WORLD rank = comm.Get_rank() From 26fc52e04d3e69ded1f4eee57fe5e573cbfbda4b Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 22 Oct 2025 06:03:54 +0000 Subject: [PATCH 156/292] Resolve "Rename struphy array imports to xp" --- src/struphy/bsplines/bsplines.py | 58 +- .../bsplines/tests/test_bsplines_kernels.py | 50 +- .../bsplines/tests/test_eval_spline_mpi.py | 158 ++--- src/struphy/console/profile.py | 10 +- src/struphy/diagnostics/console_diagn.py | 6 +- src/struphy/diagnostics/continuous_spectra.py | 32 +- src/struphy/diagnostics/diagn_tools.py | 132 ++-- .../diagnostics/paraview/mesh_creator.py | 66 +- src/struphy/dispersion_relations/analytic.py | 86 +-- src/struphy/dispersion_relations/base.py | 10 +- src/struphy/dispersion_relations/utilities.py | 4 +- src/struphy/eigenvalue_solvers/derivatives.py | 8 +- .../legacy/MHD_eigenvalues_cylinder_1D.py | 372 +++++------ .../control_variates/control_variate.py | 42 +- .../fB_massless_control_variate.py | 4 +- .../fnB_massless_control_variate.py | 6 +- .../massless_control_variate.py | 6 +- .../legacy/emw_operators.py | 8 +- .../legacy/inner_products_1d.py | 10 +- .../legacy/inner_products_2d.py | 38 +- .../legacy/inner_products_3d.py | 38 +- .../eigenvalue_solvers/legacy/l2_error_1d.py | 10 +- .../eigenvalue_solvers/legacy/l2_error_2d.py | 50 +- .../eigenvalue_solvers/legacy/l2_error_3d.py | 50 +- .../legacy/mass_matrices_3d_pre.py | 74 +- .../legacy/massless_operators/fB_arrays.py | 206 +++--- .../fB_massless_linear_operators.py | 102 +-- .../legacy/massless_operators/fB_vv_kernel.py | 2 +- .../legacy/mhd_operators_MF.py | 226 +++---- .../pro_local/mhd_operators_3d_local.py | 500 +++++++------- .../pro_local/projectors_local.py | 632 +++++++++--------- .../shape_function_projectors_L2.py | 122 ++-- .../shape_function_projectors_local.py | 220 +++--- .../eigenvalue_solvers/mass_matrices_1d.py | 26 +- .../eigenvalue_solvers/mass_matrices_2d.py | 92 +-- .../eigenvalue_solvers/mass_matrices_3d.py | 92 +-- .../mhd_axisymmetric_main.py | 8 +- .../mhd_axisymmetric_pproc.py | 10 +- .../eigenvalue_solvers/mhd_operators.py | 40 +- .../eigenvalue_solvers/mhd_operators_core.py | 374 +++++------ .../eigenvalue_solvers/projectors_global.py | 270 ++++---- .../eigenvalue_solvers/spline_space.py | 342 +++++----- src/struphy/examples/_draw_parallel.py | 10 +- .../examples/restelli2018/callables.py | 66 +- src/struphy/feec/basis_projection_ops.py | 42 +- src/struphy/feec/linear_operators.py | 32 +- src/struphy/feec/mass.py | 62 +- src/struphy/feec/preconditioner.py | 42 +- src/struphy/feec/projectors.py | 114 ++-- src/struphy/feec/psydac_derham.py | 310 ++++----- src/struphy/feec/tests/test_basis_ops.py | 60 +- src/struphy/feec/tests/test_derham.py | 18 +- src/struphy/feec/tests/test_eval_field.py | 154 ++--- src/struphy/feec/tests/test_field_init.py | 192 +++--- src/struphy/feec/tests/test_l2_projectors.py | 46 +- .../feec/tests/test_local_projectors.py | 312 ++++----- .../feec/tests/test_lowdim_nel_is_1.py | 54 +- src/struphy/feec/tests/test_mass_matrices.py | 92 +-- .../feec/tests/test_toarray_struphy.py | 28 +- .../feec/tests/test_tosparse_struphy.py | 16 +- src/struphy/feec/tests/xx_test_preconds.py | 14 +- src/struphy/feec/utilities.py | 28 +- .../feec/utilities_local_projectors.py | 104 +-- src/struphy/feec/variational_utilities.py | 118 ++-- src/struphy/fields_background/base.py | 84 +-- .../fields_background/coil_fields/base.py | 6 +- .../coil_fields/coil_fields.py | 6 +- src/struphy/fields_background/equils.py | 208 +++--- .../tests/test_desc_equil.py | 68 +- .../tests/test_generic_equils.py | 28 +- .../tests/test_mhd_equils.py | 82 +-- .../tests/test_numerical_mhd_equil.py | 72 +- src/struphy/geometry/base.py | 226 +++---- src/struphy/geometry/domains.py | 34 +- src/struphy/geometry/evaluation_kernels.py | 24 +- src/struphy/geometry/tests/test_domain.py | 118 ++-- src/struphy/geometry/utilities.py | 44 +- src/struphy/geometry/utilities_kernels.py | 8 +- src/struphy/initial/eigenfunctions.py | 38 +- src/struphy/initial/perturbations.py | 190 +++--- .../initial/tests/test_init_perturbations.py | 18 +- src/struphy/initial/utilities.py | 4 +- src/struphy/io/options.py | 6 +- src/struphy/io/output_handling.py | 6 +- src/struphy/io/setup.py | 2 +- src/struphy/kinetic_background/base.py | 68 +- src/struphy/kinetic_background/maxwellians.py | 54 +- .../kinetic_background/moment_functions.py | 6 +- .../kinetic_background/tests/test_base.py | 38 +- .../tests/test_maxwellians.py | 380 +++++------ src/struphy/linear_algebra/linalg_kron.py | 8 +- src/struphy/linear_algebra/saddle_point.py | 104 +-- .../tests/test_saddlepoint_massmatrices.py | 58 +- .../tests/test_stencil_dot_kernels.py | 36 +- .../tests/test_stencil_transpose_kernels.py | 30 +- src/struphy/main.py | 44 +- src/struphy/models/base.py | 64 +- src/struphy/models/fluid.py | 18 +- src/struphy/models/hybrid.py | 30 +- src/struphy/models/kinetic.py | 24 +- src/struphy/models/species.py | 4 +- .../models/tests/test_verif_EulerSPH.py | 6 +- .../models/tests/test_verif_LinearMHD.py | 18 +- .../models/tests/test_verif_Maxwell.py | 36 +- .../models/tests/test_verif_Poisson.py | 12 +- .../test_verif_VlasovAmpereOneSpecies.py | 18 +- src/struphy/models/tests/verification.py | 76 +-- src/struphy/models/toy.py | 16 +- src/struphy/models/variables.py | 10 +- src/struphy/ode/solvers.py | 2 +- src/struphy/ode/tests/test_ode_feec.py | 16 +- src/struphy/ode/utils.py | 8 +- src/struphy/pic/accumulation/filter.py | 22 +- .../pic/accumulation/particles_to_grid.py | 10 +- src/struphy/pic/base.py | 420 ++++++------ src/struphy/pic/particles.py | 2 +- src/struphy/pic/pushing/pusher.py | 16 +- src/struphy/pic/pushing/pusher_kernels.py | 6 +- src/struphy/pic/sampling_kernels.py | 6 +- src/struphy/pic/sobol_seq.py | 54 +- src/struphy/pic/tests/test_accum_vec_H1.py | 22 +- src/struphy/pic/tests/test_accumulation.py | 36 +- src/struphy/pic/tests/test_binning.py | 134 ++-- src/struphy/pic/tests/test_draw_parallel.py | 12 +- src/struphy/pic/tests/test_mat_vec_filler.py | 70 +- .../test_pic_legacy_files/accumulation.py | 36 +- .../pic/tests/test_pic_legacy_files/pusher.py | 2 +- .../spline_evaluation_2d.py | 2 +- .../spline_evaluation_3d.py | 2 +- src/struphy/pic/tests/test_pushers.py | 76 +-- src/struphy/pic/tests/test_sorting.py | 20 +- src/struphy/pic/tests/test_sph.py | 160 ++--- src/struphy/pic/tests/test_tesselation.py | 46 +- src/struphy/pic/utilities.py | 32 +- src/struphy/polar/basic.py | 14 +- src/struphy/polar/extraction_operators.py | 238 +++---- src/struphy/polar/linear_operators.py | 10 +- .../polar/tests/test_legacy_polar_splines.py | 20 +- src/struphy/polar/tests/test_polar.py | 44 +- .../likwid/plot_likwidproject.py | 8 +- .../likwid/plot_time_traces.py | 14 +- .../likwid/roofline_plotter.py | 12 +- .../post_processing/orbits/orbits_tools.py | 48 +- .../post_processing/post_processing_tools.py | 68 +- src/struphy/post_processing/pproc_struphy.py | 4 +- .../post_processing/profile_struphy.py | 8 +- src/struphy/profiling/profiling.py | 22 +- src/struphy/propagators/base.py | 14 +- .../propagators/propagators_coupling.py | 70 +- src/struphy/propagators/propagators_fields.py | 224 +++---- .../propagators/propagators_markers.py | 34 +- .../tests/test_gyrokinetic_poisson.py | 76 +-- src/struphy/propagators/tests/test_poisson.py | 82 +-- src/struphy/utils/clone_config.py | 8 +- src/struphy/utils/utils.py | 4 +- 155 files changed, 5585 insertions(+), 5585 deletions(-) diff --git a/src/struphy/bsplines/bsplines.py b/src/struphy/bsplines/bsplines.py index a659e4c61..ab56afd78 100644 --- a/src/struphy/bsplines/bsplines.py +++ b/src/struphy/bsplines/bsplines.py @@ -16,7 +16,7 @@ """ -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp __all__ = [ "find_span", @@ -105,7 +105,7 @@ def scaling_vector(knots, degree, span): Scaling vector with elements (p + 1)/(t[i + p + 1] - t[i]) """ - x = np.zeros(degree + 1, dtype=float) + x = xp.zeros(degree + 1, dtype=float) for il in range(degree + 1): i = span - il @@ -148,9 +148,9 @@ def basis_funs(knots, degree, x, span, normalize=False): by using 'left' and 'right' temporary arrays that are one element shorter. """ - left = np.empty(degree, dtype=float) - right = np.empty(degree, dtype=float) - values = np.empty(degree + 1, dtype=float) + left = xp.empty(degree, dtype=float) + right = xp.empty(degree, dtype=float) + values = xp.empty(degree + 1, dtype=float) values[0] = 1.0 @@ -205,7 +205,7 @@ def basis_funs_1st_der(knots, degree, x, span): # Compute derivatives at x using formula based on difference of splines of degree deg - 1 # ------- # j = 0 - ders = np.empty(degree + 1, dtype=float) + ders = xp.empty(degree + 1, dtype=float) saved = degree * values[0] / (knots[span + 1] - knots[span + 1 - degree]) ders[0] = -saved @@ -261,11 +261,11 @@ def basis_funs_all_ders(knots, degree, x, span, n): - innermost loops are replaced with vector operations on slices. """ - left = np.empty(degree) - right = np.empty(degree) - ndu = np.empty((degree + 1, degree + 1)) - a = np.empty((2, degree + 1)) - ders = np.zeros((n + 1, degree + 1)) # output array + left = xp.empty(degree) + right = xp.empty(degree) + ndu = xp.empty((degree + 1, degree + 1)) + a = xp.empty((2, degree + 1)) + ders = xp.zeros((n + 1, degree + 1)) # output array # Number of derivatives that need to be effectively computed # Derivatives higher than degree are = 0. @@ -304,7 +304,7 @@ def basis_funs_all_ders(knots, degree, x, span, n): j1 = 1 if (rk > -1) else -rk j2 = k - 1 if (r - 1 <= pk) else degree - r a[s2, j1 : j2 + 1] = (a[s1, j1 : j2 + 1] - a[s1, j1 - 1 : j2]) * ndu[pk + 1, rk + j1 : rk + j2 + 1] - d += np.dot(a[s2, j1 : j2 + 1], ndu[rk + j1 : rk + j2 + 1, pk]) + d += xp.dot(a[s2, j1 : j2 + 1], ndu[rk + j1 : rk + j2 + 1, pk]) if r <= pk: a[s2, k] = -a[s1, k - 1] * ndu[pk + 1, r] d += a[s2, k] * ndu[r, pk] @@ -362,7 +362,7 @@ def collocation_matrix(knots, degree, xgrid, periodic, normalize=False): nx = len(xgrid) # Collocation matrix as 2D Numpy array (dense storage) - mat = np.zeros((nx, nb), dtype=float) + mat = xp.zeros((nx, nb), dtype=float) # Indexing of basis functions (periodic or not) for a given span if periodic: @@ -418,12 +418,12 @@ def histopolation_matrix(knots, degree, xgrid, periodic): # Number of integrals if periodic: el_b = breakpoints(knots, degree) - xgrid = np.array([el_b[0]] + list(xgrid) + [el_b[-1]]) + xgrid = xp.array([el_b[0]] + list(xgrid) + [el_b[-1]]) ni = len(xgrid) - 1 # Histopolation matrix of M-splines as 2D Numpy array (dense storage) - his = np.zeros((ni, nbD), dtype=float) + his = xp.zeros((ni, nbD), dtype=float) # Collocation matrix of B-splines col = collocation_matrix(knots, degree, xgrid, False, normalize=False) @@ -434,7 +434,7 @@ def histopolation_matrix(knots, degree, xgrid, periodic): for k in range(j + 1): his[i, j % nbD] += col[i, k] - col[i + 1, k] - if np.abs(his[i, j % nbD]) < 1e-14: + if xp.abs(his[i, j % nbD]) < 1e-14: his[i, j % nbD] = 0.0 # add first to last integration interval in case of periodic splines @@ -470,7 +470,7 @@ def breakpoints(knots, degree): else: endsl = -degree - return np.unique(knots[slice(degree, endsl)]) + return xp.unique(knots[slice(degree, endsl)]) # ============================================================================== @@ -501,13 +501,13 @@ def greville(knots, degree, periodic): n = len(T) - 2 * p - 1 if periodic else len(T) - p - 1 # Compute greville abscissas as average of p consecutive knot values - xg = np.around([sum(T[i : i + p]) / p for i in range(s, s + n)], decimals=15) + xg = xp.around([sum(T[i : i + p]) / p for i in range(s, s + n)], decimals=15) # If needed apply periodic boundary conditions if periodic: a = T[p] b = T[-p] - xg = np.around((xg - a) % (b - a) + a, decimals=15) + xg = xp.around((xg - a) % (b - a) + a, decimals=15) return xg @@ -537,7 +537,7 @@ def elements_spans(knots, degree): >>> from psydac.core.bsplines import make_knots, elements_spans >>> p = 3 ; n = 8 - >>> grid = np.arange( n-p+1 ) + >>> grid = xp.arange( n-p+1 ) >>> knots = make_knots( breaks=grid, degree=p, periodic=False ) >>> spans = elements_spans( knots=knots, degree=p ) >>> spans @@ -549,13 +549,13 @@ def elements_spans(knots, degree): 2) This function could be written in two lines: breaks = breakpoints( knots, degree ) - spans = np.searchsorted( knots, breaks[:-1], side='right' ) - 1 + spans = xp.searchsorted( knots, breaks[:-1], side='right' ) - 1 """ breaks = breakpoints(knots, degree) nk = len(knots) ne = len(breaks) - 1 - spans = np.zeros(ne, dtype=int) + spans = xp.zeros(ne, dtype=int) ie = 0 for ik in range(degree, nk - degree): @@ -600,13 +600,13 @@ def make_knots(breaks, degree, periodic): # Consistency checks assert len(breaks) > 1 - assert all(np.diff(breaks) > 0) + assert all(xp.diff(breaks) > 0) assert degree > 0 if periodic: assert len(breaks) > degree p = degree - T = np.zeros(len(breaks) + 2 * p, dtype=float) + T = xp.zeros(len(breaks) + 2 * p, dtype=float) T[p:-p] = breaks if periodic: @@ -671,13 +671,13 @@ def quadrature_grid(breaks, quad_rule_x, quad_rule_w): assert min(quad_rule_x) >= -1 assert max(quad_rule_x) <= +1 - quad_rule_x = np.asarray(quad_rule_x) - quad_rule_w = np.asarray(quad_rule_w) + quad_rule_x = xp.asarray(quad_rule_x) + quad_rule_w = xp.asarray(quad_rule_w) ne = len(breaks) - 1 nq = len(quad_rule_x) - quad_x = np.zeros((ne, nq), dtype=float) - quad_w = np.zeros((ne, nq), dtype=float) + quad_x = xp.zeros((ne, nq), dtype=float) + quad_w = xp.zeros((ne, nq), dtype=float) # Compute location and weight of quadrature points from basic rule for ie, (a, b) in enumerate(zip(breaks[:-1], breaks[1:])): @@ -724,7 +724,7 @@ def basis_ders_on_quad_grid(knots, degree, quad_grid, nders, normalize=False): # TODO: check if it is safe to compute span only once for each element ne, nq = quad_grid.shape - basis = np.zeros((ne, degree + 1, nders + 1, nq), dtype=float) + basis = xp.zeros((ne, degree + 1, nders + 1, nq), dtype=float) # Loop over elements for ie in range(ne): diff --git a/src/struphy/bsplines/tests/test_bsplines_kernels.py b/src/struphy/bsplines/tests/test_bsplines_kernels.py index 1a16712c1..3800c8653 100644 --- a/src/struphy/bsplines/tests/test_bsplines_kernels.py +++ b/src/struphy/bsplines/tests/test_bsplines_kernels.py @@ -3,7 +3,7 @@ import pytest from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @@ -34,9 +34,9 @@ def test_bsplines_span_and_basis(Nel, p, spl_kind): # Random points in domain of process n_pts = 100 dom = derham.domain_array[rank] - eta1s = np.random.rand(n_pts) * (dom[1] - dom[0]) + dom[0] - eta2s = np.random.rand(n_pts) * (dom[4] - dom[3]) + dom[3] - eta3s = np.random.rand(n_pts) * (dom[7] - dom[6]) + dom[6] + eta1s = xp.random.rand(n_pts) * (dom[1] - dom[0]) + dom[0] + eta2s = xp.random.rand(n_pts) * (dom[4] - dom[3]) + dom[3] + eta3s = xp.random.rand(n_pts) * (dom[7] - dom[6]) + dom[6] # struphy find_span t0 = time.time() @@ -60,18 +60,18 @@ def test_bsplines_span_and_basis(Nel, p, spl_kind): if rank == 0: print(f"psydac find_span_p : {t1 - t0}") - assert np.allclose(span1s, span1s_psy) - assert np.allclose(span2s, span2s_psy) - assert np.allclose(span3s, span3s_psy) + assert xp.allclose(span1s, span1s_psy) + assert xp.allclose(span2s, span2s_psy) + assert xp.allclose(span3s, span3s_psy) # allocate tmps - bn1 = np.empty(derham.p[0] + 1, dtype=float) - bn2 = np.empty(derham.p[1] + 1, dtype=float) - bn3 = np.empty(derham.p[2] + 1, dtype=float) + bn1 = xp.empty(derham.p[0] + 1, dtype=float) + bn2 = xp.empty(derham.p[1] + 1, dtype=float) + bn3 = xp.empty(derham.p[2] + 1, dtype=float) - bd1 = np.empty(derham.p[0], dtype=float) - bd2 = np.empty(derham.p[1], dtype=float) - bd3 = np.empty(derham.p[2], dtype=float) + bd1 = xp.empty(derham.p[0], dtype=float) + bd2 = xp.empty(derham.p[1], dtype=float) + bd3 = xp.empty(derham.p[2], dtype=float) # struphy b_splines_slim val1s, val2s, val3s = [], [], [] @@ -103,13 +103,13 @@ def test_bsplines_span_and_basis(Nel, p, spl_kind): # compare for val1, val1_psy in zip(val1s, val1s_psy): - assert np.allclose(val1, val1_psy) + assert xp.allclose(val1, val1_psy) for val2, val2_psy in zip(val2s, val2s_psy): - assert np.allclose(val2, val2_psy) + assert xp.allclose(val2, val2_psy) for val3, val3_psy in zip(val3s, val3s_psy): - assert np.allclose(val3, val3_psy) + assert xp.allclose(val3, val3_psy) # struphy b_d_splines_slim val1s_n, val2s_n, val3s_n = [], [], [] @@ -131,13 +131,13 @@ def test_bsplines_span_and_basis(Nel, p, spl_kind): # compare for val1, val1_psy in zip(val1s_n, val1s_psy): - assert np.allclose(val1, val1_psy) + assert xp.allclose(val1, val1_psy) for val2, val2_psy in zip(val2s_n, val2s_psy): - assert np.allclose(val2, val2_psy) + assert xp.allclose(val2, val2_psy) for val3, val3_psy in zip(val3s_n, val3s_psy): - assert np.allclose(val3, val3_psy) + assert xp.allclose(val3, val3_psy) # struphy d_splines_slim span1s, span2s, span3s = [], [], [] @@ -175,22 +175,22 @@ def test_bsplines_span_and_basis(Nel, p, spl_kind): # compare for val1, val1_psy in zip(val1s, val1s_psy): - assert np.allclose(val1, val1_psy) + assert xp.allclose(val1, val1_psy) for val2, val2_psy in zip(val2s, val2s_psy): - assert np.allclose(val2, val2_psy) + assert xp.allclose(val2, val2_psy) for val3, val3_psy in zip(val3s, val3s_psy): - assert np.allclose(val3, val3_psy) + assert xp.allclose(val3, val3_psy) for val1, val1_psy in zip(val1s_d, val1s_psy): - assert np.allclose(val1, val1_psy) + assert xp.allclose(val1, val1_psy) for val2, val2_psy in zip(val2s_d, val2s_psy): - assert np.allclose(val2, val2_psy) + assert xp.allclose(val2, val2_psy) for val3, val3_psy in zip(val3s_d, val3s_psy): - assert np.allclose(val3, val3_psy) + assert xp.allclose(val3, val3_psy) if __name__ == "__main__": diff --git a/src/struphy/bsplines/tests/test_eval_spline_mpi.py b/src/struphy/bsplines/tests/test_eval_spline_mpi.py index e40af4124..0202963ce 100644 --- a/src/struphy/bsplines/tests/test_eval_spline_mpi.py +++ b/src/struphy/bsplines/tests/test_eval_spline_mpi.py @@ -4,7 +4,7 @@ import pytest from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @@ -38,9 +38,9 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): # Random points in domain of process dom = derham.domain_array[rank] - eta1s = np.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] - eta2s = np.random.rand(n_markers) * (dom[4] - dom[3]) + dom[3] - eta3s = np.random.rand(n_markers) * (dom[7] - dom[6]) + dom[6] + eta1s = xp.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] + eta2s = xp.random.rand(n_markers) * (dom[4] - dom[3]) + dom[3] + eta3s = xp.random.rand(n_markers) * (dom[7] - dom[6]) + dom[6] for eta1, eta2, eta3 in zip(eta1s, eta2s, eta3s): comm.Barrier() @@ -56,13 +56,13 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): span3 = bsp.find_span(tn3, derham.p[2], eta3) # non-zero spline values at eta - bn1 = np.empty(derham.p[0] + 1, dtype=float) - bn2 = np.empty(derham.p[1] + 1, dtype=float) - bn3 = np.empty(derham.p[2] + 1, dtype=float) + bn1 = xp.empty(derham.p[0] + 1, dtype=float) + bn2 = xp.empty(derham.p[1] + 1, dtype=float) + bn3 = xp.empty(derham.p[2] + 1, dtype=float) - bd1 = np.empty(derham.p[0], dtype=float) - bd2 = np.empty(derham.p[1], dtype=float) - bd3 = np.empty(derham.p[2], dtype=float) + bd1 = xp.empty(derham.p[0], dtype=float) + bd2 = xp.empty(derham.p[1], dtype=float) + bd3 = xp.empty(derham.p[2], dtype=float) bsp.b_d_splines_slim(tn1, derham.p[0], eta1, span1, bn1, bd1) bsp.b_d_splines_slim(tn2, derham.p[1], eta2, span2, bn2, bd2) @@ -83,8 +83,8 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): # compare spline evaluation routines in V0 val = eval3d(*derham.p, bn1, bn2, bn3, ind_n1, ind_n2, ind_n3, x0[0]) - val_mpi = eval3d_mpi(*derham.p, bn1, bn2, bn3, span1, span2, span3, x0_psy._data, np.array(x0_psy.starts)) - assert np.allclose(val, val_mpi) + val_mpi = eval3d_mpi(*derham.p, bn1, bn2, bn3, span1, span2, span3, x0_psy._data, xp.array(x0_psy.starts)) + assert xp.allclose(val, val_mpi) # compare spline evaluation routines in V1 val = eval3d(derham.p[0] - 1, derham.p[1], derham.p[2], bd1, bn2, bn3, ind_d1, ind_n2, ind_n3, x1[0]) @@ -99,9 +99,9 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): span2, span3, x1_psy[0]._data, - np.array(x1_psy[0].starts), + xp.array(x1_psy[0].starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) val = eval3d(derham.p[0], derham.p[1] - 1, derham.p[2], bn1, bd2, bn3, ind_n1, ind_d2, ind_n3, x1[1]) val_mpi = eval3d_mpi( @@ -115,9 +115,9 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): span2, span3, x1_psy[1]._data, - np.array(x1_psy[1].starts), + xp.array(x1_psy[1].starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) val = eval3d(derham.p[0], derham.p[1], derham.p[2] - 1, bn1, bn2, bd3, ind_n1, ind_n2, ind_d3, x1[2]) val_mpi = eval3d_mpi( @@ -131,9 +131,9 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): span2, span3, x1_psy[2]._data, - np.array(x1_psy[2].starts), + xp.array(x1_psy[2].starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) # compare spline evaluation routines in V2 val = eval3d(derham.p[0], derham.p[1] - 1, derham.p[2] - 1, bn1, bd2, bd3, ind_n1, ind_d2, ind_d3, x2[0]) @@ -148,9 +148,9 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): span2, span3, x2_psy[0]._data, - np.array(x2_psy[0].starts), + xp.array(x2_psy[0].starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) val = eval3d(derham.p[0] - 1, derham.p[1], derham.p[2] - 1, bd1, bn2, bd3, ind_d1, ind_n2, ind_d3, x2[1]) val_mpi = eval3d_mpi( @@ -164,9 +164,9 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): span2, span3, x2_psy[1]._data, - np.array(x2_psy[1].starts), + xp.array(x2_psy[1].starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) val = eval3d(derham.p[0] - 1, derham.p[1] - 1, derham.p[2], bd1, bd2, bn3, ind_d1, ind_d2, ind_n3, x2[2]) val_mpi = eval3d_mpi( @@ -180,9 +180,9 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): span2, span3, x2_psy[2]._data, - np.array(x2_psy[2].starts), + xp.array(x2_psy[2].starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) # compare spline evaluation routines in V3 val = eval3d(derham.p[0] - 1, derham.p[1] - 1, derham.p[2] - 1, bd1, bd2, bd3, ind_d1, ind_d2, ind_d3, x3[0]) @@ -197,9 +197,9 @@ def test_eval_kernels(Nel, p, spl_kind, n_markers=10): span2, span3, x3_psy._data, - np.array(x3_psy.starts), + xp.array(x3_psy.starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @@ -230,9 +230,9 @@ def test_eval_pointwise(Nel, p, spl_kind, n_markers=10): # Random points in domain of process dom = derham.domain_array[rank] - eta1s = np.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] - eta2s = np.random.rand(n_markers) * (dom[4] - dom[3]) + dom[3] - eta3s = np.random.rand(n_markers) * (dom[7] - dom[6]) + dom[6] + eta1s = xp.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] + eta2s = xp.random.rand(n_markers) * (dom[4] - dom[3]) + dom[3] + eta3s = xp.random.rand(n_markers) * (dom[7] - dom[6]) + dom[6] for eta1, eta2, eta3 in zip(eta1s, eta2s, eta3s): comm.Barrier() @@ -251,14 +251,14 @@ def test_eval_pointwise(Nel, p, spl_kind, n_markers=10): eta3, x0_psy._data, derham.spline_types_pyccel["0"], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) # compare spline evaluation routines in V1 # 1st component @@ -287,14 +287,14 @@ def test_eval_pointwise(Nel, p, spl_kind, n_markers=10): eta3, x1_psy[0]._data, derham.spline_types_pyccel["1"][0], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) # 2nd component val = evaluate_3d( @@ -322,14 +322,14 @@ def test_eval_pointwise(Nel, p, spl_kind, n_markers=10): eta3, x1_psy[1]._data, derham.spline_types_pyccel["1"][1], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) # 3rd component val = evaluate_3d( @@ -357,14 +357,14 @@ def test_eval_pointwise(Nel, p, spl_kind, n_markers=10): eta3, x1_psy[2]._data, derham.spline_types_pyccel["1"][2], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) # compare spline evaluation routines in V2 # 1st component @@ -393,14 +393,14 @@ def test_eval_pointwise(Nel, p, spl_kind, n_markers=10): eta3, x2_psy[0]._data, derham.spline_types_pyccel["2"][0], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) # 2nd component val = evaluate_3d( @@ -428,14 +428,14 @@ def test_eval_pointwise(Nel, p, spl_kind, n_markers=10): eta3, x2_psy[1]._data, derham.spline_types_pyccel["2"][1], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) # 3rd component val = evaluate_3d( @@ -463,14 +463,14 @@ def test_eval_pointwise(Nel, p, spl_kind, n_markers=10): eta3, x2_psy[2]._data, derham.spline_types_pyccel["2"][2], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) # compare spline evaluation routines in V3 val = evaluate_3d( @@ -496,14 +496,14 @@ def test_eval_pointwise(Nel, p, spl_kind, n_markers=10): eta3, x3_psy._data, derham.spline_types_pyccel["3"], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), ) - assert np.allclose(val, val_mpi) + assert xp.allclose(val, val_mpi) @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @@ -544,13 +544,13 @@ def test_eval_tensor_product(Nel, p, spl_kind, n_markers=10): # Random points in domain of process dom = derham.domain_array[rank] - eta1s = np.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] - eta2s = np.random.rand(n_markers + 1) * (dom[4] - dom[3]) + dom[3] - eta3s = np.random.rand(n_markers + 2) * (dom[7] - dom[6]) + dom[6] + eta1s = xp.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] + eta2s = xp.random.rand(n_markers + 1) * (dom[4] - dom[3]) + dom[3] + eta3s = xp.random.rand(n_markers + 2) * (dom[7] - dom[6]) + dom[6] - vals = np.zeros((n_markers, n_markers + 1, n_markers + 2), dtype=float) - vals_mpi = np.zeros((n_markers, n_markers + 1, n_markers + 2), dtype=float) - vals_mpi_fast = np.zeros((n_markers, n_markers + 1, n_markers + 2), dtype=float) + vals = xp.zeros((n_markers, n_markers + 1, n_markers + 2), dtype=float) + vals_mpi = xp.zeros((n_markers, n_markers + 1, n_markers + 2), dtype=float) + vals_mpi_fast = xp.zeros((n_markers, n_markers + 1, n_markers + 2), dtype=float) comm.Barrier() sleep(0.02 * (rank + 1)) @@ -573,11 +573,11 @@ def test_eval_tensor_product(Nel, p, spl_kind, n_markers=10): eta3s, x0_psy._data, derham.spline_types_pyccel["0"], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), vals_mpi, ) t1 = time.time() @@ -591,19 +591,19 @@ def test_eval_tensor_product(Nel, p, spl_kind, n_markers=10): eta3s, x0_psy._data, derham.spline_types_pyccel["0"], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), vals_mpi_fast, ) t1 = time.time() if rank == 0: print("v0 eval_spline_mpi_tensor_product_fast:".ljust(40), t1 - t0) - assert np.allclose(vals, vals_mpi) - assert np.allclose(vals, vals_mpi_fast) + assert xp.allclose(vals, vals_mpi) + assert xp.allclose(vals, vals_mpi_fast) # compare spline evaluation routines in V3 t0 = time.time() @@ -633,11 +633,11 @@ def test_eval_tensor_product(Nel, p, spl_kind, n_markers=10): eta3s, x3_psy._data, derham.spline_types_pyccel["3"], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), vals_mpi, ) t1 = time.time() @@ -651,19 +651,19 @@ def test_eval_tensor_product(Nel, p, spl_kind, n_markers=10): eta3s, x3_psy._data, derham.spline_types_pyccel["3"], - np.array(derham.p), + xp.array(derham.p), tn1, tn2, tn3, - np.array(x0_psy.starts), + xp.array(x0_psy.starts), vals_mpi_fast, ) t1 = time.time() if rank == 0: print("v3 eval_spline_mpi_tensor_product_fast:".ljust(40), t1 - t0) - assert np.allclose(vals, vals_mpi) - assert np.allclose(vals, vals_mpi_fast) + assert xp.allclose(vals, vals_mpi) + assert xp.allclose(vals, vals_mpi_fast) @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @@ -710,9 +710,9 @@ def test_eval_tensor_product_grid(Nel, p, spl_kind, n_markers=10): spans_f, bns_f, bds_f = derham.prepare_eval_tp_fixed([eta1s, eta2s, eta3s]) # output arrays - vals = np.zeros((eta1s.size, eta2s.size, eta3s.size), dtype=float) - vals_mpi_fixed = np.zeros((eta1s.size, eta2s.size, eta3s.size), dtype=float) - vals_mpi_grid = np.zeros((eta1s.size, eta2s.size, eta3s.size), dtype=float) + vals = xp.zeros((eta1s.size, eta2s.size, eta3s.size), dtype=float) + vals_mpi_fixed = xp.zeros((eta1s.size, eta2s.size, eta3s.size), dtype=float) + vals_mpi_grid = xp.zeros((eta1s.size, eta2s.size, eta3s.size), dtype=float) comm.Barrier() sleep(0.02 * (rank + 1)) @@ -748,20 +748,20 @@ def test_eval_tensor_product_grid(Nel, p, spl_kind, n_markers=10): *bds_f, x3_psy._data, derham.spline_types_pyccel["3"], - np.array(derham.p), - np.array(x0_psy.starts), + xp.array(derham.p), + xp.array(x0_psy.starts), vals_mpi_fixed, ) t1 = time.time() if rank == 0: print("v3 eval_spline_mpi_tensor_product_fixed:".ljust(40), t1 - t0) - assert np.allclose(vals, vals_mpi_fixed) + assert xp.allclose(vals, vals_mpi_fixed) field = derham.create_spline_function("test", "L2") field.vector = x3_psy - assert np.allclose(field.vector._data, x3_psy._data) + assert xp.allclose(field.vector._data, x3_psy._data) t0 = time.time() field.eval_tp_fixed_loc(spans_f, bds_f, out=vals_mpi_fixed) @@ -769,7 +769,7 @@ def test_eval_tensor_product_grid(Nel, p, spl_kind, n_markers=10): if rank == 0: print("v3 field.eval_tp_fixed:".ljust(40), t1 - t0) - assert np.allclose(vals, vals_mpi_fixed) + assert xp.allclose(vals, vals_mpi_fixed) if __name__ == "__main__": diff --git a/src/struphy/console/profile.py b/src/struphy/console/profile.py index aa6a36d56..11c753a15 100644 --- a/src/struphy/console/profile.py +++ b/src/struphy/console/profile.py @@ -11,7 +11,7 @@ def struphy_profile(dirs, replace, all, n_lines, print_callers, savefig): import struphy.utils.utils as utils from struphy.post_processing.cprofile_analyser import get_cprofile_data, replace_keys - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # Read struphy state file state = utils.read_state() @@ -167,10 +167,10 @@ def struphy_profile(dirs, replace, all, n_lines, print_callers, savefig): ratio.append(str(int(float(t) / runtime * 100)) + "%") # strong scaling plot - if np.all([Nel == val["Nel"][0] for Nel in val["Nel"]]): + if xp.all([Nel == val["Nel"][0] for Nel in val["Nel"]]): # ideal scaling if n == 0: - ax.loglog(val["mpi_size"], 1 / 2 ** np.arange(len(val["time"])), "k--", alpha=0.3, label="ideal") + ax.loglog(val["mpi_size"], 1 / 2 ** xp.arange(len(val["time"])), "k--", alpha=0.3, label="ideal") # print average time per one time step if "integrate" in key: @@ -206,11 +206,11 @@ def struphy_profile(dirs, replace, all, n_lines, print_callers, savefig): ax.set_ylabel("time [s]") ax.set( title="Weak scaling for cells/mpi_size=" - + str(np.prod(val["Nel"][0]) / val["mpi_size"][0]) + + str(xp.prod(val["Nel"][0]) / val["mpi_size"][0]) + "=const." ) ax.legend(loc="upper left") - # ax.loglog(val['mpi_size'], val['time'][0]*np.ones_like(val['time']), 'k--', alpha=0.3) + # ax.loglog(val['mpi_size'], val['time'][0]*xp.ones_like(val['time']), 'k--', alpha=0.3) ax.set_xscale("log") if savefig is None: diff --git a/src/struphy/diagnostics/console_diagn.py b/src/struphy/diagnostics/console_diagn.py index 35294b37a..3fa1f11ca 100644 --- a/src/struphy/diagnostics/console_diagn.py +++ b/src/struphy/diagnostics/console_diagn.py @@ -11,7 +11,7 @@ import struphy import struphy.utils.utils as utils from struphy.diagnostics.diagn_tools import plot_distr_fun, plot_scalars, plots_videos_2d -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def main(): @@ -301,7 +301,7 @@ def main(): bckgr_fun = getattr(maxwellians, default_bckgr_type)() # Get values of background shifts in velocity space - positions = [np.array([grid_slices["e" + str(k)]]) for k in range(1, 4)] + positions = [xp.array([grid_slices["e" + str(k)]]) for k in range(1, 4)] u = bckgr_fun.u(*positions) eval_params = {"u" + str(k + 1): u[k][0] for k in range(3)} @@ -315,7 +315,7 @@ def main(): # Plot the distribution function if "plot_distr" in actions: # Get index of where to plot in time - time_idx = np.argmin(np.abs(time - saved_time)) + time_idx = xp.argmin(xp.abs(time - saved_time)) plot_distr_fun( path=os.path.join( diff --git a/src/struphy/diagnostics/continuous_spectra.py b/src/struphy/diagnostics/continuous_spectra.py index cbe42dc94..ce805a45d 100644 --- a/src/struphy/diagnostics/continuous_spectra.py +++ b/src/struphy/diagnostics/continuous_spectra.py @@ -38,7 +38,7 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, """ import struphy.bsplines.bsplines as bsp - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # greville points in radial direction (s) gN_1 = bsp.greville(space.T[0], space.p[0], space.spl_kind[0]) @@ -49,7 +49,7 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, gD_2 = bsp.greville(space.t[1], space.p[1] - 1, space.spl_kind[1]) # poloidal mode numbers - ms = np.arange(m_range[1] - m_range[0] + 1) + m_range[0] + ms = xp.arange(m_range[1] - m_range[0] + 1) + m_range[0] # grid for normalized Jacobian determinant det_df = domain.jacobian_det(gD_1, gD_2, 0.0) @@ -65,7 +65,7 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, s_spec = [[[], [], []] for m in ms] # only consider eigenmodes in range omega^2/omega_A^2 = [0, 1] - modes_ind = np.where((np.real(omega2) / omega_A**2 < 1.0) & (np.real(omega2) / omega_A**2 > 0.0))[0] + modes_ind = xp.where((xp.real(omega2) / omega_A**2 < 1.0) & (xp.real(omega2) / omega_A**2 > 0.0))[0] for i in range(modes_ind.size): # determine whether it's an Alfvén branch or sound branch by checking DIV(U) @@ -85,14 +85,14 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, U2_1_coeff = (U2_1_coeff[:, :, 0] - 1j * U2_1_coeff[:, :, 1]) / 2 # determine radial location of singularity by looking for a peak in eigenfunction U2_1 - s_ind = np.unravel_index(np.argmax(abs(U2_1_coeff)), U2_1_coeff.shape)[0] + s_ind = xp.unravel_index(xp.argmax(abs(U2_1_coeff)), U2_1_coeff.shape)[0] s = gN_1[s_ind] # perform fft to determine m - U2_1_fft = np.fft.fft(U2_1_coeff) + U2_1_fft = xp.fft.fft(U2_1_coeff) # determine m by looking for peak in Fourier spectrum at singularity - m = int((np.fft.fftfreq(U2_1_fft[s_ind].size) * U2_1_fft[s_ind].size)[np.argmax(abs(U2_1_fft[s_ind]))]) + m = int((xp.fft.fftfreq(U2_1_fft[s_ind].size) * U2_1_fft[s_ind].size)[xp.argmax(abs(U2_1_fft[s_ind]))]) ## perform shift for negative m # if m >= (space.Nel[1] + 1)//2: @@ -102,7 +102,7 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, for j in range(ms.size): if ms[j] == m: a_spec[j][0].append(s) - a_spec[j][1].append(np.real(omega2[modes_ind[i]])) + a_spec[j][1].append(xp.real(omega2[modes_ind[i]])) a_spec[j][2].append(modes_ind[i]) # Sound branch @@ -116,14 +116,14 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, U2_coeff = (U2_coeff[:, :, 0] - 1j * U2_coeff[:, :, 1]) / 2 # determine radial location of singularity by looking for a peak in eigenfunction (U2_2 or U2_3) - s_ind = np.unravel_index(np.argmax(abs(U2_coeff)), U2_coeff.shape)[0] + s_ind = xp.unravel_index(xp.argmax(abs(U2_coeff)), U2_coeff.shape)[0] s = gD_1[s_ind] # perform fft to determine m - U2_fft = np.fft.fft(U2_coeff) + U2_fft = xp.fft.fft(U2_coeff) # determine m by looking for peak in Fourier spectrum at singularity - m = int((np.fft.fftfreq(U2_fft[s_ind].size) * U2_fft[s_ind].size)[np.argmax(abs(U2_fft[s_ind]))]) + m = int((xp.fft.fftfreq(U2_fft[s_ind].size) * U2_fft[s_ind].size)[xp.argmax(abs(U2_fft[s_ind]))]) ## perform shift for negative m # if m >= (space.Nel[1] + 1)//2: @@ -133,13 +133,13 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, for j in range(ms.size): if ms[j] == m: s_spec[j][0].append(s) - s_spec[j][1].append(np.real(omega2[modes_ind[i]])) + s_spec[j][1].append(xp.real(omega2[modes_ind[i]])) s_spec[j][2].append(modes_ind[i]) # convert to array for j in range(ms.size): - a_spec[j] = np.array(a_spec[j]) - s_spec[j] = np.array(s_spec[j]) + a_spec[j] = xp.array(a_spec[j]) + s_spec[j] = xp.array(s_spec[j]) return a_spec, s_spec @@ -153,7 +153,7 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, import yaml - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # parse arguments parser = argparse.ArgumentParser( @@ -256,7 +256,7 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, ) # load and analyze spectrum - omega2, U2_eig = np.split(np.load(spec_path), [1], axis=0) + omega2, U2_eig = xp.split(xp.load(spec_path), [1], axis=0) omega2 = omega2.flatten() m_range_alfven = [args.m_l_alfvén, args.m_u_alfvén] @@ -282,7 +282,7 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, fig.set_figheight(12) fig.set_figwidth(14) - etaplot = [np.linspace(0.0, 1.0, 201), np.linspace(0.0, 1.0, 101)] + etaplot = [xp.linspace(0.0, 1.0, 201), xp.linspace(0.0, 1.0, 101)] etaplot[0][0] += 1e-5 diff --git a/src/struphy/diagnostics/diagn_tools.py b/src/struphy/diagnostics/diagn_tools.py index 3f3e52fe8..88c81409c 100644 --- a/src/struphy/diagnostics/diagn_tools.py +++ b/src/struphy/diagnostics/diagn_tools.py @@ -10,7 +10,7 @@ from tqdm import tqdm from struphy.dispersion_relations import analytic -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def power_spectrum_2d( @@ -37,7 +37,7 @@ def power_spectrum_2d( Parameters ---------- values : dict - Dictionary holding values of a B-spline FemField on the grid as 3d np.arrays: + Dictionary holding values of a B-spline FemField on the grid as 3d xp.arrays: values[n] contains the values at time step n, where n = 0:Nt-1:step with 0 0: @@ -162,9 +162,9 @@ def power_spectrum_2d( for n in range(fit_branches): omega_fit[n] = [] for k, f_of_omega in zip(kvec[k_start:k_end], dispersion[:, k_start:k_end].T): - threshold = np.max(f_of_omega) * noise_level - extrms = argrelextrema(f_of_omega, np.greater, order=extr_order)[0] - above_noise = np.nonzero(f_of_omega > threshold)[0] + threshold = xp.max(f_of_omega) * noise_level + extrms = argrelextrema(f_of_omega, xp.greater, order=extr_order)[0] + above_noise = xp.nonzero(f_of_omega > threshold)[0] intersec = list(set(extrms) & set(above_noise)) # intersec = list(set(extrms)) if not intersec: @@ -183,14 +183,14 @@ def power_spectrum_2d( # fit coeffs = [] for m, om in omega_fit.items(): - coeffs += [np.polyfit(k_fit, om, deg=fit_degree[n])] + coeffs += [xp.polyfit(k_fit, om, deg=fit_degree[n])] print(f"\nFitted {coeffs = }") if do_plot: _, ax = plt.subplots(1, 1, figsize=(10, 10)) colormap = "plasma" - K, W = np.meshgrid(kvec, omega) - lvls = np.logspace(-15, -1, 27) + K, W = xp.meshgrid(kvec, omega) + lvls = xp.logspace(-15, -1, 27) disp_plot = ax.contourf( K, W, @@ -214,7 +214,7 @@ def power_spectrum_2d( def fun(k): out = k * 0.0 - for i, c in enumerate(np.flip(cs)): + for i, c in enumerate(xp.flip(cs)): out += c * k**i return out @@ -230,12 +230,12 @@ def fun(k): set_min = 0.0 set_max = 0.0 for key, branch in branches.items(): - vals = np.real(branch) + vals = xp.real(branch) ax.plot(kvec, vals, "--", label=key) - tmp = np.min(vals) + tmp = xp.min(vals) if tmp < set_min: set_min = tmp - tmp = np.max(vals) + tmp = xp.max(vals) if tmp > set_max: set_max = tmp @@ -331,8 +331,8 @@ def plot_scalars( plt.figure("en_tot_rel_err") plt.plot( time[1:], - np.divide( - np.abs(en_tot[1:] - en_tot[0]), + xp.divide( + xp.abs(en_tot[1:] - en_tot[0]), en_tot[0], ), ) @@ -363,9 +363,9 @@ def plot_scalars( for key, plot_quantity in plot_quantities.items(): # Get the indices of the extrema if do_fit: - inds_exs = argrelextrema(plot_quantity, np.greater, order=order) + inds_exs = argrelextrema(plot_quantity, xp.greater, order=order) elif fit_minima: - inds_exs = argrelextrema(plot_quantity, np.less, order=order) + inds_exs = argrelextrema(plot_quantity, xp.less, order=order) else: inds_exs = None @@ -376,10 +376,10 @@ def plot_scalars( # for plotting take a bit more time at start and end if len(inds_exs[0]) >= 2: - time_start_idx = np.max( + time_start_idx = xp.max( [0, 2 * inds_exs[0][start_extremum] - inds_exs[0][start_extremum + 1]], ) - time_end_idx = np.min( + time_end_idx = xp.min( [ len(time) - 1, 2 * inds_exs[0][start_extremum + no_extrema - 1] - inds_exs[0][start_extremum + no_extrema - 2], @@ -395,9 +395,9 @@ def plot_scalars( if inds_exs is not None: # do the fitting - coeffs = np.polyfit( + coeffs = xp.polyfit( times_extrema, - np.log( + xp.log( quantity_extrema, ), deg=degree, @@ -410,15 +410,15 @@ def plot_scalars( ) plt.plot( time_cut, - np.exp(coeffs[0] * time_cut + coeffs[1]), - label=r"$a * \exp(m x)$ with" + f"\na={np.round(np.exp(coeffs[1]), 3)} m={np.round(coeffs[0], 3)}", + xp.exp(coeffs[0] * time_cut + coeffs[1]), + label=r"$a * \exp(m x)$ with" + f"\na={xp.round(xp.exp(coeffs[1]), 3)} m={xp.round(coeffs[0], 3)}", ) else: plt.plot(time, plot_quantity[:], ".", label=key, markersize=2) if inds_exs is not None: # do the fitting - coeffs = np.polyfit( + coeffs = xp.polyfit( times_extrema, quantity_extrema, deg=degree, @@ -433,8 +433,8 @@ def plot_scalars( ) plt.plot( time_cut, - np.exp(coeffs[0] * time_cut + coeffs[1]), - label=r"$a x + b$ with" + f"\na={np.round(coeffs[1], 3)} b={np.round(coeffs[0], 3)}", + xp.exp(coeffs[0] * time_cut + coeffs[1]), + label=r"$a x + b$ with" + f"\na={xp.round(coeffs[1], 3)} b={xp.round(coeffs[0], 3)}", ) plt.legend() @@ -496,11 +496,11 @@ def plot_distr_fun( # load full distribution functions if filename == "f_binned.npy": - f = np.load(filepath) + f = xp.load(filepath) # load delta f elif filename == "delta_f_binned.npy": - delta_f = np.load(filepath) + delta_f = xp.load(filepath) assert f is not None, "No distribution function file found!" @@ -508,7 +508,7 @@ def plot_distr_fun( directions = folder.split("_") for direction in directions: grids += [ - np.load( + xp.load( os.path.join( subpath, "grid_" + direction + ".npy", @@ -519,8 +519,8 @@ def plot_distr_fun( # Get indices of where to plot in other directions grid_idxs = {} for k in range(f.ndim - 1): - grid_idxs[directions[k]] = np.argmin( - np.abs(grids[k] - grid_slices[directions[k]]), + grid_idxs[directions[k]] = xp.argmin( + xp.abs(grids[k] - grid_slices[directions[k]]), ) for k in range(f.ndim - 1): @@ -655,17 +655,17 @@ def plots_videos_2d( grid_idxs = {} for k in range(df_data.ndim - 1): direc = directions[k] - grid_idxs[direc] = np.argmin( - np.abs(grids[direc] - grid_slices[direc]), + grid_idxs[direc] = xp.argmin( + xp.abs(grids[direc] - grid_slices[direc]), ) - grid_1 = np.load( + grid_1 = xp.load( os.path.join( data_path, "grid_" + label_1 + ".npy", ), ) - grid_2 = np.load( + grid_2 = xp.load( os.path.join( data_path, "grid_" + label_2 + ".npy", @@ -696,9 +696,9 @@ def plots_videos_2d( var *= polar_params["r_max"] - polar_params["r_min"] var += polar_params["r_min"] elif polar_params["angular_coord"] == sl: - var *= 2 * np.pi + var *= 2 * xp.pi - grid_1_mesh, grid_2_mesh = np.meshgrid(grid_1, grid_2, indexing="ij") + grid_1_mesh, grid_2_mesh = xp.meshgrid(grid_1, grid_2, indexing="ij") if output == "video": plots_2d_video( @@ -745,7 +745,7 @@ def video_2d(slc, diagn_path, images_path): Parameters ---------- - t_grid : np.ndarray + t_grid : xp.ndarray 1D-array containing all the times grid_slices : dict @@ -833,15 +833,15 @@ def plots_2d_video( # Get parameters for time and labelling for it nt = len(t_grid) - log_nt = int(np.log10(nt)) + 1 + log_nt = int(xp.log10(nt)) + 1 len_dt = len(str(t_grid[1]).split(".")[1]) # Get the correct scale for the plots - vmin += [np.min(df_binned[:]) / 3] - vmax += [np.max(df_binned[:]) / 3] - vmin = np.min(vmin) - vmax = np.max(vmax) - vscale = np.max(np.abs([vmin, vmax])) + vmin += [xp.min(df_binned[:]) / 3] + vmax += [xp.max(df_binned[:]) / 3] + vmin = xp.min(vmin) + vmax = xp.max(vmax) + vscale = xp.max(xp.abs([vmin, vmax])) # Set up the figure and axis once if do_polar: @@ -939,18 +939,18 @@ def plots_2d_overview( fig_height = 8.5 else: n_cols = 3 - n_rows = int(np.ceil(n_times / n_cols)) + n_rows = int(xp.ceil(n_times / n_cols)) fig_height = 4 * n_rows fig_size = (4 * n_cols, fig_height) # Get the correct scale for the plots for time in times: - vmin += [np.min(df_binned[time]) / 3] - vmax += [np.max(df_binned[time]) / 3] - vmin = np.min(vmin) - vmax = np.max(vmax) - vscale = np.max(np.abs([vmin, vmax])) + vmin += [xp.min(df_binned[time]) / 3] + vmax += [xp.max(df_binned[time]) / 3] + vmin = xp.min(vmin) + vmax = xp.max(vmax) + vscale = xp.max(xp.abs([vmin, vmax])) # Plot options for polar plots subplot_kw = dict(projection="polar") if do_polar else None @@ -959,8 +959,8 @@ def plots_2d_overview( fig, axes = plt.subplots(n_rows, n_cols, figsize=fig_size, subplot_kw=subplot_kw) # So we an use .flatten() even for just 1 plot - if not isinstance(axes, np.ndarray): - axes = np.array([axes]) + if not isinstance(axes, xp.ndarray): + axes = xp.array([axes]) # fig.tight_layout(h_pad=5.0, w_pad=5.0) # fig.tight_layout(pad=5.0) @@ -976,7 +976,7 @@ def plots_2d_overview( # Set the suptitle fig.suptitle(f"Struphy model '{model_name}'") - for k in np.arange(n_times): + for k in xp.arange(n_times): obj = axes.flatten()[k] n = times[k] t = f"%.{len_dt}f" % t_grid[n] @@ -1048,13 +1048,13 @@ def get_slices_grids_directions_and_df_data(plot_full_f, grid_slices, data_path, slices_2d : list[string] A list of all the slicings - grids : list[np.ndarray] + grids : list[xp.ndarray] A list of all grids according to the slices directions : list[string] A list of the directions that appear in all slices - df_data : np.ndarray + df_data : xp.ndarray The data of delta-f (in case of full-f: distribution function minus background) """ @@ -1063,7 +1063,7 @@ def get_slices_grids_directions_and_df_data(plot_full_f, grid_slices, data_path, # Load all the grids grids = {} for direction in directions: - grids[direction] = np.load( + grids[direction] = xp.load( os.path.join(data_path, "grid_" + direction + ".npy"), ) @@ -1072,7 +1072,7 @@ def get_slices_grids_directions_and_df_data(plot_full_f, grid_slices, data_path, _name = "f_binned.npy" else: _name = "delta_f_binned.npy" - _data = np.load(os.path.join(data_path, _name)) + _data = xp.load(os.path.join(data_path, _name)) # Check how many slicings have been given and make slices_2d for all # combinations of spatial and velocity dimensions diff --git a/src/struphy/diagnostics/paraview/mesh_creator.py b/src/struphy/diagnostics/paraview/mesh_creator.py index 0a8a35903..2446d0391 100644 --- a/src/struphy/diagnostics/paraview/mesh_creator.py +++ b/src/struphy/diagnostics/paraview/mesh_creator.py @@ -4,7 +4,7 @@ from vtkmodules.util.numpy_support import vtk_to_numpy as vtk2np from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def make_ugrid_and_write_vtu(filename: str, writer, vtk_dir, gvec, s_range, u_range, v_range, periodic): @@ -81,43 +81,43 @@ def gen_vtk_points(gvec, s_range, u_range, v_range, point_data, cell_data): pt_idx = 0 vtk_points = vtk.vtkPoints() - suv_points = np.zeros((s_range.shape[0], u_range.shape[0], v_range.shape[0], 3)) - xyz_points = np.zeros((s_range.shape[0], u_range.shape[0], v_range.shape[0], 3)) - point_indices = np.zeros((s_range.shape[0], u_range.shape[0], v_range.shape[0]), dtype=np.int_) + suv_points = xp.zeros((s_range.shape[0], u_range.shape[0], v_range.shape[0], 3)) + xyz_points = xp.zeros((s_range.shape[0], u_range.shape[0], v_range.shape[0], 3)) + point_indices = xp.zeros((s_range.shape[0], u_range.shape[0], v_range.shape[0]), dtype=xp.int_) # Add metadata to grid. num_pts = s_range.shape[0] * u_range.shape[0] * v_range.shape[0] - point_data["s"] = np.zeros(num_pts, dtype=np.float_) - point_data["u"] = np.zeros(num_pts, dtype=np.float_) - point_data["v"] = np.zeros(num_pts, dtype=np.float_) - point_data["x"] = np.zeros(num_pts, dtype=np.float_) - point_data["y"] = np.zeros(num_pts, dtype=np.float_) - point_data["z"] = np.zeros(num_pts, dtype=np.float_) - point_data["theta"] = np.zeros(num_pts, dtype=np.float_) - point_data["zeta"] = np.zeros(num_pts, dtype=np.float_) - point_data["Point ID"] = np.zeros(num_pts, dtype=np.int_) - point_data["pressure"] = np.zeros(num_pts, dtype=np.float_) - point_data["phi"] = np.zeros(num_pts, dtype=np.float_) - point_data["chi"] = np.zeros(num_pts, dtype=np.float_) - point_data["iota"] = np.zeros(num_pts, dtype=np.float_) - point_data["q"] = np.zeros(num_pts, dtype=np.float_) - point_data["det"] = np.zeros(num_pts, dtype=np.float_) - point_data["det/(2pi)^2"] = np.zeros(num_pts, dtype=np.float_) - point_data["A"] = np.zeros((num_pts, 3), dtype=np.float_) - point_data["A_vec"] = np.zeros((num_pts, 3), dtype=np.float_) - point_data["A_1"] = np.zeros((num_pts, 3), dtype=np.float_) - point_data["A_2"] = np.zeros((num_pts, 3), dtype=np.float_) - point_data["B"] = np.zeros((num_pts, 3), dtype=np.float_) - point_data["B_vec"] = np.zeros((num_pts, 3), dtype=np.float_) - point_data["B_1"] = np.zeros((num_pts, 3), dtype=np.float_) - point_data["B_2"] = np.zeros((num_pts, 3), dtype=np.float_) + point_data["s"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["u"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["v"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["x"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["y"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["z"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["theta"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["zeta"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["Point ID"] = xp.zeros(num_pts, dtype=xp.int_) + point_data["pressure"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["phi"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["chi"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["iota"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["q"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["det"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["det/(2pi)^2"] = xp.zeros(num_pts, dtype=xp.float_) + point_data["A"] = xp.zeros((num_pts, 3), dtype=xp.float_) + point_data["A_vec"] = xp.zeros((num_pts, 3), dtype=xp.float_) + point_data["A_1"] = xp.zeros((num_pts, 3), dtype=xp.float_) + point_data["A_2"] = xp.zeros((num_pts, 3), dtype=xp.float_) + point_data["B"] = xp.zeros((num_pts, 3), dtype=xp.float_) + point_data["B_vec"] = xp.zeros((num_pts, 3), dtype=xp.float_) + point_data["B_1"] = xp.zeros((num_pts, 3), dtype=xp.float_) + point_data["B_2"] = xp.zeros((num_pts, 3), dtype=xp.float_) # pbar = tqdm(total=num_pts) for s_idx, s in enumerate(s_range): for u_idx, u in enumerate(u_range): for v_idx, v in enumerate(v_range): point = gvec.f(s, u, v) - suv_points[s_idx, u_idx, v_idx, :] = np.array([s, u, v]) + suv_points[s_idx, u_idx, v_idx, :] = xp.array([s, u, v]) xyz_points[s_idx, u_idx, v_idx, :] = point point_indices[s_idx, u_idx, v_idx] = pt_idx vtk_points.InsertPoint(pt_idx, point) @@ -149,10 +149,10 @@ def gen_vtk_points(gvec, s_range, u_range, v_range, point_data, cell_data): pt_idx += 1 # pbar.close() - point_data["theta"] = 2 * np.pi * point_data["u"] - point_data["zeta"] = 2 * np.pi * point_data["v"] + point_data["theta"] = 2 * xp.pi * point_data["u"] + point_data["zeta"] = 2 * xp.pi * point_data["v"] point_data["q"] = 1 / point_data["iota"] - point_data["det/(2pi)^2"] = point_data["det"] / (2 * np.pi) ** 2 + point_data["det/(2pi)^2"] = point_data["det"] / (2 * xp.pi) ** 2 return vtk_points, suv_points, xyz_points, point_indices @@ -312,4 +312,4 @@ def connect_cell(s_range, u_range, v_range, point_indices, ugrid, point_data, ce cell_data["Cell ID"].append(cell_idx) cell_idx += 1 - cell_data["Cell ID"] = np.array(cell_data["Cell ID"], dtype=np.int_) + cell_data["Cell ID"] = xp.array(cell_data["Cell ID"], dtype=xp.int_) diff --git a/src/struphy/dispersion_relations/analytic.py b/src/struphy/dispersion_relations/analytic.py index 8a87a65bb..756840f9a 100644 --- a/src/struphy/dispersion_relations/analytic.py +++ b/src/struphy/dispersion_relations/analytic.py @@ -6,7 +6,7 @@ from struphy.dispersion_relations.base import ContinuousSpectra1D, DispersionRelations1D from struphy.dispersion_relations.utilities import Zplasma from struphy.fields_background.equils import set_defaults -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Maxwell1D(DispersionRelations1D): @@ -108,18 +108,18 @@ def __call__(self, k): Bsquare = self.params["B0x"] ** 2 + self.params["B0y"] ** 2 + self.params["B0z"] ** 2 # Alfvén velocity and speed of sound - vA = np.sqrt(Bsquare / self.params["n0"]) + vA = xp.sqrt(Bsquare / self.params["n0"]) - cS = np.sqrt(self.params["gamma"] * self.params["p0"] / self.params["n0"]) + cS = xp.sqrt(self.params["gamma"] * self.params["p0"] / self.params["n0"]) # shear Alfvén branch - self._branches["shear Alfvén"] = vA * k * self.params["B0z"] / np.sqrt(Bsquare) + self._branches["shear Alfvén"] = vA * k * self.params["B0z"] / xp.sqrt(Bsquare) # slow/fast magnetosonic branch delta = (4 * self.params["B0z"] ** 2 * cS**2 * vA**2) / ((cS**2 + vA**2) ** 2 * Bsquare) - self._branches["slow magnetosonic"] = np.sqrt(1 / 2 * k**2 * (cS**2 + vA**2) * (1 - np.sqrt(1 - delta))) - self._branches["fast magnetosonic"] = np.sqrt(1 / 2 * k**2 * (cS**2 + vA**2) * (1 + np.sqrt(1 - delta))) + self._branches["slow magnetosonic"] = xp.sqrt(1 / 2 * k**2 * (cS**2 + vA**2) * (1 - xp.sqrt(1 - delta))) + self._branches["fast magnetosonic"] = xp.sqrt(1 / 2 * k**2 * (cS**2 + vA**2) * (1 + xp.sqrt(1 - delta))) return self.branches @@ -186,14 +186,14 @@ def __call__(self, k): Bsquare = self.params["B0x"] ** 2 + self.params["B0y"] ** 2 + self.params["B0z"] ** 2 - cos_theta = self.params["B0z"] / np.sqrt(Bsquare) + cos_theta = self.params["B0z"] / xp.sqrt(Bsquare) # Alfvén velocity, speed of sound and cyclotron frequency - vA = np.sqrt(Bsquare / self.params["n0"]) + vA = xp.sqrt(Bsquare / self.params["n0"]) - cS = np.sqrt(self.params["gamma"] * self.params["p0"] / self.params["n0"]) + cS = xp.sqrt(self.params["gamma"] * self.params["p0"] / self.params["n0"]) - Omega_i = np.sqrt(Bsquare) / self.params["eps"] + Omega_i = xp.sqrt(Bsquare) / self.params["eps"] # auxiliary functions def omega_0(k): @@ -218,7 +218,7 @@ def discriminant(k): ) # solve - out = np.zeros((k.size, 4), dtype=complex) + out = xp.zeros((k.size, 4), dtype=complex) for i, ki in enumerate(k): p0 = Polynomial([-(omega_0(ki) ** 2), 1.0]) p1 = Polynomial([d(ki), c(ki), b(ki), 1.0]) @@ -302,7 +302,7 @@ def discriminant(k): return -4.0 * p**3 - 27.0 * q(k) ** 2 # solve - out = np.zeros((k.size, 3), dtype=complex) + out = xp.zeros((k.size, 3), dtype=complex) for i, ki in enumerate(k): poly = Polynomial([q(ki), p, 0.0, 1.0]) out[i] = poly.roots() @@ -342,17 +342,17 @@ def __call__(self, kvec): # One complex array for each branch tmps = [] for n in range(self.nbranches): - tmps += [np.zeros_like(kvec, dtype=complex)] + tmps += [xp.zeros_like(kvec, dtype=complex)] ########### Model specific part ############################## # angle between k and magnetic field if self.params["B0z"] == 0: - theta = np.pi / 2 + theta = xp.pi / 2 else: - theta = np.arctan(np.sqrt(self.params["B0x"] ** 2 + self.params["B0y"] ** 2) / self.params["B0z"]) + theta = xp.arctan(xp.sqrt(self.params["B0x"] ** 2 + self.params["B0y"] ** 2) / self.params["B0z"]) print(theta) - cos2 = np.cos(theta) ** 2 + cos2 = xp.cos(theta) ** 2 neq = self.params["n0"] @@ -393,10 +393,10 @@ def __call__(self, kvec): e = eps6 # determinant in polynomial form - det = np.polynomial.Polynomial([a, b, c, d, e]) + det = xp.polynomial.Polynomial([a, b, c, d, e]) # solutions - sol = np.sqrt(np.abs(det.roots())) + sol = xp.sqrt(xp.abs(det.roots())) # Ion-cyclotron branch tmps[0][n] = sol[0] # Electron-cyclotron branch @@ -489,7 +489,7 @@ def __init__(self, **params): ee = 1.602176634e-19 # calculate coupling parameter alpha_c from bulk number density and mass number - self._kappa = ee * np.sqrt(mu * self.params["Ab"] * self.params["nb"] * 1e20 / mp) + self._kappa = ee * xp.sqrt(mu * self.params["Ab"] * self.params["nb"] * 1e20 / mp) def __call__(self, k, method="newton", tol=1e-10, max_it=100): """ @@ -518,7 +518,7 @@ def __call__(self, k, method="newton", tol=1e-10, max_it=100): # One complex array for each branch tmps = [] for _ in range(self.nbranches): - tmps += [np.zeros_like(k, dtype=complex)] + tmps += [xp.zeros_like(k, dtype=complex)] ########### Model specific part ############################## @@ -532,8 +532,8 @@ def __call__(self, k, method="newton", tol=1e-10, max_it=100): wR = [self.params["B0"] * ki, 0.0] wL = [self.params["B0"] * ki, 0.0] else: - wR = [np.real(tmps[0][i - 1]), np.imag(tmps[0][i - 1])] - wL = [np.real(tmps[1][i - 1]), np.imag(tmps[1][i - 1])] + wR = [xp.real(tmps[0][i - 1]), xp.imag(tmps[0][i - 1])] + wL = [xp.real(tmps[1][i - 1]), xp.imag(tmps[1][i - 1])] # apply solver if method == "newton": @@ -542,13 +542,13 @@ def __call__(self, k, method="newton", tol=1e-10, max_it=100): Dr, Di = self.D_RL(wR, ki, +1) - while np.abs(Dr + Di * 1j) > tol or counter == max_it: + while xp.abs(Dr + Di * 1j) > tol or counter == max_it: # derivative Drp, Dip = self.D_RL(wR, ki, +1, 1) # update - wR[0] = wR[0] - np.real((Dr + Di * 1j) / (Drp + Dip * 1j)) - wR[1] = wR[1] - np.imag((Dr + Di * 1j) / (Drp + Dip * 1j)) + wR[0] = wR[0] - xp.real((Dr + Di * 1j) / (Drp + Dip * 1j)) + wR[1] = wR[1] - xp.imag((Dr + Di * 1j) / (Drp + Dip * 1j)) Dr, Di = self.D_RL(wR, ki, +1) counter += 1 @@ -558,13 +558,13 @@ def __call__(self, k, method="newton", tol=1e-10, max_it=100): Dr, Di = self.D_RL(wL, ki, -1) - while np.abs(Dr + Di * 1j) > tol or counter == max_it: + while xp.abs(Dr + Di * 1j) > tol or counter == max_it: # derivative Drp, Dip = self.D_RL(wL, ki, -1, 1) # update - wL[0] = wL[0] - np.real((Dr + Di * 1j) / (Drp + Dip * 1j)) - wL[1] = wL[1] - np.imag((Dr + Di * 1j) / (Drp + Dip * 1j)) + wL[0] = wL[0] - xp.real((Dr + Di * 1j) / (Drp + Dip * 1j)) + wL[1] = wL[1] - xp.imag((Dr + Di * 1j) / (Drp + Dip * 1j)) Dr, Di = self.D_RL(wL, ki, -1) counter += 1 @@ -651,7 +651,7 @@ def D_RL(self, w, k, pol, der=0): * (Zplasma(xi, 0) + (w - k * v0) * Zplasma(xi, 1) * xip) ) - return np.real(out), np.imag(out) + return xp.real(out), xp.imag(out) class PressureCouplingFull6DParallel(DispersionRelations1D): @@ -723,7 +723,7 @@ def __call__(self, k, tol=1e-10): # One complex array for each branch tmps = [] for n in range(self.nbranches): - tmps += [np.zeros_like(k, dtype=complex)] + tmps += [xp.zeros_like(k, dtype=complex)] ########### Model specific part ############################## @@ -735,9 +735,9 @@ def __call__(self, k, tol=1e-10): wL = [1 * ki, 0.0] # TODO: use vA wS = [1 * ki, 0.0] # TODO: use cS else: - wR = [np.real(tmps[0][i - 1]), np.imag(tmps[0][i - 1])] - wL = [np.real(tmps[1][i - 1]), np.imag(tmps[1][i - 1])] - wS = [np.real(tmps[2][i - 1]), np.imag(tmps[2][i - 1])] + wR = [xp.real(tmps[0][i - 1]), xp.imag(tmps[0][i - 1])] + wL = [xp.real(tmps[1][i - 1]), xp.imag(tmps[1][i - 1])] + wS = [xp.real(tmps[2][i - 1]), xp.imag(tmps[2][i - 1])] # R/L shear Alfvén wave sol_R = fsolve(self.D_RL, x0=wR, args=(ki, +1), xtol=tol) @@ -796,8 +796,8 @@ def D_RL(self, w, k, pol): vperp = 1.0 # TODO vth = 1.0 - vA = np.sqrt((self.params["B0x"] ** 2 + self.params["B0y"] ** 2 + self.params["B0z"] ** 2) / self.params["n0"]) - # cS = np.sqrt(self.params['beta']*vA) + vA = xp.sqrt((self.params["B0x"] ** 2 + self.params["B0y"] ** 2 + self.params["B0z"] ** 2) / self.params["n0"]) + # cS = xp.sqrt(self.params['beta']*vA) cS = 1.0 a0 = u0 / vpara # TODO @@ -840,7 +840,7 @@ def D_RL(self, w, k, pol): ) ) - return np.real(c1), np.imag(c1) + return xp.real(c1), xp.imag(c1) def D_sonic(self, w, k): r""" @@ -873,8 +873,8 @@ def D_sonic(self, w, k): vperp = 1.0 # TODO vth = 1.0 - vA = np.sqrt((self.params["B0x"] ** 2 + self.params["B0y"] ** 2 + self.params["B0z"] ** 2) / self.params["n0"]) - # cS = np.sqrt(self.params['beta']*vA) + vA = xp.sqrt((self.params["B0x"] ** 2 + self.params["B0y"] ** 2 + self.params["B0z"] ** 2) / self.params["n0"]) + # cS = xp.sqrt(self.params['beta']*vA) cS = 1.0 a0 = u0 / vpara # TODO @@ -885,7 +885,7 @@ def D_sonic(self, w, k): c1 = w**2 - k**2 * cS**2 + 2 * w * k * nu * vpara * x4 - return np.real(c1), np.imag(c1) + return xp.real(c1), xp.imag(c1) # private methods: # ---------------- @@ -1014,10 +1014,10 @@ def __call__(self, x, m, n): specs = {} # shear Alfvén continuum - specs["shear_Alfvén"] = np.sqrt(F(x, m, n) ** 2 / rho(x)) + specs["shear_Alfvén"] = xp.sqrt(F(x, m, n) ** 2 / rho(x)) # slow sound continuum - specs["slow_sound"] = np.sqrt( + specs["slow_sound"] = xp.sqrt( gamma * p(x) * F(x, m, n) ** 2 / (rho(x) * (gamma * p(x) + By(x) ** 2 + Bz(x) ** 2)) ) @@ -1121,10 +1121,10 @@ def __call__(self, r, m, n): specs = {} # shear Alfvén continuum - specs["shear_Alfvén"] = np.sqrt(F(r, m, n) ** 2 / rho(r)) + specs["shear_Alfvén"] = xp.sqrt(F(r, m, n) ** 2 / rho(r)) # slow sound continuum - specs["slow_sound"] = np.sqrt( + specs["slow_sound"] = xp.sqrt( gamma * p(r) * F(r, m, n) ** 2 / (rho(r) * (gamma * p(r) + Bt(r) ** 2 + Bz(r) ** 2)) ) diff --git a/src/struphy/dispersion_relations/base.py b/src/struphy/dispersion_relations/base.py index 31a237a90..dd7293867 100644 --- a/src/struphy/dispersion_relations/base.py +++ b/src/struphy/dispersion_relations/base.py @@ -4,7 +4,7 @@ from matplotlib import pyplot as plt -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class DispersionRelations1D(metaclass=ABCMeta): @@ -100,18 +100,18 @@ def plot(self, k): plt.ylabel(rf"Im($\omega$) [{unit_om}]") for name, omega in self.branches.items(): plt.subplot(2, 1, 1) - plt.plot(k, np.real(omega), label=name) + plt.plot(k, xp.real(omega), label=name) plt.subplot(2, 1, 2) - plt.plot(k, np.imag(omega), label=name) + plt.plot(k, xp.imag(omega), label=name) plt.subplot(2, 1, 1) for lab, kc in self.k_crit.items(): - if kc > np.min(k) and kc < np.max(k): + if kc > xp.min(k) and kc < xp.max(k): plt.axvline(kc, color="k", linestyle="--", linewidth=0.5, label=lab) plt.legend() plt.subplot(2, 1, 2) for lab, kc in self.k_crit.items(): - if kc > np.min(k) and kc < np.max(k): + if kc > xp.min(k) and kc < xp.max(k): plt.axvline(kc, color="k", linestyle="--", linewidth=0.5, label=lab) diff --git a/src/struphy/dispersion_relations/utilities.py b/src/struphy/dispersion_relations/utilities.py index b0f74fbb4..f84aa1acf 100644 --- a/src/struphy/dispersion_relations/utilities.py +++ b/src/struphy/dispersion_relations/utilities.py @@ -1,6 +1,6 @@ from scipy.special import erfi -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def Zplasma(xi, der=0): @@ -24,7 +24,7 @@ def Zplasma(xi, der=0): assert der == 0 or der == 1, 'Parameter "der" must be either 0 or 1' if der == 0: - z = np.sqrt(np.pi) * np.exp(-(xi**2)) * (1j - erfi(xi)) + z = xp.sqrt(xp.pi) * xp.exp(-(xi**2)) * (1j - erfi(xi)) else: z = -2 * (1 + xi * Zplasma(xi, 0)) diff --git a/src/struphy/eigenvalue_solvers/derivatives.py b/src/struphy/eigenvalue_solvers/derivatives.py index 45d8e2d94..738b1537a 100644 --- a/src/struphy/eigenvalue_solvers/derivatives.py +++ b/src/struphy/eigenvalue_solvers/derivatives.py @@ -8,7 +8,7 @@ import scipy.sparse as spa -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ================== 1d incident matrix ======================= @@ -32,7 +32,7 @@ def grad_1d_matrix(spl_kind, NbaseN): NbaseD = NbaseN - 1 + spl_kind - grad = np.zeros((NbaseD, NbaseN), dtype=float) + grad = xp.zeros((NbaseD, NbaseN), dtype=float) for i in range(NbaseD): grad[i, i] = -1.0 @@ -80,9 +80,9 @@ def discrete_derivatives_3d(space): grad_1d_3 = 0 * spa.identity(1, format="csr") else: if space.basis_tor == "r": - grad_1d_3 = 2 * np.pi * space.n_tor * spa.csr_matrix(np.array([[0.0, 1.0], [-1.0, 0.0]])) + grad_1d_3 = 2 * xp.pi * space.n_tor * spa.csr_matrix(xp.array([[0.0, 1.0], [-1.0, 0.0]])) else: - grad_1d_3 = 1j * 2 * np.pi * space.n_tor * spa.identity(1, format="csr") + grad_1d_3 = 1j * 2 * xp.pi * space.n_tor * spa.identity(1, format="csr") # standard tensor-product derivatives if space.ck == -1: diff --git a/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py b/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py index e60d565fb..b7d0d30d7 100644 --- a/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py +++ b/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py @@ -9,7 +9,7 @@ import struphy.eigenvalue_solvers.mass_matrices_1d as mass import struphy.eigenvalue_solvers.projectors_global as pro import struphy.eigenvalue_solvers.spline_space as spl -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # numerical solution of the general ideal MHD eigenvalue problem in a cylinder using 1d B-splines in radial direction @@ -21,7 +21,7 @@ def solve_ev_problem(rho, B_phi, dB_phi, B_z, p, gamma, a, k, m, num_params, bcZ r = lambda eta: a * eta # jacobian for integration - jac = lambda eta1: a * np.ones(eta1.shape, dtype=float) + jac = lambda eta1: a * xp.ones(eta1.shape, dtype=float) # ========================== kinetic energy functional ============================== # integrands (multiplied by -2/omega**2) @@ -46,11 +46,11 @@ def solve_ev_problem(rho, B_phi, dB_phi, B_z, p, gamma, a, k, m, num_params, bcZ # Bspline_A = Bsp.Bspline(splines.T, splines.p ) # Bspline_B = Bsp.Bspline(splines.t, splines.p - 1) # - # K_11_scipy = np.zeros((splines.NbaseN, splines.NbaseN), dtype=float) - # K_22_scipy = np.zeros((splines.NbaseD, splines.NbaseD), dtype=float) - # K_33_scipy = np.zeros((splines.NbaseD, splines.NbaseD), dtype=float) - # K_23_scipy = np.zeros((splines.NbaseD, splines.NbaseD), dtype=float) - # K_32_scipy = np.zeros((splines.NbaseD, splines.NbaseD), dtype=float) + # K_11_scipy = xp.zeros((splines.NbaseN, splines.NbaseN), dtype=float) + # K_22_scipy = xp.zeros((splines.NbaseD, splines.NbaseD), dtype=float) + # K_33_scipy = xp.zeros((splines.NbaseD, splines.NbaseD), dtype=float) + # K_23_scipy = xp.zeros((splines.NbaseD, splines.NbaseD), dtype=float) + # K_32_scipy = xp.zeros((splines.NbaseD, splines.NbaseD), dtype=float) # # for i in range(1, Bspline_A.N - 1): # for j in range(1, Bspline_A.N - 1): @@ -76,11 +76,11 @@ def solve_ev_problem(rho, B_phi, dB_phi, B_z, p, gamma, a, k, m, num_params, bcZ # integrand = lambda eta : a*K_ZV(eta)*Bspline_B(eta, i)*Bspline_B(eta, j) # K_32_scipy[i, j] = integrate.quad(integrand, 0., 1.)[0] - # assert np.allclose(K_11.toarray(), K_11_scipy[1:-1, 1:-1]) - # assert np.allclose(K_22.toarray(), K_22_scipy ) - # assert np.allclose(K_33.toarray(), K_33_scipy[bcZ:, bcZ:]) - # assert np.allclose(K_23.toarray(), K_23_scipy[ : , bcZ:]) - # assert np.allclose(K_32.toarray(), K_32_scipy[bcZ:, :]) + # assert xp.allclose(K_11.toarray(), K_11_scipy[1:-1, 1:-1]) + # assert xp.allclose(K_22.toarray(), K_22_scipy ) + # assert xp.allclose(K_33.toarray(), K_33_scipy[bcZ:, bcZ:]) + # assert xp.allclose(K_23.toarray(), K_23_scipy[ : , bcZ:]) + # assert xp.allclose(K_32.toarray(), K_32_scipy[bcZ:, :]) # ========================== potential energy functional =========================== # integrands (multiplied by 2) @@ -163,15 +163,15 @@ def solve_ev_problem(rho, B_phi, dB_phi, B_z, p, gamma, a, k, m, num_params, bcZ # return W_22 ## test correct computation - # W_11_scipy = np.zeros((splines.NbaseN, splines.NbaseN), dtype=float) - # W_22_scipy = np.zeros((splines.NbaseD, splines.NbaseD), dtype=float) - # W_33_scipy = np.zeros((splines.NbaseD, splines.NbaseD), dtype=float) - # W_12_scipy = np.zeros((splines.NbaseN, splines.NbaseD), dtype=float) - # W_21_scipy = np.zeros((splines.NbaseD, splines.NbaseN), dtype=float) - # W_13_scipy = np.zeros((splines.NbaseN, splines.NbaseD), dtype=float) - # W_31_scipy = np.zeros((splines.NbaseD, splines.NbaseN), dtype=float) - # W_23_scipy = np.zeros((splines.NbaseD, splines.NbaseD), dtype=float) - # W_32_scipy = np.zeros((splines.NbaseD, splines.NbaseD), dtype=float) + # W_11_scipy = xp.zeros((splines.NbaseN, splines.NbaseN), dtype=float) + # W_22_scipy = xp.zeros((splines.NbaseD, splines.NbaseD), dtype=float) + # W_33_scipy = xp.zeros((splines.NbaseD, splines.NbaseD), dtype=float) + # W_12_scipy = xp.zeros((splines.NbaseN, splines.NbaseD), dtype=float) + # W_21_scipy = xp.zeros((splines.NbaseD, splines.NbaseN), dtype=float) + # W_13_scipy = xp.zeros((splines.NbaseN, splines.NbaseD), dtype=float) + # W_31_scipy = xp.zeros((splines.NbaseD, splines.NbaseN), dtype=float) + # W_23_scipy = xp.zeros((splines.NbaseD, splines.NbaseD), dtype=float) + # W_32_scipy = xp.zeros((splines.NbaseD, splines.NbaseD), dtype=float) # # for i in range(1, Bspline_A.N - 1): # for j in range(1, Bspline_A.N - 1): @@ -187,15 +187,15 @@ def solve_ev_problem(rho, B_phi, dB_phi, B_z, p, gamma, a, k, m, num_params, bcZ # integrand = lambda eta : W_XdX(eta) * Bspline_A(eta, i, 0) * Bspline_A(eta, j, 1) # W_11_scipy[i, j] += integrate.quad(integrand, 0., 1.)[0] # - # assert np.allclose(W_11.toarray(), W_11_scipy[1:-1, 1:-1]) + # assert xp.allclose(W_11.toarray(), W_11_scipy[1:-1, 1:-1]) - # print(np.allclose(K, K.T)) - # print(np.allclose(W, W.T)) + # print(xp.allclose(K, K.T)) + # print(xp.allclose(W, W.T)) # solve eigenvalue problem omega**2*K*xi = W*xi - A = np.linalg.inv(K).dot(W) + A = xp.linalg.inv(K).dot(W) - omega2, XVZ_eig = np.linalg.eig(A) + omega2, XVZ_eig = xp.linalg.eig(A) # extract components X_eig = XVZ_eig[: (splines.NbaseN - 2), :] @@ -203,11 +203,11 @@ def solve_ev_problem(rho, B_phi, dB_phi, B_z, p, gamma, a, k, m, num_params, bcZ Z_eig = XVZ_eig[(splines.NbaseN - 2 + splines.NbaseD) :, :] # add boundary conditions X(0) = X(1) = 0 - X_eig = np.vstack((np.zeros(X_eig.shape[1], dtype=float), X_eig, np.zeros(X_eig.shape[1], dtype=float))) + X_eig = xp.vstack((xp.zeros(X_eig.shape[1], dtype=float), X_eig, xp.zeros(X_eig.shape[1], dtype=float))) # add boundary condition Z(0) = 0 if bcZ == 1: - Z_eig = np.vstack((np.zeros(Z_eig.shape[1], dtype=float), Z_eig)) + Z_eig = xp.vstack((xp.zeros(Z_eig.shape[1], dtype=float), Z_eig)) return omega2, X_eig, V_eig, Z_eig @@ -225,43 +225,43 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, # components of metric tensor and Jacobian determinant G_r = a**2 - G_phi = lambda eta: 4 * np.pi**2 * r(eta) ** 2 - G_z = 4 * np.pi**2 * R0**2 - J = lambda eta: 4 * np.pi**2 * R0 * a * r(eta) + G_phi = lambda eta: 4 * xp.pi**2 * r(eta) ** 2 + G_z = 4 * xp.pi**2 * R0**2 + J = lambda eta: 4 * xp.pi**2 * R0 * a * r(eta) # 2-from components of equilibrium magnetic field and its projection - B2_phi = lambda eta: 2 * np.pi * R0 * a * B_phi(r(eta)) - B2_z = lambda eta: 2 * np.pi * a * r(eta) * B_z(r(eta)) + B2_phi = lambda eta: 2 * xp.pi * R0 * a * B_phi(r(eta)) + B2_z = lambda eta: 2 * xp.pi * a * r(eta) * B_z(r(eta)) - b2_eq_phi = np.linalg.solve(proj.D.toarray(), proj.rhs_1(B2_phi)) - b2_eq_z = np.append(np.array([0.0]), np.linalg.solve(proj.D.toarray()[1:, 1:], proj.rhs_1(B2_z)[1:])) + b2_eq_phi = xp.linalg.solve(proj.D.toarray(), proj.rhs_1(B2_phi)) + b2_eq_z = xp.append(xp.array([0.0]), xp.linalg.solve(proj.D.toarray()[1:, 1:], proj.rhs_1(B2_z)[1:])) # 3-form components of equilibrium density and pessure and its projection Rho3 = lambda eta: J(eta) * Rho(r(eta)) P3 = lambda eta: J(eta) * P(r(eta)) - rho3_eq = np.append(np.array([0.0]), np.linalg.solve(proj.D.toarray()[1:, 1:], proj.rhs_1(Rho3)[1:])) - p3_eq = np.append(np.array([0.0]), np.linalg.solve(proj.D.toarray()[1:, 1:], proj.rhs_1(P3)[1:])) + rho3_eq = xp.append(xp.array([0.0]), xp.linalg.solve(proj.D.toarray()[1:, 1:], proj.rhs_1(Rho3)[1:])) + p3_eq = xp.append(xp.array([0.0]), xp.linalg.solve(proj.D.toarray()[1:, 1:], proj.rhs_1(P3)[1:])) # 2-form components of initial velocity and its projection U2_r = lambda eta: J(eta) * eta * (1 - eta) u2_r = proj.pi_0(U2_r) - u2_phi = -1 / (2 * np.pi * m) * GRAD_all.dot(u2_r) - u2_z = np.zeros(len(u2_phi), dtype=float) + u2_phi = -1 / (2 * xp.pi * m) * GRAD_all.dot(u2_r) + u2_z = xp.zeros(len(u2_phi), dtype=float) - b2_r = np.zeros(len(u2_r), dtype=float) - b2_phi = np.zeros(len(u2_phi), dtype=float) - b2_z = np.zeros(len(u2_z), dtype=float) + b2_r = xp.zeros(len(u2_r), dtype=float) + b2_phi = xp.zeros(len(u2_phi), dtype=float) + b2_z = xp.zeros(len(u2_z), dtype=float) - p3 = np.zeros(len(u2_z), dtype=float) + p3 = xp.zeros(len(u2_z), dtype=float) # projection matrices pi0_N_i, pi0_D_i, pi1_N_i, pi1_D_i = proj.projection_matrices_1d_reduced() pi0_NN_i, pi0_DN_i, pi0_ND_i, pi0_DD_i, pi1_NN_i, pi1_DN_i, pi1_ND_i, pi1_DD_i = proj.projection_matrices_1d() # 1D collocation matrices for interpolation in format (point, global basis function) - x_int = np.copy(proj.x_int) + x_int = xp.copy(proj.x_int) kind_splines = [False, True] @@ -270,22 +270,22 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, # 1D integration sub-intervals, quadrature points and weights if splines.p % 2 == 0: - x_his = np.union1d(x_int, splines.el_b) + x_his = xp.union1d(x_int, splines.el_b) else: - x_his = np.copy(x_int) + x_his = xp.copy(x_int) pts, wts = bsp.quadrature_grid(x_his, proj.pts_loc, proj.wts_loc) # compute number of sub-intervals for integrations (even degree) if splines.p % 2 == 0: - subs = 2 * np.ones(proj.pts.shape[0], dtype=int) + subs = 2 * xp.ones(proj.pts.shape[0], dtype=int) subs[: splines.p // 2] = 1 subs[-splines.p // 2 :] = 1 # compute number of sub-intervals for integrations (odd degree) else: - subs = np.ones(proj.pts.shape[0], dtype=int) + subs = xp.ones(proj.pts.shape[0], dtype=int) # evaluate basis functions on quadrature points in format (interval, local quad. point, global basis function) basis_his_N = bsp.collocation_matrix(splines.T, splines.p, pts.flatten(), False, normalize=kind_splines[0]).reshape( @@ -308,26 +308,26 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, M2_z = mass.get_M1(splines, mapping=lambda eta: J(eta) / G_z).toarray() # === matrices for curl of equilibrium field (with integration by parts) ========== - MB_12_eq = np.empty((splines.NbaseN, splines.NbaseD), dtype=float) - MB_13_eq = np.empty((splines.NbaseN, splines.NbaseD), dtype=float) + MB_12_eq = xp.empty((splines.NbaseN, splines.NbaseD), dtype=float) + MB_13_eq = xp.empty((splines.NbaseN, splines.NbaseD), dtype=float) - MB_21_eq = np.empty((splines.NbaseD, splines.NbaseN), dtype=float) - MB_31_eq = np.empty((splines.NbaseD, splines.NbaseN), dtype=float) + MB_21_eq = xp.empty((splines.NbaseD, splines.NbaseN), dtype=float) + MB_31_eq = xp.empty((splines.NbaseD, splines.NbaseN), dtype=float) - f_phi = np.linalg.inv(proj.N.toarray()).T.dot(GRAD_all.T.dot(M2_phi.dot(b2_eq_phi))) - f_z = np.linalg.inv(proj.N.toarray()).T.dot(GRAD_all.T.dot(M2_z.dot(b2_eq_z))) + f_phi = xp.linalg.inv(proj.N.toarray()).T.dot(GRAD_all.T.dot(M2_phi.dot(b2_eq_phi))) + f_z = xp.linalg.inv(proj.N.toarray()).T.dot(GRAD_all.T.dot(M2_z.dot(b2_eq_z))) - pi0_ND_phi = np.empty(pi0_ND_i[3].max() + 1, dtype=float) - pi0_ND_z = np.empty(pi0_ND_i[3].max() + 1, dtype=float) + pi0_ND_phi = xp.empty(pi0_ND_i[3].max() + 1, dtype=float) + pi0_ND_z = xp.empty(pi0_ND_i[3].max() + 1, dtype=float) - row_ND = np.empty(pi0_ND_i[3].max() + 1, dtype=int) - col_ND = np.empty(pi0_ND_i[3].max() + 1, dtype=int) + row_ND = xp.empty(pi0_ND_i[3].max() + 1, dtype=int) + col_ND = xp.empty(pi0_ND_i[3].max() + 1, dtype=int) - pi0_DN_phi = np.empty(pi0_DN_i[3].max() + 1, dtype=float) - pi0_DN_z = np.empty(pi0_DN_i[3].max() + 1, dtype=float) + pi0_DN_phi = xp.empty(pi0_DN_i[3].max() + 1, dtype=float) + pi0_DN_z = xp.empty(pi0_DN_i[3].max() + 1, dtype=float) - row_DN = np.empty(pi0_DN_i[3].max() + 1, dtype=int) - col_DN = np.empty(pi0_DN_i[3].max() + 1, dtype=int) + row_DN = xp.empty(pi0_DN_i[3].max() + 1, dtype=int) + col_DN = xp.empty(pi0_DN_i[3].max() + 1, dtype=int) ker.rhs0_f_1d(pi0_ND_i, basis_int_N, basis_int_D, 1 / J(x_int), f_phi, pi0_ND_phi, row_ND, col_ND) ker.rhs0_f_1d(pi0_ND_i, basis_int_N, basis_int_D, 1 / J(x_int), f_z, pi0_ND_z, row_ND, col_ND) @@ -348,23 +348,23 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, MB_31_eq[:, :] = -pi0_DN_z # === matrices for curl of equilibrium field (without integration by parts) ====== - MB_12_eq = np.empty((splines.NbaseN, splines.NbaseD), dtype=float) - MB_13_eq = np.empty((splines.NbaseN, splines.NbaseD), dtype=float) + MB_12_eq = xp.empty((splines.NbaseN, splines.NbaseD), dtype=float) + MB_13_eq = xp.empty((splines.NbaseN, splines.NbaseD), dtype=float) - MB_21_eq = np.empty((splines.NbaseD, splines.NbaseN), dtype=float) - MB_31_eq = np.empty((splines.NbaseD, splines.NbaseN), dtype=float) + MB_21_eq = xp.empty((splines.NbaseD, splines.NbaseN), dtype=float) + MB_31_eq = xp.empty((splines.NbaseD, splines.NbaseN), dtype=float) - cN = np.empty(splines.NbaseN, dtype=float) - cD = np.empty(splines.NbaseD, dtype=float) + cN = xp.empty(splines.NbaseN, dtype=float) + cD = xp.empty(splines.NbaseD, dtype=float) for j in range(splines.NbaseD): cD[:] = 0.0 cD[j] = 1.0 integrand2 = ( - lambda eta: splines.evaluate_D(eta, cD) / J(eta) * 2 * np.pi * a * (B_phi(r(eta)) + r(eta) * dB_phi(r(eta))) + lambda eta: splines.evaluate_D(eta, cD) / J(eta) * 2 * xp.pi * a * (B_phi(r(eta)) + r(eta) * dB_phi(r(eta))) ) - integrand3 = lambda eta: splines.evaluate_D(eta, cD) / J(eta) * 2 * np.pi * a * R0 * dB_z(r(eta)) + integrand3 = lambda eta: splines.evaluate_D(eta, cD) / J(eta) * 2 * xp.pi * a * R0 * dB_z(r(eta)) MB_12_eq[:, j] = inner.inner_prod_V0(splines, integrand2) MB_13_eq[:, j] = inner.inner_prod_V0(splines, integrand3) @@ -374,39 +374,39 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, cN[j] = 1.0 integrand2 = ( - lambda eta: splines.evaluate_N(eta, cN) / J(eta) * 2 * np.pi * a * (B_phi(r(eta)) + r(eta) * dB_phi(r(eta))) + lambda eta: splines.evaluate_N(eta, cN) / J(eta) * 2 * xp.pi * a * (B_phi(r(eta)) + r(eta) * dB_phi(r(eta))) ) - integrand3 = lambda eta: splines.evaluate_N(eta, cN) / J(eta) * 2 * np.pi * a * R0 * dB_z(r(eta)) + integrand3 = lambda eta: splines.evaluate_N(eta, cN) / J(eta) * 2 * xp.pi * a * R0 * dB_z(r(eta)) MB_21_eq[:, j] = inner.inner_prod_V1(splines, integrand2) MB_31_eq[:, j] = inner.inner_prod_V1(splines, integrand3) # ===== right-hand sides of projection matrices =============== - rhs0_N_phi = np.empty(pi0_N_i[0].size, dtype=float) - rhs0_N_z = np.empty(pi0_N_i[0].size, dtype=float) + rhs0_N_phi = xp.empty(pi0_N_i[0].size, dtype=float) + rhs0_N_z = xp.empty(pi0_N_i[0].size, dtype=float) - rhs1_D_phi = np.empty(pi1_D_i[0].size, dtype=float) - rhs1_D_z = np.empty(pi1_D_i[0].size, dtype=float) + rhs1_D_phi = xp.empty(pi1_D_i[0].size, dtype=float) + rhs1_D_z = xp.empty(pi1_D_i[0].size, dtype=float) - rhs0_N_pr = np.empty(pi0_N_i[0].size, dtype=float) - rhs1_D_pr = np.empty(pi1_D_i[0].size, dtype=float) + rhs0_N_pr = xp.empty(pi0_N_i[0].size, dtype=float) + rhs1_D_pr = xp.empty(pi1_D_i[0].size, dtype=float) - rhs0_N_rho = np.empty(pi0_N_i[0].size, dtype=float) - rhs1_D_rho = np.empty(pi1_D_i[0].size, dtype=float) + rhs0_N_rho = xp.empty(pi0_N_i[0].size, dtype=float) + rhs1_D_rho = xp.empty(pi1_D_i[0].size, dtype=float) # ker.rhs0_1d(pi0_N_i[0], pi0_N_i[1], basis_int_N, splines.evaluate_D(x_int, b2_eq_phi)/J(x_int), rhs0_N_phi) # ker.rhs0_1d(pi0_N_i[0], pi0_N_i[1], basis_int_N, splines.evaluate_D(x_int, b2_eq_z )/J(x_int), rhs0_N_z ) # - # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, np.append(0, np.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), b2_eq_z )/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_z) - # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, np.append(0, np.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), b2_eq_phi)/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_phi) + # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), b2_eq_z )/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_z) + # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), b2_eq_phi)/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_phi) # # ker.rhs0_1d(pi0_N_i[0], pi0_N_i[1], basis_int_N, splines.evaluate_D(x_int, p3_eq)/J(x_int), rhs0_N_pr) - # temp = np.empty(pi0_N_i[0].size, dtype=float) + # temp = xp.empty(pi0_N_i[0].size, dtype=float) # temp[:] = rhs0_N_pr - # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, np.append(0, np.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), p3_eq)/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_pr) + # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), p3_eq)/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_pr) # # ker.rhs0_1d(pi0_N_i[0], pi0_N_i[1], basis_int_N, splines.evaluate_D(x_int, rho3)/J(x_int), rhs0_N_rho) - # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, np.append(0, np.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), rho3)/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_rho) + # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), rho3)/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_rho) ker.rhs0_1d(pi0_N_i[0], pi0_N_i[1], basis_int_N, B2_phi(x_int) / J(x_int), rhs0_N_phi) ker.rhs0_1d(pi0_N_i[0], pi0_N_i[1], basis_int_N, B2_z(x_int) / J(x_int), rhs0_N_z) @@ -415,7 +415,7 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, pi1_D_i[0], pi1_D_i[1], subs, - np.append(0, np.cumsum(subs - 1)[:-1]), + xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, (B2_phi(pts.flatten()) / J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), @@ -426,20 +426,20 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, pi1_D_i[0], pi1_D_i[1], subs, - np.append(0, np.cumsum(subs - 1)[:-1]), + xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, (B2_z(pts.flatten()) / J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_z, ) - # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, np.append(0, np.cumsum(subs - 1)[:-1]), wts, basis_his_D, np.ones(pts.shape, dtype=float), rhs1_D_z) + # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, xp.ones(pts.shape, dtype=float), rhs1_D_z) ker.rhs0_1d(pi0_N_i[0], pi0_N_i[1], basis_int_N, P3(x_int) / J(x_int), rhs0_N_pr) ker.rhs1_1d( pi1_D_i[0], pi1_D_i[1], subs, - np.append(0, np.cumsum(subs - 1)[:-1]), + xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, (P3(pts.flatten()) / J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), @@ -451,7 +451,7 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, pi1_D_i[0], pi1_D_i[1], subs, - np.append(0, np.cumsum(subs - 1)[:-1]), + xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, (Rho3(pts.flatten()) / J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), @@ -480,38 +480,38 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, (rhs1_D_rho, (pi1_D_i[0], pi1_D_i[1])), shape=(splines.NbaseD, splines.NbaseD) ).toarray() - pi0_N_phi = np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_phi[1:-1, 1:-1]) - pi0_N_z = np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_z[1:-1, 1:-1]) + pi0_N_phi = xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_phi[1:-1, 1:-1]) + pi0_N_z = xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_z[1:-1, 1:-1]) - pi1_D_phi = np.linalg.inv(proj.D.toarray()).dot(rhs1_D_phi) - pi1_D_z = np.linalg.inv(proj.D.toarray()).dot(rhs1_D_z) + pi1_D_phi = xp.linalg.inv(proj.D.toarray()).dot(rhs1_D_phi) + pi1_D_z = xp.linalg.inv(proj.D.toarray()).dot(rhs1_D_z) - pi0_N_pr = np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_pr[1:-1, 1:-1]) - pi1_D_pr = np.linalg.inv(proj.D.toarray()).dot(rhs1_D_pr) + pi0_N_pr = xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_pr[1:-1, 1:-1]) + pi1_D_pr = xp.linalg.inv(proj.D.toarray()).dot(rhs1_D_pr) - pi0_N_rho = np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_rho[1:-1, 1:-1]) - pi1_D_rho = np.linalg.inv(proj.D.toarray()).dot(rhs1_D_rho) + pi0_N_rho = xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_rho[1:-1, 1:-1]) + pi1_D_rho = xp.linalg.inv(proj.D.toarray()).dot(rhs1_D_rho) # ======= matrices in strong induction equation ================ # 11 block - I_11 = -2 * np.pi * m * pi0_N_phi - 2 * np.pi * n * pi0_N_z + I_11 = -2 * xp.pi * m * pi0_N_phi - 2 * xp.pi * n * pi0_N_z # 21 block and 31 block I_21 = -GRAD.dot(pi0_N_phi) I_31 = -GRAD.dot(pi0_N_z) # 22 block and 32 block - I_22 = 2 * np.pi * n * pi1_D_z - I_32 = -2 * np.pi * m * pi1_D_z + I_22 = 2 * xp.pi * n * pi1_D_z + I_32 = -2 * xp.pi * m * pi1_D_z # 23 block and 33 block - I_23 = -2 * np.pi * n * pi1_D_phi - I_33 = 2 * np.pi * m * pi1_D_phi + I_23 = -2 * xp.pi * n * pi1_D_phi + I_33 = 2 * xp.pi * m * pi1_D_phi # total - I_all = np.block( + I_all = xp.block( [ - [I_11, np.zeros((len(u2_r) - 2, len(u2_phi))), np.zeros((len(u2_r) - 2, len(u2_z) - 1))], + [I_11, xp.zeros((len(u2_r) - 2, len(u2_phi))), xp.zeros((len(u2_r) - 2, len(u2_z) - 1))], [I_21, I_22, I_23[:, 1:]], [I_31[1:, :], I_32[1:, :], I_33[1:, 1:]], ] @@ -519,97 +519,97 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, # ======= matrices in strong pressure equation ================ P_1 = -GRAD.dot(pi0_N_pr) - (gamma - 1) * pi1_D_pr.dot(GRAD) - P_2 = -2 * np.pi * m * gamma * pi1_D_pr - P_3 = -2 * np.pi * n * gamma * pi1_D_pr + P_2 = -2 * xp.pi * m * gamma * pi1_D_pr + P_3 = -2 * xp.pi * n * gamma * pi1_D_pr - P_all = np.block([[P_1[1:, :], P_2[1:, :], P_3[1:, 1:]]]) + P_all = xp.block([[P_1[1:, :], P_2[1:, :], P_3[1:, 1:]]]) # ========== matrices in weak momentum balance equation ====== A_1 = 1 / 2 * (pi0_N_rho.T.dot(M2_r) + M2_r.dot(pi0_N_rho)) A_2 = 1 / 2 * (pi1_D_rho.T.dot(M2_phi) + M2_phi.dot(pi1_D_rho)) A_3 = 1 / 2 * (pi1_D_rho.T.dot(M2_z) + M2_z.dot(pi1_D_rho))[:, :] - A_all = np.block( + A_all = xp.block( [ - [A_1, np.zeros((A_1.shape[0], A_2.shape[1])), np.zeros((A_1.shape[0], A_3.shape[1]))], - [np.zeros((A_2.shape[0], A_1.shape[1])), A_2, np.zeros((A_2.shape[0], A_3.shape[1]))], - [np.zeros((A_3.shape[0], A_1.shape[1])), np.zeros((A_3.shape[0], A_2.shape[1])), A_3], + [A_1, xp.zeros((A_1.shape[0], A_2.shape[1])), xp.zeros((A_1.shape[0], A_3.shape[1]))], + [xp.zeros((A_2.shape[0], A_1.shape[1])), A_2, xp.zeros((A_2.shape[0], A_3.shape[1]))], + [xp.zeros((A_3.shape[0], A_1.shape[1])), xp.zeros((A_3.shape[0], A_2.shape[1])), A_3], ] ) - MB_11 = 2 * np.pi * n * pi0_N_z.T.dot(M2_r) + 2 * np.pi * m * pi0_N_phi.T.dot(M2_r) + MB_11 = 2 * xp.pi * n * pi0_N_z.T.dot(M2_r) + 2 * xp.pi * m * pi0_N_phi.T.dot(M2_r) MB_12 = pi0_N_phi.T.dot(GRAD.T.dot(M2_phi)) - MB_12_eq[1:-1, :] MB_13 = pi0_N_z.T.dot(GRAD.T.dot(M2_z)) - MB_13_eq[1:-1, :] MB_14 = GRAD.T.dot(M3) MB_21 = MB_21_eq[:, 1:-1] - MB_22 = -2 * np.pi * n * pi1_D_z.T.dot(M2_phi) - MB_23 = 2 * np.pi * m * pi1_D_z.T.dot(M2_z) - MB_24 = 2 * np.pi * m * M3 + MB_22 = -2 * xp.pi * n * pi1_D_z.T.dot(M2_phi) + MB_23 = 2 * xp.pi * m * pi1_D_z.T.dot(M2_z) + MB_24 = 2 * xp.pi * m * M3 MB_31 = MB_31_eq[:, 1:-1] - MB_32 = 2 * np.pi * n * pi1_D_phi.T.dot(M2_phi) - MB_33 = -2 * np.pi * m * pi1_D_phi.T.dot(M2_z) - MB_34 = 2 * np.pi * n * M3 + MB_32 = 2 * xp.pi * n * pi1_D_phi.T.dot(M2_phi) + MB_33 = -2 * xp.pi * m * pi1_D_phi.T.dot(M2_z) + MB_34 = 2 * xp.pi * n * M3 - MB_b_all = np.block( + MB_b_all = xp.block( [[MB_11, MB_12, MB_13[:, 1:]], [MB_21, MB_22, MB_23[:, 1:]], [MB_31[1:, :], MB_32[1:, :], MB_33[1:, 1:]]] ) - MB_p_all = np.block([[MB_14[:, 1:]], [MB_24[:, 1:]], [MB_34[1:, 1:]]]) + MB_p_all = xp.block([[MB_14[:, 1:]], [MB_24[:, 1:]], [MB_34[1:, 1:]]]) ## ======= matrices in strong induction equation ================ ## 11 block - # I_11 = np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(-2*np.pi*m*rhs0_N_phi[1:-1, 1:-1] - 2*np.pi*n*rhs0_N_z[1:-1, 1:-1]) + # I_11 = xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(-2*xp.pi*m*rhs0_N_phi[1:-1, 1:-1] - 2*xp.pi*n*rhs0_N_z[1:-1, 1:-1]) # ## 21 block and 31 block - # I_21 = -GRAD[: , 1:-1].dot(np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_phi[1:-1, 1:-1])) - # I_31 = -GRAD[1:, 1:-1].dot(np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_z[1:-1, 1:-1])) + # I_21 = -GRAD[: , 1:-1].dot(xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_phi[1:-1, 1:-1])) + # I_31 = -GRAD[1:, 1:-1].dot(xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_z[1:-1, 1:-1])) # ## 22 block and 32 block - # I_22 = 2*np.pi*n*np.linalg.inv(proj.D.toarray()[ :, :]).dot(rhs1_D_z[ :, :]) - # I_32 = -2*np.pi*m*np.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_z[1:, :]) + # I_22 = 2*xp.pi*n*xp.linalg.inv(proj.D.toarray()[ :, :]).dot(rhs1_D_z[ :, :]) + # I_32 = -2*xp.pi*m*xp.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_z[1:, :]) # ## 23 block and 33 block - # I_23 = -2*np.pi*n*np.linalg.inv(proj.D.toarray()[ :, :]).dot(rhs1_D_phi[ :, 1:]) - # I_33 = 2*np.pi*m*np.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_phi[1:, 1:]) + # I_23 = -2*xp.pi*n*xp.linalg.inv(proj.D.toarray()[ :, :]).dot(rhs1_D_phi[ :, 1:]) + # I_33 = 2*xp.pi*m*xp.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_phi[1:, 1:]) # # ## ======= matrices in strong pressure equation ================ - # P_1 = -GRAD[1:, 1:-1].dot(np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_pr[1:-1, 1:-1])) - (gamma - 1)*np.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_pr[1:, :].dot(GRAD[:, 1:-1])) - # P_2 = -2*np.pi*m*gamma*np.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_pr[1:, :]) - # P_3 = -2*np.pi*n*gamma*np.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_pr[1:, 1:]) + # P_1 = -GRAD[1:, 1:-1].dot(xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_pr[1:-1, 1:-1])) - (gamma - 1)*xp.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_pr[1:, :].dot(GRAD[:, 1:-1])) + # P_2 = -2*xp.pi*m*gamma*xp.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_pr[1:, :]) + # P_3 = -2*xp.pi*n*gamma*xp.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_pr[1:, 1:]) # # ## ========== matrices in weak momentum balance equation ====== - # rhs0_N_rho = np.empty(pi0_N_i[0].size, dtype=float) + # rhs0_N_rho = xp.empty(pi0_N_i[0].size, dtype=float) # ker.rhs0_1d(pi0_N_i[0], pi0_N_i[1], basis_int_N, splines.evaluate_D(x_int, rho3)/J(x_int), rhs0_N_rho) # # - # rhs1_D_rho = np.empty(pi1_D_i[0].size, dtype=float) - # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, np.append(0, np.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), rho3)/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_rho) + # rhs1_D_rho = xp.empty(pi1_D_i[0].size, dtype=float) + # ker.rhs1_1d(pi1_D_i[0], pi1_D_i[1], subs, xp.append(0, xp.cumsum(subs - 1)[:-1]), wts, basis_his_D, (splines.evaluate_D(pts.flatten(), rho3)/J(pts.flatten())).reshape(pts.shape[0], pts.shape[1]), rhs1_D_rho) # # # - # A_1 = 1/2*(rhs0_N_rho[1:-1, 1:-1].T.dot(np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(M2_r[1:-1, 1:-1])) + M2_r[1:-1, 1:-1].dot(np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_rho[1:-1, 1:-1]))) - # A_2 = 1/2*(rhs1_D_rho.T.dot(np.linalg.inv(proj.D.toarray()[:, :]).T.dot(M2_phi)) + M2_phi.dot(np.linalg.inv(proj.D.toarray()[:, :]).dot(rhs1_D_rho))) - # A_3 = 1/2*(rhs1_D_rho[1:, 1:].T.dot(np.linalg.inv(proj.D.toarray()[1:, 1:]).T.dot(M2_z[1:, 1:])) + M2_z[1:, 1:].dot(np.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_rho[1:, 1:]))) + # A_1 = 1/2*(rhs0_N_rho[1:-1, 1:-1].T.dot(xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(M2_r[1:-1, 1:-1])) + M2_r[1:-1, 1:-1].dot(xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_rho[1:-1, 1:-1]))) + # A_2 = 1/2*(rhs1_D_rho.T.dot(xp.linalg.inv(proj.D.toarray()[:, :]).T.dot(M2_phi)) + M2_phi.dot(xp.linalg.inv(proj.D.toarray()[:, :]).dot(rhs1_D_rho))) + # A_3 = 1/2*(rhs1_D_rho[1:, 1:].T.dot(xp.linalg.inv(proj.D.toarray()[1:, 1:]).T.dot(M2_z[1:, 1:])) + M2_z[1:, 1:].dot(xp.linalg.inv(proj.D.toarray()[1:, 1:]).dot(rhs1_D_rho[1:, 1:]))) # # - # MB_11 = 2*np.pi*n*rhs0_N_z[1:-1, 1:-1].T.dot(np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(M2_r[1:-1, 1:-1])) + 2*np.pi*m*rhs0_N_phi[1:-1, 1:-1].T.dot(np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(M2_r[1:-1, 1:-1])) + # MB_11 = 2*xp.pi*n*rhs0_N_z[1:-1, 1:-1].T.dot(xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(M2_r[1:-1, 1:-1])) + 2*xp.pi*m*rhs0_N_phi[1:-1, 1:-1].T.dot(xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(M2_r[1:-1, 1:-1])) # - # MB_12 = rhs0_N_phi[1:-1, 1:-1].T.dot(np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(GRAD[:, 1:-1].T.dot(M2_phi))) - # MB_13 = rhs0_N_z[1:-1, 1:-1].T.dot(np.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(GRAD[1:, 1:-1].T.dot(M2_z[1:, 1:]))) + # MB_12 = rhs0_N_phi[1:-1, 1:-1].T.dot(xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(GRAD[:, 1:-1].T.dot(M2_phi))) + # MB_13 = rhs0_N_z[1:-1, 1:-1].T.dot(xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).T.dot(GRAD[1:, 1:-1].T.dot(M2_z[1:, 1:]))) # # MB_14 = GRAD[1:, 1:-1].T.dot(M3[1:, 1:]) # # - # MB_22 = -2*np.pi*n*rhs1_D_z.T.dot(np.linalg.inv(proj.D.toarray()).T.dot(M2_phi)) - # MB_23 = 2*np.pi*m*rhs1_D_z[1:, :].T.dot(np.linalg.inv(proj.D.toarray()[1:, 1:]).T.dot(M2_z[1:, 1:])) - # MB_24 = 2*np.pi*m*M3[ :, 1:] + # MB_22 = -2*xp.pi*n*rhs1_D_z.T.dot(xp.linalg.inv(proj.D.toarray()).T.dot(M2_phi)) + # MB_23 = 2*xp.pi*m*rhs1_D_z[1:, :].T.dot(xp.linalg.inv(proj.D.toarray()[1:, 1:]).T.dot(M2_z[1:, 1:])) + # MB_24 = 2*xp.pi*m*M3[ :, 1:] # - # MB_32 = 2*np.pi*n*rhs1_D_phi[:, 1:].T.dot(np.linalg.inv(proj.D.toarray()).T.dot(M2_phi)) - # MB_33 = -2*np.pi*m*rhs1_D_phi[1:, 1:].T.dot(np.linalg.inv(proj.D.toarray()[1:, 1:]).T.dot(M2_z[1:, 1:])) - # MB_34 = 2*np.pi*n*M3[1:, 1:] + # MB_32 = 2*xp.pi*n*rhs1_D_phi[:, 1:].T.dot(xp.linalg.inv(proj.D.toarray()).T.dot(M2_phi)) + # MB_33 = -2*xp.pi*m*rhs1_D_phi[1:, 1:].T.dot(xp.linalg.inv(proj.D.toarray()[1:, 1:]).T.dot(M2_z[1:, 1:])) + # MB_34 = 2*xp.pi*n*M3[1:, 1:] # # # ==== matrices in eigenvalue problem ======== @@ -625,17 +625,17 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, W_32 = MB_32.dot(I_22) + MB_33.dot(I_32) + MB_34.dot(P_2) W_33 = MB_32.dot(I_23) + MB_33.dot(I_33) + MB_34.dot(P_3) - # W = np.block([[W_11, W_12, W_13[:, 1:]], [W_21, W_22, W_23[:, 1:]], [W_31[1:, :], W_32[1:, :], W_33[1:, 1:]]]) - W = np.block([[W_11, W_12, W_13[:, :]], [W_21, W_22, W_23[:, :]], [W_31[:, :], W_32[:, :], W_33[:, :]]]) + # W = xp.block([[W_11, W_12, W_13[:, 1:]], [W_21, W_22, W_23[:, 1:]], [W_31[1:, :], W_32[1:, :], W_33[1:, 1:]]]) + W = xp.block([[W_11, W_12, W_13[:, :]], [W_21, W_22, W_23[:, :]], [W_31[:, :], W_32[:, :], W_33[:, :]]]) - # print(np.allclose(K, K.T)) - # print(np.allclose(W, W.T)) + # print(xp.allclose(K, K.T)) + # print(xp.allclose(W, W.T)) # solve eigenvalue problem omega**2*K*xi = W*xi - MAT = np.linalg.inv(-A_all).dot(W) + MAT = xp.linalg.inv(-A_all).dot(W) - omega2, XYZ_eig = np.linalg.eig(MAT) - # omega2, XYZ_eig = np.linalg.eig(np.linalg.inv(-A_all).dot(MB_b_all.dot(I_all) + MB_p_all.dot(P_all))) + omega2, XYZ_eig = xp.linalg.eig(MAT) + # omega2, XYZ_eig = xp.linalg.eig(xp.linalg.inv(-A_all).dot(MB_b_all.dot(I_all) + MB_p_all.dot(P_all))) # extract components X_eig = XYZ_eig[: (splines.NbaseN - 2), :] @@ -643,35 +643,35 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, Z_eig = XYZ_eig[(splines.NbaseN - 2 + splines.NbaseD) :, :] # add boundary conditions X(0) = X(1) = 0 - X_eig = np.vstack((np.zeros(X_eig.shape[1], dtype=float), X_eig, np.zeros(X_eig.shape[1], dtype=float))) + X_eig = xp.vstack((xp.zeros(X_eig.shape[1], dtype=float), X_eig, xp.zeros(X_eig.shape[1], dtype=float))) # add boundary condition Z(0) = 0 - Z_eig = np.vstack((np.zeros(Z_eig.shape[1], dtype=float), Z_eig)) + Z_eig = xp.vstack((xp.zeros(Z_eig.shape[1], dtype=float), Z_eig)) return omega2, X_eig, Y_eig, Z_eig ## ========== matrices in initial value problem === - LHS = np.block( + LHS = xp.block( [ - [A_all, np.zeros((A_all.shape[0], A_all.shape[1])), np.zeros((A_all.shape[0], len(p3) - 1))], + [A_all, xp.zeros((A_all.shape[0], A_all.shape[1])), xp.zeros((A_all.shape[0], len(p3) - 1))], [ - np.zeros((A_all.shape[0], A_all.shape[1])), - np.identity(A_all.shape[0]), - np.zeros((A_all.shape[0], len(p3) - 1)), + xp.zeros((A_all.shape[0], A_all.shape[1])), + xp.identity(A_all.shape[0]), + xp.zeros((A_all.shape[0], len(p3) - 1)), ], [ - np.zeros((len(p3) - 1, A_all.shape[1])), - np.zeros((len(p3) - 1, A_all.shape[1])), - np.identity(len(p3) - 1), + xp.zeros((len(p3) - 1, A_all.shape[1])), + xp.zeros((len(p3) - 1, A_all.shape[1])), + xp.identity(len(p3) - 1), ], ] ) - RHS = np.block( + RHS = xp.block( [ - [np.zeros((MB_b_all.shape[0], I_all.shape[1])), MB_b_all, MB_p_all], - [I_all, np.zeros((I_all.shape[0], MB_b_all.shape[1])), np.zeros((I_all.shape[0], MB_p_all.shape[1]))], - [P_all, np.zeros((P_all.shape[0], MB_b_all.shape[1])), np.zeros((P_all.shape[0], MB_p_all.shape[1]))], + [xp.zeros((MB_b_all.shape[0], I_all.shape[1])), MB_b_all, MB_p_all], + [I_all, xp.zeros((I_all.shape[0], MB_b_all.shape[1])), xp.zeros((I_all.shape[0], MB_p_all.shape[1]))], + [P_all, xp.zeros((P_all.shape[0], MB_b_all.shape[1])), xp.zeros((P_all.shape[0], MB_p_all.shape[1]))], ] ) @@ -679,35 +679,35 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, T = 200.0 Nt = int(T / dt) - UPDATE = np.linalg.inv(LHS - dt / 2 * RHS).dot(LHS + dt / 2 * RHS) - ##UPDATE = np.linalg.inv(LHS).dot(LHS + dt*RHS) + UPDATE = xp.linalg.inv(LHS - dt / 2 * RHS).dot(LHS + dt / 2 * RHS) + ##UPDATE = xp.linalg.inv(LHS).dot(LHS + dt*RHS) # - # lambdas, eig_vecs = np.linalg.eig(UPDATE) + # lambdas, eig_vecs = xp.linalg.eig(UPDATE) # return lambdas # # return lambdas # - u2_r_all = np.zeros((Nt + 1, len(u2_r)), dtype=float) - u2_phi_all = np.zeros((Nt + 1, len(u2_phi)), dtype=float) - u2_z_all = np.zeros((Nt + 1, len(u2_z)), dtype=float) + u2_r_all = xp.zeros((Nt + 1, len(u2_r)), dtype=float) + u2_phi_all = xp.zeros((Nt + 1, len(u2_phi)), dtype=float) + u2_z_all = xp.zeros((Nt + 1, len(u2_z)), dtype=float) - b2_r_all = np.zeros((Nt + 1, len(b2_r)), dtype=float) - b2_phi_all = np.zeros((Nt + 1, len(b2_phi)), dtype=float) - b2_z_all = np.zeros((Nt + 1, len(b2_z)), dtype=float) + b2_r_all = xp.zeros((Nt + 1, len(b2_r)), dtype=float) + b2_phi_all = xp.zeros((Nt + 1, len(b2_phi)), dtype=float) + b2_z_all = xp.zeros((Nt + 1, len(b2_z)), dtype=float) - p3_all = np.zeros((Nt + 1, len(p3)), dtype=float) + p3_all = xp.zeros((Nt + 1, len(p3)), dtype=float) # initialization # u2_r_all[0, :] = u2_r # u2_phi_all[0, :] = u2_phi - u2_r_all[0, 1:-1] = np.random.rand(len(u2_r) - 2) - p3_all[0, 1:] = np.random.rand(len(p3) - 1) + u2_r_all[0, 1:-1] = xp.random.rand(len(u2_r) - 2) + p3_all[0, 1:] = xp.random.rand(len(p3) - 1) # time integration for n in range(Nt): - old = np.concatenate( + old = xp.concatenate( ( u2_r_all[n, 1:-1], u2_phi_all[n, :], @@ -721,18 +721,18 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, new = UPDATE.dot(old) # extract components - unew, bnew, pnew = np.split( + unew, bnew, pnew = xp.split( new, [len(u2_r) - 2 + len(u2_phi) + len(u2_z) - 1, 2 * (len(u2_r) - 2 + len(u2_phi) + len(u2_z) - 1)] ) - u2_r_all[n + 1, :] = np.array([0.0] + list(unew[: (splines.NbaseN - 2)]) + [0.0]) + u2_r_all[n + 1, :] = xp.array([0.0] + list(unew[: (splines.NbaseN - 2)]) + [0.0]) u2_phi_all[n + 1, :] = unew[(splines.NbaseN - 2) : (splines.NbaseN - 2 + splines.NbaseD)] - u2_z_all[n + 1, :] = np.array([0.0] + list(unew[(splines.NbaseN - 2 + splines.NbaseD) :])) + u2_z_all[n + 1, :] = xp.array([0.0] + list(unew[(splines.NbaseN - 2 + splines.NbaseD) :])) - b2_r_all[n + 1, :] = np.array([0.0] + list(bnew[: (splines.NbaseN - 2)]) + [0.0]) + b2_r_all[n + 1, :] = xp.array([0.0] + list(bnew[: (splines.NbaseN - 2)]) + [0.0]) b2_phi_all[n + 1, :] = bnew[(splines.NbaseN - 2) : (splines.NbaseN - 2 + splines.NbaseD)] - b2_z_all[n + 1, :] = np.array([0.0] + list(bnew[(splines.NbaseN - 2 + splines.NbaseD) :])) + b2_z_all[n + 1, :] = xp.array([0.0] + list(bnew[(splines.NbaseN - 2 + splines.NbaseD) :])) - p3_all[n + 1, :] = np.array([0.0] + list(pnew)) + p3_all[n + 1, :] = xp.array([0.0] + list(pnew)) return u2_r_all, u2_phi_all, u2_z_all, b2_r_all, b2_phi_all, b2_z_all, p3_all, omega2 diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py index 37940a8d6..448298451 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py @@ -10,7 +10,7 @@ import struphy.feec.basics.kernels_3d as ker import struphy.feec.control_variates.kernels_control_variate as ker_cv -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class terms_control_variate: @@ -40,7 +40,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): kind_fun_eq = [11, 12, 13, 14] # ========= evaluation of DF^(-1) * jh_eq_phys * |det(DF)| at quadrature points ========= - self.mat_jh1 = np.empty( + self.mat_jh1 = xp.empty( ( self.space.Nel[0], self.space.n_quad[0], @@ -51,7 +51,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ), dtype=float, ) - self.mat_jh2 = np.empty( + self.mat_jh2 = xp.empty( ( self.space.Nel[0], self.space.n_quad[0], @@ -62,7 +62,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ), dtype=float, ) - self.mat_jh3 = np.empty( + self.mat_jh3 = xp.empty( ( self.space.Nel[0], self.space.n_quad[0], @@ -133,7 +133,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ) # ========= evaluation of nh_eq_phys * |det(DF)| at quadrature points =================== - self.mat_nh = np.empty( + self.mat_nh = xp.empty( ( self.space.Nel[0], self.space.n_quad[0], @@ -166,7 +166,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ) # =========== 2-form magnetic field at quadrature points ================================= - self.B2_1 = np.empty( + self.B2_1 = xp.empty( ( self.space.Nel[0], self.space.n_quad[0], @@ -177,7 +177,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ), dtype=float, ) - self.B2_2 = np.empty( + self.B2_2 = xp.empty( ( self.space.Nel[0], self.space.n_quad[0], @@ -188,7 +188,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ), dtype=float, ) - self.B2_3 = np.empty( + self.B2_3 = xp.empty( ( self.space.Nel[0], self.space.n_quad[0], @@ -202,7 +202,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): # ================== correction matrices in step 1 ======================== if self.basis_u == 0: - self.M12 = np.empty( + self.M12 = xp.empty( ( self.space.NbaseN[0], self.space.NbaseN[1], @@ -213,7 +213,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ), dtype=float, ) - self.M13 = np.empty( + self.M13 = xp.empty( ( self.space.NbaseN[0], self.space.NbaseN[1], @@ -224,7 +224,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ), dtype=float, ) - self.M23 = np.empty( + self.M23 = xp.empty( ( self.space.NbaseN[0], self.space.NbaseN[1], @@ -237,7 +237,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ) elif self.basis_u == 2: - self.M12 = np.empty( + self.M12 = xp.empty( ( self.space.NbaseN[0], self.space.NbaseD[1], @@ -248,7 +248,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ), dtype=float, ) - self.M13 = np.empty( + self.M13 = xp.empty( ( self.space.NbaseN[0], self.space.NbaseD[1], @@ -259,7 +259,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u): ), dtype=float, ) - self.M23 = np.empty( + self.M23 = xp.empty( ( self.space.NbaseD[0], self.space.NbaseN[1], @@ -273,14 +273,14 @@ def __init__(self, tensor_space_FEM, domain, basis_u): # ==================== correction vectors in step 3 ======================= if self.basis_u == 0: - self.F1 = np.empty((self.space.NbaseN[0], self.space.NbaseN[1], self.space.NbaseN[2]), dtype=float) - self.F2 = np.empty((self.space.NbaseN[0], self.space.NbaseN[1], self.space.NbaseN[2]), dtype=float) - self.F3 = np.empty((self.space.NbaseN[0], self.space.NbaseN[1], self.space.NbaseN[2]), dtype=float) + self.F1 = xp.empty((self.space.NbaseN[0], self.space.NbaseN[1], self.space.NbaseN[2]), dtype=float) + self.F2 = xp.empty((self.space.NbaseN[0], self.space.NbaseN[1], self.space.NbaseN[2]), dtype=float) + self.F3 = xp.empty((self.space.NbaseN[0], self.space.NbaseN[1], self.space.NbaseN[2]), dtype=float) elif self.basis_u == 2: - self.F1 = np.empty((self.space.NbaseN[0], self.space.NbaseD[1], self.space.NbaseD[2]), dtype=float) - self.F2 = np.empty((self.space.NbaseD[0], self.space.NbaseN[1], self.space.NbaseD[2]), dtype=float) - self.F3 = np.empty((self.space.NbaseD[0], self.space.NbaseD[1], self.space.NbaseN[2]), dtype=float) + self.F1 = xp.empty((self.space.NbaseN[0], self.space.NbaseD[1], self.space.NbaseD[2]), dtype=float) + self.F2 = xp.empty((self.space.NbaseD[0], self.space.NbaseN[1], self.space.NbaseD[2]), dtype=float) + self.F3 = xp.empty((self.space.NbaseD[0], self.space.NbaseD[1], self.space.NbaseN[2]), dtype=float) # ===== inner product in V0^3 resp. V2 of (B x jh_eq) - term ========== def inner_prod_jh_eq(self, b1, b2, b3): @@ -511,7 +511,7 @@ def inner_prod_jh_eq(self, b1, b2, b3): self.B2_1 * self.mat_jh2 - self.B2_2 * self.mat_jh1, ) - return np.concatenate((self.F1.flatten(), self.F2.flatten(), self.F3.flatten())) + return xp.concatenate((self.F1.flatten(), self.F2.flatten(), self.F3.flatten())) # ===== mass matrix in V0^3 resp. V2 of -(rhoh_eq * (B x U)) - term ======= def mass_nh_eq(self, b1, b2, b3): diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py index e39a463ab..d4713c12a 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py @@ -3,7 +3,7 @@ import struphy.feec.basics.kernels_3d as ker import struphy.feec.control_variates.kinetic_extended.fB_massless_kernels_control_variate as ker_cv import struphy.feec.control_variates.kinetic_extended.fnB_massless_cv_kernel_2 as ker_cv2 -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def bv_right( @@ -204,7 +204,7 @@ def bv_right( ) # ========================= C.T =========================== return tensor_space_FEM.C.T.dot( - np.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())) + xp.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())) ) diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py index cb8877c6f..b8673c837 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py @@ -2,7 +2,7 @@ import hylife.utilitis_FEEC.control_variates.fnB_massless_kernels_control_variate as ker_cv import scipy.sparse as spa -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def bv_pre(tol, n, LO_inv, tensor_space_FEM, p, Nel, idnx, idny, idnz): @@ -249,7 +249,7 @@ def bv_right( ) # ========================= C.T =========================== return tensor_space_FEM.C.T.dot( - np.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())) + xp.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())) ) @@ -430,7 +430,7 @@ def uv_right( ) # ========================= C.T =========================== temp_final = temp_final_0.flatten() + tensor_space_FEM.G.T.dot( - np.concatenate((temp_final_1.flatten(), temp_final_2.flatten(), temp_final_3.flatten())) + xp.concatenate((temp_final_1.flatten(), temp_final_2.flatten(), temp_final_3.flatten())) ) return temp_final diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py index 1ea567314..3d09c5bce 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py @@ -2,7 +2,7 @@ import hylife.utilitis_FEEC.control_variates.massless_kernels_control_variate as ker_cv import scipy.sparse as spa -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def bv_pre(u, uvalue, tensor_space_FEM, p, Nel, idnx, idny, idnz): @@ -248,7 +248,7 @@ def bv_right( ) # ========================= C.T =========================== return tensor_space_FEM.C.T.dot( - np.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())) + xp.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())) ) @@ -431,7 +431,7 @@ def uv_right( ) # ========================= C.T =========================== temp_final = temp_final_0.flatten() + tensor_space_FEM.G.T.dot( - np.concatenate((temp_final_1.flatten(), temp_final_2.flatten(), temp_final_3.flatten())) + xp.concatenate((temp_final_1.flatten(), temp_final_2.flatten(), temp_final_3.flatten())) ) return temp_final diff --git a/src/struphy/eigenvalue_solvers/legacy/emw_operators.py b/src/struphy/eigenvalue_solvers/legacy/emw_operators.py index 3187ac649..0a0f16106 100755 --- a/src/struphy/eigenvalue_solvers/legacy/emw_operators.py +++ b/src/struphy/eigenvalue_solvers/legacy/emw_operators.py @@ -10,7 +10,7 @@ import struphy.eigenvalue_solvers.kernels_3d as ker import struphy.eigenvalue_solvers.legacy.mass_matrices_3d_pre as mass_3d_pre -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class EMW_operators: @@ -134,7 +134,7 @@ def __assemble_M1_cross(self, weight): Ni = self.SPACES.Nbase_1form[a] Nj = self.SPACES.Nbase_1form[b] - M[a][b] = np.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) + M[a][b] = xp.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) # evaluate metric tensor at quadrature points if a == 1 and b == 2: @@ -185,9 +185,9 @@ def __assemble_M1_cross(self, weight): mat_w, ) # convert to sparse matrix - indices = np.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) + indices = xp.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() diff --git a/src/struphy/eigenvalue_solvers/legacy/inner_products_1d.py b/src/struphy/eigenvalue_solvers/legacy/inner_products_1d.py index 5cae935fb..469798bb5 100644 --- a/src/struphy/eigenvalue_solvers/legacy/inner_products_1d.py +++ b/src/struphy/eigenvalue_solvers/legacy/inner_products_1d.py @@ -8,7 +8,7 @@ import scipy.sparse as spa -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ======= inner product in V0 ==================== @@ -40,7 +40,7 @@ def inner_prod_V0(spline_space, fun, mapping=None): # evaluation of mapping at quadrature points if mapping == None: - mat_map = np.ones(pts.shape, dtype=float) + mat_map = xp.ones(pts.shape, dtype=float) else: mat_map = mapping(pts.flatten()).reshape(pts.shape) @@ -48,7 +48,7 @@ def inner_prod_V0(spline_space, fun, mapping=None): mat_f = fun(pts.flatten()).reshape(pts.shape) # assembly - F = np.zeros(NbaseN, dtype=float) + F = xp.zeros(NbaseN, dtype=float) for ie in range(Nel): for il in range(p + 1): @@ -91,7 +91,7 @@ def inner_prod_V1(spline_space, fun, mapping=None): # evaluation of mapping at quadrature points if mapping == None: - mat_map = np.ones(pts.shape, dtype=float) + mat_map = xp.ones(pts.shape, dtype=float) else: mat_map = 1 / mapping(pts.flatten()).reshape(pts.shape) @@ -99,7 +99,7 @@ def inner_prod_V1(spline_space, fun, mapping=None): mat_f = fun(pts.flatten()).reshape(pts.shape) # assembly - F = np.zeros(NbaseD, dtype=float) + F = xp.zeros(NbaseD, dtype=float) for ie in range(Nel): for il in range(p): diff --git a/src/struphy/eigenvalue_solvers/legacy/inner_products_2d.py b/src/struphy/eigenvalue_solvers/legacy/inner_products_2d.py index fd7ecabd4..58ea01f2f 100644 --- a/src/struphy/eigenvalue_solvers/legacy/inner_products_2d.py +++ b/src/struphy/eigenvalue_solvers/legacy/inner_products_2d.py @@ -9,7 +9,7 @@ import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_2d as ker -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ================ inner product in V0 =========================== @@ -25,7 +25,7 @@ def inner_prod_V0(tensor_space_FEM, domain, fun): domain : domain domain object defining the geometry - fun : callable or np.ndarray + fun : callable or xp.ndarray the 0-form with which the inner products shall be computed (either callable or 2D array with values at quadrature points) """ @@ -46,10 +46,10 @@ def inner_prod_V0(tensor_space_FEM, domain, fun): det_df = det_df.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1]) # evaluation of given 0-form at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size), dtype=float) if callable(fun): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") mat_f[:, :] = fun(quad_mesh[0], quad_mesh[1], 0.0) else: mat_f[:, :] = fun @@ -57,7 +57,7 @@ def inner_prod_V0(tensor_space_FEM, domain, fun): # assembly Ni = tensor_space_FEM.Nbase_0form - F = np.zeros((Ni[0], Ni[1]), dtype=float) + F = xp.zeros((Ni[0], Ni[1]), dtype=float) mat_f = mat_f.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1]) @@ -94,7 +94,7 @@ def inner_prod_V1(tensor_space_FEM, domain, fun): domain : domain domain object defining the geometry - fun : list of callables or np.ndarrays + fun : list of callables or xp.ndarrays the 1-form components with which the inner products shall be computed (either list of 3 callables or 2D arrays with values at quadrature points) """ @@ -127,10 +127,10 @@ def inner_prod_V1(tensor_space_FEM, domain, fun): g_inv = domain.metric_inv(pts[0].flatten(), pts[1].flatten(), 0.0) # 1-form components at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size), dtype=float) if callable(fun[0]): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") # components of global inner product F = [0, 0, 0] @@ -138,7 +138,7 @@ def inner_prod_V1(tensor_space_FEM, domain, fun): # assembly for a in range(3): Ni = tensor_space_FEM.Nbase_1form[a] - F[a] = np.zeros((Ni[0], Ni[1]), dtype=float) + F[a] = xp.zeros((Ni[0], Ni[1]), dtype=float) mat_f[:, :] = 0.0 @@ -170,7 +170,7 @@ def inner_prod_V1(tensor_space_FEM, domain, fun): mat_f * det_df, ) - F1 = tensor_space_FEM.E1_pol_0.dot(np.concatenate((F[0].flatten(), F[1].flatten()))) + F1 = tensor_space_FEM.E1_pol_0.dot(xp.concatenate((F[0].flatten(), F[1].flatten()))) F2 = tensor_space_FEM.E0_pol_0.dot(F[2].flatten()) return F1, F2 @@ -187,7 +187,7 @@ def inner_prod_V2(tensor_space_FEM, domain, fun): domain : domain domain object defining the geometry - fun : list of callables or np.ndarrays + fun : list of callables or xp.ndarrays the 2-form components with which the inner products shall be computed (either list of 3 callables or 2D arrays with values at quadrature points) """ @@ -220,10 +220,10 @@ def inner_prod_V2(tensor_space_FEM, domain, fun): g = domain.metric(pts[0].flatten(), pts[1].flatten(), 0.0) # 2-form components at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size), dtype=float) if callable(fun[0]): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") # components of global inner product F = [0, 0, 0] @@ -231,7 +231,7 @@ def inner_prod_V2(tensor_space_FEM, domain, fun): # assembly for a in range(3): Ni = tensor_space_FEM.Nbase_2form[a] - F[a] = np.zeros((Ni[0], Ni[1]), dtype=float) + F[a] = xp.zeros((Ni[0], Ni[1]), dtype=float) mat_f[:, :] = 0.0 @@ -263,7 +263,7 @@ def inner_prod_V2(tensor_space_FEM, domain, fun): mat_f / det_df, ) - F1 = tensor_space_FEM.E2_pol_0.dot(np.concatenate((F[0].flatten(), F[1].flatten()))) + F1 = tensor_space_FEM.E2_pol_0.dot(xp.concatenate((F[0].flatten(), F[1].flatten()))) F2 = tensor_space_FEM.E3_pol_0.dot(F[2].flatten()) return F1, F2 @@ -280,7 +280,7 @@ def inner_prod_V3(tensor_space_FEM, domain, fun): domain : domain domain object defining the geometry - fun : callable or np.ndarray + fun : callable or xp.ndarray the 3-form component with which the inner products shall be computed (either callable or 2D array with values at quadrature points) """ @@ -301,10 +301,10 @@ def inner_prod_V3(tensor_space_FEM, domain, fun): det_df = det_df.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1]) # evaluation of given 3-form at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size), dtype=float) if callable(fun): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") mat_f[:, :] = fun(quad_mesh[0], quad_mesh[1], 0.0) else: mat_f[:, :] = fun @@ -312,7 +312,7 @@ def inner_prod_V3(tensor_space_FEM, domain, fun): # assembly Ni = tensor_space_FEM.Nbase_3form - F = np.zeros((Ni[0], Ni[1]), dtype=float) + F = xp.zeros((Ni[0], Ni[1]), dtype=float) mat_f = mat_f.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1]) diff --git a/src/struphy/eigenvalue_solvers/legacy/inner_products_3d.py b/src/struphy/eigenvalue_solvers/legacy/inner_products_3d.py index 5aa9f710a..2616857b6 100644 --- a/src/struphy/eigenvalue_solvers/legacy/inner_products_3d.py +++ b/src/struphy/eigenvalue_solvers/legacy/inner_products_3d.py @@ -9,7 +9,7 @@ import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_3d as ker -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ================ inner product in V0 =========================== @@ -25,7 +25,7 @@ def inner_prod_V0(tensor_space_FEM, domain, fun): domain : domain domain object defining the geometry - fun : callable or np.ndarray + fun : callable or xp.ndarray the 0-form with which the inner products shall be computed (either callable or 3D array with values at quadrature points) """ @@ -46,10 +46,10 @@ def inner_prod_V0(tensor_space_FEM, domain, fun): det_df = det_df.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]) # evaluation of given 0-form at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) if callable(fun): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") mat_f[:, :, :] = fun(quad_mesh[0], quad_mesh[1], quad_mesh[2]) else: mat_f[:, :, :] = fun @@ -57,7 +57,7 @@ def inner_prod_V0(tensor_space_FEM, domain, fun): # assembly Ni = tensor_space.Nbase_0form - F = np.zeros((Ni[0], Ni[1], Ni[2]), dtype=float) + F = xp.zeros((Ni[0], Ni[1], Ni[2]), dtype=float) mat_f = mat_f.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]) @@ -101,7 +101,7 @@ def inner_prod_V1(tensor_space_FEM, domain, fun): domain : domain domain object defining the geometry - fun : list of callables or np.ndarrays + fun : list of callables or xp.ndarrays the 1-form components with which the inner products shall be computed (either list of 3 callables or 3D arrays with values at quadrature points) """ @@ -134,10 +134,10 @@ def inner_prod_V1(tensor_space_FEM, domain, fun): g_inv = domain.metric_inv(pts[0].flatten(), pts[1].flatten(), pts[2].flatten()) # 1-form components at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) if callable(fun[0]): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") # components of global inner product F = [0, 0, 0] @@ -146,7 +146,7 @@ def inner_prod_V1(tensor_space_FEM, domain, fun): for a in range(3): Ni = tensor_space_FEM.Nbase_1form[a] - F[a] = np.zeros((Ni[0], Ni[1], Ni[2]), dtype=float) + F[a] = xp.zeros((Ni[0], Ni[1], Ni[2]), dtype=float) mat_f[:, :, :] = 0.0 @@ -185,7 +185,7 @@ def inner_prod_V1(tensor_space_FEM, domain, fun): mat_f * det_df, ) - return tensor_space_FEM.E1_0.dot(np.concatenate((F[0].flatten(), F[1].flatten(), F[2].flatten()))) + return tensor_space_FEM.E1_0.dot(xp.concatenate((F[0].flatten(), F[1].flatten(), F[2].flatten()))) # ================ inner product in V2 =========================== @@ -199,7 +199,7 @@ def inner_prod_V2(tensor_space_FEM, domain, fun): domain : domain domain object defining the geometry - fun : list of callables or np.ndarrays + fun : list of callables or xp.ndarrays the 2-form components with which the inner products shall be computed (either list of 3 callables or 3D arrays with values at quadrature points) """ @@ -232,10 +232,10 @@ def inner_prod_V2(tensor_space_FEM, domain, fun): g = domain.metric(pts[0].flatten(), pts[1].flatten(), pts[2].flatten()) # 2-form components at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) if callable(fun[0]): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") # components of global inner product F = [0, 0, 0] @@ -244,7 +244,7 @@ def inner_prod_V2(tensor_space_FEM, domain, fun): for a in range(3): Ni = tensor_space_FEM.Nbase_2form[a] - F[a] = np.zeros((Ni[0], Ni[1], Ni[2]), dtype=float) + F[a] = xp.zeros((Ni[0], Ni[1], Ni[2]), dtype=float) mat_f[:, :, :] = 0.0 @@ -283,7 +283,7 @@ def inner_prod_V2(tensor_space_FEM, domain, fun): mat_f / det_df, ) - return tensor_space_FEM.E2_0.dot(np.concatenate((F[0].flatten(), F[1].flatten(), F[2].flatten()))) + return tensor_space_FEM.E2_0.dot(xp.concatenate((F[0].flatten(), F[1].flatten(), F[2].flatten()))) # ================ inner product in V3 =========================== @@ -297,7 +297,7 @@ def inner_prod_V3(tensor_space_FEM, domain, fun): domain : domain domain object defining the geometry - fun : callable or np.ndarray + fun : callable or xp.ndarray the 3-form component with which the inner products shall be computed (either callable or 3D array with values at quadrature points) """ @@ -318,10 +318,10 @@ def inner_prod_V3(tensor_space_FEM, domain, fun): det_df = det_df.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]) # evaluation of given 3-form at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) if callable(fun): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") mat_f[:, :, :] = fun(quad_mesh[0], quad_mesh[1], quad_mesh[2]) else: mat_f[:, :, :] = fun @@ -329,7 +329,7 @@ def inner_prod_V3(tensor_space_FEM, domain, fun): # assembly Ni = tensor_space.Nbase_3form - F = np.zeros((Ni[0], Ni[1], Ni[2]), dtype=float) + F = xp.zeros((Ni[0], Ni[1], Ni[2]), dtype=float) ker.kernel_inner( Nel[0], diff --git a/src/struphy/eigenvalue_solvers/legacy/l2_error_1d.py b/src/struphy/eigenvalue_solvers/legacy/l2_error_1d.py index f8544c1cf..708d0352e 100644 --- a/src/struphy/eigenvalue_solvers/legacy/l2_error_1d.py +++ b/src/struphy/eigenvalue_solvers/legacy/l2_error_1d.py @@ -8,7 +8,7 @@ import scipy.sparse as spa -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ======= error in V0 ==================== @@ -48,7 +48,7 @@ def l2_error_V0(spline_space, mapping, coeff, fun): mat_f = fun(pts) # assembly - error = np.zeros(Nel, dtype=float) + error = xp.zeros(Nel, dtype=float) for ie in range(Nel): for q in range(n_quad): @@ -59,7 +59,7 @@ def l2_error_V0(spline_space, mapping, coeff, fun): error[ie] += wts[ie, q] * (bi - mat_f[ie, q]) ** 2 - return np.sqrt(error.sum()) + return xp.sqrt(error.sum()) # ======= error in V1 ==================== @@ -99,7 +99,7 @@ def l2_error_V1(spline_space, mapping, coeff, fun): mat_f = fun(pts) # assembly - error = np.zeros(Nel, dtype=float) + error = xp.zeros(Nel, dtype=float) for ie in range(Nel): for q in range(n_quad): @@ -110,4 +110,4 @@ def l2_error_V1(spline_space, mapping, coeff, fun): error[ie] += wts[ie, q] * (bi - mat_f[ie, q]) ** 2 - return np.sqrt(error.sum()) + return xp.sqrt(error.sum()) diff --git a/src/struphy/eigenvalue_solvers/legacy/l2_error_2d.py b/src/struphy/eigenvalue_solvers/legacy/l2_error_2d.py index 818d0d7c2..ce2142ced 100644 --- a/src/struphy/eigenvalue_solvers/legacy/l2_error_2d.py +++ b/src/struphy/eigenvalue_solvers/legacy/l2_error_2d.py @@ -9,7 +9,7 @@ import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_2d as ker -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ======= error in V0 ==================== @@ -25,7 +25,7 @@ def l2_error_V0(tensor_space_FEM, domain, f0, c0, method="standard"): domain : domain domain object defining the geometry - f0 : callable or np.ndarray + f0 : callable or xp.ndarray the 0-form with which the error shall be computed c0 : array_like @@ -63,12 +63,12 @@ def l2_error_V0(tensor_space_FEM, domain, f0, c0, method="standard"): # evaluation of exact 0-form at quadrature points if callable(f0): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") f0 = f0(quad_mesh[0], quad_mesh[1], 0.0) if method == "standard": # evaluation of discrete 0-form at quadrature points - f0_h = tensor_space_FEM.evaluate_NN(pts[0].flatten(), pts[1].flatten(), np.array([0.0]), c0, "V0")[:, :, 0] + f0_h = tensor_space_FEM.evaluate_NN(pts[0].flatten(), pts[1].flatten(), xp.array([0.0]), c0, "V0")[:, :, 0] # compute error error = 0.0 @@ -78,7 +78,7 @@ def l2_error_V0(tensor_space_FEM, domain, f0, c0, method="standard"): else: # compute error in each element - error = np.zeros(Nel[:2], dtype=float) + error = xp.zeros(Nel[:2], dtype=float) ker.kernel_l2error( Nel, @@ -106,7 +106,7 @@ def l2_error_V0(tensor_space_FEM, domain, f0, c0, method="standard"): error = error.sum() - return np.sqrt(error) + return xp.sqrt(error) # ======= error in V1 ==================== @@ -122,7 +122,7 @@ def l2_error_V1(tensor_space_FEM, domain, f1, c1, method="standard"): domain : domain domain object defining the geometry - f1 : list of callables or np.ndarrays + f1 : list of callables or xp.ndarrays the three 1-form components with which the error shall be computed c1 : list of array_like @@ -162,16 +162,16 @@ def l2_error_V1(tensor_space_FEM, domain, f1, c1, method="standard"): # evaluation of exact 1-form components at quadrature points if callable(f1[0]): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") f1_1 = f1[0](quad_mesh[0], quad_mesh[1], 0.0) f1_2 = f1[1](quad_mesh[0], quad_mesh[1], 0.0) f1_3 = f1[2](quad_mesh[0], quad_mesh[1], 0.0) if method == "standard": # evaluation of discrete 1-form components at quadrature points - f1_h_1 = tensor_space_FEM.evaluate_DN(pts[0].flatten(), pts[1].flatten(), np.array([0.0]), c1_1, "V1")[:, :, 0] - f1_h_2 = tensor_space_FEM.evaluate_ND(pts[0].flatten(), pts[1].flatten(), np.array([0.0]), c1_2, "V1")[:, :, 0] - f1_h_3 = tensor_space_FEM.evaluate_NN(pts[0].flatten(), pts[1].flatten(), np.array([0.0]), c1_3, "V1")[:, :, 0] + f1_h_1 = tensor_space_FEM.evaluate_DN(pts[0].flatten(), pts[1].flatten(), xp.array([0.0]), c1_1, "V1")[:, :, 0] + f1_h_2 = tensor_space_FEM.evaluate_ND(pts[0].flatten(), pts[1].flatten(), xp.array([0.0]), c1_2, "V1")[:, :, 0] + f1_h_3 = tensor_space_FEM.evaluate_NN(pts[0].flatten(), pts[1].flatten(), xp.array([0.0]), c1_3, "V1")[:, :, 0] # compute error error = 0.0 @@ -194,7 +194,7 @@ def l2_error_V1(tensor_space_FEM, domain, f1, c1, method="standard"): else: # compute error in each element - error = np.zeros(Nel[:2], dtype=float) + error = xp.zeros(Nel[:2], dtype=float) # 1 * d_f1 * G^11 * |det(DF)| * d_f1 ker.kernel_l2error( @@ -298,7 +298,7 @@ def l2_error_V1(tensor_space_FEM, domain, f1, c1, method="standard"): error = error.sum() - return np.sqrt(error) + return xp.sqrt(error) # ======= error in V2 ==================== @@ -314,7 +314,7 @@ def l2_error_V2(tensor_space_FEM, domain, f2, c2, method="standard"): domain : domain domain object defining the geometry - f2 : list of callables or np.ndarrays + f2 : list of callables or xp.ndarrays the three 2-form components with which the error shall be computed c2 : list of array_like @@ -354,16 +354,16 @@ def l2_error_V2(tensor_space_FEM, domain, f2, c2, method="standard"): # evaluation of exact 2-form components at quadrature points if callable(f2[0]): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") f2_1 = f2[0](quad_mesh[0], quad_mesh[1], 0.0) f2_2 = f2[1](quad_mesh[0], quad_mesh[1], 0.0) f2_3 = f2[2](quad_mesh[0], quad_mesh[1], 0.0) if method == "standard": # evaluation of discrete 2-form components at quadrature points - f2_h_1 = tensor_space_FEM.evaluate_ND(pts[0].flatten(), pts[1].flatten(), np.array([0.0]), c2_1, "V2")[:, :, 0] - f2_h_2 = tensor_space_FEM.evaluate_DN(pts[0].flatten(), pts[1].flatten(), np.array([0.0]), c2_2, "V2")[:, :, 0] - f2_h_3 = tensor_space_FEM.evaluate_DD(pts[0].flatten(), pts[1].flatten(), np.array([0.0]), c2_3, "V2")[:, :, 0] + f2_h_1 = tensor_space_FEM.evaluate_ND(pts[0].flatten(), pts[1].flatten(), xp.array([0.0]), c2_1, "V2")[:, :, 0] + f2_h_2 = tensor_space_FEM.evaluate_DN(pts[0].flatten(), pts[1].flatten(), xp.array([0.0]), c2_2, "V2")[:, :, 0] + f2_h_3 = tensor_space_FEM.evaluate_DD(pts[0].flatten(), pts[1].flatten(), xp.array([0.0]), c2_3, "V2")[:, :, 0] # compute error error = 0.0 @@ -386,7 +386,7 @@ def l2_error_V2(tensor_space_FEM, domain, f2, c2, method="standard"): else: # compute error in each element - error = np.zeros(Nel[:2], dtype=float) + error = xp.zeros(Nel[:2], dtype=float) # 1 * d_f1 * G_11 / |det(DF)| * d_f1 ker.kernel_l2error( @@ -490,7 +490,7 @@ def l2_error_V2(tensor_space_FEM, domain, f2, c2, method="standard"): error = error.sum() - return np.sqrt(error) + return xp.sqrt(error) # ======= error in V3 ==================== @@ -506,7 +506,7 @@ def l2_error_V3(tensor_space_FEM, domain, f3, c3, method="standard"): domain : domain domain object defining the geometry - f3 : callable or np.ndarray + f3 : callable or xp.ndarray the 3-form component with which the error shall be computed c3 : array_like @@ -544,12 +544,12 @@ def l2_error_V3(tensor_space_FEM, domain, f3, c3, method="standard"): # evaluation of exact 3-form at quadrature points if callable(f3): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), indexing="ij") f3 = f3(quad_mesh[0], quad_mesh[1], 0.0) if method == "standard": # evaluation of discrete 3-form at quadrature points - f3_h = tensor_space_FEM.evaluate_DD(pts[0].flatten(), pts[1].flatten(), np.array([0.0]), c3, "V3")[:, :, 0] + f3_h = tensor_space_FEM.evaluate_DD(pts[0].flatten(), pts[1].flatten(), xp.array([0.0]), c3, "V3")[:, :, 0] # compute error error = 0.0 @@ -559,7 +559,7 @@ def l2_error_V3(tensor_space_FEM, domain, f3, c3, method="standard"): else: # compute error in each element - error = np.zeros(Nel[:2], dtype=float) + error = xp.zeros(Nel[:2], dtype=float) ker.kernel_l2error( Nel, @@ -587,4 +587,4 @@ def l2_error_V3(tensor_space_FEM, domain, f3, c3, method="standard"): error = error.sum() - return np.sqrt(error) + return xp.sqrt(error) diff --git a/src/struphy/eigenvalue_solvers/legacy/l2_error_3d.py b/src/struphy/eigenvalue_solvers/legacy/l2_error_3d.py index 39eac0b66..93dcc4053 100644 --- a/src/struphy/eigenvalue_solvers/legacy/l2_error_3d.py +++ b/src/struphy/eigenvalue_solvers/legacy/l2_error_3d.py @@ -9,7 +9,7 @@ import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_3d as ker -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ======= error in V0 ==================== @@ -25,7 +25,7 @@ def l2_error_V0(tensor_space_FEM, domain, fun, coeff): domain : domain domain object defining the geometry - fun : callable or np.ndarray + fun : callable or xp.ndarray the 0-form with which the error shall be computed coeff : array_like @@ -54,16 +54,16 @@ def l2_error_V0(tensor_space_FEM, domain, fun, coeff): det_df = abs(domain.jacobian_det(pts[0].flatten(), pts[1].flatten(), pts[2].flatten())) # evaluation of given 0-form at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) if callable(fun): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") mat_f[:, :, :] = fun(quad_mesh[0], quad_mesh[1], quad_mesh[2]) else: mat_f[:, :, :] = fun # compute error - error = np.zeros(Nel, dtype=float) + error = xp.zeros(Nel, dtype=float) ker.kernel_l2error( Nel, @@ -94,7 +94,7 @@ def l2_error_V0(tensor_space_FEM, domain, fun, coeff): det_df.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]), ) - return np.sqrt(error.sum()) + return xp.sqrt(error.sum()) # ======= error in V1 ==================== @@ -110,7 +110,7 @@ def l2_error_V1(tensor_space_FEM, domain, fun, coeff): domain : domain domain object defining the geometry - fun : list of callables or np.ndarrays + fun : list of callables or xp.ndarrays the three 1-form components with which the error shall be computed coeff : list of array_like @@ -141,12 +141,12 @@ def l2_error_V1(tensor_space_FEM, domain, fun, coeff): metric_coeffs *= abs(domain.jacobian_det(pts[0].flatten(), pts[1].flatten(), pts[2].flatten())) # evaluation of given 1-form components at quadrature points - mat_f1 = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) - mat_f2 = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) - mat_f3 = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f1 = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f2 = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f3 = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) if callable(fun[0]): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") mat_f1[:, :, :] = fun[0](quad_mesh[0], quad_mesh[1], quad_mesh[2]) mat_f2[:, :, :] = fun[1](quad_mesh[0], quad_mesh[1], quad_mesh[2]) mat_f3[:, :, :] = fun[2](quad_mesh[0], quad_mesh[1], quad_mesh[2]) @@ -156,7 +156,7 @@ def l2_error_V1(tensor_space_FEM, domain, fun, coeff): mat_f3[:, :, :] = fun[2] # compute error - error = np.zeros(Nel, dtype=float) + error = xp.zeros(Nel, dtype=float) # 1 * f1 * G^11 * |det(DF)| * f1 ker.kernel_l2error( @@ -314,7 +314,7 @@ def l2_error_V1(tensor_space_FEM, domain, fun, coeff): 1 * metric_coeffs[2, 2].reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]), ) - return np.sqrt(error.sum()) + return xp.sqrt(error.sum()) # ======= error in V2 ==================== @@ -330,7 +330,7 @@ def l2_error_V2(tensor_space_FEM, domain, fun, coeff): domain : domain domain object defining the geometry - fun : list of callables or np.ndarrays + fun : list of callables or xp.ndarrays the three 2-form components with which the error shall be computed coeff : list of array_like @@ -361,12 +361,12 @@ def l2_error_V2(tensor_space_FEM, domain, fun, coeff): metric_coeffs /= abs(domain.jacobian_det(pts[0].flatten(), pts[1].flatten(), pts[2].flatten())) # evaluation of given 2-form components at quadrature points - mat_f1 = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) - mat_f2 = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) - mat_f3 = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f1 = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f2 = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f3 = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) if callable(fun[0]): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") mat_f1[:, :, :] = fun[0](quad_mesh[0], quad_mesh[1], quad_mesh[2]) mat_f2[:, :, :] = fun[1](quad_mesh[0], quad_mesh[1], quad_mesh[2]) mat_f3[:, :, :] = fun[2](quad_mesh[0], quad_mesh[1], quad_mesh[2]) @@ -376,7 +376,7 @@ def l2_error_V2(tensor_space_FEM, domain, fun, coeff): mat_f3[:, :, :] = fun[2] # compute error - error = np.zeros(Nel, dtype=float) + error = xp.zeros(Nel, dtype=float) # 1 * f1 * G_11 / |det(DF)| * f1 ker.kernel_l2error( @@ -534,7 +534,7 @@ def l2_error_V2(tensor_space_FEM, domain, fun, coeff): 1 * metric_coeffs[2, 2].reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]), ) - return np.sqrt(error.sum()) + return xp.sqrt(error.sum()) # ======= error in V3 ==================== @@ -550,7 +550,7 @@ def l2_error_V3(tensor_space_FEM, domain, fun, coeff): domain : domain domain object defining the geometry - fun : callable or np.ndarray + fun : callable or xp.ndarray the 3-form component with which the error shall be computed coeff : array_like @@ -579,16 +579,16 @@ def l2_error_V3(tensor_space_FEM, domain, fun, coeff): det_df = abs(domain.jacobian_det(pts[0].flatten(), pts[1].flatten(), pts[2].flatten())) # evaluation of given 3-form component at quadrature points - mat_f = np.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) + mat_f = xp.empty((pts[0].size, pts[1].size, pts[2].size), dtype=float) if callable(fun): - quad_mesh = np.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") + quad_mesh = xp.meshgrid(pts[0].flatten(), pts[1].flatten(), pts[2].flatten(), indexing="ij") mat_f[:, :, :] = fun(quad_mesh[0], quad_mesh[1], quad_mesh[2]) else: mat_f[:, :, :] = fun # compute error - error = np.zeros(Nel, dtype=float) + error = xp.zeros(Nel, dtype=float) ker.kernel_l2error( Nel, @@ -619,4 +619,4 @@ def l2_error_V3(tensor_space_FEM, domain, fun, coeff): 1 / det_df.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]), ) - return np.sqrt(error.sum()) + return xp.sqrt(error.sum()) diff --git a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py index 673069c0e..fedff3c90 100644 --- a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py +++ b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py @@ -10,7 +10,7 @@ import struphy.eigenvalue_solvers.spline_space as spl import struphy.linear_algebra.linalg_kron as linkron -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ================ inverse mass matrix in V0 =========================== @@ -32,9 +32,9 @@ def get_M0_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M0(lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M0(lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M0(lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M0(lambda eta: (domain.params[1] - domain.params[0]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: (domain.params[3] - domain.params[2]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: (domain.params[5] - domain.params[4]) * xp.ones(eta.shape, dtype=float)) c_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] @@ -63,20 +63,20 @@ def get_M1_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M0(lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M0(lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M0(lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M0(lambda eta: (domain.params[1] - domain.params[0]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: (domain.params[3] - domain.params[2]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: (domain.params[5] - domain.params[4]) * xp.ones(eta.shape, dtype=float)) - spaces_pre[0].assemble_M1(lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M1(lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M1(lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M1(lambda eta: 1 / (domain.params[1] - domain.params[0]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M1(lambda eta: 1 / (domain.params[3] - domain.params[2]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M1(lambda eta: 1 / (domain.params[5] - domain.params[4]) * xp.ones(eta.shape, dtype=float)) c11_pre = [spaces_pre[0].M1.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] c22_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M1.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] c33_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M1.toarray()[:, 0]] def solve(x): - x1, x2, x3 = np.split(x, 3) + x1, x2, x3 = xp.split(x, 3) x1 = x1.reshape(Nel_pre[0], Nel_pre[1], Nel_pre[2]) x2 = x2.reshape(Nel_pre[0], Nel_pre[1], Nel_pre[2]) @@ -86,7 +86,7 @@ def solve(x): r2 = linkron.kron_fftsolve_3d(c22_pre, x2).flatten() r3 = linkron.kron_fftsolve_3d(c33_pre, x3).flatten() - return np.concatenate((r1, r2, r3)) + return xp.concatenate((r1, r2, r3)) return spa.linalg.LinearOperator(shape=tensor_space_FEM.M1.shape, matvec=solve) @@ -110,20 +110,20 @@ def get_M2_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M0(lambda eta: (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M0(lambda eta: (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M0(lambda eta: (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M0(lambda eta: (domain.params[1] - domain.params[0]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: (domain.params[3] - domain.params[2]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: (domain.params[5] - domain.params[4]) * xp.ones(eta.shape, dtype=float)) - spaces_pre[0].assemble_M1(lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M1(lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M1(lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M1(lambda eta: 1 / (domain.params[1] - domain.params[0]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M1(lambda eta: 1 / (domain.params[3] - domain.params[2]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M1(lambda eta: 1 / (domain.params[5] - domain.params[4]) * xp.ones(eta.shape, dtype=float)) c11_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M1.toarray()[:, 0], spaces_pre[2].M1.toarray()[:, 0]] c22_pre = [spaces_pre[0].M1.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M1.toarray()[:, 0]] c33_pre = [spaces_pre[0].M1.toarray()[:, 0], spaces_pre[1].M1.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] def solve(x): - x1, x2, x3 = np.split(x, 3) + x1, x2, x3 = xp.split(x, 3) x1 = x1.reshape(Nel_pre[0], Nel_pre[1], Nel_pre[2]) x2 = x2.reshape(Nel_pre[0], Nel_pre[1], Nel_pre[2]) @@ -133,7 +133,7 @@ def solve(x): r2 = linkron.kron_fftsolve_3d(c22_pre, x2).flatten() r3 = linkron.kron_fftsolve_3d(c33_pre, x3).flatten() - return np.concatenate((r1, r2, r3)) + return xp.concatenate((r1, r2, r3)) return spa.linalg.LinearOperator(shape=tensor_space_FEM.M2.shape, matvec=solve) @@ -157,9 +157,9 @@ def get_M3_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M1(lambda eta: 1 / (domain.params[1] - domain.params[0]) * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M1(lambda eta: 1 / (domain.params[3] - domain.params[2]) * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M1(lambda eta: 1 / (domain.params[5] - domain.params[4]) * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M1(lambda eta: 1 / (domain.params[1] - domain.params[0]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M1(lambda eta: 1 / (domain.params[3] - domain.params[2]) * xp.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M1(lambda eta: 1 / (domain.params[5] - domain.params[4]) * xp.ones(eta.shape, dtype=float)) c_pre = [spaces_pre[0].M1.toarray()[:, 0], spaces_pre[1].M1.toarray()[:, 0], spaces_pre[2].M1.toarray()[:, 0]] @@ -188,26 +188,26 @@ def get_Mv_PRE(tensor_space_FEM, domain): # spaces_pre[1].set_extraction_operators() # spaces_pre[2].set_extraction_operators() - spaces_pre[0].assemble_M0(lambda eta: domain.params[0] ** 3 * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M0(lambda eta: domain.params[1] * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M0(lambda eta: domain.params[2] * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M0(lambda eta: domain.params[0] ** 3 * xp.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: domain.params[1] * xp.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: domain.params[2] * xp.ones(eta.shape, dtype=float)) c11_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] - spaces_pre[0].assemble_M0(lambda eta: domain.params[0] * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M0(lambda eta: domain.params[1] ** 3 * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M0(lambda eta: domain.params[2] * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M0(lambda eta: domain.params[0] * xp.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: domain.params[1] ** 3 * xp.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: domain.params[2] * xp.ones(eta.shape, dtype=float)) c22_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] - spaces_pre[0].assemble_M0(lambda eta: domain.params[0] * np.ones(eta.shape, dtype=float)) - spaces_pre[1].assemble_M0(lambda eta: domain.params[1] * np.ones(eta.shape, dtype=float)) - spaces_pre[2].assemble_M0(lambda eta: domain.params[2] ** 3 * np.ones(eta.shape, dtype=float)) + spaces_pre[0].assemble_M0(lambda eta: domain.params[0] * xp.ones(eta.shape, dtype=float)) + spaces_pre[1].assemble_M0(lambda eta: domain.params[1] * xp.ones(eta.shape, dtype=float)) + spaces_pre[2].assemble_M0(lambda eta: domain.params[2] ** 3 * xp.ones(eta.shape, dtype=float)) c33_pre = [spaces_pre[0].M0.toarray()[:, 0], spaces_pre[1].M0.toarray()[:, 0], spaces_pre[2].M0.toarray()[:, 0]] def solve(x): - x1, x2, x3 = np.split(x, 3) + x1, x2, x3 = xp.split(x, 3) x1 = x1.reshape(Nel_pre[0], Nel_pre[1], Nel_pre[2]) x2 = x2.reshape(Nel_pre[0], Nel_pre[1], Nel_pre[2]) @@ -217,7 +217,7 @@ def solve(x): r2 = linkron.kron_fftsolve_3d(c22_pre, x2).flatten() r3 = linkron.kron_fftsolve_3d(c33_pre, x3).flatten() - return np.concatenate((r1, r2, r3)) + return xp.concatenate((r1, r2, r3)) return spa.linalg.LinearOperator(shape=tensor_space_FEM.Mv.shape, matvec=solve) @@ -282,7 +282,7 @@ def solve(x): r1 = linkron.kron_fftsolve_2d(M1_pol_0_11_LU, tor_vec0, x1).flatten() r2 = linkron.kron_fftsolve_2d(M1_pol_0_22_LU, tor_vec1, x2).flatten() - return np.concatenate((r1, r2)) + return xp.concatenate((r1, r2)) return spa.linalg.LinearOperator(shape=tensor_space_FEM.M1_0.shape, matvec=solve) @@ -320,7 +320,7 @@ def solve(x): r1 = linkron.kron_fftsolve_2d(M2_pol_0_11_LU, tor_vec1, x1).flatten() r2 = linkron.kron_fftsolve_2d(M2_pol_0_22_LU, tor_vec0, x2).flatten() - return np.concatenate((r1, r2)) + return xp.concatenate((r1, r2)) return spa.linalg.LinearOperator(shape=tensor_space_FEM.M2_0.shape, matvec=solve) @@ -382,6 +382,6 @@ def solve(x): r1 = linkron.kron_fftsolve_2d(Mv_pol_0_11_LU, tor_vec0, x1).flatten() r2 = linkron.kron_fftsolve_2d(Mv_pol_0_22_LU, tor_vec0, x2).flatten() - return np.concatenate((r1, r2)) + return xp.concatenate((r1, r2)) return spa.linalg.LinearOperator(shape=tensor_space_FEM.Mv_0.shape, matvec=solve) diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py index 8514d25fc..193d51a11 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py @@ -7,7 +7,7 @@ import struphy.geometry.mappings_3d as mapping3d import struphy.geometry.mappings_3d_fast as mapping_fast import struphy.linear_algebra.linalg_kernels as linalg -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Temp_arrays: @@ -39,67 +39,67 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): self.Ntot_1form = TENSOR_SPACE_FEM.Ntot_1form self.Ntot_2form = TENSOR_SPACE_FEM.Ntot_2form - self.b1_old = np.empty(TENSOR_SPACE_FEM.Nbase_1form[0], dtype=float) - self.b2_old = np.empty(TENSOR_SPACE_FEM.Nbase_1form[1], dtype=float) - self.b3_old = np.empty(TENSOR_SPACE_FEM.Nbase_1form[2], dtype=float) + self.b1_old = xp.empty(TENSOR_SPACE_FEM.Nbase_1form[0], dtype=float) + self.b2_old = xp.empty(TENSOR_SPACE_FEM.Nbase_1form[1], dtype=float) + self.b3_old = xp.empty(TENSOR_SPACE_FEM.Nbase_1form[2], dtype=float) - self.b1_iter = np.empty(TENSOR_SPACE_FEM.Nbase_1form[0], dtype=float) - self.b2_iter = np.empty(TENSOR_SPACE_FEM.Nbase_1form[1], dtype=float) - self.b3_iter = np.empty(TENSOR_SPACE_FEM.Nbase_1form[2], dtype=float) + self.b1_iter = xp.empty(TENSOR_SPACE_FEM.Nbase_1form[0], dtype=float) + self.b2_iter = xp.empty(TENSOR_SPACE_FEM.Nbase_1form[1], dtype=float) + self.b3_iter = xp.empty(TENSOR_SPACE_FEM.Nbase_1form[2], dtype=float) - self.temp_dft = np.empty((3, 3), dtype=float) - self.temp_generate_weight1 = np.empty(3, dtype=float) - self.temp_generate_weight2 = np.empty(3, dtype=float) - self.temp_generate_weight3 = np.empty(3, dtype=float) + self.temp_dft = xp.empty((3, 3), dtype=float) + self.temp_generate_weight1 = xp.empty(3, dtype=float) + self.temp_generate_weight2 = xp.empty(3, dtype=float) + self.temp_generate_weight3 = xp.empty(3, dtype=float) - self.zerosform_temp_long = np.empty(TENSOR_SPACE_FEM.Ntot_0form, dtype=float) - self.oneform_temp1_long = np.empty(TENSOR_SPACE_FEM.Ntot_1form[0], dtype=float) - self.oneform_temp2_long = np.empty(TENSOR_SPACE_FEM.Ntot_1form[1], dtype=float) - self.oneform_temp3_long = np.empty(TENSOR_SPACE_FEM.Ntot_1form[2], dtype=float) + self.zerosform_temp_long = xp.empty(TENSOR_SPACE_FEM.Ntot_0form, dtype=float) + self.oneform_temp1_long = xp.empty(TENSOR_SPACE_FEM.Ntot_1form[0], dtype=float) + self.oneform_temp2_long = xp.empty(TENSOR_SPACE_FEM.Ntot_1form[1], dtype=float) + self.oneform_temp3_long = xp.empty(TENSOR_SPACE_FEM.Ntot_1form[2], dtype=float) - self.oneform_temp_long = np.empty( + self.oneform_temp_long = xp.empty( TENSOR_SPACE_FEM.Ntot_1form[0] + TENSOR_SPACE_FEM.Ntot_1form[1] + TENSOR_SPACE_FEM.Ntot_1form[2], dtype=float, ) - self.twoform_temp1_long = np.empty(TENSOR_SPACE_FEM.Ntot_2form[0], dtype=float) - self.twoform_temp2_long = np.empty(TENSOR_SPACE_FEM.Ntot_2form[1], dtype=float) - self.twoform_temp3_long = np.empty(TENSOR_SPACE_FEM.Ntot_2form[2], dtype=float) + self.twoform_temp1_long = xp.empty(TENSOR_SPACE_FEM.Ntot_2form[0], dtype=float) + self.twoform_temp2_long = xp.empty(TENSOR_SPACE_FEM.Ntot_2form[1], dtype=float) + self.twoform_temp3_long = xp.empty(TENSOR_SPACE_FEM.Ntot_2form[2], dtype=float) - self.twoform_temp_long = np.empty( + self.twoform_temp_long = xp.empty( TENSOR_SPACE_FEM.Ntot_2form[0] + TENSOR_SPACE_FEM.Ntot_2form[1] + TENSOR_SPACE_FEM.Ntot_2form[2], dtype=float, ) - self.temp_twoform1 = np.empty(TENSOR_SPACE_FEM.Nbase_2form[0], dtype=float) - self.temp_twoform2 = np.empty(TENSOR_SPACE_FEM.Nbase_2form[1], dtype=float) - self.temp_twoform3 = np.empty(TENSOR_SPACE_FEM.Nbase_2form[2], dtype=float) + self.temp_twoform1 = xp.empty(TENSOR_SPACE_FEM.Nbase_2form[0], dtype=float) + self.temp_twoform2 = xp.empty(TENSOR_SPACE_FEM.Nbase_2form[1], dtype=float) + self.temp_twoform3 = xp.empty(TENSOR_SPACE_FEM.Nbase_2form[2], dtype=float) # arrays used to store intermidaite values - self.form_0_flatten = np.empty(self.Ntot_0form, dtype=float) + self.form_0_flatten = xp.empty(self.Ntot_0form, dtype=float) - self.form_1_1_flatten = np.empty(self.Ntot_1form[0], dtype=float) - self.form_1_2_flatten = np.empty(self.Ntot_1form[1], dtype=float) - self.form_1_3_flatten = np.empty(self.Ntot_1form[2], dtype=float) + self.form_1_1_flatten = xp.empty(self.Ntot_1form[0], dtype=float) + self.form_1_2_flatten = xp.empty(self.Ntot_1form[1], dtype=float) + self.form_1_3_flatten = xp.empty(self.Ntot_1form[2], dtype=float) - self.form_1_tot_flatten = np.empty(self.Ntot_1form[0] + self.Ntot_1form[1] + self.Ntot_1form[2], dtype=float) + self.form_1_tot_flatten = xp.empty(self.Ntot_1form[0] + self.Ntot_1form[1] + self.Ntot_1form[2], dtype=float) - self.form_2_1_flatten = np.empty(self.Ntot_2form[0], dtype=float) - self.form_2_2_flatten = np.empty(self.Ntot_2form[1], dtype=float) - self.form_2_3_flatten = np.empty(self.Ntot_2form[2], dtype=float) + self.form_2_1_flatten = xp.empty(self.Ntot_2form[0], dtype=float) + self.form_2_2_flatten = xp.empty(self.Ntot_2form[1], dtype=float) + self.form_2_3_flatten = xp.empty(self.Ntot_2form[2], dtype=float) - self.form_2_tot_flatten = np.empty(self.Ntot_2form[0] + self.Ntot_2form[1] + self.Ntot_2form[2], dtype=float) + self.form_2_tot_flatten = xp.empty(self.Ntot_2form[0] + self.Ntot_2form[1] + self.Ntot_2form[2], dtype=float) - self.bulkspeed_loc = np.zeros((3, self.Nel[0], self.Nel[1], self.Nel[2]), dtype=float) - self.temperature_loc = np.zeros((3, self.Nel[0], self.Nel[1], self.Nel[2]), dtype=float) - self.bulkspeed = np.zeros((3, self.Nel[0], self.Nel[1], self.Nel[2]), dtype=float) + self.bulkspeed_loc = xp.zeros((3, self.Nel[0], self.Nel[1], self.Nel[2]), dtype=float) + self.temperature_loc = xp.zeros((3, self.Nel[0], self.Nel[1], self.Nel[2]), dtype=float) + self.bulkspeed = xp.zeros((3, self.Nel[0], self.Nel[1], self.Nel[2]), dtype=float) if self.mpi_rank == 0: - temperature = np.zeros((3, self.Nel[0], self.Nel[1], self.Nel[2]), dtype=float) + temperature = xp.zeros((3, self.Nel[0], self.Nel[1], self.Nel[2]), dtype=float) else: temperature = None # values of magnetic fields at all quadrature points - self.LO_inv = np.empty( + self.LO_inv = xp.empty( ( self.Nel[0], self.Nel[1], @@ -111,7 +111,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) - self.LO_b1 = np.empty( + self.LO_b1 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -122,7 +122,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.LO_b2 = np.empty( + self.LO_b2 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -133,7 +133,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.LO_b3 = np.empty( + self.LO_b3 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -145,7 +145,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) # values of weights (used in the linear operators) - self.LO_w1 = np.empty( + self.LO_w1 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -156,7 +156,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.LO_w2 = np.empty( + self.LO_w2 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -167,7 +167,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.LO_w3 = np.empty( + self.LO_w3 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -179,7 +179,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) # values of a function (given its finite element coefficients) at all quadrature points - self.LO_r1 = np.empty( + self.LO_r1 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -190,7 +190,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.LO_r2 = np.empty( + self.LO_r2 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -201,7 +201,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.LO_r3 = np.empty( + self.LO_r3 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -213,7 +213,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) # values of determinant of Jacobi matrix of the map at all quadrature points - self.df_det = np.empty( + self.df_det = xp.empty( ( self.Nel[0], self.Nel[1], @@ -226,7 +226,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ) # when using delta f method, the values of current equilibrium at all quadrature points if control == True: - self.Jeqx = np.empty( + self.Jeqx = xp.empty( ( self.Nel[0], self.Nel[1], @@ -237,7 +237,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.Jeqy = np.empty( + self.Jeqy = xp.empty( ( self.Nel[0], self.Nel[1], @@ -248,7 +248,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.Jeqz = np.empty( + self.Jeqz = xp.empty( ( self.Nel[0], self.Nel[1], @@ -260,7 +260,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) # values of DF and inverse of DF at all quadrature points - self.DF_11 = np.empty( + self.DF_11 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -271,7 +271,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DF_12 = np.empty( + self.DF_12 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -282,7 +282,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DF_13 = np.empty( + self.DF_13 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -293,7 +293,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DF_21 = np.empty( + self.DF_21 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -304,7 +304,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DF_22 = np.empty( + self.DF_22 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -315,7 +315,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DF_23 = np.empty( + self.DF_23 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -326,7 +326,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DF_31 = np.empty( + self.DF_31 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -337,7 +337,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DF_32 = np.empty( + self.DF_32 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -348,7 +348,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DF_33 = np.empty( + self.DF_33 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -360,7 +360,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) - self.DFI_11 = np.empty( + self.DFI_11 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -371,7 +371,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFI_12 = np.empty( + self.DFI_12 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -382,7 +382,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFI_13 = np.empty( + self.DFI_13 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -393,7 +393,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFI_21 = np.empty( + self.DFI_21 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -404,7 +404,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFI_22 = np.empty( + self.DFI_22 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -415,7 +415,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFI_23 = np.empty( + self.DFI_23 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -426,7 +426,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFI_31 = np.empty( + self.DFI_31 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -437,7 +437,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFI_32 = np.empty( + self.DFI_32 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -448,7 +448,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFI_33 = np.empty( + self.DFI_33 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -460,7 +460,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) - self.DFIT_11 = np.empty( + self.DFIT_11 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -471,7 +471,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFIT_12 = np.empty( + self.DFIT_12 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -482,7 +482,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFIT_13 = np.empty( + self.DFIT_13 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -493,7 +493,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFIT_21 = np.empty( + self.DFIT_21 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -504,7 +504,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFIT_22 = np.empty( + self.DFIT_22 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -515,7 +515,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFIT_23 = np.empty( + self.DFIT_23 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -526,7 +526,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFIT_31 = np.empty( + self.DFIT_31 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -537,7 +537,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFIT_32 = np.empty( + self.DFIT_32 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -548,7 +548,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.DFIT_33 = np.empty( + self.DFIT_33 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -560,7 +560,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) - self.G_inv_11 = np.empty( + self.G_inv_11 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -571,7 +571,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.G_inv_12 = np.empty( + self.G_inv_12 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -582,7 +582,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.G_inv_13 = np.empty( + self.G_inv_13 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -594,7 +594,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) - self.G_inv_22 = np.empty( + self.G_inv_22 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -605,7 +605,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): ), dtype=float, ) - self.G_inv_23 = np.empty( + self.G_inv_23 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -617,7 +617,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) - self.G_inv_33 = np.empty( + self.G_inv_33 = xp.empty( ( self.Nel[0], self.Nel[1], @@ -629,7 +629,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) - self.temp_particle = np.empty(3, dtype=float) + self.temp_particle = xp.empty(3, dtype=float) # initialization of DF and its inverse # ================ for mapping evaluation ================== # spline degrees @@ -638,34 +638,34 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): pf3 = DOMAIN.p[2] # pf + 1 non-vanishing basis functions up tp degree pf - b1f = np.empty((pf1 + 1, pf1 + 1), dtype=float) - b2f = np.empty((pf2 + 1, pf2 + 1), dtype=float) - b3f = np.empty((pf3 + 1, pf3 + 1), dtype=float) + b1f = xp.empty((pf1 + 1, pf1 + 1), dtype=float) + b2f = xp.empty((pf2 + 1, pf2 + 1), dtype=float) + b3f = xp.empty((pf3 + 1, pf3 + 1), dtype=float) # left and right values for spline evaluation - l1f = np.empty(pf1, dtype=float) - l2f = np.empty(pf2, dtype=float) - l3f = np.empty(pf3, dtype=float) + l1f = xp.empty(pf1, dtype=float) + l2f = xp.empty(pf2, dtype=float) + l3f = xp.empty(pf3, dtype=float) - r1f = np.empty(pf1, dtype=float) - r2f = np.empty(pf2, dtype=float) - r3f = np.empty(pf3, dtype=float) + r1f = xp.empty(pf1, dtype=float) + r2f = xp.empty(pf2, dtype=float) + r3f = xp.empty(pf3, dtype=float) # scaling arrays for M-splines - d1f = np.empty(pf1, dtype=float) - d2f = np.empty(pf2, dtype=float) - d3f = np.empty(pf3, dtype=float) + d1f = xp.empty(pf1, dtype=float) + d2f = xp.empty(pf2, dtype=float) + d3f = xp.empty(pf3, dtype=float) # pf + 1 derivatives - der1f = np.empty(pf1 + 1, dtype=float) - der2f = np.empty(pf2 + 1, dtype=float) - der3f = np.empty(pf3 + 1, dtype=float) + der1f = xp.empty(pf1 + 1, dtype=float) + der2f = xp.empty(pf2 + 1, dtype=float) + der3f = xp.empty(pf3 + 1, dtype=float) # needed mapping quantities - df = np.empty((3, 3), dtype=float) - fx = np.empty(3, dtype=float) - ginv = np.empty((3, 3), dtype=float) - dfinv = np.empty((3, 3), dtype=float) + df = xp.empty((3, 3), dtype=float) + fx = xp.empty(3, dtype=float) + ginv = xp.empty((3, 3), dtype=float) + dfinv = xp.empty((3, 3), dtype=float) for ie1 in range(self.Nel[0]): for ie2 in range(self.Nel[1]): diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py index 843952a7b..c8b04f877 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py @@ -5,7 +5,7 @@ import struphy.feec.massless_operators.fB_bb_kernel as bb_kernel import struphy.feec.massless_operators.fB_bv_kernel as bv_kernel import struphy.feec.massless_operators.fB_vv_kernel as vv_kernel -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Massless_linear_operators: @@ -49,15 +49,15 @@ def linearoperator_step_vv(self, M2_PRE, M2, M1_PRE, M1, TEMP, ACC_VV): This function is used in substep vv with L2 projector. """ - dft = np.empty((3, 3), dtype=float) - generate_weight1 = np.zeros(3, dtype=float) - generate_weight2 = np.zeros(3, dtype=float) - generate_weight3 = np.zeros(3, dtype=float) + dft = xp.empty((3, 3), dtype=float) + generate_weight1 = xp.zeros(3, dtype=float) + generate_weight2 = xp.zeros(3, dtype=float) + generate_weight3 = xp.zeros(3, dtype=float) # =========================inverse of M1 =========================== - ACC_VV.temp1[:], ACC_VV.temp2[:], ACC_VV.temp3[:] = np.split( + ACC_VV.temp1[:], ACC_VV.temp2[:], ACC_VV.temp3[:] = xp.split( spa.linalg.cg( M1, - 1.0 / self.Np * np.concatenate((ACC_VV.vec1.flatten(), ACC_VV.vec2.flatten(), ACC_VV.vec3.flatten())), + 1.0 / self.Np * xp.concatenate((ACC_VV.vec1.flatten(), ACC_VV.vec2.flatten(), ACC_VV.vec3.flatten())), tol=10 ** (-14), M=M1_PRE, )[0], @@ -153,10 +153,10 @@ def linearoperator_step_vv(self, M2_PRE, M2, M1_PRE, M1, TEMP, ACC_VV): ) # =========================inverse of M1 =========================== - ACC_VV.temp1[:], ACC_VV.temp2[:], ACC_VV.temp3[:] = np.split( + ACC_VV.temp1[:], ACC_VV.temp2[:], ACC_VV.temp3[:] = xp.split( spa.linalg.cg( M1, - np.concatenate((ACC_VV.one_form1.flatten(), ACC_VV.one_form2.flatten(), ACC_VV.one_form3.flatten())), + xp.concatenate((ACC_VV.one_form1.flatten(), ACC_VV.one_form2.flatten(), ACC_VV.one_form3.flatten())), tol=10 ** (-14), M=M1_PRE, )[0], @@ -310,10 +310,10 @@ def linearoperator_pre_step_vv( indN = tensor_space_FEM.indN indD = tensor_space_FEM.indD - dft = np.empty((3, 3), dtype=float) - generate_weight1 = np.zeros(3, dtype=float) - generate_weight2 = np.zeros(3, dtype=float) - generate_weight3 = np.zeros(3, dtype=float) + dft = xp.empty((3, 3), dtype=float) + generate_weight1 = xp.zeros(3, dtype=float) + generate_weight2 = xp.zeros(3, dtype=float) + generate_weight3 = xp.zeros(3, dtype=float) vv_kernel.prepre( indN[0], @@ -716,15 +716,15 @@ def linearoperator_step3( Ntot_2form = tensor_space_FEM.Ntot_2form Nbase_2form = tensor_space_FEM.Nbase_2form - dft = np.empty((3, 3), dtype=float) - generate_weight1 = np.empty(3, dtype=float) - generate_weight2 = np.empty(3, dtype=float) - generate_weight3 = np.empty(3, dtype=float) + dft = xp.empty((3, 3), dtype=float) + generate_weight1 = xp.empty(3, dtype=float) + generate_weight2 = xp.empty(3, dtype=float) + generate_weight3 = xp.empty(3, dtype=float) # ================================================================== # ========================= C =========================== # time1 = time.time() - twoform_temp1_long[:], twoform_temp2_long[:], twoform_temp3_long[:] = np.split( + twoform_temp1_long[:], twoform_temp2_long[:], twoform_temp3_long[:] = xp.split( tensor_space_FEM.C.dot(input_vector), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]] ) temp_vector_1[:, :, :] = twoform_temp1_long.reshape(Nbase_2form[0]) @@ -826,7 +826,7 @@ def linearoperator_step3( # ========================= C.T =========================== # time1 = time.time() temp_final = tensor_space_FEM.M1.dot(input_vector) - dt / 2.0 * tensor_space_FEM.C.T.dot( - np.concatenate((temp_vector_1.flatten(), temp_vector_2.flatten(), temp_vector_3.flatten())) + xp.concatenate((temp_vector_1.flatten(), temp_vector_2.flatten(), temp_vector_3.flatten())) ) # time2 = time.time() # print('second_curl_time', time2 - time1) @@ -921,14 +921,14 @@ def linearoperator_right_step3( Ntot_2form = tensor_space_FEM.Ntot_2form Nbase_2form = tensor_space_FEM.Nbase_2form - dft = np.empty((3, 3), dtype=float) - generate_weight1 = np.empty(3, dtype=float) - generate_weight2 = np.empty(3, dtype=float) - generate_weight3 = np.empty(3, dtype=float) + dft = xp.empty((3, 3), dtype=float) + generate_weight1 = xp.empty(3, dtype=float) + generate_weight2 = xp.empty(3, dtype=float) + generate_weight3 = xp.empty(3, dtype=float) # ================================================================== # ========================= C =========================== - twoform_temp1_long[:], twoform_temp2_long[:], twoform_temp3_long[:] = np.split( + twoform_temp1_long[:], twoform_temp2_long[:], twoform_temp3_long[:] = xp.split( tensor_space_FEM.C.dot(input_vector), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]] ) temp_vector_1[:, :, :] = twoform_temp1_long.reshape(Nbase_2form[0]) @@ -1082,7 +1082,7 @@ def linearoperator_right_step3( # print('final_bb', time2 - time1) # ========================= C.T =========================== temp_final = tensor_space_FEM.M1.dot(input_vector) + dt / 2.0 * tensor_space_FEM.C.T.dot( - np.concatenate((temp_vector_1.flatten(), temp_vector_2.flatten(), temp_vector_3.flatten())) + xp.concatenate((temp_vector_1.flatten(), temp_vector_2.flatten(), temp_vector_3.flatten())) ) return temp_final @@ -1145,7 +1145,7 @@ def substep4_linear_operator( wts = tensor_space_FEM.wts # global quadrature weights # ========================================== - acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = np.split( + acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = xp.split( tensor_space_FEM.C.dot(input), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]] ) acc.twoform_temp1[:, :, :] = acc.twoform_temp1_long.reshape(Nbase_2form[0]) @@ -1249,12 +1249,12 @@ def substep4_linear_operator( ) acc.oneform_temp_long[:] = spa.linalg.gmres( M1, - np.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())), + xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())), tol=10 ** (-10), M=M1_PRE, )[0] - acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = np.split( + acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = xp.split( spa.linalg.gmres(M1, mat.dot(acc.oneform_temp_long), tol=10 ** (-10), M=M1_PRE)[0], [Ntot_1form[0], Ntot_1form[0] + Ntot_1form[1]], ) @@ -1359,7 +1359,7 @@ def substep4_linear_operator( ) return M1.dot(input) + dt**2 / 4.0 * tensor_space_FEM.C.T.dot( - np.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) + xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) ) # ========================================================================================================== @@ -1422,8 +1422,8 @@ def substep4_linear_operator_right( wts = tensor_space_FEM.wts # global quadrature weights # ========================================== - acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = np.split( - CURL.dot(np.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))), + acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = xp.split( + CURL.dot(xp.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]], ) acc.twoform_temp1[:, :, :] = acc.twoform_temp1_long.reshape(Nbase_2form[0]) @@ -1529,13 +1529,13 @@ def substep4_linear_operator_right( acc.oneform_temp_long[:] = mat.dot( spa.linalg.gmres( M1, - np.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())), + xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())), tol=10 ** (-10), M=M1_PRE, )[0] ) - acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = np.split( + acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = xp.split( spa.linalg.gmres(M1, dt**2.0 / 4.0 * acc.oneform_temp_long + dt * vec, tol=10 ** (-10), M=M1_PRE)[0], [Ntot_1form[0], Ntot_1form[0] + Ntot_1form[1]], ) @@ -1639,8 +1639,8 @@ def substep4_linear_operator_right( tensor_space_FEM.basisD[2], ) - return M1.dot(np.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))) - CURL.T.dot( - np.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) + return M1.dot(xp.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))) - CURL.T.dot( + xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) ) # ========================================================================================================== @@ -1792,8 +1792,8 @@ def substep4_pusher_field( wts = tensor_space_FEM.wts # global quadrature weights # ========================================== - acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = np.split( - CURL.dot(np.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))), + acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = xp.split( + CURL.dot(xp.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]], ) acc.twoform_temp1[:, :, :] = acc.twoform_temp1_long.reshape(Nbase_2form[0]) @@ -1898,7 +1898,7 @@ def substep4_pusher_field( return spa.linalg.cg( M1, - np.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())), + xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())), tol=10 ** (-13), M=M1_PRE, )[0] @@ -1960,7 +1960,7 @@ def substep4_localproj_linear_operator( wts = tensor_space_FEM.wts # global quadrature weights # ========================================== - acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = np.split( + acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = xp.split( tensor_space_FEM.C.dot(input), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]] ) acc.twoform_temp1[:, :, :] = acc.twoform_temp1_long.reshape(Nbase_2form[0]) @@ -2063,9 +2063,9 @@ def substep4_localproj_linear_operator( tensor_space_FEM.basisD[2], ) - acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = np.split( + acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = xp.split( mat.dot( - np.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())) + xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())) ), [Ntot_1form[0], Ntot_1form[0] + Ntot_1form[1]], ) @@ -2171,7 +2171,7 @@ def substep4_localproj_linear_operator( ) return M1.dot(input) + dt**2 / 4.0 * tensor_space_FEM.C.T.dot( - np.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) + xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) ) # ========================================================================================================== @@ -2234,8 +2234,8 @@ def substep4_localproj_linear_operator_right( wts = tensor_space_FEM.wts # global quadrature weights # ========================================== - acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = np.split( - CURL.dot(np.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))), + acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = xp.split( + CURL.dot(xp.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]], ) acc.twoform_temp1[:, :, :] = acc.twoform_temp1_long.reshape(Nbase_2form[0]) @@ -2339,10 +2339,10 @@ def substep4_localproj_linear_operator_right( tensor_space_FEM.basisD[2], ) acc.oneform_temp_long[:] = mat.dot( - np.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())) + xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())) ) - acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = np.split( + acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = xp.split( (dt**2.0 / 4.0 * acc.oneform_temp_long + dt * vec), [Ntot_1form[0], Ntot_1form[0] + Ntot_1form[1]] ) @@ -2446,8 +2446,8 @@ def substep4_localproj_linear_operator_right( tensor_space_FEM.basisD[2], ) - return M1.dot(np.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))) - CURL.T.dot( - np.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) + return M1.dot(xp.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))) - CURL.T.dot( + xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) ) # ========================================================================================================== @@ -2509,8 +2509,8 @@ def substep4_localproj_pusher_field( wts = tensor_space_FEM.wts # global quadrature weights # ========================================== - acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = np.split( - CURL.dot(np.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))), + acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = xp.split( + CURL.dot(xp.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]], ) acc.twoform_temp1[:, :, :] = acc.twoform_temp1_long.reshape(Nbase_2form[0]) @@ -2613,4 +2613,4 @@ def substep4_localproj_pusher_field( tensor_space_FEM.basisD[2], ) - return np.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())) + return xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())) diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py index 216103640..449e31e50 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py @@ -3,7 +3,7 @@ import struphy.bsplines.bsplines_kernels as bsp import struphy.geometry.mappings_kernels as mapping_fast import struphy.linear_algebra.linalg_kernels as linalg -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ========================================================================================== diff --git a/src/struphy/eigenvalue_solvers/legacy/mhd_operators_MF.py b/src/struphy/eigenvalue_solvers/legacy/mhd_operators_MF.py index 9591fe06b..7f3775389 100644 --- a/src/struphy/eigenvalue_solvers/legacy/mhd_operators_MF.py +++ b/src/struphy/eigenvalue_solvers/legacy/mhd_operators_MF.py @@ -3,7 +3,7 @@ from struphy.eigenvalue_solvers.projectors_global import Projectors_tensor_3d from struphy.eigenvalue_solvers.spline_space import Tensor_spline_space from struphy.linear_algebra.linalg_kron import kron_matvec_3d, kron_solve_3d -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ================================================================================================= @@ -107,9 +107,9 @@ def __init__(self, space, eq_MHD): self.pts1_D_2 = self.space.spaces[1].projectors.D_pts self.pts1_D_3 = self.space.spaces[2].projectors.D_pts - # assert np.allclose(self.N_1.toarray(), self.pts0_N_1.toarray(), atol=1e-14) - # assert np.allclose(self.N_2.toarray(), self.pts0_N_2.toarray(), atol=1e-14) - # assert np.allclose(self.N_3.toarray(), self.pts0_N_3.toarray(), atol=1e-14) + # assert xp.allclose(self.N_1.toarray(), self.pts0_N_1.toarray(), atol=1e-14) + # assert xp.allclose(self.N_2.toarray(), self.pts0_N_2.toarray(), atol=1e-14) + # assert xp.allclose(self.N_3.toarray(), self.pts0_N_3.toarray(), atol=1e-14) # ===== call equilibrium_mhd values at the projection points ===== # projection points @@ -210,11 +210,11 @@ def __init__(self, space, eq_MHD): # # Operator A # if self.basis_u == 1: # self.A = spa.linalg.LinearOperator((self.dim_1, self.dim_1), matvec = lambda x : (self.M1.dot(self.W1_dot(x)) + self.transpose_W1_dot(self.M1.dot(x))) / 2 ) - # self.A_mat = spa.csc_matrix(self.A.dot(np.identity(self.dim_1))) + # self.A_mat = spa.csc_matrix(self.A.dot(xp.identity(self.dim_1))) # elif self.basis_u == 2: # self.A = spa.linalg.LinearOperator((self.dim_2, self.dim_2), matvec = lambda x : (self.M2.dot(self.Q2_dot(x)) + self.transpose_Q2_dot(self.M2.dot(x))) / 2 ) - # self.A_mat = spa.csc_matrix(self.A.dot(np.identity(self.dim_2))) + # self.A_mat = spa.csc_matrix(self.A.dot(xp.identity(self.dim_2))) # self.A_inv = spa.linalg.inv(self.A_mat) @@ -228,12 +228,12 @@ def Q1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^1} Returns ---------- - res : np.array + res : xp.array dim R^{N^2} Notes @@ -320,7 +320,7 @@ def Q1_dot(self, x): # xi3 : histo(xi1)-histo(xi2)-inter(xi3)-polation. res_3 = self.space.projectors.PI_mat("23", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ================================================================== def transpose_Q1_dot(self, x): @@ -329,12 +329,12 @@ def transpose_Q1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^2} Returns ---------- - res : np.array + res : xp.array dim R^{N^1} Notes @@ -403,7 +403,7 @@ def transpose_Q1_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # =================================================================== def W1_dot(self, x): @@ -412,12 +412,12 @@ def W1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^1} Returns ---------- - res : np.array + res : xp.array dim R^{N^1} Notes @@ -482,7 +482,7 @@ def W1_dot(self, x): # xi3 : inter(xi1)-inter(xi2)-histo(xi3)-polation. res_3 = self.space.projectors.PI_mat("13", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # =================================================================== def transpose_W1_dot(self, x): @@ -491,12 +491,12 @@ def transpose_W1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^1} Returns ---------- - res : np.array + res : xp.array dim R{N^1} Notes @@ -545,7 +545,7 @@ def transpose_W1_dot(self, x): res_2 = kron_matvec_3d([self.pts0_N_1.T, self.pts1_D_2.T, self.pts0_N_3.T], mat_f_2_c) res_3 = kron_matvec_3d([self.pts0_N_1.T, self.pts0_N_2.T, self.pts1_D_3.T], mat_f_3_c) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def U1_dot(self, x): @@ -554,12 +554,12 @@ def U1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^1} Returns ---------- - res : np.array + res : xp.array dim R^{N^2} Notes @@ -645,7 +645,7 @@ def U1_dot(self, x): # xi3 : histo(xi1)-histo(xi2)-inter(xi3)-polation. res_3 = self.space.projectors.PI_mat("23", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def transpose_U1_dot(self, x): @@ -654,12 +654,12 @@ def transpose_U1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^2} Returns ---------- - res : np.array + res : xp.array dim R{N^1} Notes @@ -728,7 +728,7 @@ def transpose_U1_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def P1_dot(self, x): @@ -737,12 +737,12 @@ def P1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^2} Returns ---------- - res : np.array + res : xp.array dim R^{N^1} Notes @@ -831,7 +831,7 @@ def P1_dot(self, x): # xi3 : inter(xi1)-inter(xi2)-histo(xi3)-polation. res_3 = self.space.projectors.PI_mat("13", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def transpose_P1_dot(self, x): @@ -840,12 +840,12 @@ def transpose_P1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^1} Returns ---------- - res : np.array + res : xp.array dim R{N^2} Notes @@ -914,7 +914,7 @@ def transpose_P1_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def S1_dot(self, x): @@ -923,12 +923,12 @@ def S1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^1} Returns ---------- - res : np.array + res : xp.array dim R^{N^2} Notes @@ -1015,7 +1015,7 @@ def S1_dot(self, x): # xi3 : histo(xi1)-histo(xi2)-inter(xi3)-polation. res_3 = self.space.projectors.PI_mat("23", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def transpose_S1_dot(self, x): @@ -1024,12 +1024,12 @@ def transpose_S1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^2} Returns ---------- - res : np.array + res : xp.array dim R{N^1} Notes @@ -1098,7 +1098,7 @@ def transpose_S1_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # =================================================================== def S10_dot(self, x): @@ -1107,12 +1107,12 @@ def S10_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^1} Returns ---------- - res : np.array + res : xp.array dim R^{N^1} Notes @@ -1178,7 +1178,7 @@ def S10_dot(self, x): # xi3 : inter(xi1)-inter(xi2)-histo(xi3)-polation. res_3 = self.space.projectors.PI_mat("13", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # =================================================================== def transpose_S10_dot(self, x): @@ -1187,12 +1187,12 @@ def transpose_S10_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^1} Returns ---------- - res : np.array + res : xp.array dim R{N^1} Notes @@ -1241,7 +1241,7 @@ def transpose_S10_dot(self, x): res_2 = kron_matvec_3d([self.pts0_N_1.T, self.pts1_D_2.T, self.pts0_N_3.T], mat_f_2_c) res_3 = kron_matvec_3d([self.pts0_N_1.T, self.pts0_N_2.T, self.pts1_D_3.T], mat_f_3_c) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ================================================================= def K1_dot(self, x): @@ -1250,12 +1250,12 @@ def K1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^3} Returns ---------- - res : np.array + res : xp.array dim R^{N^3} Notes @@ -1307,7 +1307,7 @@ def transpose_K1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^3} Returns @@ -1350,12 +1350,12 @@ def K10_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^0} Returns ---------- - res : np.array + res : xp.array dim R^{N^0} Notes @@ -1406,7 +1406,7 @@ def transpose_K10_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^0} Returns @@ -1449,12 +1449,12 @@ def T1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^1} Returns ---------- - res : np.array + res : xp.array dim R^{N^1} Notes @@ -1543,7 +1543,7 @@ def T1_dot(self, x): # xi3 : inter(xi1)-inter(xi2)-histo(xi3)-polation. res_3 = self.space.projectors.PI_mat("13", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ================================================================= def transpose_T1_dot(self, x): @@ -1552,12 +1552,12 @@ def transpose_T1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^1} Returns ---------- - res : np.array + res : xp.array dim R{N^1} Notes @@ -1626,7 +1626,7 @@ def transpose_T1_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ================================================================= def X1_dot(self, x): @@ -1635,13 +1635,13 @@ def X1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^1} Returns ---------- res : list - 3 np.arrays of dim R^{N^0} + 3 xp.arrays of dim R^{N^0} Notes ----- @@ -1718,12 +1718,12 @@ def transpose_X1_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^0 x 3} Returns ---------- - res : np.array + res : xp.array dim R{N^1} Notes @@ -1738,9 +1738,9 @@ def transpose_X1_dot(self, x): # x dim check # x should be R{N^0 * 3} # assert len(x) == self.space.Ntot_0form * 3 - # x_loc_1 = self.space.extract_0(np.split(x,3)[0]) - # x_loc_2 = self.space.extract_0(np.split(x,3)[1]) - # x_loc_3 = self.space.extract_0(np.split(x,3)[2]) + # x_loc_1 = self.space.extract_0(xp.split(x,3)[0]) + # x_loc_2 = self.space.extract_0(xp.split(x,3)[1]) + # x_loc_3 = self.space.extract_0(xp.split(x,3)[2]) # x_loc = list((x_loc_1, x_loc_2, x_loc_3)) x_loc_1 = self.space.extract_0(x[0]) @@ -1794,7 +1794,7 @@ def transpose_X1_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) ######################################## ########## 2-form formulation ########## @@ -1806,12 +1806,12 @@ def Q2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^2} Returns ---------- - res : np.array + res : xp.array dim R^{N^2} Notes @@ -1882,7 +1882,7 @@ def Q2_dot(self, x): # xi3 : histo(xi1)-histo(xi2)-inter(xi3)-polation. res_3 = self.space.projectors.PI_mat("23", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def transpose_Q2_dot(self, x): @@ -1891,12 +1891,12 @@ def transpose_Q2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^2} Returns ---------- - res : np.array + res : xp.array dim R{N^2} Notes @@ -1946,7 +1946,7 @@ def transpose_Q2_dot(self, x): res_2 = kron_matvec_3d([self.pts1_D_1.T, self.pts0_N_2.T, self.pts1_D_3.T], mat_f_2_c) res_3 = kron_matvec_3d([self.pts1_D_1.T, self.pts1_D_2.T, self.pts0_N_3.T], mat_f_3_c) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def T2_dot(self, x): @@ -1955,12 +1955,12 @@ def T2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^2} Returns ---------- - res : np.array + res : xp.array dim R^{N^1} Notes @@ -2049,7 +2049,7 @@ def T2_dot(self, x): # xi3 : inter(xi1)-inter(xi2)-histo(xi3)-polation. res_3 = self.space.projectors.PI_mat("13", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def transpose_T2_dot(self, x): @@ -2058,12 +2058,12 @@ def transpose_T2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^1} Returns ---------- - res : np.array + res : xp.array dim R{N^2} Notes @@ -2132,7 +2132,7 @@ def transpose_T2_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def P2_dot(self, x): @@ -2141,12 +2141,12 @@ def P2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^2} Returns ---------- - res : np.array + res : xp.array dim R^{N^2} Notes @@ -2235,7 +2235,7 @@ def P2_dot(self, x): # xi3 : histo(xi1)-histo(xi2)-inter(xi3)-polation. res_3 = self.space.projectors.PI_mat("23", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def transpose_P2_dot(self, x): @@ -2244,12 +2244,12 @@ def transpose_P2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^2} Returns ---------- - res : np.array + res : xp.array dim R{N^2} Notes @@ -2317,7 +2317,7 @@ def transpose_P2_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def S2_dot(self, x): @@ -2326,12 +2326,12 @@ def S2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^2} Returns ---------- - res : np.array + res : xp.array dim R^{N^2} Notes @@ -2402,7 +2402,7 @@ def S2_dot(self, x): # xi3 : histo(xi1)-histo(xi2)-inter(xi3)-polation. res_3 = self.space.projectors.PI_mat("23", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def transpose_S2_dot(self, x): @@ -2411,12 +2411,12 @@ def transpose_S2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^2} Returns ---------- - res : np.array + res : xp.array dim R{N^2} Notes @@ -2466,7 +2466,7 @@ def transpose_S2_dot(self, x): res_2 = kron_matvec_3d([self.pts1_D_1.T, self.pts0_N_2.T, self.pts1_D_3.T], mat_f_2_c) res_3 = kron_matvec_3d([self.pts1_D_1.T, self.pts1_D_2.T, self.pts0_N_3.T], mat_f_3_c) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def K2_dot(self, x): @@ -2475,12 +2475,12 @@ def K2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^3} Returns ---------- - res : np.array + res : xp.array dim R^{N^3} Notes @@ -2532,7 +2532,7 @@ def transpose_K2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^3} Returns @@ -2575,13 +2575,13 @@ def X2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^2} Returns ---------- res : list - 3 np.arrays of dim R^{N^0} + 3 xp.arrays of dim R^{N^0} Notes ----- @@ -2658,12 +2658,12 @@ def transpose_X2_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^0 x 3} Returns ---------- - res : np.array + res : xp.array dim R{N^2} Notes @@ -2678,9 +2678,9 @@ def transpose_X2_dot(self, x): # x dim check # x should be R{N^0 * 3} # assert len(x) == self.space.Ntot_0form * 3 - # x_loc_1 = self.space.extract_0(np.split(x,3)[0]) - # x_loc_2 = self.space.extract_0(np.split(x,3)[1]) - # x_loc_3 = self.space.extract_0(np.split(x,3)[2]) + # x_loc_1 = self.space.extract_0(xp.split(x,3)[0]) + # x_loc_2 = self.space.extract_0(xp.split(x,3)[1]) + # x_loc_3 = self.space.extract_0(xp.split(x,3)[2]) # x_loc = list((x_loc_1, x_loc_2, x_loc_3)) x_loc_1 = self.space.extract_0(x[0]) @@ -2734,7 +2734,7 @@ def transpose_X2_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def Z20_dot(self, x): @@ -2743,12 +2743,12 @@ def Z20_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^2} Returns ---------- - res : np.array + res : xp.array dim R^{N^1} Notes @@ -2835,7 +2835,7 @@ def Z20_dot(self, x): # xi3 : inter(xi1)-inter(xi2)-histo(xi3)-polation. res_3 = self.space.projectors.PI_mat("13", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def transpose_Z20_dot(self, x): @@ -2844,12 +2844,12 @@ def transpose_Z20_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^2} Returns ---------- - res : np.array + res : xp.array dim R{N^1} Notes @@ -2918,7 +2918,7 @@ def transpose_Z20_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def Y20_dot(self, x): @@ -2927,12 +2927,12 @@ def Y20_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^0} Returns ---------- - res : np.array + res : xp.array dim R^{N^3} Notes @@ -2984,12 +2984,12 @@ def transpose_Y20_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^3} Returns ---------- - res : np.array + res : xp.array dim R{N^0} Notes @@ -3027,12 +3027,12 @@ def S20_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R^{N^2} Returns ---------- - res : np.array + res : xp.array dim R^{N^1} Notes @@ -3119,7 +3119,7 @@ def S20_dot(self, x): # xi3 : inter(xi1)-inter(xi2)-histo(xi3)-polation. res_3 = self.space.projectors.PI_mat("13", DOF_3) - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) # ==================================================================== def transpose_S20_dot(self, x): @@ -3128,12 +3128,12 @@ def transpose_S20_dot(self, x): Parameters ---------- - x : np.array + x : xp.array dim R{N^1} Returns ---------- - res : np.array + res : xp.array dim R{N^2} Notes @@ -3202,4 +3202,4 @@ def transpose_S20_dot(self, x): res_2 = res_12 + res_22 + res_32 res_3 = res_13 + res_23 + res_33 - return np.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) + return xp.concatenate((res_1.flatten(), res_2.flatten(), res_3.flatten())) diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py index bd840b3f0..5d6c133e2 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py @@ -14,7 +14,7 @@ import struphy.feec.basics.kernels_3d as ker_loc_3d import struphy.feec.bsplines as bsp import struphy.feec.projectors.pro_local.kernels_projectors_local_mhd as ker_loc -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class projectors_local_mhd: @@ -44,8 +44,8 @@ def __init__(self, tensor_space, n_quad): self.n_quad = n_quad # number of quadrature point per integration interval # Gauss - Legendre quadrature points and weights in (-1, 1) - self.pts_loc = [np.polynomial.legendre.leggauss(n_quad)[0] for n_quad in self.n_quad] - self.wts_loc = [np.polynomial.legendre.leggauss(n_quad)[1] for n_quad in self.n_quad] + self.pts_loc = [xp.polynomial.legendre.leggauss(n_quad)[0] for n_quad in self.n_quad] + self.wts_loc = [xp.polynomial.legendre.leggauss(n_quad)[1] for n_quad in self.n_quad] # set interpolation and histopolation coefficients self.coeff_i = [0, 0, 0] @@ -53,78 +53,78 @@ def __init__(self, tensor_space, n_quad): for a in range(3): if self.bc[a] == True: - self.coeff_i[a] = np.zeros((1, 2 * self.p[a] - 1), dtype=float) - self.coeff_h[a] = np.zeros((1, 2 * self.p[a]), dtype=float) + self.coeff_i[a] = xp.zeros((1, 2 * self.p[a] - 1), dtype=float) + self.coeff_h[a] = xp.zeros((1, 2 * self.p[a]), dtype=float) if self.p[a] == 1: - self.coeff_i[a][0, :] = np.array([1.0]) - self.coeff_h[a][0, :] = np.array([1.0, 1.0]) + self.coeff_i[a][0, :] = xp.array([1.0]) + self.coeff_h[a][0, :] = xp.array([1.0, 1.0]) elif self.p[a] == 2: - self.coeff_i[a][0, :] = 1 / 2 * np.array([-1.0, 4.0, -1.0]) - self.coeff_h[a][0, :] = 1 / 2 * np.array([-1.0, 3.0, 3.0, -1.0]) + self.coeff_i[a][0, :] = 1 / 2 * xp.array([-1.0, 4.0, -1.0]) + self.coeff_h[a][0, :] = 1 / 2 * xp.array([-1.0, 3.0, 3.0, -1.0]) elif self.p[a] == 3: - self.coeff_i[a][0, :] = 1 / 6 * np.array([1.0, -8.0, 20.0, -8.0, 1.0]) - self.coeff_h[a][0, :] = 1 / 6 * np.array([1.0, -7.0, 12.0, 12.0, -7.0, 1.0]) + self.coeff_i[a][0, :] = 1 / 6 * xp.array([1.0, -8.0, 20.0, -8.0, 1.0]) + self.coeff_h[a][0, :] = 1 / 6 * xp.array([1.0, -7.0, 12.0, 12.0, -7.0, 1.0]) elif self.p[a] == 4: - self.coeff_i[a][0, :] = 2 / 45 * np.array([-1.0, 16.0, -295 / 4, 140.0, -295 / 4, 16.0, -1.0]) + self.coeff_i[a][0, :] = 2 / 45 * xp.array([-1.0, 16.0, -295 / 4, 140.0, -295 / 4, 16.0, -1.0]) self.coeff_h[a][0, :] = ( - 2 / 45 * np.array([-1.0, 15.0, -231 / 4, 265 / 4, 265 / 4, -231 / 4, 15.0, -1.0]) + 2 / 45 * xp.array([-1.0, 15.0, -231 / 4, 265 / 4, 265 / 4, -231 / 4, 15.0, -1.0]) ) else: print("degree > 4 not implemented!") else: - self.coeff_i[a] = np.zeros((2 * self.p[a] - 1, 2 * self.p[a] - 1), dtype=float) - self.coeff_h[a] = np.zeros((2 * self.p[a] - 1, 2 * self.p[a]), dtype=float) + self.coeff_i[a] = xp.zeros((2 * self.p[a] - 1, 2 * self.p[a] - 1), dtype=float) + self.coeff_h[a] = xp.zeros((2 * self.p[a] - 1, 2 * self.p[a]), dtype=float) if self.p[a] == 1: - self.coeff_i[a][0, :] = np.array([1.0]) - self.coeff_h[a][0, :] = np.array([1.0, 1.0]) + self.coeff_i[a][0, :] = xp.array([1.0]) + self.coeff_h[a][0, :] = xp.array([1.0, 1.0]) elif self.p[a] == 2: - self.coeff_i[a][0, :] = 1 / 2 * np.array([2.0, 0.0, 0.0]) - self.coeff_i[a][1, :] = 1 / 2 * np.array([-1.0, 4.0, -1.0]) - self.coeff_i[a][2, :] = 1 / 2 * np.array([0.0, 0.0, 2.0]) + self.coeff_i[a][0, :] = 1 / 2 * xp.array([2.0, 0.0, 0.0]) + self.coeff_i[a][1, :] = 1 / 2 * xp.array([-1.0, 4.0, -1.0]) + self.coeff_i[a][2, :] = 1 / 2 * xp.array([0.0, 0.0, 2.0]) - self.coeff_h[a][0, :] = 1 / 2 * np.array([3.0, -1.0, 0.0, 0.0]) - self.coeff_h[a][1, :] = 1 / 2 * np.array([-1.0, 3.0, 3.0, -1.0]) - self.coeff_h[a][2, :] = 1 / 2 * np.array([0.0, 0.0, -1.0, 3.0]) + self.coeff_h[a][0, :] = 1 / 2 * xp.array([3.0, -1.0, 0.0, 0.0]) + self.coeff_h[a][1, :] = 1 / 2 * xp.array([-1.0, 3.0, 3.0, -1.0]) + self.coeff_h[a][2, :] = 1 / 2 * xp.array([0.0, 0.0, -1.0, 3.0]) elif self.p[a] == 3: - self.coeff_i[a][0, :] = 1 / 18 * np.array([18.0, 0.0, 0.0, 0.0, 0.0]) - self.coeff_i[a][1, :] = 1 / 18 * np.array([-5.0, 40.0, -24.0, 8.0, -1.0]) - self.coeff_i[a][2, :] = 1 / 18 * np.array([3.0, -24.0, 60.0, -24.0, 3.0]) - self.coeff_i[a][3, :] = 1 / 18 * np.array([-1.0, 8.0, -24.0, 40.0, -5.0]) - self.coeff_i[a][4, :] = 1 / 18 * np.array([0.0, 0.0, 0.0, 0.0, 18.0]) - - self.coeff_h[a][0, :] = 1 / 18 * np.array([23.0, -17.0, 7.0, -1.0, 0.0, 0.0]) - self.coeff_h[a][1, :] = 1 / 18 * np.array([-8.0, 56.0, -28.0, 4.0, 0.0, 0.0]) - self.coeff_h[a][2, :] = 1 / 18 * np.array([3.0, -21.0, 36.0, 36.0, -21.0, 3.0]) - self.coeff_h[a][3, :] = 1 / 18 * np.array([0.0, 0.0, 4.0, -28.0, 56.0, -8.0]) - self.coeff_h[a][4, :] = 1 / 18 * np.array([0.0, 0.0, -1.0, 7.0, -17.0, 23.0]) + self.coeff_i[a][0, :] = 1 / 18 * xp.array([18.0, 0.0, 0.0, 0.0, 0.0]) + self.coeff_i[a][1, :] = 1 / 18 * xp.array([-5.0, 40.0, -24.0, 8.0, -1.0]) + self.coeff_i[a][2, :] = 1 / 18 * xp.array([3.0, -24.0, 60.0, -24.0, 3.0]) + self.coeff_i[a][3, :] = 1 / 18 * xp.array([-1.0, 8.0, -24.0, 40.0, -5.0]) + self.coeff_i[a][4, :] = 1 / 18 * xp.array([0.0, 0.0, 0.0, 0.0, 18.0]) + + self.coeff_h[a][0, :] = 1 / 18 * xp.array([23.0, -17.0, 7.0, -1.0, 0.0, 0.0]) + self.coeff_h[a][1, :] = 1 / 18 * xp.array([-8.0, 56.0, -28.0, 4.0, 0.0, 0.0]) + self.coeff_h[a][2, :] = 1 / 18 * xp.array([3.0, -21.0, 36.0, 36.0, -21.0, 3.0]) + self.coeff_h[a][3, :] = 1 / 18 * xp.array([0.0, 0.0, 4.0, -28.0, 56.0, -8.0]) + self.coeff_h[a][4, :] = 1 / 18 * xp.array([0.0, 0.0, -1.0, 7.0, -17.0, 23.0]) elif self.p[a] == 4: - self.coeff_i[a][0, :] = 1 / 360 * np.array([360.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) - self.coeff_i[a][1, :] = 1 / 360 * np.array([-59.0, 944.0, -1000.0, 720.0, -305.0, 64.0, -4.0]) - self.coeff_i[a][2, :] = 1 / 360 * np.array([23.0, -368.0, 1580.0, -1360.0, 605.0, -128.0, 8.0]) - self.coeff_i[a][3, :] = 1 / 360 * np.array([-16.0, 256.0, -1180.0, 2240.0, -1180.0, 256.0, -16.0]) - self.coeff_i[a][4, :] = 1 / 360 * np.array([8.0, -128.0, 605.0, -1360.0, 1580.0, -368.0, 23.0]) - self.coeff_i[a][5, :] = 1 / 360 * np.array([-4.0, 64.0, -305.0, 720.0, -1000.0, 944.0, -59.0]) - self.coeff_i[a][6, :] = 1 / 360 * np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 360.0]) - - self.coeff_h[a][0, :] = 1 / 360 * np.array([419.0, -525.0, 475.0, -245.0, 60.0, -4.0, 0.0, 0.0]) - self.coeff_h[a][1, :] = 1 / 360 * np.array([-82.0, 1230.0, -1350.0, 730.0, -180.0, 12.0, 0.0, 0.0]) - self.coeff_h[a][2, :] = 1 / 360 * np.array([39.0, -585.0, 2175.0, -1425.0, 360.0, -24.0, 0.0, 0.0]) + self.coeff_i[a][0, :] = 1 / 360 * xp.array([360.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + self.coeff_i[a][1, :] = 1 / 360 * xp.array([-59.0, 944.0, -1000.0, 720.0, -305.0, 64.0, -4.0]) + self.coeff_i[a][2, :] = 1 / 360 * xp.array([23.0, -368.0, 1580.0, -1360.0, 605.0, -128.0, 8.0]) + self.coeff_i[a][3, :] = 1 / 360 * xp.array([-16.0, 256.0, -1180.0, 2240.0, -1180.0, 256.0, -16.0]) + self.coeff_i[a][4, :] = 1 / 360 * xp.array([8.0, -128.0, 605.0, -1360.0, 1580.0, -368.0, 23.0]) + self.coeff_i[a][5, :] = 1 / 360 * xp.array([-4.0, 64.0, -305.0, 720.0, -1000.0, 944.0, -59.0]) + self.coeff_i[a][6, :] = 1 / 360 * xp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 360.0]) + + self.coeff_h[a][0, :] = 1 / 360 * xp.array([419.0, -525.0, 475.0, -245.0, 60.0, -4.0, 0.0, 0.0]) + self.coeff_h[a][1, :] = 1 / 360 * xp.array([-82.0, 1230.0, -1350.0, 730.0, -180.0, 12.0, 0.0, 0.0]) + self.coeff_h[a][2, :] = 1 / 360 * xp.array([39.0, -585.0, 2175.0, -1425.0, 360.0, -24.0, 0.0, 0.0]) self.coeff_h[a][3, :] = ( - 1 / 360 * np.array([-16.0, 240.0, -924.0, 1060.0, 1060.0, -924.0, 240.0, -16.0]) + 1 / 360 * xp.array([-16.0, 240.0, -924.0, 1060.0, 1060.0, -924.0, 240.0, -16.0]) ) - self.coeff_h[a][4, :] = 1 / 360 * np.array([0.0, 0.0, -24.0, 360.0, -1425.0, 2175.0, -585.0, 39.0]) - self.coeff_h[a][5, :] = 1 / 360 * np.array([0.0, 0.0, 12.0, -180.0, 730.0, -1350.0, 1230.0, -82.0]) - self.coeff_h[a][6, :] = 1 / 360 * np.array([0.0, 0.0, -4.0, 60.0, -245.0, 475.0, -525.0, 419.0]) + self.coeff_h[a][4, :] = 1 / 360 * xp.array([0.0, 0.0, -24.0, 360.0, -1425.0, 2175.0, -585.0, 39.0]) + self.coeff_h[a][5, :] = 1 / 360 * xp.array([0.0, 0.0, 12.0, -180.0, 730.0, -1350.0, 1230.0, -82.0]) + self.coeff_h[a][6, :] = 1 / 360 * xp.array([0.0, 0.0, -4.0, 60.0, -245.0, 475.0, -525.0, 419.0]) else: print("degree > 4 not implemented!") @@ -150,31 +150,31 @@ def __init__(self, tensor_space, n_quad): ) # number of non-vanishing D bf in interpolation interval (1, 2, 4, 6) self.x_int = [ - np.zeros((n_lambda_int, n_int), dtype=float) for n_lambda_int, n_int in zip(n_lambda_int, self.n_int) + xp.zeros((n_lambda_int, n_int), dtype=float) for n_lambda_int, n_int in zip(n_lambda_int, self.n_int) ] self.int_global_N = [ - np.zeros((n_lambda_int, n_int_locbf_N), dtype=int) + xp.zeros((n_lambda_int, n_int_locbf_N), dtype=int) for n_lambda_int, n_int_locbf_N in zip(n_lambda_int, self.n_int_locbf_N) ] self.int_global_D = [ - np.zeros((n_lambda_int, n_int_locbf_D), dtype=int) + xp.zeros((n_lambda_int, n_int_locbf_D), dtype=int) for n_lambda_int, n_int_locbf_D in zip(n_lambda_int, self.n_int_locbf_D) ] self.int_loccof_N = [ - np.zeros((n_lambda_int, n_int_locbf_N), dtype=int) + xp.zeros((n_lambda_int, n_int_locbf_N), dtype=int) for n_lambda_int, n_int_locbf_N in zip(n_lambda_int, self.n_int_locbf_N) ] self.int_loccof_D = [ - np.zeros((n_lambda_int, n_int_locbf_D), dtype=int) + xp.zeros((n_lambda_int, n_int_locbf_D), dtype=int) for n_lambda_int, n_int_locbf_D in zip(n_lambda_int, self.n_int_locbf_D) ] self.x_int_indices = [ - np.zeros((n_lambda_int, n_int), dtype=int) for n_lambda_int, n_int in zip(n_lambda_int, self.n_int) + xp.zeros((n_lambda_int, n_int), dtype=int) for n_lambda_int, n_int in zip(n_lambda_int, self.n_int) ] - self.coeffi_indices = [np.zeros(n_lambda_int, dtype=int) for n_lambda_int in n_lambda_int] + self.coeffi_indices = [xp.zeros(n_lambda_int, dtype=int) for n_lambda_int in n_lambda_int] self.n_int_nvcof_D = [None, None, None] self.n_int_nvcof_N = [None, None, None] @@ -197,39 +197,39 @@ def __init__(self, tensor_space, n_quad): self.n_int_nvcof_N[a] = 3 * self.p[a] - 2 # shift in local coefficient indices at right boundary (only for non-periodic boundary conditions) - self.int_add_D[a] = np.arange(self.n_int[a] - 2) + 1 - self.int_add_N[a] = np.arange(self.n_int[a] - 1) + 1 + self.int_add_D[a] = xp.arange(self.n_int[a] - 2) + 1 + self.int_add_N[a] = xp.arange(self.n_int[a] - 1) + 1 counter_D = 0 counter_N = 0 # shift local coefficients --> global coefficients (D) if self.p[a] == 1: - self.int_shift_D[a] = np.arange(self.NbaseD[a]) + self.int_shift_D[a] = xp.arange(self.NbaseD[a]) else: - self.int_shift_D[a] = np.arange(self.NbaseD[a]) - (self.p[a] - 2) + self.int_shift_D[a] = xp.arange(self.NbaseD[a]) - (self.p[a] - 2) self.int_shift_D[a][: 2 * self.p[a] - 2] = 0 self.int_shift_D[a][-(2 * self.p[a] - 2) :] = self.int_shift_D[a][-(2 * self.p[a] - 2)] # shift local coefficients --> global coefficients (N) if self.p[a] == 1: - self.int_shift_N[a] = np.arange(self.NbaseN[a]) + self.int_shift_N[a] = xp.arange(self.NbaseN[a]) self.int_shift_N[a][-1] = self.int_shift_N[a][-2] else: - self.int_shift_N[a] = np.arange(self.NbaseN[a]) - (self.p[a] - 1) + self.int_shift_N[a] = xp.arange(self.NbaseN[a]) - (self.p[a] - 1) self.int_shift_N[a][: 2 * self.p[a] - 1] = 0 self.int_shift_N[a][-(2 * self.p[a] - 1) :] = self.int_shift_N[a][-(2 * self.p[a] - 1)] - counter_coeffi = np.copy(self.p[a]) + counter_coeffi = xp.copy(self.p[a]) for i in range(n_lambda_int[a]): # left boundary region if i < self.p[a] - 1: - self.int_global_N[a][i] = np.arange(self.n_int_locbf_N[a]) - self.int_global_D[a][i] = np.arange(self.n_int_locbf_D[a]) + self.int_global_N[a][i] = xp.arange(self.n_int_locbf_N[a]) + self.int_global_D[a][i] = xp.arange(self.n_int_locbf_D[a]) - self.x_int_indices[a][i] = np.arange(self.n_int[a]) + self.x_int_indices[a][i] = xp.arange(self.n_int[a]) self.coeffi_indices[a][i] = i for j in range(2 * (self.p[a] - 1) + 1): xi = self.p[a] - 1 @@ -240,13 +240,13 @@ def __init__(self, tensor_space, n_quad): # right boundary region elif i > n_lambda_int[a] - self.p[a]: self.int_global_N[a][i] = ( - np.arange(self.n_int_locbf_N[a]) + n_lambda_int[a] - self.p[a] - (self.p[a] - 1) + xp.arange(self.n_int_locbf_N[a]) + n_lambda_int[a] - self.p[a] - (self.p[a] - 1) ) self.int_global_D[a][i] = ( - np.arange(self.n_int_locbf_D[a]) + n_lambda_int[a] - self.p[a] - (self.p[a] - 1) + xp.arange(self.n_int_locbf_D[a]) + n_lambda_int[a] - self.p[a] - (self.p[a] - 1) ) - self.x_int_indices[a][i] = np.arange(self.n_int[a]) + 2 * ( + self.x_int_indices[a][i] = xp.arange(self.n_int[a]) + 2 * ( n_lambda_int[a] - self.p[a] - (self.p[a] - 1) ) self.coeffi_indices[a][i] = counter_coeffi @@ -260,20 +260,20 @@ def __init__(self, tensor_space, n_quad): # interior else: if self.p[a] == 1: - self.int_global_N[a][i] = np.arange(self.n_int_locbf_N[a]) + i - self.int_global_D[a][i] = np.arange(self.n_int_locbf_D[a]) + i + self.int_global_N[a][i] = xp.arange(self.n_int_locbf_N[a]) + i + self.int_global_D[a][i] = xp.arange(self.n_int_locbf_D[a]) + i self.int_global_N[a][-1] = self.int_global_N[a][-2] self.int_global_D[a][-1] = self.int_global_D[a][-2] else: - self.int_global_N[a][i] = np.arange(self.n_int_locbf_N[a]) + i - (self.p[a] - 1) - self.int_global_D[a][i] = np.arange(self.n_int_locbf_D[a]) + i - (self.p[a] - 1) + self.int_global_N[a][i] = xp.arange(self.n_int_locbf_N[a]) + i - (self.p[a] - 1) + self.int_global_D[a][i] = xp.arange(self.n_int_locbf_D[a]) + i - (self.p[a] - 1) if self.p[a] == 1: self.x_int_indices[a][i] = i else: - self.x_int_indices[a][i] = np.arange(self.n_int[a]) + 2 * (i - (self.p[a] - 1)) + self.x_int_indices[a][i] = xp.arange(self.n_int[a]) + 2 * (i - (self.p[a] - 1)) self.coeffi_indices[a][i] = self.p[a] - 1 @@ -284,8 +284,8 @@ def __init__(self, tensor_space, n_quad): # local coefficient index if self.p[a] == 1: - self.int_loccof_N[a][i] = np.array([0, 1]) - self.int_loccof_D[a][-1] = np.array([1]) + self.int_loccof_N[a][i] = xp.array([0, 1]) + self.int_loccof_D[a][-1] = xp.array([1]) else: if i > 0: @@ -293,8 +293,8 @@ def __init__(self, tensor_space, n_quad): k_glob_new = self.int_global_D[a][i, il] bol = k_glob_new == self.int_global_D[a][i - 1] - if np.any(bol): - self.int_loccof_D[a][i, il] = self.int_loccof_D[a][i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.int_loccof_D[a][i, il] = self.int_loccof_D[a][i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_int[a] - self.p[a] - (self.p[a] - 2)) and ( self.int_loccof_D[a][i, il] == 0 @@ -306,8 +306,8 @@ def __init__(self, tensor_space, n_quad): k_glob_new = self.int_global_N[a][i, il] bol = k_glob_new == self.int_global_N[a][i - 1] - if np.any(bol): - self.int_loccof_N[a][i, il] = self.int_loccof_N[a][i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.int_loccof_N[a][i, il] = self.int_loccof_N[a][i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_int[a] - self.p[a] - (self.p[a] - 2)) and ( self.int_loccof_N[a][i, il] == 0 @@ -327,24 +327,24 @@ def __init__(self, tensor_space, n_quad): # shift local coefficients --> global coefficients if self.p[a] == 1: - self.int_shift_D[a] = np.arange(self.NbaseN[a]) - (self.p[a] - 1) - self.int_shift_N[a] = np.arange(self.NbaseN[a]) - (self.p[a]) + self.int_shift_D[a] = xp.arange(self.NbaseN[a]) - (self.p[a] - 1) + self.int_shift_N[a] = xp.arange(self.NbaseN[a]) - (self.p[a]) else: - self.int_shift_D[a] = np.arange(self.NbaseN[a]) - (self.p[a] - 2) - self.int_shift_N[a] = np.arange(self.NbaseN[a]) - (self.p[a] - 1) + self.int_shift_D[a] = xp.arange(self.NbaseN[a]) - (self.p[a] - 2) + self.int_shift_N[a] = xp.arange(self.NbaseN[a]) - (self.p[a] - 1) for i in range(n_lambda_int[a]): # global indices of non-vanishing basis functions and position of coefficients in final matrix - self.int_global_N[a][i] = (np.arange(self.n_int_locbf_N[a]) + i - (self.p[a] - 1)) % self.NbaseN[a] - self.int_global_D[a][i] = (np.arange(self.n_int_locbf_D[a]) + i - (self.p[a] - 1)) % self.NbaseD[a] + self.int_global_N[a][i] = (xp.arange(self.n_int_locbf_N[a]) + i - (self.p[a] - 1)) % self.NbaseN[a] + self.int_global_D[a][i] = (xp.arange(self.n_int_locbf_D[a]) + i - (self.p[a] - 1)) % self.NbaseD[a] - self.int_loccof_N[a][i] = np.arange(self.n_int_locbf_N[a] - 1, -1, -1) - self.int_loccof_D[a][i] = np.arange(self.n_int_locbf_D[a] - 1, -1, -1) + self.int_loccof_N[a][i] = xp.arange(self.n_int_locbf_N[a] - 1, -1, -1) + self.int_loccof_D[a][i] = xp.arange(self.n_int_locbf_D[a] - 1, -1, -1) if self.p[a] == 1: self.x_int_indices[a][i] = i else: - self.x_int_indices[a][i] = (np.arange(self.n_int[a]) + 2 * (i - (self.p[a] - 1))) % ( + self.x_int_indices[a][i] = (xp.arange(self.n_int[a]) + 2 * (i - (self.p[a] - 1))) % ( 2 * self.Nel[a] ) @@ -356,41 +356,41 @@ def __init__(self, tensor_space, n_quad): ) % 1.0 # identify unique interpolation points to save memory - self.x_int[a] = np.unique(self.x_int[a].flatten()) + self.x_int[a] = xp.unique(self.x_int[a].flatten()) # set histopolation points, quadrature points and weights - n_lambda_his = [np.copy(NbaseD) for NbaseD in self.NbaseD] # number of coefficients in space V1 + n_lambda_his = [xp.copy(NbaseD) for NbaseD in self.NbaseD] # number of coefficients in space V1 self.n_his = [2 * p for p in self.p] # number of histopolation intervals self.n_his_locbf_N = [2 * p for p in self.p] # number of non-vanishing N bf in histopolation interval self.n_his_locbf_D = [2 * p - 1 for p in self.p] # number of non-vanishing D bf in histopolation interval self.x_his = [ - np.zeros((n_lambda_his, n_his + 1), dtype=float) for n_lambda_his, n_his in zip(n_lambda_his, self.n_his) + xp.zeros((n_lambda_his, n_his + 1), dtype=float) for n_lambda_his, n_his in zip(n_lambda_his, self.n_his) ] self.his_global_N = [ - np.zeros((n_lambda_his, n_his_locbf_N), dtype=int) + xp.zeros((n_lambda_his, n_his_locbf_N), dtype=int) for n_lambda_his, n_his_locbf_N in zip(n_lambda_his, self.n_his_locbf_N) ] self.his_global_D = [ - np.zeros((n_lambda_his, n_his_locbf_D), dtype=int) + xp.zeros((n_lambda_his, n_his_locbf_D), dtype=int) for n_lambda_his, n_his_locbf_D in zip(n_lambda_his, self.n_his_locbf_D) ] self.his_loccof_N = [ - np.zeros((n_lambda_his, n_his_locbf_N), dtype=int) + xp.zeros((n_lambda_his, n_his_locbf_N), dtype=int) for n_lambda_his, n_his_locbf_N in zip(n_lambda_his, self.n_his_locbf_N) ] self.his_loccof_D = [ - np.zeros((n_lambda_his, n_his_locbf_D), dtype=int) + xp.zeros((n_lambda_his, n_his_locbf_D), dtype=int) for n_lambda_his, n_his_locbf_D in zip(n_lambda_his, self.n_his_locbf_D) ] self.x_his_indices = [ - np.zeros((n_lambda_his, n_his), dtype=int) for n_lambda_his, n_his in zip(n_lambda_his, self.n_his) + xp.zeros((n_lambda_his, n_his), dtype=int) for n_lambda_his, n_his in zip(n_lambda_his, self.n_his) ] - self.coeffh_indices = [np.zeros(n_lambda_his, dtype=int) for n_lambda_his in n_lambda_his] + self.coeffh_indices = [xp.zeros(n_lambda_his, dtype=int) for n_lambda_his in n_lambda_his] self.pts = [0, 0, 0] self.wts = [0, 0, 0] @@ -411,31 +411,31 @@ def __init__(self, tensor_space, n_quad): self.n_his_nvcof_N[a] = 3 * self.p[a] - 1 # shift in local coefficient indices at right boundary (only for non-periodic boundary conditions) - self.his_add_D[a] = np.arange(self.n_his[a] - 2) + 1 - self.his_add_N[a] = np.arange(self.n_his[a] - 1) + 1 + self.his_add_D[a] = xp.arange(self.n_his[a] - 2) + 1 + self.his_add_N[a] = xp.arange(self.n_his[a] - 1) + 1 counter_D = 0 counter_N = 0 # shift local coefficients --> global coefficients (D) - self.his_shift_D[a] = np.arange(self.NbaseD[a]) - (self.p[a] - 1) + self.his_shift_D[a] = xp.arange(self.NbaseD[a]) - (self.p[a] - 1) self.his_shift_D[a][: 2 * self.p[a] - 1] = 0 self.his_shift_D[a][-(2 * self.p[a] - 1) :] = self.his_shift_D[a][-(2 * self.p[a] - 1)] # shift local coefficients --> global coefficients (N) - self.his_shift_N[a] = np.arange(self.NbaseN[a]) - self.p[a] + self.his_shift_N[a] = xp.arange(self.NbaseN[a]) - self.p[a] self.his_shift_N[a][: 2 * self.p[a]] = 0 self.his_shift_N[a][-2 * self.p[a] :] = self.his_shift_N[a][-2 * self.p[a]] - counter_coeffh = np.copy(self.p[a]) + counter_coeffh = xp.copy(self.p[a]) for i in range(n_lambda_his[a]): # left boundary region if i < self.p[a] - 1: - self.his_global_N[a][i] = np.arange(self.n_his_locbf_N[a]) - self.his_global_D[a][i] = np.arange(self.n_his_locbf_D[a]) + self.his_global_N[a][i] = xp.arange(self.n_his_locbf_N[a]) + self.his_global_D[a][i] = xp.arange(self.n_his_locbf_D[a]) - self.x_his_indices[a][i] = np.arange(self.n_his[a]) + self.x_his_indices[a][i] = xp.arange(self.n_his[a]) self.coeffh_indices[a][i] = i for j in range(2 * self.p[a] + 1): xi = self.p[a] - 1 @@ -446,13 +446,13 @@ def __init__(self, tensor_space, n_quad): # right boundary region elif i > n_lambda_his[a] - self.p[a]: self.his_global_N[a][i] = ( - np.arange(self.n_his_locbf_N[a]) + n_lambda_his[a] - self.p[a] - (self.p[a] - 1) + xp.arange(self.n_his_locbf_N[a]) + n_lambda_his[a] - self.p[a] - (self.p[a] - 1) ) self.his_global_D[a][i] = ( - np.arange(self.n_his_locbf_D[a]) + n_lambda_his[a] - self.p[a] - (self.p[a] - 1) + xp.arange(self.n_his_locbf_D[a]) + n_lambda_his[a] - self.p[a] - (self.p[a] - 1) ) - self.x_his_indices[a][i] = np.arange(self.n_his[a]) + 2 * ( + self.x_his_indices[a][i] = xp.arange(self.n_his[a]) + 2 * ( n_lambda_his[a] - self.p[a] - (self.p[a] - 1) ) self.coeffh_indices[a][i] = counter_coeffh @@ -465,10 +465,10 @@ def __init__(self, tensor_space, n_quad): # interior else: - self.his_global_N[a][i] = np.arange(self.n_his_locbf_N[a]) + i - (self.p[a] - 1) - self.his_global_D[a][i] = np.arange(self.n_his_locbf_D[a]) + i - (self.p[a] - 1) + self.his_global_N[a][i] = xp.arange(self.n_his_locbf_N[a]) + i - (self.p[a] - 1) + self.his_global_D[a][i] = xp.arange(self.n_his_locbf_D[a]) + i - (self.p[a] - 1) - self.x_his_indices[a][i] = np.arange(self.n_his[a]) + 2 * (i - (self.p[a] - 1)) + self.x_his_indices[a][i] = xp.arange(self.n_his[a]) + 2 * (i - (self.p[a] - 1)) self.coeffh_indices[a][i] = self.p[a] - 1 for j in range(2 * self.p[a] + 1): self.x_his[a][i, j] = ( @@ -481,8 +481,8 @@ def __init__(self, tensor_space, n_quad): k_glob_new = self.his_global_D[a][i, il] bol = k_glob_new == self.his_global_D[a][i - 1] - if np.any(bol): - self.his_loccof_D[a][i, il] = self.his_loccof_D[a][i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.his_loccof_D[a][i, il] = self.his_loccof_D[a][i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_his[a] - self.p[a] - (self.p[a] - 2)) and ( self.his_loccof_D[a][i, il] == 0 @@ -494,8 +494,8 @@ def __init__(self, tensor_space, n_quad): k_glob_new = self.his_global_N[a][i, il] bol = k_glob_new == self.his_global_N[a][i - 1] - if np.any(bol): - self.his_loccof_N[a][i, il] = self.his_loccof_N[a][i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.his_loccof_N[a][i, il] = self.his_loccof_N[a][i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_his[a] - self.p[a] - (self.p[a] - 2)) and ( self.his_loccof_N[a][i, il] == 0 @@ -505,7 +505,7 @@ def __init__(self, tensor_space, n_quad): # quadrature points and weights self.pts[a], self.wts[a] = bsp.quadrature_grid( - np.unique(self.x_his[a].flatten()), self.pts_loc[a], self.wts_loc[a] + xp.unique(self.x_his[a].flatten()), self.pts_loc[a], self.wts_loc[a] ) else: @@ -514,18 +514,18 @@ def __init__(self, tensor_space, n_quad): self.n_his_nvcof_N[a] = 2 * self.p[a] # shift local coefficients --> global coefficients (D) - self.his_shift_D[a] = np.arange(self.NbaseD[a]) - (self.p[a] - 1) + self.his_shift_D[a] = xp.arange(self.NbaseD[a]) - (self.p[a] - 1) # shift local coefficients --> global coefficients (N) - self.his_shift_N[a] = np.arange(self.NbaseD[a]) - self.p[a] + self.his_shift_N[a] = xp.arange(self.NbaseD[a]) - self.p[a] for i in range(n_lambda_his[a]): - self.his_global_N[a][i] = (np.arange(self.n_his_locbf_N[a]) + i - (self.p[a] - 1)) % self.NbaseN[a] - self.his_global_D[a][i] = (np.arange(self.n_his_locbf_D[a]) + i - (self.p[a] - 1)) % self.NbaseD[a] - self.his_loccof_N[a][i] = np.arange(self.n_his_locbf_N[a] - 1, -1, -1) - self.his_loccof_D[a][i] = np.arange(self.n_his_locbf_D[a] - 1, -1, -1) + self.his_global_N[a][i] = (xp.arange(self.n_his_locbf_N[a]) + i - (self.p[a] - 1)) % self.NbaseN[a] + self.his_global_D[a][i] = (xp.arange(self.n_his_locbf_D[a]) + i - (self.p[a] - 1)) % self.NbaseD[a] + self.his_loccof_N[a][i] = xp.arange(self.n_his_locbf_N[a] - 1, -1, -1) + self.his_loccof_D[a][i] = xp.arange(self.n_his_locbf_D[a] - 1, -1, -1) - self.x_his_indices[a][i] = (np.arange(self.n_his[a]) + 2 * (i - (self.p[a] - 1))) % ( + self.x_his_indices[a][i] = (xp.arange(self.n_his[a]) + 2 * (i - (self.p[a] - 1))) % ( 2 * self.Nel[a] ) self.coeffh_indices[a][i] = 0 @@ -535,7 +535,7 @@ def __init__(self, tensor_space, n_quad): # quadrature points and weights self.pts[a], self.wts[a] = bsp.quadrature_grid( - np.append(np.unique(self.x_his[a].flatten() % 1.0), 1.0), self.pts_loc[a], self.wts_loc[a] + xp.append(xp.unique(self.x_his[a].flatten() % 1.0), 1.0), self.pts_loc[a], self.wts_loc[a] ) # evaluate N basis functions at interpolation and quadrature points @@ -586,7 +586,7 @@ def projection_Q_0form(self, domain): """ # non-vanishing coefficients - Q11 = np.empty( + Q11 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -597,7 +597,7 @@ def projection_Q_0form(self, domain): ), dtype=float, ) - Q22 = np.empty( + Q22 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -608,7 +608,7 @@ def projection_Q_0form(self, domain): ), dtype=float, ) - Q33 = np.empty( + Q33 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -626,7 +626,7 @@ def projection_Q_0form(self, domain): n_unique3 = [self.pts[0].flatten().size, self.pts[1].flatten().size, self.x_int[2].size] # ========= assembly of 1 - component (pi2_1 : int, his, his) ============ - mat_eq = np.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) + mat_eq = xp.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -682,7 +682,7 @@ def projection_Q_0form(self, domain): ) # ========= assembly of 2 - component (pi2_2 : his, int, his) ============ - mat_eq = np.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) + mat_eq = xp.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -738,7 +738,7 @@ def projection_Q_0form(self, domain): ) # ========= assembly of 3 - component (pi2_3 : his, his, int) ============ - mat_eq = np.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) + mat_eq = xp.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -794,7 +794,7 @@ def projection_Q_0form(self, domain): ) # ========= conversion to sparse matrices (1 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -819,7 +819,7 @@ def projection_Q_0form(self, domain): Q11.eliminate_zeros() # ========= conversion to sparse matrices (2 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -844,7 +844,7 @@ def projection_Q_0form(self, domain): Q22.eliminate_zeros() # ========= conversion to sparse matrices (3 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -895,7 +895,7 @@ def projection_Q_2form(self, domain): """ # non-vanishing coefficients - Q11 = np.empty( + Q11 = xp.empty( ( self.NbaseN[0], self.NbaseD[1], @@ -906,7 +906,7 @@ def projection_Q_2form(self, domain): ), dtype=float, ) - Q22 = np.empty( + Q22 = xp.empty( ( self.NbaseD[0], self.NbaseN[1], @@ -917,7 +917,7 @@ def projection_Q_2form(self, domain): ), dtype=float, ) - Q33 = np.empty( + Q33 = xp.empty( ( self.NbaseD[0], self.NbaseD[1], @@ -935,7 +935,7 @@ def projection_Q_2form(self, domain): n_unique3 = [self.pts[0].flatten().size, self.pts[1].flatten().size, self.x_int[2].size] # ========= assembly of 1 - component (pi2_1 : int, his, his) ============ - mat_eq = np.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) + mat_eq = xp.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -991,7 +991,7 @@ def projection_Q_2form(self, domain): ) # ========= assembly of 2 - component (pi2_2 : his, int, his) ============ - mat_eq = np.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) + mat_eq = xp.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -1047,7 +1047,7 @@ def projection_Q_2form(self, domain): ) # ========= assembly of 3 - component (pi2_3 : his, his, int) ============ - mat_eq = np.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) + mat_eq = xp.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -1103,7 +1103,7 @@ def projection_Q_2form(self, domain): ) # ========= conversion to sparse matrices (1 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseD[1], @@ -1128,7 +1128,7 @@ def projection_Q_2form(self, domain): Q11.eliminate_zeros() # ========= conversion to sparse matrices (2 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseN[1], @@ -1153,7 +1153,7 @@ def projection_Q_2form(self, domain): Q22.eliminate_zeros() # ========= conversion to sparse matrices (3 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseD[1], @@ -1204,7 +1204,7 @@ def projection_W_0form(self, domain): """ # non-vanishing coefficients - W1 = np.empty( + W1 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -1215,14 +1215,14 @@ def projection_W_0form(self, domain): ), dtype=float, ) - # W2 = np.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2], self.n_int_nvcof_N[0], self.n_int_nvcof_N[1], self.n_int_nvcof_N[2]), dtype=float) - # W3 = np.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2], self.n_int_nvcof_N[0], self.n_int_nvcof_N[1], self.n_int_nvcof_N[2]), dtype=float) + # W2 = xp.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2], self.n_int_nvcof_N[0], self.n_int_nvcof_N[1], self.n_int_nvcof_N[2]), dtype=float) + # W3 = xp.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2], self.n_int_nvcof_N[0], self.n_int_nvcof_N[1], self.n_int_nvcof_N[2]), dtype=float) # size of interpolation/quadrature points of the 3 components n_unique = [self.x_int[0].size, self.x_int[1].size, self.x_int[2].size] # assembly - mat_eq = np.empty((n_unique[0], n_unique[1], n_unique[2]), dtype=float) + mat_eq = xp.empty((n_unique[0], n_unique[1], n_unique[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -1284,7 +1284,7 @@ def projection_W_0form(self, domain): """ # conversion to sparse matrix - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -1345,7 +1345,7 @@ def projection_T_0form(self, domain): """ # non-vanishing coefficients - T12 = np.empty( + T12 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -1356,7 +1356,7 @@ def projection_T_0form(self, domain): ), dtype=float, ) - T13 = np.empty( + T13 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -1368,7 +1368,7 @@ def projection_T_0form(self, domain): dtype=float, ) - T21 = np.empty( + T21 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -1379,7 +1379,7 @@ def projection_T_0form(self, domain): ), dtype=float, ) - T23 = np.empty( + T23 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -1391,7 +1391,7 @@ def projection_T_0form(self, domain): dtype=float, ) - T31 = np.empty( + T31 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -1402,7 +1402,7 @@ def projection_T_0form(self, domain): ), dtype=float, ) - T32 = np.empty( + T32 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -1420,7 +1420,7 @@ def projection_T_0form(self, domain): n_unique3 = [self.x_int[0].size, self.x_int[1].size, self.pts[2].flatten().size] # ================= assembly of 1 - component (pi1_1 : his, int, int) ============ - mat_eq = np.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) + mat_eq = xp.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -1515,7 +1515,7 @@ def projection_T_0form(self, domain): ) # ================= assembly of 2 - component (PI_1_2 : int, his, int) ============ - mat_eq = np.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) + mat_eq = xp.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -1621,7 +1621,7 @@ def projection_T_0form(self, domain): ) # ================= assembly of 3 - component (PI_1_3 : int, int, his) ============ - mat_eq = np.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) + mat_eq = xp.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -1727,7 +1727,7 @@ def projection_T_0form(self, domain): ) # conversion to sparse matrices (1 - component) - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -1751,7 +1751,7 @@ def projection_T_0form(self, domain): ) T12.eliminate_zeros() - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -1776,7 +1776,7 @@ def projection_T_0form(self, domain): T13.eliminate_zeros() # conversion to sparse matrices (2 - component) - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -1800,7 +1800,7 @@ def projection_T_0form(self, domain): ) T21.eliminate_zeros() - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -1825,7 +1825,7 @@ def projection_T_0form(self, domain): T23.eliminate_zeros() # conversion to sparse matrices (3 - component) - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -1849,7 +1849,7 @@ def projection_T_0form(self, domain): ) T31.eliminate_zeros() - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -1900,7 +1900,7 @@ def projection_T_1form(self, domain): """ # non-vanishing coefficients - T12 = np.empty( + T12 = xp.empty( ( self.NbaseN[0], self.NbaseD[1], @@ -1911,7 +1911,7 @@ def projection_T_1form(self, domain): ), dtype=float, ) - T13 = np.empty( + T13 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -1923,7 +1923,7 @@ def projection_T_1form(self, domain): dtype=float, ) - T21 = np.empty( + T21 = xp.empty( ( self.NbaseD[0], self.NbaseN[1], @@ -1934,7 +1934,7 @@ def projection_T_1form(self, domain): ), dtype=float, ) - T23 = np.empty( + T23 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -1946,7 +1946,7 @@ def projection_T_1form(self, domain): dtype=float, ) - T31 = np.empty( + T31 = xp.empty( ( self.NbaseD[0], self.NbaseN[1], @@ -1957,7 +1957,7 @@ def projection_T_1form(self, domain): ), dtype=float, ) - T32 = np.empty( + T32 = xp.empty( ( self.NbaseN[0], self.NbaseD[1], @@ -1975,7 +1975,7 @@ def projection_T_1form(self, domain): n_unique3 = [self.x_int[0].size, self.x_int[1].size, self.pts[2].flatten().size] # ================= assembly of 1 - component (pi1_1 : his, int, int) ============ - mat_eq = np.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) + mat_eq = xp.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -2070,7 +2070,7 @@ def projection_T_1form(self, domain): ) # ================= assembly of 2 - component (PI_1_2 : int, his, int) ============ - mat_eq = np.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) + mat_eq = xp.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -2165,7 +2165,7 @@ def projection_T_1form(self, domain): ) # ================= assembly of 3 - component (PI_1_3 : int, int, his) ============ - mat_eq = np.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) + mat_eq = xp.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -2260,7 +2260,7 @@ def projection_T_1form(self, domain): ) # conversion to sparse matrices (1 - component) - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseD[1], @@ -2284,7 +2284,7 @@ def projection_T_1form(self, domain): ) T12.eliminate_zeros() - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -2309,7 +2309,7 @@ def projection_T_1form(self, domain): T13.eliminate_zeros() # conversion to sparse matrices (2 - component) - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseN[1], @@ -2333,7 +2333,7 @@ def projection_T_1form(self, domain): ) T21.eliminate_zeros() - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -2358,7 +2358,7 @@ def projection_T_1form(self, domain): T23.eliminate_zeros() # conversion to sparse matrices (3 - component) - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseN[1], @@ -2382,7 +2382,7 @@ def projection_T_1form(self, domain): ) T31.eliminate_zeros() - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseD[1], @@ -2433,7 +2433,7 @@ def projection_T_2form(self, domain): """ # non-vanishing coefficients - T12 = np.empty( + T12 = xp.empty( ( self.NbaseD[0], self.NbaseN[1], @@ -2444,7 +2444,7 @@ def projection_T_2form(self, domain): ), dtype=float, ) - T13 = np.empty( + T13 = xp.empty( ( self.NbaseD[0], self.NbaseD[1], @@ -2456,7 +2456,7 @@ def projection_T_2form(self, domain): dtype=float, ) - T21 = np.empty( + T21 = xp.empty( ( self.NbaseN[0], self.NbaseD[1], @@ -2467,7 +2467,7 @@ def projection_T_2form(self, domain): ), dtype=float, ) - T23 = np.empty( + T23 = xp.empty( ( self.NbaseD[0], self.NbaseD[1], @@ -2479,7 +2479,7 @@ def projection_T_2form(self, domain): dtype=float, ) - T31 = np.empty( + T31 = xp.empty( ( self.NbaseN[0], self.NbaseD[1], @@ -2490,7 +2490,7 @@ def projection_T_2form(self, domain): ), dtype=float, ) - T32 = np.empty( + T32 = xp.empty( ( self.NbaseD[0], self.NbaseN[1], @@ -2508,7 +2508,7 @@ def projection_T_2form(self, domain): n_unique3 = [self.x_int[0].size, self.x_int[1].size, self.pts[2].flatten().size] # ================= assembly of 1 - component (pi1_1 : his, int, int) ============ - mat_eq = np.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) + mat_eq = xp.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -2603,7 +2603,7 @@ def projection_T_2form(self, domain): ) # ================= assembly of 2 - component (PI_1_2 : int, his, int) ============ - mat_eq = np.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) + mat_eq = xp.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -2698,7 +2698,7 @@ def projection_T_2form(self, domain): ) # ================= assembly of 3 - component (PI_1_3 : int, int, his) ============ - mat_eq = np.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) + mat_eq = xp.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -2793,7 +2793,7 @@ def projection_T_2form(self, domain): ) # ============== conversion to sparse matrices (1 - component) ============== - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseN[1], @@ -2817,7 +2817,7 @@ def projection_T_2form(self, domain): ) T12.eliminate_zeros() - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseD[1], @@ -2842,7 +2842,7 @@ def projection_T_2form(self, domain): T13.eliminate_zeros() # ============== conversion to sparse matrices (2 - component) ============== - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseD[1], @@ -2866,7 +2866,7 @@ def projection_T_2form(self, domain): ) T21.eliminate_zeros() - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseD[1], @@ -2891,7 +2891,7 @@ def projection_T_2form(self, domain): T23.eliminate_zeros() # ============== conversion to sparse matrices (3 - component) ============== - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseD[1], @@ -2915,7 +2915,7 @@ def projection_T_2form(self, domain): ) T31.eliminate_zeros() - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseN[1], @@ -2966,7 +2966,7 @@ def projection_S_0form(self, domain): """ # non-vanishing coefficients - S11 = np.empty( + S11 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -2977,7 +2977,7 @@ def projection_S_0form(self, domain): ), dtype=float, ) - S22 = np.empty( + S22 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -2988,7 +2988,7 @@ def projection_S_0form(self, domain): ), dtype=float, ) - S33 = np.empty( + S33 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -3006,7 +3006,7 @@ def projection_S_0form(self, domain): n_unique3 = [self.pts[0].flatten().size, self.pts[1].flatten().size, self.x_int[2].size] # ========= assembly of 1 - component (pi2_1 : int, his, his) ============ - mat_eq = np.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) + mat_eq = xp.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -3062,7 +3062,7 @@ def projection_S_0form(self, domain): ) # ========= assembly of 2 - component (pi2_2 : his, int, his) ============ - mat_eq = np.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) + mat_eq = xp.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -3118,7 +3118,7 @@ def projection_S_0form(self, domain): ) # ========= assembly of 3 - component (pi2_3 : his, his, int) ============ - mat_eq = np.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) + mat_eq = xp.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -3174,7 +3174,7 @@ def projection_S_0form(self, domain): ) # ========= conversion to sparse matrices (1 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -3199,7 +3199,7 @@ def projection_S_0form(self, domain): S11.eliminate_zeros() # ========= conversion to sparse matrices (2 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -3224,7 +3224,7 @@ def projection_S_0form(self, domain): S22.eliminate_zeros() # ========= conversion to sparse matrices (3 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -3275,7 +3275,7 @@ def projection_S_2form(self, domain): """ # non-vanishing coefficients - S11 = np.empty( + S11 = xp.empty( ( self.NbaseN[0], self.NbaseD[1], @@ -3286,7 +3286,7 @@ def projection_S_2form(self, domain): ), dtype=float, ) - S22 = np.empty( + S22 = xp.empty( ( self.NbaseD[0], self.NbaseN[1], @@ -3297,7 +3297,7 @@ def projection_S_2form(self, domain): ), dtype=float, ) - S33 = np.empty( + S33 = xp.empty( ( self.NbaseD[0], self.NbaseD[1], @@ -3315,7 +3315,7 @@ def projection_S_2form(self, domain): n_unique3 = [self.pts[0].flatten().size, self.pts[1].flatten().size, self.x_int[2].size] # ========= assembly of 1 - component (pi2_1 : int, his, his) ============ - mat_eq = np.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) + mat_eq = xp.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -3371,7 +3371,7 @@ def projection_S_2form(self, domain): ) # ========= assembly of 2 - component (pi2_2 : his, int, his) ============ - mat_eq = np.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) + mat_eq = xp.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -3427,7 +3427,7 @@ def projection_S_2form(self, domain): ) # ========= assembly of 3 - component (pi2_3 : his, his, int) ============ - mat_eq = np.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) + mat_eq = xp.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -3483,7 +3483,7 @@ def projection_S_2form(self, domain): ) # ========= conversion to sparse matrices (1 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseD[1], @@ -3508,7 +3508,7 @@ def projection_S_2form(self, domain): S11.eliminate_zeros() # ========= conversion to sparse matrices (2 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseN[1], @@ -3533,7 +3533,7 @@ def projection_S_2form(self, domain): S22.eliminate_zeros() # ========= conversion to sparse matrices (3 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseD[1], @@ -3582,7 +3582,7 @@ def projection_K_3form(self, domain): """ # non-vanishing coefficients - K = np.zeros( + K = xp.zeros( ( self.NbaseD[0], self.NbaseD[1], @@ -3597,7 +3597,7 @@ def projection_K_3form(self, domain): # evaluation of equilibrium pressure at interpolation points n_unique = [self.pts[0].flatten().size, self.pts[1].flatten().size, self.pts[2].flatten().size] - mat_eq = np.zeros((n_unique[0], n_unique[1], n_unique[2]), dtype=float) + mat_eq = xp.zeros((n_unique[0], n_unique[1], n_unique[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -3656,7 +3656,7 @@ def projection_K_3form(self, domain): ) # conversion to sparse matrix - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseD[1], @@ -3711,7 +3711,7 @@ def projection_N_0form(self, domain): """ # non-vanishing coefficients - N11 = np.empty( + N11 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -3722,7 +3722,7 @@ def projection_N_0form(self, domain): ), dtype=float, ) - N22 = np.empty( + N22 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -3733,7 +3733,7 @@ def projection_N_0form(self, domain): ), dtype=float, ) - N33 = np.empty( + N33 = xp.empty( ( self.NbaseN[0], self.NbaseN[1], @@ -3751,7 +3751,7 @@ def projection_N_0form(self, domain): n_unique3 = [self.pts[0].flatten().size, self.pts[1].flatten().size, self.x_int[2].size] # ========= assembly of 1 - component (pi2_1 : int, his, his) ============ - mat_eq = np.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) + mat_eq = xp.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -3807,7 +3807,7 @@ def projection_N_0form(self, domain): ) # ========= assembly of 2 - component (pi2_2 : his, int, his) ============ - mat_eq = np.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) + mat_eq = xp.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -3863,7 +3863,7 @@ def projection_N_0form(self, domain): ) # ========= assembly of 3 - component (pi2_3 : his, his, int) ============ - mat_eq = np.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) + mat_eq = xp.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -3919,7 +3919,7 @@ def projection_N_0form(self, domain): ) # ========= conversion to sparse matrices (1 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -3944,7 +3944,7 @@ def projection_N_0form(self, domain): N11.eliminate_zeros() # ========= conversion to sparse matrices (2 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -3969,7 +3969,7 @@ def projection_N_0form(self, domain): N22.eliminate_zeros() # ========= conversion to sparse matrices (3 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseN[1], @@ -4020,7 +4020,7 @@ def projection_N_2form(self, domain): """ # non-vanishing coefficients - N11 = np.empty( + N11 = xp.empty( ( self.NbaseN[0], self.NbaseD[1], @@ -4031,7 +4031,7 @@ def projection_N_2form(self, domain): ), dtype=float, ) - N22 = np.empty( + N22 = xp.empty( ( self.NbaseD[0], self.NbaseN[1], @@ -4042,7 +4042,7 @@ def projection_N_2form(self, domain): ), dtype=float, ) - N33 = np.empty( + N33 = xp.empty( ( self.NbaseD[0], self.NbaseD[1], @@ -4060,7 +4060,7 @@ def projection_N_2form(self, domain): n_unique3 = [self.pts[0].flatten().size, self.pts[1].flatten().size, self.x_int[2].size] # ========= assembly of 1 - component (pi2_1 : int, his, his) ============ - mat_eq = np.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) + mat_eq = xp.empty((n_unique1[0], n_unique1[1], n_unique1[2]), dtype=float) ker_eva.kernel_eva( self.x_int[0], @@ -4108,7 +4108,7 @@ def projection_N_2form(self, domain): ) # ========= assembly of 2 - component (pi2_2 : his, int, his) ============ - mat_eq = np.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) + mat_eq = xp.empty((n_unique2[0], n_unique2[1], n_unique2[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -4156,7 +4156,7 @@ def projection_N_2form(self, domain): ) # ========= assembly of 3 - component (pi2_3 : his, his, int) ============ - mat_eq = np.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) + mat_eq = xp.empty((n_unique3[0], n_unique3[1], n_unique3[2]), dtype=float) ker_eva.kernel_eva( self.pts[0].flatten(), @@ -4204,7 +4204,7 @@ def projection_N_2form(self, domain): ) # ========= conversion to sparse matrices (1 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseN[0], self.NbaseD[1], @@ -4229,7 +4229,7 @@ def projection_N_2form(self, domain): N11.eliminate_zeros() # ========= conversion to sparse matrices (2 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseN[1], @@ -4254,7 +4254,7 @@ def projection_N_2form(self, domain): N22.eliminate_zeros() # ========= conversion to sparse matrices (3 - component) ================= - indices = np.indices( + indices = xp.indices( ( self.NbaseD[0], self.NbaseD[1], @@ -4356,13 +4356,13 @@ def __init__( self.cz = cz # ============= evaluation of background magnetic field at quadrature points ========= - self.mat_curl_beq_1 = np.empty( + self.mat_curl_beq_1 = xp.empty( (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float ) - self.mat_curl_beq_2 = np.empty( + self.mat_curl_beq_2 = xp.empty( (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float ) - self.mat_curl_beq_3 = np.empty( + self.mat_curl_beq_3 = xp.empty( (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float ) @@ -4454,20 +4454,20 @@ def __init__( ) # ====================== perturbed magnetic field at quadrature points ========== - self.B1 = np.empty( + self.B1 = xp.empty( (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float ) - self.B2 = np.empty( + self.B2 = xp.empty( (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float ) - self.B3 = np.empty( + self.B3 = xp.empty( (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float ) # ========================== inner products ===================================== - self.F1 = np.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) - self.F2 = np.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) - self.F3 = np.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) + self.F1 = xp.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) + self.F2 = xp.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) + self.F3 = xp.empty((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) # ============================================================ def inner_curl_beq(self, b1, b2, b3): @@ -4598,7 +4598,7 @@ def inner_curl_beq(self, b1, b2, b3): # ker_loc_3d.kernel_inner_2(self.Nel[0], self.Nel[1], self.Nel[2], self.p[0], self.p[1], self.p[2], self.n_quad[0], self.n_quad[1], self.n_quad[2], 0, 0, 0, self.wts[0], self.wts[1], self.wts[2], self.basisN[0], self.basisN[1], self.basisN[2], self.NbaseN[0], self.NbaseN[1], self.NbaseN[2], self.F3, self.mat_curl_beq_3) # convert to 1d array and return - return np.concatenate((self.F1.flatten(), self.F2.flatten(), self.F3.flatten())) + return xp.concatenate((self.F1.flatten(), self.F2.flatten(), self.F3.flatten())) # ================ mass matrix in V1 =========================== @@ -4641,9 +4641,9 @@ def mass_curl(tensor_space, kind_map, params_map): Nbj3 = [NbaseD[2], NbaseN[2], NbaseD[2], NbaseN[2], NbaseD[2], NbaseD[2]] # ============= evaluation of background magnetic field at quadrature points ========= - mat_curl_beq_1 = np.empty((Nel[0], Nel[1], Nel[2], n_quad[0], n_quad[1], n_quad[2]), dtype=float) - mat_curl_beq_2 = np.empty((Nel[0], Nel[1], Nel[2], n_quad[0], n_quad[1], n_quad[2]), dtype=float) - mat_curl_beq_3 = np.empty((Nel[0], Nel[1], Nel[2], n_quad[0], n_quad[1], n_quad[2]), dtype=float) + mat_curl_beq_1 = xp.empty((Nel[0], Nel[1], Nel[2], n_quad[0], n_quad[1], n_quad[2]), dtype=float) + mat_curl_beq_2 = xp.empty((Nel[0], Nel[1], Nel[2], n_quad[0], n_quad[1], n_quad[2]), dtype=float) + mat_curl_beq_3 = xp.empty((Nel[0], Nel[1], Nel[2], n_quad[0], n_quad[1], n_quad[2]), dtype=float) ker_eva.kernel_eva_quad(Nel, n_quad, pts[0], pts[1], pts[2], mat_curl_beq_1, 61, kind_map, params_map) ker_eva.kernel_eva_quad(Nel, n_quad, pts[0], pts[1], pts[2], mat_curl_beq_2, 62, kind_map, params_map) @@ -4652,7 +4652,7 @@ def mass_curl(tensor_space, kind_map, params_map): # blocks of global mass matrix M = [ - np.zeros((Nbi1, Nbi2, Nbi3, 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) + xp.zeros((Nbi1, Nbi2, Nbi3, 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) for Nbi1, Nbi2, Nbi3 in zip(Nbi1, Nbi2, Nbi3) ] @@ -4858,11 +4858,11 @@ def mass_curl(tensor_space, kind_map, params_map): counter = 0 for i in range(6): - indices = np.indices((Nbi1[counter], Nbi2[counter], Nbi3[counter], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) + indices = xp.indices((Nbi1[counter], Nbi2[counter], Nbi3[counter], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) - shift1 = np.arange(Nbi1[counter]) - p[0] - shift2 = np.arange(Nbi2[counter]) - p[1] - shift3 = np.arange(Nbi3[counter]) - p[2] + shift1 = xp.arange(Nbi1[counter]) - p[0] + shift2 = xp.arange(Nbi2[counter]) - p[1] + shift3 = xp.arange(Nbi3[counter]) - p[2] row = (Nbi2[counter] * Nbi3[counter] * indices[0] + Nbi3[counter] * indices[1] + indices[2]).flatten() diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py index 72c5babfd..f7d366c99 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py @@ -10,7 +10,7 @@ import struphy.feec.bsplines as bsp import struphy.feec.projectors.pro_local.kernels_projectors_local as ker_loc -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ======================= 1d ==================================== @@ -41,85 +41,85 @@ def __init__(self, spline_space, n_quad): self.n_quad = n_quad # number of quadrature point per integration interval # Gauss - Legendre quadrature points and weights in (-1, 1) - self.pts_loc = np.polynomial.legendre.leggauss(self.n_quad)[0] - self.wts_loc = np.polynomial.legendre.leggauss(self.n_quad)[1] + self.pts_loc = xp.polynomial.legendre.leggauss(self.n_quad)[0] + self.wts_loc = xp.polynomial.legendre.leggauss(self.n_quad)[1] # set interpolation and histopolation coefficients if self.bc == True: - self.coeff_i = np.zeros((1, 2 * self.p - 1), dtype=float) - self.coeff_h = np.zeros((1, 2 * self.p), dtype=float) + self.coeff_i = xp.zeros((1, 2 * self.p - 1), dtype=float) + self.coeff_h = xp.zeros((1, 2 * self.p), dtype=float) if self.p == 1: - self.coeff_i[0, :] = np.array([1.0]) - self.coeff_h[0, :] = np.array([1.0, 1.0]) + self.coeff_i[0, :] = xp.array([1.0]) + self.coeff_h[0, :] = xp.array([1.0, 1.0]) elif self.p == 2: - self.coeff_i[0, :] = 1 / 2 * np.array([-1.0, 4.0, -1.0]) - self.coeff_h[0, :] = 1 / 2 * np.array([-1.0, 3.0, 3.0, -1.0]) + self.coeff_i[0, :] = 1 / 2 * xp.array([-1.0, 4.0, -1.0]) + self.coeff_h[0, :] = 1 / 2 * xp.array([-1.0, 3.0, 3.0, -1.0]) elif self.p == 3: - self.coeff_i[0, :] = 1 / 6 * np.array([1.0, -8.0, 20.0, -8.0, 1.0]) - self.coeff_h[0, :] = 1 / 6 * np.array([1.0, -7.0, 12.0, 12.0, -7.0, 1.0]) + self.coeff_i[0, :] = 1 / 6 * xp.array([1.0, -8.0, 20.0, -8.0, 1.0]) + self.coeff_h[0, :] = 1 / 6 * xp.array([1.0, -7.0, 12.0, 12.0, -7.0, 1.0]) elif self.p == 4: - self.coeff_i[0, :] = 2 / 45 * np.array([-1.0, 16.0, -295 / 4, 140.0, -295 / 4, 16.0, -1.0]) - self.coeff_h[0, :] = 2 / 45 * np.array([-1.0, 15.0, -231 / 4, 265 / 4, 265 / 4, -231 / 4, 15.0, -1.0]) + self.coeff_i[0, :] = 2 / 45 * xp.array([-1.0, 16.0, -295 / 4, 140.0, -295 / 4, 16.0, -1.0]) + self.coeff_h[0, :] = 2 / 45 * xp.array([-1.0, 15.0, -231 / 4, 265 / 4, 265 / 4, -231 / 4, 15.0, -1.0]) else: print("degree > 4 not implemented!") else: - self.coeff_i = np.zeros((2 * self.p - 1, 2 * self.p - 1), dtype=float) - self.coeff_h = np.zeros((2 * self.p - 1, 2 * self.p), dtype=float) + self.coeff_i = xp.zeros((2 * self.p - 1, 2 * self.p - 1), dtype=float) + self.coeff_h = xp.zeros((2 * self.p - 1, 2 * self.p), dtype=float) if self.p == 1: - self.coeff_i[0, :] = np.array([1.0]) - self.coeff_h[0, :] = np.array([1.0, 1.0]) + self.coeff_i[0, :] = xp.array([1.0]) + self.coeff_h[0, :] = xp.array([1.0, 1.0]) elif self.p == 2: - self.coeff_i[0, :] = 1 / 2 * np.array([2.0, 0.0, 0.0]) - self.coeff_i[1, :] = 1 / 2 * np.array([-1.0, 4.0, -1.0]) - self.coeff_i[2, :] = 1 / 2 * np.array([0.0, 0.0, 2.0]) + self.coeff_i[0, :] = 1 / 2 * xp.array([2.0, 0.0, 0.0]) + self.coeff_i[1, :] = 1 / 2 * xp.array([-1.0, 4.0, -1.0]) + self.coeff_i[2, :] = 1 / 2 * xp.array([0.0, 0.0, 2.0]) - self.coeff_h[0, :] = 1 / 2 * np.array([3.0, -1.0, 0.0, 0.0]) - self.coeff_h[1, :] = 1 / 2 * np.array([-1.0, 3.0, 3.0, -1.0]) - self.coeff_h[2, :] = 1 / 2 * np.array([0.0, 0.0, -1.0, 3.0]) + self.coeff_h[0, :] = 1 / 2 * xp.array([3.0, -1.0, 0.0, 0.0]) + self.coeff_h[1, :] = 1 / 2 * xp.array([-1.0, 3.0, 3.0, -1.0]) + self.coeff_h[2, :] = 1 / 2 * xp.array([0.0, 0.0, -1.0, 3.0]) elif self.p == 3: - self.coeff_i[0, :] = 1 / 18 * np.array([18.0, 0.0, 0.0, 0.0, 0.0]) - self.coeff_i[1, :] = 1 / 18 * np.array([-5.0, 40.0, -24.0, 8.0, -1.0]) - self.coeff_i[2, :] = 1 / 18 * np.array([3.0, -24.0, 60.0, -24.0, 3.0]) - self.coeff_i[3, :] = 1 / 18 * np.array([-1.0, 8.0, -24.0, 40.0, -5.0]) - self.coeff_i[4, :] = 1 / 18 * np.array([0.0, 0.0, 0.0, 0.0, 18.0]) - - self.coeff_h[0, :] = 1 / 18 * np.array([23.0, -17.0, 7.0, -1.0, 0.0, 0.0]) - self.coeff_h[1, :] = 1 / 18 * np.array([-8.0, 56.0, -28.0, 4.0, 0.0, 0.0]) - self.coeff_h[2, :] = 1 / 18 * np.array([3.0, -21.0, 36.0, 36.0, -21.0, 3.0]) - self.coeff_h[3, :] = 1 / 18 * np.array([0.0, 0.0, 4.0, -28.0, 56.0, -8.0]) - self.coeff_h[4, :] = 1 / 18 * np.array([0.0, 0.0, -1.0, 7.0, -17.0, 23.0]) + self.coeff_i[0, :] = 1 / 18 * xp.array([18.0, 0.0, 0.0, 0.0, 0.0]) + self.coeff_i[1, :] = 1 / 18 * xp.array([-5.0, 40.0, -24.0, 8.0, -1.0]) + self.coeff_i[2, :] = 1 / 18 * xp.array([3.0, -24.0, 60.0, -24.0, 3.0]) + self.coeff_i[3, :] = 1 / 18 * xp.array([-1.0, 8.0, -24.0, 40.0, -5.0]) + self.coeff_i[4, :] = 1 / 18 * xp.array([0.0, 0.0, 0.0, 0.0, 18.0]) + + self.coeff_h[0, :] = 1 / 18 * xp.array([23.0, -17.0, 7.0, -1.0, 0.0, 0.0]) + self.coeff_h[1, :] = 1 / 18 * xp.array([-8.0, 56.0, -28.0, 4.0, 0.0, 0.0]) + self.coeff_h[2, :] = 1 / 18 * xp.array([3.0, -21.0, 36.0, 36.0, -21.0, 3.0]) + self.coeff_h[3, :] = 1 / 18 * xp.array([0.0, 0.0, 4.0, -28.0, 56.0, -8.0]) + self.coeff_h[4, :] = 1 / 18 * xp.array([0.0, 0.0, -1.0, 7.0, -17.0, 23.0]) elif self.p == 4: - self.coeff_i[0, :] = 1 / 360 * np.array([360.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) - self.coeff_i[1, :] = 1 / 360 * np.array([-59.0, 944.0, -1000.0, 720.0, -305.0, 64.0, -4.0]) - self.coeff_i[2, :] = 1 / 360 * np.array([23.0, -368.0, 1580.0, -1360.0, 605.0, -128.0, 8.0]) - self.coeff_i[3, :] = 1 / 360 * np.array([-16.0, 256.0, -1180.0, 2240.0, -1180.0, 256.0, -16.0]) - self.coeff_i[4, :] = 1 / 360 * np.array([8.0, -128.0, 605.0, -1360.0, 1580.0, -368.0, 23.0]) - self.coeff_i[5, :] = 1 / 360 * np.array([-4.0, 64.0, -305.0, 720.0, -1000.0, 944.0, -59.0]) - self.coeff_i[6, :] = 1 / 360 * np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 360.0]) - - self.coeff_h[0, :] = 1 / 360 * np.array([419.0, -525.0, 475.0, -245.0, 60.0, -4.0, 0.0, 0.0]) - self.coeff_h[1, :] = 1 / 360 * np.array([-82.0, 1230.0, -1350.0, 730.0, -180.0, 12.0, 0.0, 0.0]) - self.coeff_h[2, :] = 1 / 360 * np.array([39.0, -585.0, 2175.0, -1425.0, 360.0, -24.0, 0.0, 0.0]) - self.coeff_h[3, :] = 1 / 360 * np.array([-16.0, 240.0, -924.0, 1060.0, 1060.0, -924.0, 240.0, -16.0]) - self.coeff_h[4, :] = 1 / 360 * np.array([0.0, 0.0, -24.0, 360.0, -1425.0, 2175.0, -585.0, 39.0]) - self.coeff_h[5, :] = 1 / 360 * np.array([0.0, 0.0, 12.0, -180.0, 730.0, -1350.0, 1230.0, -82.0]) - self.coeff_h[6, :] = 1 / 360 * np.array([0.0, 0.0, -4.0, 60.0, -245.0, 475.0, -525.0, 419.0]) + self.coeff_i[0, :] = 1 / 360 * xp.array([360.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + self.coeff_i[1, :] = 1 / 360 * xp.array([-59.0, 944.0, -1000.0, 720.0, -305.0, 64.0, -4.0]) + self.coeff_i[2, :] = 1 / 360 * xp.array([23.0, -368.0, 1580.0, -1360.0, 605.0, -128.0, 8.0]) + self.coeff_i[3, :] = 1 / 360 * xp.array([-16.0, 256.0, -1180.0, 2240.0, -1180.0, 256.0, -16.0]) + self.coeff_i[4, :] = 1 / 360 * xp.array([8.0, -128.0, 605.0, -1360.0, 1580.0, -368.0, 23.0]) + self.coeff_i[5, :] = 1 / 360 * xp.array([-4.0, 64.0, -305.0, 720.0, -1000.0, 944.0, -59.0]) + self.coeff_i[6, :] = 1 / 360 * xp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 360.0]) + + self.coeff_h[0, :] = 1 / 360 * xp.array([419.0, -525.0, 475.0, -245.0, 60.0, -4.0, 0.0, 0.0]) + self.coeff_h[1, :] = 1 / 360 * xp.array([-82.0, 1230.0, -1350.0, 730.0, -180.0, 12.0, 0.0, 0.0]) + self.coeff_h[2, :] = 1 / 360 * xp.array([39.0, -585.0, 2175.0, -1425.0, 360.0, -24.0, 0.0, 0.0]) + self.coeff_h[3, :] = 1 / 360 * xp.array([-16.0, 240.0, -924.0, 1060.0, 1060.0, -924.0, 240.0, -16.0]) + self.coeff_h[4, :] = 1 / 360 * xp.array([0.0, 0.0, -24.0, 360.0, -1425.0, 2175.0, -585.0, 39.0]) + self.coeff_h[5, :] = 1 / 360 * xp.array([0.0, 0.0, 12.0, -180.0, 730.0, -1350.0, 1230.0, -82.0]) + self.coeff_h[6, :] = 1 / 360 * xp.array([0.0, 0.0, -4.0, 60.0, -245.0, 475.0, -525.0, 419.0]) else: print("degree > 4 not implemented!") # set interpolation points - n_lambda_int = np.copy(self.NbaseN) # number of coefficients in space V0 + n_lambda_int = xp.copy(self.NbaseN) # number of coefficients in space V0 self.n_int = 2 * self.p - 1 # number of local interpolation points (1, 3, 5, 7, ...) if self.p == 1: @@ -134,21 +134,21 @@ def __init__(self, spline_space, n_quad): 2 * self.p - 2 ) # number of non-vanishing D bf in interpolation interval (1, 2, 4, 6, ...) - self.x_int = np.zeros((n_lambda_int, self.n_int), dtype=float) # interpolation points for each coeff. + self.x_int = xp.zeros((n_lambda_int, self.n_int), dtype=float) # interpolation points for each coeff. - self.int_global_N = np.zeros( + self.int_global_N = xp.zeros( (n_lambda_int, self.n_int_locbf_N), dtype=int ) # global indices of non-vanishing N bf - self.int_global_D = np.zeros( + self.int_global_D = xp.zeros( (n_lambda_int, self.n_int_locbf_D), dtype=int ) # global indices of non-vanishing D bf - self.int_loccof_N = np.zeros((n_lambda_int, self.n_int_locbf_N), dtype=int) # index of non-vanishing coeff. (N) - self.int_loccof_D = np.zeros((n_lambda_int, self.n_int_locbf_D), dtype=int) # index of non-vanishing coeff. (D) + self.int_loccof_N = xp.zeros((n_lambda_int, self.n_int_locbf_N), dtype=int) # index of non-vanishing coeff. (N) + self.int_loccof_D = xp.zeros((n_lambda_int, self.n_int_locbf_D), dtype=int) # index of non-vanishing coeff. (D) - self.x_int_indices = np.zeros((n_lambda_int, self.n_int), dtype=int) + self.x_int_indices = xp.zeros((n_lambda_int, self.n_int), dtype=int) - self.coeffi_indices = np.zeros(n_lambda_int, dtype=int) + self.coeffi_indices = xp.zeros(n_lambda_int, dtype=int) if self.bc == False: # maximum number of non-vanishing coefficients @@ -160,39 +160,39 @@ def __init__(self, spline_space, n_quad): self.n_int_nvcof_N = 3 * self.p - 2 # shift in local coefficient indices at right boundary (only for non-periodic boundary conditions) - self.int_add_D = np.arange(self.n_int - 2) + 1 - self.int_add_N = np.arange(self.n_int - 1) + 1 + self.int_add_D = xp.arange(self.n_int - 2) + 1 + self.int_add_N = xp.arange(self.n_int - 1) + 1 counter_D = 0 counter_N = 0 # shift local coefficients --> global coefficients (D) if self.p == 1: - self.int_shift_D = np.arange(self.NbaseD) + self.int_shift_D = xp.arange(self.NbaseD) else: - self.int_shift_D = np.arange(self.NbaseD) - (self.p - 2) + self.int_shift_D = xp.arange(self.NbaseD) - (self.p - 2) self.int_shift_D[: 2 * self.p - 2] = 0 self.int_shift_D[-(2 * self.p - 2) :] = self.int_shift_D[-(2 * self.p - 2)] # shift local coefficients --> global coefficients (N) if self.p == 1: - self.int_shift_N = np.arange(self.NbaseN) + self.int_shift_N = xp.arange(self.NbaseN) self.int_shift_N[-1] = self.int_shift_N[-2] else: - self.int_shift_N = np.arange(self.NbaseN) - (self.p - 1) + self.int_shift_N = xp.arange(self.NbaseN) - (self.p - 1) self.int_shift_N[: 2 * self.p - 1] = 0 self.int_shift_N[-(2 * self.p - 1) :] = self.int_shift_N[-(2 * self.p - 1)] - counter_coeffi = np.copy(self.p) + counter_coeffi = xp.copy(self.p) for i in range(n_lambda_int): # left boundary region if i < self.p - 1: - self.int_global_N[i] = np.arange(self.n_int_locbf_N) - self.int_global_D[i] = np.arange(self.n_int_locbf_D) + self.int_global_N[i] = xp.arange(self.n_int_locbf_N) + self.int_global_D[i] = xp.arange(self.n_int_locbf_D) - self.x_int_indices[i] = np.arange(self.n_int) + self.x_int_indices[i] = xp.arange(self.n_int) self.coeffi_indices[i] = i for j in range(2 * (self.p - 1) + 1): xi = self.p - 1 @@ -200,10 +200,10 @@ def __init__(self, spline_space, n_quad): # right boundary region elif i > n_lambda_int - self.p: - self.int_global_N[i] = np.arange(self.n_int_locbf_N) + n_lambda_int - self.p - (self.p - 1) - self.int_global_D[i] = np.arange(self.n_int_locbf_D) + n_lambda_int - self.p - (self.p - 1) + self.int_global_N[i] = xp.arange(self.n_int_locbf_N) + n_lambda_int - self.p - (self.p - 1) + self.int_global_D[i] = xp.arange(self.n_int_locbf_D) + n_lambda_int - self.p - (self.p - 1) - self.x_int_indices[i] = np.arange(self.n_int) + 2 * (n_lambda_int - self.p - (self.p - 1)) + self.x_int_indices[i] = xp.arange(self.n_int) + 2 * (n_lambda_int - self.p - (self.p - 1)) self.coeffi_indices[i] = counter_coeffi counter_coeffi += 1 for j in range(2 * (self.p - 1) + 1): @@ -213,20 +213,20 @@ def __init__(self, spline_space, n_quad): # interior else: if self.p == 1: - self.int_global_N[i] = np.arange(self.n_int_locbf_N) + i - self.int_global_D[i] = np.arange(self.n_int_locbf_D) + i + self.int_global_N[i] = xp.arange(self.n_int_locbf_N) + i + self.int_global_D[i] = xp.arange(self.n_int_locbf_D) + i self.int_global_N[-1] = self.int_global_N[-2] self.int_global_D[-1] = self.int_global_D[-2] else: - self.int_global_N[i] = np.arange(self.n_int_locbf_N) + i - (self.p - 1) - self.int_global_D[i] = np.arange(self.n_int_locbf_D) + i - (self.p - 1) + self.int_global_N[i] = xp.arange(self.n_int_locbf_N) + i - (self.p - 1) + self.int_global_D[i] = xp.arange(self.n_int_locbf_D) + i - (self.p - 1) if self.p == 1: self.x_int_indices[i] = i else: - self.x_int_indices[i] = np.arange(self.n_int) + 2 * (i - (self.p - 1)) + self.x_int_indices[i] = xp.arange(self.n_int) + 2 * (i - (self.p - 1)) self.coeffi_indices[i] = self.p - 1 for j in range(2 * (self.p - 1) + 1): @@ -234,8 +234,8 @@ def __init__(self, spline_space, n_quad): # local coefficient index if self.p == 1: - self.int_loccof_N[i] = np.array([0, 1]) - self.int_loccof_D[-1] = np.array([1]) + self.int_loccof_N[i] = xp.array([0, 1]) + self.int_loccof_D[-1] = xp.array([1]) else: if i > 0: @@ -243,8 +243,8 @@ def __init__(self, spline_space, n_quad): k_glob_new = self.int_global_D[i, il] bol = k_glob_new == self.int_global_D[i - 1] - if np.any(bol): - self.int_loccof_D[i, il] = self.int_loccof_D[i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.int_loccof_D[i, il] = self.int_loccof_D[i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_int - self.p - (self.p - 2)) and (self.int_loccof_D[i, il] == 0): self.int_loccof_D[i, il] = self.int_add_D[counter_D] @@ -254,8 +254,8 @@ def __init__(self, spline_space, n_quad): k_glob_new = self.int_global_N[i, il] bol = k_glob_new == self.int_global_N[i - 1] - if np.any(bol): - self.int_loccof_N[i, il] = self.int_loccof_N[i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.int_loccof_N[i, il] = self.int_loccof_N[i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_int - self.p - (self.p - 2)) and (self.int_loccof_N[i, il] == 0): self.int_loccof_N[i, il] = self.int_add_N[counter_N] @@ -273,24 +273,24 @@ def __init__(self, spline_space, n_quad): # shift local coefficients --> global coefficients if self.p == 1: - self.int_shift_D = np.arange(self.NbaseN) - (self.p - 1) - self.int_shift_N = np.arange(self.NbaseN) - (self.p) + self.int_shift_D = xp.arange(self.NbaseN) - (self.p - 1) + self.int_shift_N = xp.arange(self.NbaseN) - (self.p) else: - self.int_shift_D = np.arange(self.NbaseN) - (self.p - 2) - self.int_shift_N = np.arange(self.NbaseN) - (self.p - 1) + self.int_shift_D = xp.arange(self.NbaseN) - (self.p - 2) + self.int_shift_N = xp.arange(self.NbaseN) - (self.p - 1) for i in range(n_lambda_int): # global indices of non-vanishing basis functions and position of coefficients in final matrix - self.int_global_D[i] = (np.arange(self.n_int_locbf_D) + i - (self.p - 1)) % self.NbaseD - self.int_loccof_D[i] = np.arange(self.n_int_locbf_D - 1, -1, -1) + self.int_global_D[i] = (xp.arange(self.n_int_locbf_D) + i - (self.p - 1)) % self.NbaseD + self.int_loccof_D[i] = xp.arange(self.n_int_locbf_D - 1, -1, -1) - self.int_global_N[i] = (np.arange(self.n_int_locbf_N) + i - (self.p - 1)) % self.NbaseN - self.int_loccof_N[i] = np.arange(self.n_int_locbf_N - 1, -1, -1) + self.int_global_N[i] = (xp.arange(self.n_int_locbf_N) + i - (self.p - 1)) % self.NbaseN + self.int_loccof_N[i] = xp.arange(self.n_int_locbf_N - 1, -1, -1) if self.p == 1: self.x_int_indices[i] = i else: - self.x_int_indices[i] = np.arange(self.n_int) + 2 * (i - (self.p - 1)) + self.x_int_indices[i] = xp.arange(self.n_int) + 2 * (i - (self.p - 1)) self.coeffi_indices[i] = 0 @@ -298,23 +298,23 @@ def __init__(self, spline_space, n_quad): self.x_int[i, j] = ((self.T[i + 1 + int(j / 2)] + self.T[i + 1 + int((j + 1) / 2)]) / 2) % 1.0 # set histopolation points, quadrature points and weights - n_lambda_his = np.copy(self.NbaseD) # number of coefficients in space V1 + n_lambda_his = xp.copy(self.NbaseD) # number of coefficients in space V1 self.n_his = 2 * self.p # number of histopolation intervals (2, 4, 6, 8, ...) self.n_his_locbf_N = 2 * self.p # number of non-vanishing N bf in histopolation interval (2, 4, 6, 8, ...) self.n_his_locbf_D = 2 * self.p - 1 # number of non-vanishing D bf in histopolation interval (2, 4, 6, 8, ...) - self.x_his = np.zeros((n_lambda_his, self.n_his + 1), dtype=float) # histopolation boundaries + self.x_his = xp.zeros((n_lambda_his, self.n_his + 1), dtype=float) # histopolation boundaries - self.his_global_N = np.zeros((n_lambda_his, self.n_his_locbf_N), dtype=int) - self.his_global_D = np.zeros((n_lambda_his, self.n_his_locbf_D), dtype=int) + self.his_global_N = xp.zeros((n_lambda_his, self.n_his_locbf_N), dtype=int) + self.his_global_D = xp.zeros((n_lambda_his, self.n_his_locbf_D), dtype=int) - self.his_loccof_N = np.zeros((n_lambda_his, self.n_his_locbf_N), dtype=int) - self.his_loccof_D = np.zeros((n_lambda_his, self.n_his_locbf_D), dtype=int) + self.his_loccof_N = xp.zeros((n_lambda_his, self.n_his_locbf_N), dtype=int) + self.his_loccof_D = xp.zeros((n_lambda_his, self.n_his_locbf_D), dtype=int) - self.x_his_indices = np.zeros((n_lambda_his, self.n_his), dtype=int) + self.x_his_indices = xp.zeros((n_lambda_his, self.n_his), dtype=int) - self.coeffh_indices = np.zeros(n_lambda_his, dtype=int) + self.coeffh_indices = xp.zeros(n_lambda_his, dtype=int) if self.bc == False: # maximum number of non-vanishing coefficients @@ -322,31 +322,31 @@ def __init__(self, spline_space, n_quad): self.n_his_nvcof_N = 3 * self.p - 1 # shift in local coefficient indices at right boundary (only for non-periodic boundary conditions) - self.his_add_D = np.arange(self.n_his - 2) + 1 - self.his_add_N = np.arange(self.n_his - 1) + 1 + self.his_add_D = xp.arange(self.n_his - 2) + 1 + self.his_add_N = xp.arange(self.n_his - 1) + 1 counter_D = 0 counter_N = 0 # shift local coefficients --> global coefficients (D) - self.his_shift_D = np.arange(self.NbaseD) - (self.p - 1) + self.his_shift_D = xp.arange(self.NbaseD) - (self.p - 1) self.his_shift_D[: 2 * self.p - 1] = 0 self.his_shift_D[-(2 * self.p - 1) :] = self.his_shift_D[-(2 * self.p - 1)] # shift local coefficients --> global coefficients (N) - self.his_shift_N = np.arange(self.NbaseN) - self.p + self.his_shift_N = xp.arange(self.NbaseN) - self.p self.his_shift_N[: 2 * self.p] = 0 self.his_shift_N[-2 * self.p :] = self.his_shift_N[-2 * self.p] - counter_coeffh = np.copy(self.p) + counter_coeffh = xp.copy(self.p) for i in range(n_lambda_his): # left boundary region if i < self.p - 1: - self.his_global_N[i] = np.arange(self.n_his_locbf_N) - self.his_global_D[i] = np.arange(self.n_his_locbf_D) + self.his_global_N[i] = xp.arange(self.n_his_locbf_N) + self.his_global_D[i] = xp.arange(self.n_his_locbf_D) - self.x_his_indices[i] = np.arange(self.n_his) + self.x_his_indices[i] = xp.arange(self.n_his) self.coeffh_indices[i] = i for j in range(2 * self.p + 1): xi = self.p - 1 @@ -354,10 +354,10 @@ def __init__(self, spline_space, n_quad): # right boundary region elif i > n_lambda_his - self.p: - self.his_global_N[i] = np.arange(self.n_his_locbf_N) + n_lambda_his - self.p - (self.p - 1) - self.his_global_D[i] = np.arange(self.n_his_locbf_D) + n_lambda_his - self.p - (self.p - 1) + self.his_global_N[i] = xp.arange(self.n_his_locbf_N) + n_lambda_his - self.p - (self.p - 1) + self.his_global_D[i] = xp.arange(self.n_his_locbf_D) + n_lambda_his - self.p - (self.p - 1) - self.x_his_indices[i] = np.arange(self.n_his) + 2 * (n_lambda_his - self.p - (self.p - 1)) + self.x_his_indices[i] = xp.arange(self.n_his) + 2 * (n_lambda_his - self.p - (self.p - 1)) self.coeffh_indices[i] = counter_coeffh counter_coeffh += 1 for j in range(2 * self.p + 1): @@ -366,10 +366,10 @@ def __init__(self, spline_space, n_quad): # interior else: - self.his_global_N[i] = np.arange(self.n_his_locbf_N) + i - (self.p - 1) - self.his_global_D[i] = np.arange(self.n_his_locbf_D) + i - (self.p - 1) + self.his_global_N[i] = xp.arange(self.n_his_locbf_N) + i - (self.p - 1) + self.his_global_D[i] = xp.arange(self.n_his_locbf_D) + i - (self.p - 1) - self.x_his_indices[i] = np.arange(self.n_his) + 2 * (i - (self.p - 1)) + self.x_his_indices[i] = xp.arange(self.n_his) + 2 * (i - (self.p - 1)) self.coeffh_indices[i] = self.p - 1 for j in range(2 * self.p + 1): self.x_his[i, j] = (self.T[i + 1 + int(j / 2)] + self.T[i + 1 + int((j + 1) / 2)]) / 2 @@ -380,8 +380,8 @@ def __init__(self, spline_space, n_quad): k_glob_new = self.his_global_D[i, il] bol = k_glob_new == self.his_global_D[i - 1] - if np.any(bol): - self.his_loccof_D[i, il] = self.his_loccof_D[i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.his_loccof_D[i, il] = self.his_loccof_D[i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_his - self.p - (self.p - 2)) and (self.his_loccof_D[i, il] == 0): self.his_loccof_D[i, il] = self.his_add_D[counter_D] @@ -391,15 +391,15 @@ def __init__(self, spline_space, n_quad): k_glob_new = self.his_global_N[i, il] bol = k_glob_new == self.his_global_N[i - 1] - if np.any(bol): - self.his_loccof_N[i, il] = self.his_loccof_N[i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.his_loccof_N[i, il] = self.his_loccof_N[i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_his - self.p - (self.p - 2)) and (self.his_loccof_N[i, il] == 0): self.his_loccof_N[i, il] = self.his_add_N[counter_N] counter_N += 1 # quadrature points and weights - self.pts, self.wts = bsp.quadrature_grid(np.unique(self.x_his.flatten()), self.pts_loc, self.wts_loc) + self.pts, self.wts = bsp.quadrature_grid(xp.unique(self.x_his.flatten()), self.pts_loc, self.wts_loc) else: # maximum number of non-vanishing coefficients @@ -407,31 +407,31 @@ def __init__(self, spline_space, n_quad): self.n_his_nvcof_N = 2 * self.p # shift local coefficients --> global coefficients - self.his_shift_D = np.arange(self.NbaseD) - (self.p - 1) - self.his_shift_N = np.arange(self.NbaseD) - self.p + self.his_shift_D = xp.arange(self.NbaseD) - (self.p - 1) + self.his_shift_N = xp.arange(self.NbaseD) - self.p for i in range(n_lambda_his): - self.his_global_N[i] = (np.arange(self.n_his_locbf_N) + i - (self.p - 1)) % self.NbaseN - self.his_global_D[i] = (np.arange(self.n_his_locbf_D) + i - (self.p - 1)) % self.NbaseD - self.his_loccof_N[i] = np.arange(self.n_his_locbf_N - 1, -1, -1) - self.his_loccof_D[i] = np.arange(self.n_his_locbf_D - 1, -1, -1) + self.his_global_N[i] = (xp.arange(self.n_his_locbf_N) + i - (self.p - 1)) % self.NbaseN + self.his_global_D[i] = (xp.arange(self.n_his_locbf_D) + i - (self.p - 1)) % self.NbaseD + self.his_loccof_N[i] = xp.arange(self.n_his_locbf_N - 1, -1, -1) + self.his_loccof_D[i] = xp.arange(self.n_his_locbf_D - 1, -1, -1) - self.x_his_indices[i] = np.arange(self.n_his) + 2 * (i - (self.p - 1)) + self.x_his_indices[i] = xp.arange(self.n_his) + 2 * (i - (self.p - 1)) self.coeffh_indices[i] = 0 for j in range(2 * self.p + 1): self.x_his[i, j] = (self.T[i + 1 + int(j / 2)] + self.T[i + 1 + int((j + 1) / 2)]) / 2 # quadrature points and weights self.pts, self.wts = bsp.quadrature_grid( - np.append(np.unique(self.x_his.flatten() % 1.0), 1.0), self.pts_loc, self.wts_loc + xp.append(xp.unique(self.x_his.flatten() % 1.0), 1.0), self.pts_loc, self.wts_loc ) # quasi interpolation def pi_0(self, fun): - lambdas = np.zeros(self.NbaseN, dtype=float) + lambdas = xp.zeros(self.NbaseN, dtype=float) # evaluate function at interpolation points - mat_f = fun(np.unique(self.x_int.flatten())) + mat_f = fun(xp.unique(self.x_int.flatten())) for i in range(self.NbaseN): for j in range(self.n_int): @@ -441,7 +441,7 @@ def pi_0(self, fun): # quasi histopolation def pi_1(self, fun): - lambdas = np.zeros(self.NbaseD, dtype=float) + lambdas = xp.zeros(self.NbaseD, dtype=float) # evaluate function at quadrature points mat_f = fun(self.pts) @@ -459,17 +459,17 @@ def pi_1(self, fun): # projection matrices of products of basis functions: pi0_i(A_j*B_k) and pi1_i(A_j*B_k) def projection_matrices_1d(self, bc_kind=["free", "free"]): - PI0_NN = np.empty((self.NbaseN, self.NbaseN, self.NbaseN), dtype=float) - PI0_DN = np.empty((self.NbaseN, self.NbaseD, self.NbaseN), dtype=float) - PI0_DD = np.empty((self.NbaseN, self.NbaseD, self.NbaseD), dtype=float) + PI0_NN = xp.empty((self.NbaseN, self.NbaseN, self.NbaseN), dtype=float) + PI0_DN = xp.empty((self.NbaseN, self.NbaseD, self.NbaseN), dtype=float) + PI0_DD = xp.empty((self.NbaseN, self.NbaseD, self.NbaseD), dtype=float) - PI1_NN = np.empty((self.NbaseD, self.NbaseN, self.NbaseN), dtype=float) - PI1_DN = np.empty((self.NbaseD, self.NbaseD, self.NbaseN), dtype=float) - PI1_DD = np.empty((self.NbaseD, self.NbaseD, self.NbaseD), dtype=float) + PI1_NN = xp.empty((self.NbaseD, self.NbaseN, self.NbaseN), dtype=float) + PI1_DN = xp.empty((self.NbaseD, self.NbaseD, self.NbaseN), dtype=float) + PI1_DD = xp.empty((self.NbaseD, self.NbaseD, self.NbaseD), dtype=float) # ========= PI0__NN and PI1_NN ============= - ci = np.zeros(self.NbaseN, dtype=float) - cj = np.zeros(self.NbaseN, dtype=float) + ci = xp.zeros(self.NbaseN, dtype=float) + cj = xp.zeros(self.NbaseN, dtype=float) for i in range(self.NbaseN): for j in range(self.NbaseN): @@ -485,8 +485,8 @@ def projection_matrices_1d(self, bc_kind=["free", "free"]): PI1_NN[:, i, j] = self.pi_1(fun) # ========= PI0__DN and PI1_DN ============= - ci = np.zeros(self.NbaseD, dtype=float) - cj = np.zeros(self.NbaseN, dtype=float) + ci = xp.zeros(self.NbaseD, dtype=float) + cj = xp.zeros(self.NbaseN, dtype=float) for i in range(self.NbaseD): for j in range(self.NbaseN): @@ -502,8 +502,8 @@ def projection_matrices_1d(self, bc_kind=["free", "free"]): PI1_DN[:, i, j] = self.pi_1(fun) # ========= PI0__DD and PI1_DD ============= - ci = np.zeros(self.NbaseD, dtype=float) - cj = np.zeros(self.NbaseD, dtype=float) + ci = xp.zeros(self.NbaseD, dtype=float) + cj = xp.zeros(self.NbaseD, dtype=float) for i in range(self.NbaseD): for j in range(self.NbaseD): @@ -518,8 +518,8 @@ def projection_matrices_1d(self, bc_kind=["free", "free"]): PI0_DD[:, i, j] = self.pi_0(fun) PI1_DD[:, i, j] = self.pi_1(fun) - PI0_ND = np.transpose(PI0_DN, (0, 2, 1)) - PI1_ND = np.transpose(PI1_DN, (0, 2, 1)) + PI0_ND = xp.transpose(PI0_DN, (0, 2, 1)) + PI1_ND = xp.transpose(PI1_DN, (0, 2, 1)) # remove contributions from first and last N-splines if bc_kind[0] == "dirichlet": @@ -544,25 +544,25 @@ def projection_matrices_1d(self, bc_kind=["free", "free"]): PI1_DN[:, :, -1] = 0.0 PI1_ND[:, -1, :] = 0.0 - PI0_NN_indices = np.nonzero(PI0_NN) - PI0_DN_indices = np.nonzero(PI0_DN) - PI0_ND_indices = np.nonzero(PI0_ND) - PI0_DD_indices = np.nonzero(PI0_DD) + PI0_NN_indices = xp.nonzero(PI0_NN) + PI0_DN_indices = xp.nonzero(PI0_DN) + PI0_ND_indices = xp.nonzero(PI0_ND) + PI0_DD_indices = xp.nonzero(PI0_DD) - PI1_NN_indices = np.nonzero(PI1_NN) - PI1_DN_indices = np.nonzero(PI1_DN) - PI1_ND_indices = np.nonzero(PI1_ND) - PI1_DD_indices = np.nonzero(PI1_DD) + PI1_NN_indices = xp.nonzero(PI1_NN) + PI1_DN_indices = xp.nonzero(PI1_DN) + PI1_ND_indices = xp.nonzero(PI1_ND) + PI1_DD_indices = xp.nonzero(PI1_DD) - PI0_NN_indices = np.vstack((PI0_NN_indices[0], PI0_NN_indices[1], PI0_NN_indices[2])) - PI0_DN_indices = np.vstack((PI0_DN_indices[0], PI0_DN_indices[1], PI0_DN_indices[2])) - PI0_ND_indices = np.vstack((PI0_ND_indices[0], PI0_ND_indices[1], PI0_ND_indices[2])) - PI0_DD_indices = np.vstack((PI0_DD_indices[0], PI0_DD_indices[1], PI0_DD_indices[2])) + PI0_NN_indices = xp.vstack((PI0_NN_indices[0], PI0_NN_indices[1], PI0_NN_indices[2])) + PI0_DN_indices = xp.vstack((PI0_DN_indices[0], PI0_DN_indices[1], PI0_DN_indices[2])) + PI0_ND_indices = xp.vstack((PI0_ND_indices[0], PI0_ND_indices[1], PI0_ND_indices[2])) + PI0_DD_indices = xp.vstack((PI0_DD_indices[0], PI0_DD_indices[1], PI0_DD_indices[2])) - PI1_NN_indices = np.vstack((PI1_NN_indices[0], PI1_NN_indices[1], PI1_NN_indices[2])) - PI1_DN_indices = np.vstack((PI1_DN_indices[0], PI1_DN_indices[1], PI1_DN_indices[2])) - PI1_ND_indices = np.vstack((PI1_ND_indices[0], PI1_ND_indices[1], PI1_ND_indices[2])) - PI1_DD_indices = np.vstack((PI1_DD_indices[0], PI1_DD_indices[1], PI1_DD_indices[2])) + PI1_NN_indices = xp.vstack((PI1_NN_indices[0], PI1_NN_indices[1], PI1_NN_indices[2])) + PI1_DN_indices = xp.vstack((PI1_DN_indices[0], PI1_DN_indices[1], PI1_DN_indices[2])) + PI1_ND_indices = xp.vstack((PI1_ND_indices[0], PI1_ND_indices[1], PI1_ND_indices[2])) + PI1_DD_indices = xp.vstack((PI1_DD_indices[0], PI1_DD_indices[1], PI1_DD_indices[2])) return ( PI0_NN, @@ -617,8 +617,8 @@ def __init__(self, tensor_space, n_quad): self.polar = False # local projectors for polar splines are not implemented yet # Gauss - Legendre quadrature points and weights in (-1, 1) - self.pts_loc = [np.polynomial.legendre.leggauss(n_quad)[0] for n_quad in self.n_quad] - self.wts_loc = [np.polynomial.legendre.leggauss(n_quad)[1] for n_quad in self.n_quad] + self.pts_loc = [xp.polynomial.legendre.leggauss(n_quad)[0] for n_quad in self.n_quad] + self.wts_loc = [xp.polynomial.legendre.leggauss(n_quad)[1] for n_quad in self.n_quad] # set interpolation and histopolation coefficients self.coeff_i = [0, 0, 0] @@ -626,78 +626,78 @@ def __init__(self, tensor_space, n_quad): for a in range(3): if self.bc[a] == True: - self.coeff_i[a] = np.zeros((1, 2 * self.p[a] - 1), dtype=float) - self.coeff_h[a] = np.zeros((1, 2 * self.p[a]), dtype=float) + self.coeff_i[a] = xp.zeros((1, 2 * self.p[a] - 1), dtype=float) + self.coeff_h[a] = xp.zeros((1, 2 * self.p[a]), dtype=float) if self.p[a] == 1: - self.coeff_i[a][0, :] = np.array([1.0]) - self.coeff_h[a][0, :] = np.array([1.0, 1.0]) + self.coeff_i[a][0, :] = xp.array([1.0]) + self.coeff_h[a][0, :] = xp.array([1.0, 1.0]) elif self.p[a] == 2: - self.coeff_i[a][0, :] = 1 / 2 * np.array([-1.0, 4.0, -1.0]) - self.coeff_h[a][0, :] = 1 / 2 * np.array([-1.0, 3.0, 3.0, -1.0]) + self.coeff_i[a][0, :] = 1 / 2 * xp.array([-1.0, 4.0, -1.0]) + self.coeff_h[a][0, :] = 1 / 2 * xp.array([-1.0, 3.0, 3.0, -1.0]) elif self.p[a] == 3: - self.coeff_i[a][0, :] = 1 / 6 * np.array([1.0, -8.0, 20.0, -8.0, 1.0]) - self.coeff_h[a][0, :] = 1 / 6 * np.array([1.0, -7.0, 12.0, 12.0, -7.0, 1.0]) + self.coeff_i[a][0, :] = 1 / 6 * xp.array([1.0, -8.0, 20.0, -8.0, 1.0]) + self.coeff_h[a][0, :] = 1 / 6 * xp.array([1.0, -7.0, 12.0, 12.0, -7.0, 1.0]) elif self.p[a] == 4: - self.coeff_i[a][0, :] = 2 / 45 * np.array([-1.0, 16.0, -295 / 4, 140.0, -295 / 4, 16.0, -1.0]) + self.coeff_i[a][0, :] = 2 / 45 * xp.array([-1.0, 16.0, -295 / 4, 140.0, -295 / 4, 16.0, -1.0]) self.coeff_h[a][0, :] = ( - 2 / 45 * np.array([-1.0, 15.0, -231 / 4, 265 / 4, 265 / 4, -231 / 4, 15.0, -1.0]) + 2 / 45 * xp.array([-1.0, 15.0, -231 / 4, 265 / 4, 265 / 4, -231 / 4, 15.0, -1.0]) ) else: print("degree > 4 not implemented!") else: - self.coeff_i[a] = np.zeros((2 * self.p[a] - 1, 2 * self.p[a] - 1), dtype=float) - self.coeff_h[a] = np.zeros((2 * self.p[a] - 1, 2 * self.p[a]), dtype=float) + self.coeff_i[a] = xp.zeros((2 * self.p[a] - 1, 2 * self.p[a] - 1), dtype=float) + self.coeff_h[a] = xp.zeros((2 * self.p[a] - 1, 2 * self.p[a]), dtype=float) if self.p[a] == 1: - self.coeff_i[a][0, :] = np.array([1.0]) - self.coeff_h[a][0, :] = np.array([1.0, 1.0]) + self.coeff_i[a][0, :] = xp.array([1.0]) + self.coeff_h[a][0, :] = xp.array([1.0, 1.0]) elif self.p[a] == 2: - self.coeff_i[a][0, :] = 1 / 2 * np.array([2.0, 0.0, 0.0]) - self.coeff_i[a][1, :] = 1 / 2 * np.array([-1.0, 4.0, -1.0]) - self.coeff_i[a][2, :] = 1 / 2 * np.array([0.0, 0.0, 2.0]) + self.coeff_i[a][0, :] = 1 / 2 * xp.array([2.0, 0.0, 0.0]) + self.coeff_i[a][1, :] = 1 / 2 * xp.array([-1.0, 4.0, -1.0]) + self.coeff_i[a][2, :] = 1 / 2 * xp.array([0.0, 0.0, 2.0]) - self.coeff_h[a][0, :] = 1 / 2 * np.array([3.0, -1.0, 0.0, 0.0]) - self.coeff_h[a][1, :] = 1 / 2 * np.array([-1.0, 3.0, 3.0, -1.0]) - self.coeff_h[a][2, :] = 1 / 2 * np.array([0.0, 0.0, -1.0, 3.0]) + self.coeff_h[a][0, :] = 1 / 2 * xp.array([3.0, -1.0, 0.0, 0.0]) + self.coeff_h[a][1, :] = 1 / 2 * xp.array([-1.0, 3.0, 3.0, -1.0]) + self.coeff_h[a][2, :] = 1 / 2 * xp.array([0.0, 0.0, -1.0, 3.0]) elif self.p[a] == 3: - self.coeff_i[a][0, :] = 1 / 18 * np.array([18.0, 0.0, 0.0, 0.0, 0.0]) - self.coeff_i[a][1, :] = 1 / 18 * np.array([-5.0, 40.0, -24.0, 8.0, -1.0]) - self.coeff_i[a][2, :] = 1 / 18 * np.array([3.0, -24.0, 60.0, -24.0, 3.0]) - self.coeff_i[a][3, :] = 1 / 18 * np.array([-1.0, 8.0, -24.0, 40.0, -5.0]) - self.coeff_i[a][4, :] = 1 / 18 * np.array([0.0, 0.0, 0.0, 0.0, 18.0]) - - self.coeff_h[a][0, :] = 1 / 18 * np.array([23.0, -17.0, 7.0, -1.0, 0.0, 0.0]) - self.coeff_h[a][1, :] = 1 / 18 * np.array([-8.0, 56.0, -28.0, 4.0, 0.0, 0.0]) - self.coeff_h[a][2, :] = 1 / 18 * np.array([3.0, -21.0, 36.0, 36.0, -21.0, 3.0]) - self.coeff_h[a][3, :] = 1 / 18 * np.array([0.0, 0.0, 4.0, -28.0, 56.0, -8.0]) - self.coeff_h[a][4, :] = 1 / 18 * np.array([0.0, 0.0, -1.0, 7.0, -17.0, 23.0]) + self.coeff_i[a][0, :] = 1 / 18 * xp.array([18.0, 0.0, 0.0, 0.0, 0.0]) + self.coeff_i[a][1, :] = 1 / 18 * xp.array([-5.0, 40.0, -24.0, 8.0, -1.0]) + self.coeff_i[a][2, :] = 1 / 18 * xp.array([3.0, -24.0, 60.0, -24.0, 3.0]) + self.coeff_i[a][3, :] = 1 / 18 * xp.array([-1.0, 8.0, -24.0, 40.0, -5.0]) + self.coeff_i[a][4, :] = 1 / 18 * xp.array([0.0, 0.0, 0.0, 0.0, 18.0]) + + self.coeff_h[a][0, :] = 1 / 18 * xp.array([23.0, -17.0, 7.0, -1.0, 0.0, 0.0]) + self.coeff_h[a][1, :] = 1 / 18 * xp.array([-8.0, 56.0, -28.0, 4.0, 0.0, 0.0]) + self.coeff_h[a][2, :] = 1 / 18 * xp.array([3.0, -21.0, 36.0, 36.0, -21.0, 3.0]) + self.coeff_h[a][3, :] = 1 / 18 * xp.array([0.0, 0.0, 4.0, -28.0, 56.0, -8.0]) + self.coeff_h[a][4, :] = 1 / 18 * xp.array([0.0, 0.0, -1.0, 7.0, -17.0, 23.0]) elif self.p[a] == 4: - self.coeff_i[a][0, :] = 1 / 360 * np.array([360.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) - self.coeff_i[a][1, :] = 1 / 360 * np.array([-59.0, 944.0, -1000.0, 720.0, -305.0, 64.0, -4.0]) - self.coeff_i[a][2, :] = 1 / 360 * np.array([23.0, -368.0, 1580.0, -1360.0, 605.0, -128.0, 8.0]) - self.coeff_i[a][3, :] = 1 / 360 * np.array([-16.0, 256.0, -1180.0, 2240.0, -1180.0, 256.0, -16.0]) - self.coeff_i[a][4, :] = 1 / 360 * np.array([8.0, -128.0, 605.0, -1360.0, 1580.0, -368.0, 23.0]) - self.coeff_i[a][5, :] = 1 / 360 * np.array([-4.0, 64.0, -305.0, 720.0, -1000.0, 944.0, -59.0]) - self.coeff_i[a][6, :] = 1 / 360 * np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 360.0]) - - self.coeff_h[a][0, :] = 1 / 360 * np.array([419.0, -525.0, 475.0, -245.0, 60.0, -4.0, 0.0, 0.0]) - self.coeff_h[a][1, :] = 1 / 360 * np.array([-82.0, 1230.0, -1350.0, 730.0, -180.0, 12.0, 0.0, 0.0]) - self.coeff_h[a][2, :] = 1 / 360 * np.array([39.0, -585.0, 2175.0, -1425.0, 360.0, -24.0, 0.0, 0.0]) + self.coeff_i[a][0, :] = 1 / 360 * xp.array([360.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + self.coeff_i[a][1, :] = 1 / 360 * xp.array([-59.0, 944.0, -1000.0, 720.0, -305.0, 64.0, -4.0]) + self.coeff_i[a][2, :] = 1 / 360 * xp.array([23.0, -368.0, 1580.0, -1360.0, 605.0, -128.0, 8.0]) + self.coeff_i[a][3, :] = 1 / 360 * xp.array([-16.0, 256.0, -1180.0, 2240.0, -1180.0, 256.0, -16.0]) + self.coeff_i[a][4, :] = 1 / 360 * xp.array([8.0, -128.0, 605.0, -1360.0, 1580.0, -368.0, 23.0]) + self.coeff_i[a][5, :] = 1 / 360 * xp.array([-4.0, 64.0, -305.0, 720.0, -1000.0, 944.0, -59.0]) + self.coeff_i[a][6, :] = 1 / 360 * xp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 360.0]) + + self.coeff_h[a][0, :] = 1 / 360 * xp.array([419.0, -525.0, 475.0, -245.0, 60.0, -4.0, 0.0, 0.0]) + self.coeff_h[a][1, :] = 1 / 360 * xp.array([-82.0, 1230.0, -1350.0, 730.0, -180.0, 12.0, 0.0, 0.0]) + self.coeff_h[a][2, :] = 1 / 360 * xp.array([39.0, -585.0, 2175.0, -1425.0, 360.0, -24.0, 0.0, 0.0]) self.coeff_h[a][3, :] = ( - 1 / 360 * np.array([-16.0, 240.0, -924.0, 1060.0, 1060.0, -924.0, 240.0, -16.0]) + 1 / 360 * xp.array([-16.0, 240.0, -924.0, 1060.0, 1060.0, -924.0, 240.0, -16.0]) ) - self.coeff_h[a][4, :] = 1 / 360 * np.array([0.0, 0.0, -24.0, 360.0, -1425.0, 2175.0, -585.0, 39.0]) - self.coeff_h[a][5, :] = 1 / 360 * np.array([0.0, 0.0, 12.0, -180.0, 730.0, -1350.0, 1230.0, -82.0]) - self.coeff_h[a][6, :] = 1 / 360 * np.array([0.0, 0.0, -4.0, 60.0, -245.0, 475.0, -525.0, 419.0]) + self.coeff_h[a][4, :] = 1 / 360 * xp.array([0.0, 0.0, -24.0, 360.0, -1425.0, 2175.0, -585.0, 39.0]) + self.coeff_h[a][5, :] = 1 / 360 * xp.array([0.0, 0.0, 12.0, -180.0, 730.0, -1350.0, 1230.0, -82.0]) + self.coeff_h[a][6, :] = 1 / 360 * xp.array([0.0, 0.0, -4.0, 60.0, -245.0, 475.0, -525.0, 419.0]) else: print("degree > 4 not implemented!") @@ -723,31 +723,31 @@ def __init__(self, tensor_space, n_quad): ) # number of non-vanishing D bf in interpolation interval (1, 2, 4, 6) self.x_int = [ - np.zeros((n_lambda_int, n_int), dtype=float) for n_lambda_int, n_int in zip(n_lambda_int, self.n_int) + xp.zeros((n_lambda_int, n_int), dtype=float) for n_lambda_int, n_int in zip(n_lambda_int, self.n_int) ] self.int_global_N = [ - np.zeros((n_lambda_int, n_int_locbf_N), dtype=int) + xp.zeros((n_lambda_int, n_int_locbf_N), dtype=int) for n_lambda_int, n_int_locbf_N in zip(n_lambda_int, self.n_int_locbf_N) ] self.int_global_D = [ - np.zeros((n_lambda_int, n_int_locbf_D), dtype=int) + xp.zeros((n_lambda_int, n_int_locbf_D), dtype=int) for n_lambda_int, n_int_locbf_D in zip(n_lambda_int, self.n_int_locbf_D) ] self.int_loccof_N = [ - np.zeros((n_lambda_int, n_int_locbf_N), dtype=int) + xp.zeros((n_lambda_int, n_int_locbf_N), dtype=int) for n_lambda_int, n_int_locbf_N in zip(n_lambda_int, self.n_int_locbf_N) ] self.int_loccof_D = [ - np.zeros((n_lambda_int, n_int_locbf_D), dtype=int) + xp.zeros((n_lambda_int, n_int_locbf_D), dtype=int) for n_lambda_int, n_int_locbf_D in zip(n_lambda_int, self.n_int_locbf_D) ] self.x_int_indices = [ - np.zeros((n_lambda_int, n_int), dtype=int) for n_lambda_int, n_int in zip(n_lambda_int, self.n_int) + xp.zeros((n_lambda_int, n_int), dtype=int) for n_lambda_int, n_int in zip(n_lambda_int, self.n_int) ] - self.coeffi_indices = [np.zeros(n_lambda_int, dtype=int) for n_lambda_int in n_lambda_int] + self.coeffi_indices = [xp.zeros(n_lambda_int, dtype=int) for n_lambda_int in n_lambda_int] self.n_int_nvcof_D = [None, None, None] self.n_int_nvcof_N = [None, None, None] @@ -770,39 +770,39 @@ def __init__(self, tensor_space, n_quad): self.n_int_nvcof_N[a] = 3 * self.p[a] - 2 # shift in local coefficient indices at right boundary (only for non-periodic boundary conditions) - self.int_add_D[a] = np.arange(self.n_int[a] - 2) + 1 - self.int_add_N[a] = np.arange(self.n_int[a] - 1) + 1 + self.int_add_D[a] = xp.arange(self.n_int[a] - 2) + 1 + self.int_add_N[a] = xp.arange(self.n_int[a] - 1) + 1 counter_D = 0 counter_N = 0 # shift local coefficients --> global coefficients (D) if self.p[a] == 1: - self.int_shift_D[a] = np.arange(self.NbaseD[a]) + self.int_shift_D[a] = xp.arange(self.NbaseD[a]) else: - self.int_shift_D[a] = np.arange(self.NbaseD[a]) - (self.p[a] - 2) + self.int_shift_D[a] = xp.arange(self.NbaseD[a]) - (self.p[a] - 2) self.int_shift_D[a][: 2 * self.p[a] - 2] = 0 self.int_shift_D[a][-(2 * self.p[a] - 2) :] = self.int_shift_D[a][-(2 * self.p[a] - 2)] # shift local coefficients --> global coefficients (N) if self.p[a] == 1: - self.int_shift_N[a] = np.arange(self.NbaseN[a]) + self.int_shift_N[a] = xp.arange(self.NbaseN[a]) self.int_shift_N[a][-1] = self.int_shift_N[a][-2] else: - self.int_shift_N[a] = np.arange(self.NbaseN[a]) - (self.p[a] - 1) + self.int_shift_N[a] = xp.arange(self.NbaseN[a]) - (self.p[a] - 1) self.int_shift_N[a][: 2 * self.p[a] - 1] = 0 self.int_shift_N[a][-(2 * self.p[a] - 1) :] = self.int_shift_N[a][-(2 * self.p[a] - 1)] - counter_coeffi = np.copy(self.p[a]) + counter_coeffi = xp.copy(self.p[a]) for i in range(n_lambda_int[a]): # left boundary region if i < self.p[a] - 1: - self.int_global_N[a][i] = np.arange(self.n_int_locbf_N[a]) - self.int_global_D[a][i] = np.arange(self.n_int_locbf_D[a]) + self.int_global_N[a][i] = xp.arange(self.n_int_locbf_N[a]) + self.int_global_D[a][i] = xp.arange(self.n_int_locbf_D[a]) - self.x_int_indices[a][i] = np.arange(self.n_int[a]) + self.x_int_indices[a][i] = xp.arange(self.n_int[a]) self.coeffi_indices[a][i] = i for j in range(2 * (self.p[a] - 1) + 1): xi = self.p[a] - 1 @@ -813,13 +813,13 @@ def __init__(self, tensor_space, n_quad): # right boundary region elif i > n_lambda_int[a] - self.p[a]: self.int_global_N[a][i] = ( - np.arange(self.n_int_locbf_N[a]) + n_lambda_int[a] - self.p[a] - (self.p[a] - 1) + xp.arange(self.n_int_locbf_N[a]) + n_lambda_int[a] - self.p[a] - (self.p[a] - 1) ) self.int_global_D[a][i] = ( - np.arange(self.n_int_locbf_D[a]) + n_lambda_int[a] - self.p[a] - (self.p[a] - 1) + xp.arange(self.n_int_locbf_D[a]) + n_lambda_int[a] - self.p[a] - (self.p[a] - 1) ) - self.x_int_indices[a][i] = np.arange(self.n_int[a]) + 2 * ( + self.x_int_indices[a][i] = xp.arange(self.n_int[a]) + 2 * ( n_lambda_int[a] - self.p[a] - (self.p[a] - 1) ) self.coeffi_indices[a][i] = counter_coeffi @@ -833,20 +833,20 @@ def __init__(self, tensor_space, n_quad): # interior else: if self.p[a] == 1: - self.int_global_N[a][i] = np.arange(self.n_int_locbf_N[a]) + i - self.int_global_D[a][i] = np.arange(self.n_int_locbf_D[a]) + i + self.int_global_N[a][i] = xp.arange(self.n_int_locbf_N[a]) + i + self.int_global_D[a][i] = xp.arange(self.n_int_locbf_D[a]) + i self.int_global_N[a][-1] = self.int_global_N[a][-2] self.int_global_D[a][-1] = self.int_global_D[a][-2] else: - self.int_global_N[a][i] = np.arange(self.n_int_locbf_N[a]) + i - (self.p[a] - 1) - self.int_global_D[a][i] = np.arange(self.n_int_locbf_D[a]) + i - (self.p[a] - 1) + self.int_global_N[a][i] = xp.arange(self.n_int_locbf_N[a]) + i - (self.p[a] - 1) + self.int_global_D[a][i] = xp.arange(self.n_int_locbf_D[a]) + i - (self.p[a] - 1) if self.p[a] == 1: self.x_int_indices[a][i] = i else: - self.x_int_indices[a][i] = np.arange(self.n_int[a]) + 2 * (i - (self.p[a] - 1)) + self.x_int_indices[a][i] = xp.arange(self.n_int[a]) + 2 * (i - (self.p[a] - 1)) self.coeffi_indices[a][i] = self.p[a] - 1 @@ -857,8 +857,8 @@ def __init__(self, tensor_space, n_quad): # local coefficient index if self.p[a] == 1: - self.int_loccof_N[a][i] = np.array([0, 1]) - self.int_loccof_D[a][-1] = np.array([1]) + self.int_loccof_N[a][i] = xp.array([0, 1]) + self.int_loccof_D[a][-1] = xp.array([1]) else: if i > 0: @@ -866,8 +866,8 @@ def __init__(self, tensor_space, n_quad): k_glob_new = self.int_global_D[a][i, il] bol = k_glob_new == self.int_global_D[a][i - 1] - if np.any(bol): - self.int_loccof_D[a][i, il] = self.int_loccof_D[a][i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.int_loccof_D[a][i, il] = self.int_loccof_D[a][i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_int[a] - self.p[a] - (self.p[a] - 2)) and ( self.int_loccof_D[a][i, il] == 0 @@ -879,8 +879,8 @@ def __init__(self, tensor_space, n_quad): k_glob_new = self.int_global_N[a][i, il] bol = k_glob_new == self.int_global_N[a][i - 1] - if np.any(bol): - self.int_loccof_N[a][i, il] = self.int_loccof_N[a][i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.int_loccof_N[a][i, il] = self.int_loccof_N[a][i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_int[a] - self.p[a] - (self.p[a] - 2)) and ( self.int_loccof_N[a][i, il] == 0 @@ -900,24 +900,24 @@ def __init__(self, tensor_space, n_quad): # shift local coefficients --> global coefficients if self.p[a] == 1: - self.int_shift_D[a] = np.arange(self.NbaseN[a]) - (self.p[a] - 1) - self.int_shift_N[a] = np.arange(self.NbaseN[a]) - (self.p[a]) + self.int_shift_D[a] = xp.arange(self.NbaseN[a]) - (self.p[a] - 1) + self.int_shift_N[a] = xp.arange(self.NbaseN[a]) - (self.p[a]) else: - self.int_shift_D[a] = np.arange(self.NbaseN[a]) - (self.p[a] - 2) - self.int_shift_N[a] = np.arange(self.NbaseN[a]) - (self.p[a] - 1) + self.int_shift_D[a] = xp.arange(self.NbaseN[a]) - (self.p[a] - 2) + self.int_shift_N[a] = xp.arange(self.NbaseN[a]) - (self.p[a] - 1) for i in range(n_lambda_int[a]): # global indices of non-vanishing basis functions and position of coefficients in final matrix - self.int_global_N[a][i] = (np.arange(self.n_int_locbf_N[a]) + i - (self.p[a] - 1)) % self.NbaseN[a] - self.int_global_D[a][i] = (np.arange(self.n_int_locbf_D[a]) + i - (self.p[a] - 1)) % self.NbaseD[a] + self.int_global_N[a][i] = (xp.arange(self.n_int_locbf_N[a]) + i - (self.p[a] - 1)) % self.NbaseN[a] + self.int_global_D[a][i] = (xp.arange(self.n_int_locbf_D[a]) + i - (self.p[a] - 1)) % self.NbaseD[a] - self.int_loccof_N[a][i] = np.arange(self.n_int_locbf_N[a] - 1, -1, -1) - self.int_loccof_D[a][i] = np.arange(self.n_int_locbf_D[a] - 1, -1, -1) + self.int_loccof_N[a][i] = xp.arange(self.n_int_locbf_N[a] - 1, -1, -1) + self.int_loccof_D[a][i] = xp.arange(self.n_int_locbf_D[a] - 1, -1, -1) if self.p[a] == 1: self.x_int_indices[a][i] = i else: - self.x_int_indices[a][i] = (np.arange(self.n_int[a]) + 2 * (i - (self.p[a] - 1))) % ( + self.x_int_indices[a][i] = (xp.arange(self.n_int[a]) + 2 * (i - (self.p[a] - 1))) % ( 2 * self.Nel[a] ) @@ -929,38 +929,38 @@ def __init__(self, tensor_space, n_quad): ) % 1.0 # set histopolation points, quadrature points and weights - n_lambda_his = [np.copy(NbaseD) for NbaseD in self.NbaseD] # number of coefficients in space V1 + n_lambda_his = [xp.copy(NbaseD) for NbaseD in self.NbaseD] # number of coefficients in space V1 self.n_his = [2 * p for p in self.p] # number of histopolation intervals self.n_his_locbf_N = [2 * p for p in self.p] # number of non-vanishing N bf in histopolation interval self.n_his_locbf_D = [2 * p - 1 for p in self.p] # number of non-vanishing D bf in histopolation interval self.x_his = [ - np.zeros((n_lambda_his, n_his + 1), dtype=float) for n_lambda_his, n_his in zip(n_lambda_his, self.n_his) + xp.zeros((n_lambda_his, n_his + 1), dtype=float) for n_lambda_his, n_his in zip(n_lambda_his, self.n_his) ] self.his_global_N = [ - np.zeros((n_lambda_his, n_his_locbf_N), dtype=int) + xp.zeros((n_lambda_his, n_his_locbf_N), dtype=int) for n_lambda_his, n_his_locbf_N in zip(n_lambda_his, self.n_his_locbf_N) ] self.his_global_D = [ - np.zeros((n_lambda_his, n_his_locbf_D), dtype=int) + xp.zeros((n_lambda_his, n_his_locbf_D), dtype=int) for n_lambda_his, n_his_locbf_D in zip(n_lambda_his, self.n_his_locbf_D) ] self.his_loccof_N = [ - np.zeros((n_lambda_his, n_his_locbf_N), dtype=int) + xp.zeros((n_lambda_his, n_his_locbf_N), dtype=int) for n_lambda_his, n_his_locbf_N in zip(n_lambda_his, self.n_his_locbf_N) ] self.his_loccof_D = [ - np.zeros((n_lambda_his, n_his_locbf_D), dtype=int) + xp.zeros((n_lambda_his, n_his_locbf_D), dtype=int) for n_lambda_his, n_his_locbf_D in zip(n_lambda_his, self.n_his_locbf_D) ] self.x_his_indices = [ - np.zeros((n_lambda_his, n_his), dtype=int) for n_lambda_his, n_his in zip(n_lambda_his, self.n_his) + xp.zeros((n_lambda_his, n_his), dtype=int) for n_lambda_his, n_his in zip(n_lambda_his, self.n_his) ] - self.coeffh_indices = [np.zeros(n_lambda_his, dtype=int) for n_lambda_his in n_lambda_his] + self.coeffh_indices = [xp.zeros(n_lambda_his, dtype=int) for n_lambda_his in n_lambda_his] self.pts = [0, 0, 0] self.wts = [0, 0, 0] @@ -981,31 +981,31 @@ def __init__(self, tensor_space, n_quad): self.n_his_nvcof_N[a] = 3 * self.p[a] - 1 # shift in local coefficient indices at right boundary (only for non-periodic boundary conditions) - self.his_add_D[a] = np.arange(self.n_his[a] - 2) + 1 - self.his_add_N[a] = np.arange(self.n_his[a] - 1) + 1 + self.his_add_D[a] = xp.arange(self.n_his[a] - 2) + 1 + self.his_add_N[a] = xp.arange(self.n_his[a] - 1) + 1 counter_D = 0 counter_N = 0 # shift local coefficients --> global coefficients (D) - self.his_shift_D[a] = np.arange(self.NbaseD[a]) - (self.p[a] - 1) + self.his_shift_D[a] = xp.arange(self.NbaseD[a]) - (self.p[a] - 1) self.his_shift_D[a][: 2 * self.p[a] - 1] = 0 self.his_shift_D[a][-(2 * self.p[a] - 1) :] = self.his_shift_D[a][-(2 * self.p[a] - 1)] # shift local coefficients --> global coefficients (N) - self.his_shift_N[a] = np.arange(self.NbaseN[a]) - self.p[a] + self.his_shift_N[a] = xp.arange(self.NbaseN[a]) - self.p[a] self.his_shift_N[a][: 2 * self.p[a]] = 0 self.his_shift_N[a][-2 * self.p[a] :] = self.his_shift_N[a][-2 * self.p[a]] - counter_coeffh = np.copy(self.p[a]) + counter_coeffh = xp.copy(self.p[a]) for i in range(n_lambda_his[a]): # left boundary region if i < self.p[a] - 1: - self.his_global_N[a][i] = np.arange(self.n_his_locbf_N[a]) - self.his_global_D[a][i] = np.arange(self.n_his_locbf_D[a]) + self.his_global_N[a][i] = xp.arange(self.n_his_locbf_N[a]) + self.his_global_D[a][i] = xp.arange(self.n_his_locbf_D[a]) - self.x_his_indices[a][i] = np.arange(self.n_his[a]) + self.x_his_indices[a][i] = xp.arange(self.n_his[a]) self.coeffh_indices[a][i] = i for j in range(2 * self.p[a] + 1): xi = self.p[a] - 1 @@ -1016,13 +1016,13 @@ def __init__(self, tensor_space, n_quad): # right boundary region elif i > n_lambda_his[a] - self.p[a]: self.his_global_N[a][i] = ( - np.arange(self.n_his_locbf_N[a]) + n_lambda_his[a] - self.p[a] - (self.p[a] - 1) + xp.arange(self.n_his_locbf_N[a]) + n_lambda_his[a] - self.p[a] - (self.p[a] - 1) ) self.his_global_D[a][i] = ( - np.arange(self.n_his_locbf_D[a]) + n_lambda_his[a] - self.p[a] - (self.p[a] - 1) + xp.arange(self.n_his_locbf_D[a]) + n_lambda_his[a] - self.p[a] - (self.p[a] - 1) ) - self.x_his_indices[a][i] = np.arange(self.n_his[a]) + 2 * ( + self.x_his_indices[a][i] = xp.arange(self.n_his[a]) + 2 * ( n_lambda_his[a] - self.p[a] - (self.p[a] - 1) ) self.coeffh_indices[a][i] = counter_coeffh @@ -1035,10 +1035,10 @@ def __init__(self, tensor_space, n_quad): # interior else: - self.his_global_N[a][i] = np.arange(self.n_his_locbf_N[a]) + i - (self.p[a] - 1) - self.his_global_D[a][i] = np.arange(self.n_his_locbf_D[a]) + i - (self.p[a] - 1) + self.his_global_N[a][i] = xp.arange(self.n_his_locbf_N[a]) + i - (self.p[a] - 1) + self.his_global_D[a][i] = xp.arange(self.n_his_locbf_D[a]) + i - (self.p[a] - 1) - self.x_his_indices[a][i] = np.arange(self.n_his[a]) + 2 * (i - (self.p[a] - 1)) + self.x_his_indices[a][i] = xp.arange(self.n_his[a]) + 2 * (i - (self.p[a] - 1)) self.coeffh_indices[a][i] = self.p[a] - 1 for j in range(2 * self.p[a] + 1): self.x_his[a][i, j] = ( @@ -1051,8 +1051,8 @@ def __init__(self, tensor_space, n_quad): k_glob_new = self.his_global_D[a][i, il] bol = k_glob_new == self.his_global_D[a][i - 1] - if np.any(bol): - self.his_loccof_D[a][i, il] = self.his_loccof_D[a][i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.his_loccof_D[a][i, il] = self.his_loccof_D[a][i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_his[a] - self.p[a] - (self.p[a] - 2)) and ( self.his_loccof_D[a][i, il] == 0 @@ -1064,8 +1064,8 @@ def __init__(self, tensor_space, n_quad): k_glob_new = self.his_global_N[a][i, il] bol = k_glob_new == self.his_global_N[a][i - 1] - if np.any(bol): - self.his_loccof_N[a][i, il] = self.his_loccof_N[a][i - 1, np.where(bol)[0][0]] + 1 + if xp.any(bol): + self.his_loccof_N[a][i, il] = self.his_loccof_N[a][i - 1, xp.where(bol)[0][0]] + 1 if (k_glob_new >= n_lambda_his[a] - self.p[a] - (self.p[a] - 2)) and ( self.his_loccof_N[a][i, il] == 0 @@ -1075,7 +1075,7 @@ def __init__(self, tensor_space, n_quad): # quadrature points and weights self.pts[a], self.wts[a] = bsp.quadrature_grid( - np.unique(self.x_his[a].flatten()), self.pts_loc[a], self.wts_loc[a] + xp.unique(self.x_his[a].flatten()), self.pts_loc[a], self.wts_loc[a] ) else: @@ -1084,18 +1084,18 @@ def __init__(self, tensor_space, n_quad): self.n_his_nvcof_N[a] = 2 * self.p[a] # shift local coefficients --> global coefficients (D) - self.his_shift_D[a] = np.arange(self.NbaseD[a]) - (self.p[a] - 1) + self.his_shift_D[a] = xp.arange(self.NbaseD[a]) - (self.p[a] - 1) # shift local coefficients --> global coefficients (N) - self.his_shift_N[a] = np.arange(self.NbaseD[a]) - self.p[a] + self.his_shift_N[a] = xp.arange(self.NbaseD[a]) - self.p[a] for i in range(n_lambda_his[a]): - self.his_global_N[a][i] = (np.arange(self.n_his_locbf_N[a]) + i - (self.p[a] - 1)) % self.NbaseN[a] - self.his_global_D[a][i] = (np.arange(self.n_his_locbf_D[a]) + i - (self.p[a] - 1)) % self.NbaseD[a] - self.his_loccof_N[a][i] = np.arange(self.n_his_locbf_N[a] - 1, -1, -1) - self.his_loccof_D[a][i] = np.arange(self.n_his_locbf_D[a] - 1, -1, -1) + self.his_global_N[a][i] = (xp.arange(self.n_his_locbf_N[a]) + i - (self.p[a] - 1)) % self.NbaseN[a] + self.his_global_D[a][i] = (xp.arange(self.n_his_locbf_D[a]) + i - (self.p[a] - 1)) % self.NbaseD[a] + self.his_loccof_N[a][i] = xp.arange(self.n_his_locbf_N[a] - 1, -1, -1) + self.his_loccof_D[a][i] = xp.arange(self.n_his_locbf_D[a] - 1, -1, -1) - self.x_his_indices[a][i] = (np.arange(self.n_his[a]) + 2 * (i - (self.p[a] - 1))) % ( + self.x_his_indices[a][i] = (xp.arange(self.n_his[a]) + 2 * (i - (self.p[a] - 1))) % ( 2 * self.Nel[a] ) self.coeffh_indices[a][i] = 0 @@ -1105,7 +1105,7 @@ def __init__(self, tensor_space, n_quad): # quadrature points and weights self.pts[a], self.wts[a] = bsp.quadrature_grid( - np.append(np.unique(self.x_his[a].flatten() % 1.0), 1.0), self.pts_loc[a], self.wts_loc[a] + xp.append(xp.unique(self.x_his[a].flatten() % 1.0), 1.0), self.pts_loc[a], self.wts_loc[a] ) # projector on space V0 (interpolation) @@ -1131,18 +1131,18 @@ def pi_0(self, fun, include_bc=True, eval_kind="meshgrid"): """ # interpolation points - x_int1 = np.unique(self.x_int[0].flatten()) - x_int2 = np.unique(self.x_int[1].flatten()) - x_int3 = np.unique(self.x_int[2].flatten()) + x_int1 = xp.unique(self.x_int[0].flatten()) + x_int2 = xp.unique(self.x_int[1].flatten()) + x_int3 = xp.unique(self.x_int[2].flatten()) # evaluation of function at interpolation points - mat_f = np.empty((x_int1.size, x_int2.size, x_int3.size), dtype=float) + mat_f = xp.empty((x_int1.size, x_int2.size, x_int3.size), dtype=float) # external function call if a callable is passed if callable(fun): # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": - pts1, pts2, pts3 = np.meshgrid(x_int1, x_int2, x_int3, indexing="ij") + pts1, pts2, pts3 = xp.meshgrid(x_int1, x_int2, x_int3, indexing="ij") mat_f[:, :, :] = fun(pts1, pts2, pts3) # tensor-product evaluation is done by input function @@ -1161,7 +1161,7 @@ def pi_0(self, fun, include_bc=True, eval_kind="meshgrid"): print("no internal 3D function implemented!") # coefficients - lambdas = np.zeros((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) + lambdas = xp.zeros((self.NbaseN[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) ker_loc.kernel_pi0_3d( self.NbaseN, @@ -1204,20 +1204,20 @@ def pi_1(self, fun, include_bc=True, eval_kind="meshgrid"): """ # interpolation points - x_int1 = np.unique(self.x_int[0].flatten()) - x_int2 = np.unique(self.x_int[1].flatten()) - x_int3 = np.unique(self.x_int[2].flatten()) + x_int1 = xp.unique(self.x_int[0].flatten()) + x_int2 = xp.unique(self.x_int[1].flatten()) + x_int3 = xp.unique(self.x_int[2].flatten()) # ======== 1-component ======== # evaluation of function at interpolation/quadrature points - mat_f = np.empty((self.pts[0].flatten().size, x_int2.size, x_int3.size), dtype=float) + mat_f = xp.empty((self.pts[0].flatten().size, x_int2.size, x_int3.size), dtype=float) # external function call if a callable is passed if callable(fun[0]): # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": - pts1, pts2, pts3 = np.meshgrid(self.pts[0].flatten(), x_int2, x_int3, indexing="ij") + pts1, pts2, pts3 = xp.meshgrid(self.pts[0].flatten(), x_int2, x_int3, indexing="ij") mat_f[:, :, :] = fun[0](pts1, pts2, pts3) # tensor-product evaluation is done by input function @@ -1236,7 +1236,7 @@ def pi_1(self, fun, include_bc=True, eval_kind="meshgrid"): print("no internal 3D function implemented!") # compute coefficients - lambdas1 = np.zeros((self.NbaseD[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) + lambdas1 = xp.zeros((self.NbaseD[0], self.NbaseN[1], self.NbaseN[2]), dtype=float) ker_loc.kernel_pi11_3d( [self.NbaseD[0], self.NbaseN[1], self.NbaseN[2]], @@ -1259,13 +1259,13 @@ def pi_1(self, fun, include_bc=True, eval_kind="meshgrid"): # ======== 2-component ======== # evaluation of function at interpolation/quadrature points - mat_f = np.empty((x_int1.size, self.pts[1].flatten().size, x_int3.size), dtype=float) + mat_f = xp.empty((x_int1.size, self.pts[1].flatten().size, x_int3.size), dtype=float) # external function call if a callable is passed if callable(fun[1]): # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": - pts1, pts2, pts3 = np.meshgrid(x_int1, self.pts[1].flatten(), x_int3, indexing="ij") + pts1, pts2, pts3 = xp.meshgrid(x_int1, self.pts[1].flatten(), x_int3, indexing="ij") mat_f[:, :, :] = fun[1](pts1, pts2, pts3) # tensor-product evaluation is done by input function @@ -1284,7 +1284,7 @@ def pi_1(self, fun, include_bc=True, eval_kind="meshgrid"): print("no internal 3D function implemented!") # compute coefficients - lambdas2 = np.zeros((self.NbaseN[0], self.NbaseD[1], self.NbaseN[2]), dtype=float) + lambdas2 = xp.zeros((self.NbaseN[0], self.NbaseD[1], self.NbaseN[2]), dtype=float) ker_loc.kernel_pi12_3d( [self.NbaseN[0], self.NbaseD[1], self.NbaseN[2]], @@ -1307,13 +1307,13 @@ def pi_1(self, fun, include_bc=True, eval_kind="meshgrid"): # ======== 3-component ======== # evaluation of function at interpolation/quadrature points - mat_f = np.empty((x_int1.size, x_int1.size, self.pts[2].flatten().size), dtype=float) + mat_f = xp.empty((x_int1.size, x_int1.size, self.pts[2].flatten().size), dtype=float) # external function call if a callable is passed if callable(fun[2]): # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": - pts1, pts2, pts3 = np.meshgrid(x_int1, x_int2, self.pts[2].flatten(), indexing="ij") + pts1, pts2, pts3 = xp.meshgrid(x_int1, x_int2, self.pts[2].flatten(), indexing="ij") mat_f[:, :, :] = fun[2](pts1, pts2, pts3) # tensor-product evaluation is done by input function @@ -1332,7 +1332,7 @@ def pi_1(self, fun, include_bc=True, eval_kind="meshgrid"): print("no internal 3D function implemented!") # compute coefficients - lambdas3 = np.zeros((self.NbaseN[0], self.NbaseN[1], self.NbaseD[2]), dtype=float) + lambdas3 = xp.zeros((self.NbaseN[0], self.NbaseN[1], self.NbaseD[2]), dtype=float) ker_loc.kernel_pi13_3d( [self.NbaseN[0], self.NbaseN[1], self.NbaseD[2]], @@ -1352,7 +1352,7 @@ def pi_1(self, fun, include_bc=True, eval_kind="meshgrid"): lambdas3, ) - return np.concatenate((lambdas1.flatten(), lambdas2.flatten(), lambdas3.flatten())) + return xp.concatenate((lambdas1.flatten(), lambdas2.flatten(), lambdas3.flatten())) # projector on space V1 ([inter, histo, histo], [histo, inter, histo], [histo, histo, inter]) def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): @@ -1377,20 +1377,20 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): """ # interpolation points - x_int1 = np.unique(self.x_int[0].flatten()) - x_int2 = np.unique(self.x_int[1].flatten()) - x_int3 = np.unique(self.x_int[2].flatten()) + x_int1 = xp.unique(self.x_int[0].flatten()) + x_int2 = xp.unique(self.x_int[1].flatten()) + x_int3 = xp.unique(self.x_int[2].flatten()) # ======== 1-component ======== # evaluation of function at interpolation/quadrature points - mat_f = np.empty((x_int1.size, self.pts[1].flatten().size, self.pts[2].flatten().size), dtype=float) + mat_f = xp.empty((x_int1.size, self.pts[1].flatten().size, self.pts[2].flatten().size), dtype=float) # external function call if a callable is passed if callable(fun[0]): # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": - pts1, pts2, pts3 = np.meshgrid(x_int1, self.pts[1].flatten(), self.pts[2].flatten(), indexing="ij") + pts1, pts2, pts3 = xp.meshgrid(x_int1, self.pts[1].flatten(), self.pts[2].flatten(), indexing="ij") mat_f[:, :, :] = fun[0](pts1, pts2, pts3) # tensor-product evaluation is done by input function @@ -1409,7 +1409,7 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): print("no internal 3D function implemented!") # compute coefficients - lambdas1 = np.zeros((self.NbaseN[0], self.NbaseD[1], self.NbaseD[2]), dtype=float) + lambdas1 = xp.zeros((self.NbaseN[0], self.NbaseD[1], self.NbaseD[2]), dtype=float) ker_loc.kernel_pi21_3d( [self.NbaseN[0], self.NbaseD[1], self.NbaseD[2]], @@ -1435,13 +1435,13 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): # ======== 2-component ======== # evaluation of function at interpolation/quadrature points - mat_f = np.empty((self.pts[0].flatten().size, x_int2.size, self.pts[2].flatten().size), dtype=float) + mat_f = xp.empty((self.pts[0].flatten().size, x_int2.size, self.pts[2].flatten().size), dtype=float) # external function call if a callable is passed if callable(fun[1]): # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": - pts1, pts2, pts3 = np.meshgrid(self.pts[0].flatten(), x_int2, self.pts[2].flatten(), indexing="ij") + pts1, pts2, pts3 = xp.meshgrid(self.pts[0].flatten(), x_int2, self.pts[2].flatten(), indexing="ij") mat_f[:, :, :] = fun[1](pts1, pts2, pts3) # tensor-product evaluation is done by input function @@ -1460,7 +1460,7 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): print("no internal 3D function implemented!") # compute coefficients - lambdas2 = np.zeros((self.NbaseD[0], self.NbaseN[1], self.NbaseD[2]), dtype=float) + lambdas2 = xp.zeros((self.NbaseD[0], self.NbaseN[1], self.NbaseD[2]), dtype=float) ker_loc.kernel_pi22_3d( [self.NbaseD[0], self.NbaseN[1], self.NbaseD[2]], @@ -1486,13 +1486,13 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): # ======== 3-component ======== # evaluation of function at interpolation/quadrature points - mat_f = np.empty((self.pts[0].flatten().size, self.pts[1].flatten().size, x_int3.size), dtype=float) + mat_f = xp.empty((self.pts[0].flatten().size, self.pts[1].flatten().size, x_int3.size), dtype=float) # external function call if a callable is passed if callable(fun[2]): # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": - pts1, pts2, pts3 = np.meshgrid(self.pts[0].flatten(), self.pts[1].flatten(), x_int3, indexing="ij") + pts1, pts2, pts3 = xp.meshgrid(self.pts[0].flatten(), self.pts[1].flatten(), x_int3, indexing="ij") mat_f[:, :, :] = fun[2](pts1, pts2, pts3) # tensor-product evaluation is done by input function @@ -1511,7 +1511,7 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): print("no internal 3D function implemented!") # compute coefficients - lambdas3 = np.zeros((self.NbaseD[0], self.NbaseD[1], self.NbaseN[2]), dtype=float) + lambdas3 = xp.zeros((self.NbaseD[0], self.NbaseD[1], self.NbaseN[2]), dtype=float) ker_loc.kernel_pi23_3d( [self.NbaseD[0], self.NbaseD[1], self.NbaseN[2]], @@ -1534,7 +1534,7 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): lambdas3, ) - return np.concatenate((lambdas1.flatten(), lambdas2.flatten(), lambdas3.flatten())) + return xp.concatenate((lambdas1.flatten(), lambdas2.flatten(), lambdas3.flatten())) # projector on space V3 (histopolation) def pi_3(self, fun, include_bc=True, eval_kind="meshgrid"): @@ -1559,7 +1559,7 @@ def pi_3(self, fun, include_bc=True, eval_kind="meshgrid"): """ # evaluation of function at quadrature points - mat_f = np.empty( + mat_f = xp.empty( (self.pts[0].flatten().size, self.pts[1].flatten().size, self.pts[2].flatten().size), dtype=float ) @@ -1567,7 +1567,7 @@ def pi_3(self, fun, include_bc=True, eval_kind="meshgrid"): if callable(fun): # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": - pts1, pts2, pts3 = np.meshgrid( + pts1, pts2, pts3 = xp.meshgrid( self.pts[0].flatten(), self.pts[1].flatten(), self.pts[2].flatten(), indexing="ij" ) mat_f[:, :, :] = fun(pts1, pts2, pts3) @@ -1590,7 +1590,7 @@ def pi_3(self, fun, include_bc=True, eval_kind="meshgrid"): print("no internal 3D function implemented!") # compute coefficients - lambdas = np.zeros((self.NbaseD[0], self.NbaseD[1], self.NbaseD[2]), dtype=float) + lambdas = xp.zeros((self.NbaseD[0], self.NbaseD[1], self.NbaseD[2]), dtype=float) ker_loc.kernel_pi3_3d( self.NbaseD, diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py index e85dfaeb5..04705a6a5 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py @@ -10,7 +10,7 @@ import struphy.feec.bsplines as bsp import struphy.feec.projectors.shape_pro_local.shape_L2_projector_kernel as ker_loc -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ======================= 3d ==================================== @@ -50,48 +50,48 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): self.indD = tensor_space.indD self.polar = False # local projectors for polar splines are not implemented yet - self.lambdas_0 = np.zeros((NbaseN[0], NbaseN[1], NbaseN[2]), dtype=float) - self.potential_lambdas_0 = np.zeros((NbaseN[0], NbaseN[1], NbaseN[2]), dtype=float) + self.lambdas_0 = xp.zeros((NbaseN[0], NbaseN[1], NbaseN[2]), dtype=float) + self.potential_lambdas_0 = xp.zeros((NbaseN[0], NbaseN[1], NbaseN[2]), dtype=float) - self.lambdas_1_11 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.lambdas_1_12 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_1_13 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_1_11 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.lambdas_1_12 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_1_13 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_1_21 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.lambdas_1_22 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_1_23 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_1_21 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.lambdas_1_22 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_1_23 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_1_31 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.lambdas_1_32 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_1_33 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_1_31 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.lambdas_1_32 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_1_33 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_2_11 = np.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) - self.lambdas_2_12 = np.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_2_13 = np.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_2_11 = xp.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) + self.lambdas_2_12 = xp.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_2_13 = xp.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_2_21 = np.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) - self.lambdas_2_22 = np.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_2_23 = np.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_2_21 = xp.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) + self.lambdas_2_22 = xp.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_2_23 = xp.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_2_31 = np.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) - self.lambdas_2_32 = np.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_2_33 = np.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_2_31 = xp.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) + self.lambdas_2_32 = xp.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_2_33 = xp.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_3 = np.zeros((NbaseD[0], NbaseD[1], NbaseD[2]), dtype=float) + self.lambdas_3 = xp.zeros((NbaseD[0], NbaseD[1], NbaseD[2]), dtype=float) self.p_size = p_size self.p_shape = p_shape - self.related = np.zeros(3, dtype=int) + self.related = xp.zeros(3, dtype=int) for a in range(3): - # self.related[a] = int(np.floor(NbaseN[a]/2.0)) + # self.related[a] = int(xp.floor(NbaseN[a]/2.0)) self.related[a] = int( - np.floor((3 * int((self.p_size[a] * (self.p_shape[a] + 1)) * self.Nel[a] + 1) + 3 * self.p[a]) / 2.0) + xp.floor((3 * int((self.p_size[a] * (self.p_shape[a] + 1)) * self.Nel[a] + 1) + 3 * self.p[a]) / 2.0) ) if (2 * self.related[a] + 1) > NbaseN[a]: - self.related[a] = int(np.floor(NbaseN[a] / 2.0)) + self.related[a] = int(xp.floor(NbaseN[a] / 2.0)) - self.kernel_0_loc = np.zeros( + self.kernel_0_loc = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -103,7 +103,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): dtype=float, ) - self.kernel_1_11_loc = np.zeros( + self.kernel_1_11_loc = xp.zeros( ( NbaseD[0], NbaseN[1], @@ -114,7 +114,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): ), dtype=float, ) - self.kernel_1_12_loc = np.zeros( + self.kernel_1_12_loc = xp.zeros( ( NbaseD[0], NbaseN[1], @@ -125,7 +125,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): ), dtype=float, ) - self.kernel_1_13_loc = np.zeros( + self.kernel_1_13_loc = xp.zeros( ( NbaseD[0], NbaseN[1], @@ -137,7 +137,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): dtype=float, ) - self.kernel_1_22_loc = np.zeros( + self.kernel_1_22_loc = xp.zeros( ( NbaseN[0], NbaseD[1], @@ -148,7 +148,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): ), dtype=float, ) - self.kernel_1_23_loc = np.zeros( + self.kernel_1_23_loc = xp.zeros( ( NbaseN[0], NbaseD[1], @@ -160,7 +160,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): dtype=float, ) - self.kernel_1_33_loc = np.zeros( + self.kernel_1_33_loc = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -172,12 +172,12 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): dtype=float, ) - self.right_loc_1 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.right_loc_2 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.right_loc_3 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.right_loc_1 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.right_loc_2 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.right_loc_3 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) if self.mpi_rank == 0: - self.kernel_0 = np.zeros( + self.kernel_0 = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -189,7 +189,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): dtype=float, ) - self.kernel_1_11 = np.zeros( + self.kernel_1_11 = xp.zeros( ( NbaseD[0], NbaseN[1], @@ -200,7 +200,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): ), dtype=float, ) - self.kernel_1_12 = np.zeros( + self.kernel_1_12 = xp.zeros( ( NbaseN[0], NbaseD[1], @@ -211,7 +211,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): ), dtype=float, ) - self.kernel_1_13 = np.zeros( + self.kernel_1_13 = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -223,7 +223,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): dtype=float, ) - self.kernel_1_22 = np.zeros( + self.kernel_1_22 = xp.zeros( ( NbaseN[0], NbaseD[1], @@ -234,7 +234,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): ), dtype=float, ) - self.kernel_1_23 = np.zeros( + self.kernel_1_23 = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -246,7 +246,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): dtype=float, ) - self.kernel_1_33 = np.zeros( + self.kernel_1_33 = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -258,9 +258,9 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): dtype=float, ) - self.right_1 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.right_2 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.right_3 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.right_1 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.right_2 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.right_3 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) else: self.kernel_0 = None @@ -301,11 +301,11 @@ def assemble_0_form(self, tensor_space_FEM, mpi_comm): Nj = tensor_space_FEM.Nbase_0form # conversion to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -358,11 +358,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -384,11 +384,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -410,11 +410,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -436,11 +436,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -462,11 +462,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -488,11 +488,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -510,7 +510,7 @@ def assemble_1_form(self, tensor_space_FEM): # final block matrix M = spa.bmat([[M11, M12, M13], [M12.T, M22, M23], [M13.T, M23.T, M33]], format="csr") # print('insider_check', self.kernel_1_33) - return (M, np.concatenate((self.right_1.flatten(), self.right_2.flatten(), self.right_3.flatten()))) + return (M, xp.concatenate((self.right_1.flatten(), self.right_2.flatten(), self.right_3.flatten()))) def heavy_test(self, test1, test2, test3, acc, particles_loc, Np, domain): ker_loc.kernel_1_heavy( diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py index 5951e835a..dd5a313f2 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py @@ -10,7 +10,7 @@ import struphy.feec.bsplines as bsp import struphy.feec.projectors.shape_pro_local.shape_local_projector_kernel as ker_loc -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ======================= 3d ==================================== @@ -51,48 +51,48 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co self.polar = False # local projectors for polar splines are not implemented yet - self.lambdas_0 = np.zeros((NbaseN[0], NbaseN[1], NbaseN[2]), dtype=float) - self.potential_lambdas_0 = np.zeros((NbaseN[0], NbaseN[1], NbaseN[2]), dtype=float) + self.lambdas_0 = xp.zeros((NbaseN[0], NbaseN[1], NbaseN[2]), dtype=float) + self.potential_lambdas_0 = xp.zeros((NbaseN[0], NbaseN[1], NbaseN[2]), dtype=float) - self.lambdas_1_11 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.lambdas_1_12 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_1_13 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_1_11 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.lambdas_1_12 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_1_13 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_1_21 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.lambdas_1_22 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_1_23 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_1_21 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.lambdas_1_22 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_1_23 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_1_31 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.lambdas_1_32 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_1_33 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_1_31 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.lambdas_1_32 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_1_33 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_2_11 = np.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) - self.lambdas_2_12 = np.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_2_13 = np.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_2_11 = xp.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) + self.lambdas_2_12 = xp.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_2_13 = xp.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_2_21 = np.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) - self.lambdas_2_22 = np.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_2_23 = np.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_2_21 = xp.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) + self.lambdas_2_22 = xp.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_2_23 = xp.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_2_31 = np.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) - self.lambdas_2_32 = np.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) - self.lambdas_2_33 = np.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) + self.lambdas_2_31 = xp.zeros((NbaseN[0], NbaseD[1], NbaseD[2]), dtype=float) + self.lambdas_2_32 = xp.zeros((NbaseD[0], NbaseN[1], NbaseD[2]), dtype=float) + self.lambdas_2_33 = xp.zeros((NbaseD[0], NbaseD[1], NbaseN[2]), dtype=float) - self.lambdas_3 = np.zeros((NbaseD[0], NbaseD[1], NbaseD[2]), dtype=float) + self.lambdas_3 = xp.zeros((NbaseD[0], NbaseD[1], NbaseD[2]), dtype=float) self.p_size = p_size self.p_shape = p_shape - self.related = np.zeros(3, dtype=int) + self.related = xp.zeros(3, dtype=int) for a in range(3): - # self.related[a] = int(np.floor(NbaseN[a]/2.0)) + # self.related[a] = int(xp.floor(NbaseN[a]/2.0)) self.related[a] = int( - np.floor((3 * int((self.p_size[a] * (self.p_shape[a] + 1)) * self.Nel[a] + 1) + 3 * self.p[a]) / 2.0) + xp.floor((3 * int((self.p_size[a] * (self.p_shape[a] + 1)) * self.Nel[a] + 1) + 3 * self.p[a]) / 2.0) ) if (2 * self.related[a] + 1) > NbaseN[a]: - self.related[a] = int(np.floor(NbaseN[a] / 2.0)) + self.related[a] = int(xp.floor(NbaseN[a] / 2.0)) - self.kernel_0_loc = np.zeros( + self.kernel_0_loc = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -104,7 +104,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co dtype=float, ) - self.kernel_1_11_loc = np.zeros( + self.kernel_1_11_loc = xp.zeros( ( NbaseD[0], NbaseN[1], @@ -115,7 +115,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co ), dtype=float, ) - self.kernel_1_12_loc = np.zeros( + self.kernel_1_12_loc = xp.zeros( ( NbaseD[0], NbaseN[1], @@ -126,7 +126,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co ), dtype=float, ) - self.kernel_1_13_loc = np.zeros( + self.kernel_1_13_loc = xp.zeros( ( NbaseD[0], NbaseN[1], @@ -138,7 +138,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co dtype=float, ) - self.kernel_1_22_loc = np.zeros( + self.kernel_1_22_loc = xp.zeros( ( NbaseN[0], NbaseD[1], @@ -149,7 +149,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co ), dtype=float, ) - self.kernel_1_23_loc = np.zeros( + self.kernel_1_23_loc = xp.zeros( ( NbaseN[0], NbaseD[1], @@ -161,7 +161,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co dtype=float, ) - self.kernel_1_33_loc = np.zeros( + self.kernel_1_33_loc = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -173,12 +173,12 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co dtype=float, ) - self.right_loc_1 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.right_loc_2 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.right_loc_3 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.right_loc_1 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.right_loc_2 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.right_loc_3 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) if self.mpi_rank == 0: - self.kernel_0 = np.zeros( + self.kernel_0 = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -190,7 +190,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co dtype=float, ) - self.kernel_1_11 = np.zeros( + self.kernel_1_11 = xp.zeros( ( NbaseD[0], NbaseN[1], @@ -201,7 +201,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co ), dtype=float, ) - self.kernel_1_12 = np.zeros( + self.kernel_1_12 = xp.zeros( ( NbaseN[0], NbaseD[1], @@ -212,7 +212,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co ), dtype=float, ) - self.kernel_1_13 = np.zeros( + self.kernel_1_13 = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -224,7 +224,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co dtype=float, ) - self.kernel_1_22 = np.zeros( + self.kernel_1_22 = xp.zeros( ( NbaseN[0], NbaseD[1], @@ -235,7 +235,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co ), dtype=float, ) - self.kernel_1_23 = np.zeros( + self.kernel_1_23 = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -247,7 +247,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co dtype=float, ) - self.kernel_1_33 = np.zeros( + self.kernel_1_33 = xp.zeros( ( NbaseN[0], NbaseN[1], @@ -259,9 +259,9 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co dtype=float, ) - self.right_1 = np.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) - self.right_2 = np.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) - self.right_3 = np.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) + self.right_1 = xp.zeros((NbaseD[0], NbaseN[1], NbaseN[2]), dtype=float) + self.right_2 = xp.zeros((NbaseN[0], NbaseD[1], NbaseN[2]), dtype=float) + self.right_3 = xp.zeros((NbaseN[0], NbaseN[1], NbaseD[2]), dtype=float) else: self.kernel_0 = None @@ -279,7 +279,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co self.right_2 = None self.right_3 = None - self.num_cell = np.empty(3, dtype=int) + self.num_cell = xp.empty(3, dtype=int) for i in range(3): if self.p[i] == 1: self.num_cell[i] = 1 @@ -287,8 +287,8 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co self.num_cell[i] = 2 # Gauss - Legendre quadrature points and weights in (-1, 1) - self.pts_loc = [np.polynomial.legendre.leggauss(n_quad)[0] for n_quad in self.n_quad] - self.wts_loc = [np.polynomial.legendre.leggauss(n_quad)[1] for n_quad in self.n_quad] + self.pts_loc = [xp.polynomial.legendre.leggauss(n_quad)[0] for n_quad in self.n_quad] + self.wts_loc = [xp.polynomial.legendre.leggauss(n_quad)[1] for n_quad in self.n_quad] self.pts = [0, 0, 0] self.wts = [0, 0, 0] @@ -303,78 +303,78 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co self.coeff_h = [0, 0, 0] for a in range(3): if self.bc[a] == True: - self.coeff_i[a] = np.zeros(2 * self.p[a], dtype=float) - self.coeff_h[a] = np.zeros(2 * self.p[a], dtype=float) + self.coeff_i[a] = xp.zeros(2 * self.p[a], dtype=float) + self.coeff_h[a] = xp.zeros(2 * self.p[a], dtype=float) if self.p[a] == 1: - self.coeff_i[a][:] = np.array([1.0, 0.0]) - self.coeff_h[a][:] = np.array([1.0, 1.0]) + self.coeff_i[a][:] = xp.array([1.0, 0.0]) + self.coeff_h[a][:] = xp.array([1.0, 1.0]) elif self.p[a] == 2: - self.coeff_i[a][:] = 1 / 2 * np.array([-1.0, 4.0, -1.0, 0.0]) - self.coeff_h[a][:] = 1 / 2 * np.array([-1.0, 3.0, 3.0, -1.0]) + self.coeff_i[a][:] = 1 / 2 * xp.array([-1.0, 4.0, -1.0, 0.0]) + self.coeff_h[a][:] = 1 / 2 * xp.array([-1.0, 3.0, 3.0, -1.0]) elif self.p[a] == 3: - self.coeff_i[a][:] = 1 / 6 * np.array([1.0, -8.0, 20.0, -8.0, 1.0, 0.0]) - self.coeff_h[a][:] = 1 / 6 * np.array([1.0, -7.0, 12.0, 12.0, -7.0, 1.0]) + self.coeff_i[a][:] = 1 / 6 * xp.array([1.0, -8.0, 20.0, -8.0, 1.0, 0.0]) + self.coeff_h[a][:] = 1 / 6 * xp.array([1.0, -7.0, 12.0, 12.0, -7.0, 1.0]) elif self.p[a] == 4: - self.coeff_i[a][:] = 2 / 45 * np.array([-1.0, 16.0, -295 / 4, 140.0, -295 / 4, 16.0, -1.0, 0.0]) + self.coeff_i[a][:] = 2 / 45 * xp.array([-1.0, 16.0, -295 / 4, 140.0, -295 / 4, 16.0, -1.0, 0.0]) self.coeff_h[a][:] = ( - 2 / 45 * np.array([-1.0, 15.0, -231 / 4, 265 / 4, 265 / 4, -231 / 4, 15.0, -1.0]) + 2 / 45 * xp.array([-1.0, 15.0, -231 / 4, 265 / 4, 265 / 4, -231 / 4, 15.0, -1.0]) ) else: print("degree > 4 not implemented!") else: - self.coeff_i[a] = np.zeros((2 * self.p[a] - 1, 2 * self.p[a] - 1), dtype=float) - self.coeff_h[a] = np.zeros((2 * self.p[a] - 1, 2 * self.p[a]), dtype=float) + self.coeff_i[a] = xp.zeros((2 * self.p[a] - 1, 2 * self.p[a] - 1), dtype=float) + self.coeff_h[a] = xp.zeros((2 * self.p[a] - 1, 2 * self.p[a]), dtype=float) if self.p[a] == 1: - self.coeff_i[a][0, :] = np.array([1.0]) - self.coeff_h[a][0, :] = np.array([1.0, 1.0]) + self.coeff_i[a][0, :] = xp.array([1.0]) + self.coeff_h[a][0, :] = xp.array([1.0, 1.0]) elif self.p[a] == 2: - self.coeff_i[a][0, :] = 1 / 2 * np.array([2.0, 0.0, 0.0]) - self.coeff_i[a][1, :] = 1 / 2 * np.array([-1.0, 4.0, -1.0]) - self.coeff_i[a][2, :] = 1 / 2 * np.array([0.0, 0.0, 2.0]) + self.coeff_i[a][0, :] = 1 / 2 * xp.array([2.0, 0.0, 0.0]) + self.coeff_i[a][1, :] = 1 / 2 * xp.array([-1.0, 4.0, -1.0]) + self.coeff_i[a][2, :] = 1 / 2 * xp.array([0.0, 0.0, 2.0]) - self.coeff_h[a][0, :] = 1 / 2 * np.array([3.0, -1.0, 0.0, 0.0]) - self.coeff_h[a][1, :] = 1 / 2 * np.array([-1.0, 3.0, 3.0, -1.0]) - self.coeff_h[a][2, :] = 1 / 2 * np.array([0.0, 0.0, -1.0, 3.0]) + self.coeff_h[a][0, :] = 1 / 2 * xp.array([3.0, -1.0, 0.0, 0.0]) + self.coeff_h[a][1, :] = 1 / 2 * xp.array([-1.0, 3.0, 3.0, -1.0]) + self.coeff_h[a][2, :] = 1 / 2 * xp.array([0.0, 0.0, -1.0, 3.0]) elif self.p[a] == 3: - self.coeff_i[a][0, :] = 1 / 18 * np.array([18.0, 0.0, 0.0, 0.0, 0.0]) - self.coeff_i[a][1, :] = 1 / 18 * np.array([-5.0, 40.0, -24.0, 8.0, -1.0]) - self.coeff_i[a][2, :] = 1 / 18 * np.array([3.0, -24.0, 60.0, -24.0, 3.0]) - self.coeff_i[a][3, :] = 1 / 18 * np.array([-1.0, 8.0, -24.0, 40.0, -5.0]) - self.coeff_i[a][4, :] = 1 / 18 * np.array([0.0, 0.0, 0.0, 0.0, 18.0]) - - self.coeff_h[a][0, :] = 1 / 18 * np.array([23.0, -17.0, 7.0, -1.0, 0.0, 0.0]) - self.coeff_h[a][1, :] = 1 / 18 * np.array([-8.0, 56.0, -28.0, 4.0, 0.0, 0.0]) - self.coeff_h[a][2, :] = 1 / 18 * np.array([3.0, -21.0, 36.0, 36.0, -21.0, 3.0]) - self.coeff_h[a][3, :] = 1 / 18 * np.array([0.0, 0.0, 4.0, -28.0, 56.0, -8.0]) - self.coeff_h[a][4, :] = 1 / 18 * np.array([0.0, 0.0, -1.0, 7.0, -17.0, 23.0]) + self.coeff_i[a][0, :] = 1 / 18 * xp.array([18.0, 0.0, 0.0, 0.0, 0.0]) + self.coeff_i[a][1, :] = 1 / 18 * xp.array([-5.0, 40.0, -24.0, 8.0, -1.0]) + self.coeff_i[a][2, :] = 1 / 18 * xp.array([3.0, -24.0, 60.0, -24.0, 3.0]) + self.coeff_i[a][3, :] = 1 / 18 * xp.array([-1.0, 8.0, -24.0, 40.0, -5.0]) + self.coeff_i[a][4, :] = 1 / 18 * xp.array([0.0, 0.0, 0.0, 0.0, 18.0]) + + self.coeff_h[a][0, :] = 1 / 18 * xp.array([23.0, -17.0, 7.0, -1.0, 0.0, 0.0]) + self.coeff_h[a][1, :] = 1 / 18 * xp.array([-8.0, 56.0, -28.0, 4.0, 0.0, 0.0]) + self.coeff_h[a][2, :] = 1 / 18 * xp.array([3.0, -21.0, 36.0, 36.0, -21.0, 3.0]) + self.coeff_h[a][3, :] = 1 / 18 * xp.array([0.0, 0.0, 4.0, -28.0, 56.0, -8.0]) + self.coeff_h[a][4, :] = 1 / 18 * xp.array([0.0, 0.0, -1.0, 7.0, -17.0, 23.0]) elif self.p[a] == 4: - self.coeff_i[a][0, :] = 1 / 360 * np.array([360.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) - self.coeff_i[a][1, :] = 1 / 360 * np.array([-59.0, 944.0, -1000.0, 720.0, -305.0, 64.0, -4.0]) - self.coeff_i[a][2, :] = 1 / 360 * np.array([23.0, -368.0, 1580.0, -1360.0, 605.0, -128.0, 8.0]) - self.coeff_i[a][3, :] = 1 / 360 * np.array([-16.0, 256.0, -1180.0, 2240.0, -1180.0, 256.0, -16.0]) - self.coeff_i[a][4, :] = 1 / 360 * np.array([8.0, -128.0, 605.0, -1360.0, 1580.0, -368.0, 23.0]) - self.coeff_i[a][5, :] = 1 / 360 * np.array([-4.0, 64.0, -305.0, 720.0, -1000.0, 944.0, -59.0]) - self.coeff_i[a][6, :] = 1 / 360 * np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 360.0]) - - self.coeff_h[a][0, :] = 1 / 360 * np.array([419.0, -525.0, 475.0, -245.0, 60.0, -4.0, 0.0, 0.0]) - self.coeff_h[a][1, :] = 1 / 360 * np.array([-82.0, 1230.0, -1350.0, 730.0, -180.0, 12.0, 0.0, 0.0]) - self.coeff_h[a][2, :] = 1 / 360 * np.array([39.0, -585.0, 2175.0, -1425.0, 360.0, -24.0, 0.0, 0.0]) + self.coeff_i[a][0, :] = 1 / 360 * xp.array([360.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + self.coeff_i[a][1, :] = 1 / 360 * xp.array([-59.0, 944.0, -1000.0, 720.0, -305.0, 64.0, -4.0]) + self.coeff_i[a][2, :] = 1 / 360 * xp.array([23.0, -368.0, 1580.0, -1360.0, 605.0, -128.0, 8.0]) + self.coeff_i[a][3, :] = 1 / 360 * xp.array([-16.0, 256.0, -1180.0, 2240.0, -1180.0, 256.0, -16.0]) + self.coeff_i[a][4, :] = 1 / 360 * xp.array([8.0, -128.0, 605.0, -1360.0, 1580.0, -368.0, 23.0]) + self.coeff_i[a][5, :] = 1 / 360 * xp.array([-4.0, 64.0, -305.0, 720.0, -1000.0, 944.0, -59.0]) + self.coeff_i[a][6, :] = 1 / 360 * xp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 360.0]) + + self.coeff_h[a][0, :] = 1 / 360 * xp.array([419.0, -525.0, 475.0, -245.0, 60.0, -4.0, 0.0, 0.0]) + self.coeff_h[a][1, :] = 1 / 360 * xp.array([-82.0, 1230.0, -1350.0, 730.0, -180.0, 12.0, 0.0, 0.0]) + self.coeff_h[a][2, :] = 1 / 360 * xp.array([39.0, -585.0, 2175.0, -1425.0, 360.0, -24.0, 0.0, 0.0]) self.coeff_h[a][3, :] = ( - 1 / 360 * np.array([-16.0, 240.0, -924.0, 1060.0, 1060.0, -924.0, 240.0, -16.0]) + 1 / 360 * xp.array([-16.0, 240.0, -924.0, 1060.0, 1060.0, -924.0, 240.0, -16.0]) ) - self.coeff_h[a][4, :] = 1 / 360 * np.array([0.0, 0.0, -24.0, 360.0, -1425.0, 2175.0, -585.0, 39.0]) - self.coeff_h[a][5, :] = 1 / 360 * np.array([0.0, 0.0, 12.0, -180.0, 730.0, -1350.0, 1230.0, -82.0]) - self.coeff_h[a][6, :] = 1 / 360 * np.array([0.0, 0.0, -4.0, 60.0, -245.0, 475.0, -525.0, 419.0]) + self.coeff_h[a][4, :] = 1 / 360 * xp.array([0.0, 0.0, -24.0, 360.0, -1425.0, 2175.0, -585.0, 39.0]) + self.coeff_h[a][5, :] = 1 / 360 * xp.array([0.0, 0.0, 12.0, -180.0, 730.0, -1350.0, 1230.0, -82.0]) + self.coeff_h[a][6, :] = 1 / 360 * xp.array([0.0, 0.0, -4.0, 60.0, -245.0, 475.0, -525.0, 419.0]) else: print("degree > 4 not implemented!") @@ -402,11 +402,11 @@ def assemble_0_form(self, tensor_space_FEM, mpi_comm): Nj = tensor_space_FEM.Nbase_0form # conversion to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -459,11 +459,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -485,11 +485,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -511,11 +511,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -537,11 +537,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -563,11 +563,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -589,11 +589,11 @@ def assemble_1_form(self, tensor_space_FEM): Nj = tensor_space_FEM.Nbase_1form[b] # convert to sparse matrix - indices = np.indices( + indices = xp.indices( (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) ) - shift = [np.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] + shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -611,7 +611,7 @@ def assemble_1_form(self, tensor_space_FEM): # final block matrix M = spa.bmat([[M11, M12, M13], [M12.T, M22, M23], [M13.T, M23.T, M33]], format="csr") # print('insider_check', self.kernel_1_33) - return (M, np.concatenate((self.right_1.flatten(), self.right_2.flatten(), self.right_3.flatten()))) + return (M, xp.concatenate((self.right_1.flatten(), self.right_2.flatten(), self.right_3.flatten()))) def heavy_test(self, test1, test2, test3, acc, particles_loc, Np, domain): ker_loc.kernel_1_heavy( diff --git a/src/struphy/eigenvalue_solvers/mass_matrices_1d.py b/src/struphy/eigenvalue_solvers/mass_matrices_1d.py index 99316215d..9c41ca85d 100644 --- a/src/struphy/eigenvalue_solvers/mass_matrices_1d.py +++ b/src/struphy/eigenvalue_solvers/mass_matrices_1d.py @@ -5,7 +5,7 @@ import scipy.sparse as spa import struphy.bsplines.bsplines as bsp -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ======= mass matrices in 1D ==================== @@ -47,7 +47,7 @@ def get_M(spline_space, phi_i=0, phi_j=0, fun=None): # evaluation of weight function at quadrature points (optional) if fun == None: - mat_fun = np.ones(pts.shape, dtype=float) + mat_fun = xp.ones(pts.shape, dtype=float) else: mat_fun = fun(pts.flatten()).reshape(Nel, n_quad) @@ -74,7 +74,7 @@ def get_M(spline_space, phi_i=0, phi_j=0, fun=None): bj = basisD[:, :, 0, :] # matrix assembly - M = np.zeros((Ni, 2 * p + 1), dtype=float) + M = xp.zeros((Ni, 2 * p + 1), dtype=float) for ie in range(Nel): for il in range(p + 1 - ni): @@ -86,8 +86,8 @@ def get_M(spline_space, phi_i=0, phi_j=0, fun=None): M[(ie + il) % Ni, p + jl - il] += value - indices = np.indices((Ni, 2 * p + 1)) - shift = np.arange(Ni) - p + indices = xp.indices((Ni, 2 * p + 1)) + shift = xp.arange(Ni) - p row = indices[0].flatten() col = (indices[1] + shift[:, None]) % Nj @@ -137,13 +137,13 @@ def get_M_gen(spline_space, phi_i=0, phi_j=0, fun=None, jac=None): # evaluation of weight function at quadrature points (optional) if fun == None: - mat_fun = np.ones(pts.shape, dtype=float) + mat_fun = xp.ones(pts.shape, dtype=float) else: mat_fun = fun(pts.flatten()).reshape(Nel, n_quad) # evaluation of jacobian at quadrature points if jac == None: - mat_jac = np.ones(pts.shape, dtype=float) + mat_jac = xp.ones(pts.shape, dtype=float) else: mat_jac = jac(pts.flatten()).reshape(Nel, n_quad) @@ -180,7 +180,7 @@ def get_M_gen(spline_space, phi_i=0, phi_j=0, fun=None, jac=None): bj = basis_t[:, :, 0, :] # matrix assembly - M = np.zeros((Ni, 2 * p + 1), dtype=float) + M = xp.zeros((Ni, 2 * p + 1), dtype=float) for ie in range(Nel): for il in range(p + 1 - ni): @@ -192,8 +192,8 @@ def get_M_gen(spline_space, phi_i=0, phi_j=0, fun=None, jac=None): M[(ie + il) % Ni, p + jl - il] += value - indices = np.indices((Ni, 2 * p + 1)) - shift = np.arange(Ni) - p + indices = xp.indices((Ni, 2 * p + 1)) + shift = xp.arange(Ni) - p row = indices[0].flatten() col = (indices[1] + shift[:, None]) % Nj @@ -235,11 +235,11 @@ def test_M(spline_space, phi_i=0, phi_j=0, fun=lambda eta: 1.0, jac=lambda eta: bj = lambda eta: spline_space.evaluate_D(eta, cj) / spline_space.Nel # coefficients - ci = np.zeros(Ni, dtype=float) - cj = np.zeros(Nj, dtype=float) + ci = xp.zeros(Ni, dtype=float) + cj = xp.zeros(Nj, dtype=float) # integration - M = np.zeros((Ni, Nj), dtype=float) + M = xp.zeros((Ni, Nj), dtype=float) for i in range(Ni): for j in range(Nj): diff --git a/src/struphy/eigenvalue_solvers/mass_matrices_2d.py b/src/struphy/eigenvalue_solvers/mass_matrices_2d.py index e19b31ac6..e15d5e4c7 100644 --- a/src/struphy/eigenvalue_solvers/mass_matrices_2d.py +++ b/src/struphy/eigenvalue_solvers/mass_matrices_2d.py @@ -5,7 +5,7 @@ import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_2d as ker -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ================ mass matrix in V0 =========================== @@ -46,7 +46,7 @@ def get_M0(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): # evaluation of weight function at quadrature points if weight == None: - mat_w = np.ones(det_df.shape, dtype=float) + mat_w = xp.ones(det_df.shape, dtype=float) else: mat_w = weight(pts[0].flatten(), pts[1].flatten(), 0.0) mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1]) @@ -55,14 +55,14 @@ def get_M0(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_0form Nj = tensor_space_FEM.Nbase_0form - M = np.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) + M = xp.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array([0, 0]), - np.array([0, 0]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array([0, 0]), + xp.array([0, 0]), wts[0], wts[1], basisN[0], @@ -76,9 +76,9 @@ def get_M0(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # conversion to sparse matrix - indices = np.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) + indices = xp.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * indices[0] + indices[1]).flatten() @@ -156,7 +156,7 @@ def get_M1(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_1form[a] Nj = tensor_space_FEM.Nbase_1form[b] - M[a][b] = np.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) + M[a][b] = xp.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) # evaluate inverse metric tensor at quadrature points if weight == None: @@ -167,13 +167,13 @@ def get_M1(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1]) # assemble block if weight is not zero - if np.any(mat_w): + if xp.any(mat_w): ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array(ns[a]), - np.array(ns[b]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array(ns[a]), + xp.array(ns[b]), wts[0], wts[1], basis[a][0], @@ -187,9 +187,9 @@ def get_M1(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # convert to sparse matrix - indices = np.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) + indices = xp.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * indices[0] + indices[1]).flatten() @@ -272,7 +272,7 @@ def get_M2(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_2form[a] Nj = tensor_space_FEM.Nbase_2form[b] - M[a][b] = np.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) + M[a][b] = xp.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) # evaluate metric tensor at quadrature points if weight == None: @@ -283,13 +283,13 @@ def get_M2(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1]) # assemble block if weight is not zero - if np.any(mat_w): + if xp.any(mat_w): ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array(ns[a]), - np.array(ns[b]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array(ns[a]), + xp.array(ns[b]), wts[0], wts[1], basis[a][0], @@ -303,9 +303,9 @@ def get_M2(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # convert to sparse matrix - indices = np.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) + indices = xp.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * indices[0] + indices[1]).flatten() @@ -369,7 +369,7 @@ def get_M3(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): # evaluation of weight function at quadrature points if weight == None: - mat_w = np.ones(det_df.shape, dtype=float) + mat_w = xp.ones(det_df.shape, dtype=float) else: mat_w = weight(pts[0].flatten(), pts[1].flatten(), 0.0) mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1]) @@ -378,14 +378,14 @@ def get_M3(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_3form Nj = tensor_space_FEM.Nbase_3form - M = np.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) + M = xp.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array([1, 1]), - np.array([1, 1]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array([1, 1]), + xp.array([1, 1]), wts[0], wts[1], basisD[0], @@ -399,9 +399,9 @@ def get_M3(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # conversion to sparse matrix - indices = np.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) + indices = xp.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * indices[0] + indices[1]).flatten() @@ -475,7 +475,7 @@ def get_Mv(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_0form Nj = tensor_space_FEM.Nbase_0form - M[a][b] = np.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) + M[a][b] = xp.zeros((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1), dtype=float) # evaluate metric tensor at quadrature points if weight == None: @@ -486,13 +486,13 @@ def get_Mv(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1]) # assemble block if weight is not zero - if np.any(mat_w): + if xp.any(mat_w): ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array(ns[a]), - np.array(ns[b]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array(ns[a]), + xp.array(ns[b]), wts[0], wts[1], basis[a][0], @@ -506,9 +506,9 @@ def get_Mv(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # convert to sparse matrix - indices = np.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) + indices = xp.indices((Ni[0], Ni[1], 2 * p[0] + 1, 2 * p[1] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * indices[0] + indices[1]).flatten() diff --git a/src/struphy/eigenvalue_solvers/mass_matrices_3d.py b/src/struphy/eigenvalue_solvers/mass_matrices_3d.py index ef6ee1e0c..f7846f0b2 100644 --- a/src/struphy/eigenvalue_solvers/mass_matrices_3d.py +++ b/src/struphy/eigenvalue_solvers/mass_matrices_3d.py @@ -5,7 +5,7 @@ import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_3d as ker -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ================ mass matrix in V0 =========================== @@ -46,7 +46,7 @@ def get_M0(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): # evaluation of weight function at quadrature points if weight == None: - mat_w = np.ones(det_df.shape, dtype=float) + mat_w = xp.ones(det_df.shape, dtype=float) else: mat_w = weight(pts[0].flatten(), pts[1].flatten(), pts[2].flatten()) mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]) @@ -55,14 +55,14 @@ def get_M0(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_0form Nj = tensor_space_FEM.Nbase_0form - M = np.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) + M = xp.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array([0, 0, 0]), - np.array([0, 0, 0]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array([0, 0, 0]), + xp.array([0, 0, 0]), wts[0], wts[1], wts[2], @@ -80,9 +80,9 @@ def get_M0(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # conversion to sparse matrix - indices = np.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) + indices = xp.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -161,7 +161,7 @@ def get_M1(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_1form[a] Nj = tensor_space_FEM.Nbase_1form[b] - M[a][b] = np.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) + M[a][b] = xp.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) # evaluate metric tensor at quadrature points if weight == None: @@ -172,13 +172,13 @@ def get_M1(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]) # assemble block if weight is not zero - if np.any(mat_w): + if xp.any(mat_w): ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array(ns[a]), - np.array(ns[b]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array(ns[a]), + xp.array(ns[b]), wts[0], wts[1], wts[2], @@ -196,9 +196,9 @@ def get_M1(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # convert to sparse matrix - indices = np.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) + indices = xp.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -280,7 +280,7 @@ def get_M2(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_2form[a] Nj = tensor_space_FEM.Nbase_2form[b] - M[a][b] = np.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) + M[a][b] = xp.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) # evaluate metric tensor at quadrature points if weight == None: @@ -291,13 +291,13 @@ def get_M2(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]) # assemble block if weight is not zero - if np.any(mat_w): + if xp.any(mat_w): ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array(ns[a]), - np.array(ns[b]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array(ns[a]), + xp.array(ns[b]), wts[0], wts[1], wts[2], @@ -315,9 +315,9 @@ def get_M2(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # convert to sparse matrix - indices = np.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) + indices = xp.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -381,7 +381,7 @@ def get_M3(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): # evaluation of weight function at quadrature points if weight == None: - mat_w = np.ones(det_df.shape, dtype=float) + mat_w = xp.ones(det_df.shape, dtype=float) else: mat_w = weight(pts[0].flatten(), pts[1].flatten(), pts[2].flatten()) mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]) @@ -390,14 +390,14 @@ def get_M3(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_3form Nj = tensor_space_FEM.Nbase_3form - M = np.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) + M = xp.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array([1, 1, 1]), - np.array([1, 1, 1]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array([1, 1, 1]), + xp.array([1, 1, 1]), wts[0], wts[1], wts[2], @@ -415,9 +415,9 @@ def get_M3(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # conversion to sparse matrix - indices = np.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) + indices = xp.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() @@ -515,7 +515,7 @@ def get_Mv(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): Ni = tensor_space_FEM.Nbase_2form[a] Nj = tensor_space_FEM.Nbase_2form[b] - M[a][b] = np.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) + M[a][b] = xp.zeros((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1), dtype=float) # evaluate metric tensor at quadrature points if weight == None: @@ -526,13 +526,13 @@ def get_Mv(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): mat_w = mat_w.reshape(Nel[0], n_quad[0], Nel[1], n_quad[1], Nel[2], n_quad[2]) # assemble block if weight is not zero - if np.any(mat_w): + if xp.any(mat_w): ker.kernel_mass( - np.array(Nel), - np.array(p), - np.array(n_quad), - np.array(ns[a]), - np.array(ns[b]), + xp.array(Nel), + xp.array(p), + xp.array(n_quad), + xp.array(ns[a]), + xp.array(ns[b]), wts[0], wts[1], wts[2], @@ -550,9 +550,9 @@ def get_Mv(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): ) # convert to sparse matrix - indices = np.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) + indices = xp.indices((Ni[0], Ni[1], Ni[2], 2 * p[0] + 1, 2 * p[1] + 1, 2 * p[2] + 1)) - shift = [np.arange(Ni) - p for Ni, p in zip(Ni, p)] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni, p)] row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() diff --git a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py index 009190582..e0ee63cc5 100644 --- a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py +++ b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py @@ -36,7 +36,7 @@ def solve_mhd_ev_problem_2d(num_params, eq_mhd, n_tor, basis_tor="i", path_out=N from struphy.eigenvalue_solvers.mhd_operators import MHDOperators from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp print("\nStart of eigenspectrum calculation for toroidal mode number", n_tor) print("") @@ -148,7 +148,7 @@ def solve_mhd_ev_problem_2d(num_params, eq_mhd, n_tor, basis_tor="i", path_out=N print("Assembly of final system matrix done --> start of eigenvalue calculation") - omega2, U2_eig = np.linalg.eig(MAT) + omega2, U2_eig = xp.linalg.eig(MAT) print("Eigenstates calculated") @@ -161,8 +161,8 @@ def solve_mhd_ev_problem_2d(num_params, eq_mhd, n_tor, basis_tor="i", path_out=N else: n_tor_str = "+" + str(n_tor) - np.save( - os.path.join(path_out, "spec_n_" + n_tor_str + ".npy"), np.vstack((omega2.reshape(1, omega2.size), U2_eig)) + xp.save( + os.path.join(path_out, "spec_n_" + n_tor_str + ".npy"), xp.vstack((omega2.reshape(1, omega2.size), U2_eig)) ) # or return eigenfrequencies, eigenvectors and system matrix diff --git a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py index dc6e53ddd..779b9463e 100644 --- a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py +++ b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py @@ -5,7 +5,7 @@ def main(): import yaml - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # parse arguments parser = argparse.ArgumentParser(description="Restrict a full .npy eigenspectrum to a range of eigenfrequencies.") @@ -52,18 +52,18 @@ def main(): spec_path = os.path.join(input_path, "spec_n_" + n_tor_str + ".npy") - omega2, U2_eig = np.split(np.load(spec_path), [1], axis=0) + omega2, U2_eig = xp.split(xp.load(spec_path), [1], axis=0) omega2 = omega2.flatten() - modes_ind = np.where((np.real(omega2) < args.upper) & (np.real(omega2) > args.lower))[0] + modes_ind = xp.where((xp.real(omega2) < args.upper) & (xp.real(omega2) > args.lower))[0] omega2 = omega2[modes_ind] U2_eig = U2_eig[:, modes_ind] # save restricted spectrum - np.save( + xp.save( os.path.join(input_path, "spec_" + str(args.lower) + "_" + str(args.upper) + "_n_" + n_tor_str + ".npy"), - np.vstack((omega2.reshape(1, omega2.size), U2_eig)), + xp.vstack((omega2.reshape(1, omega2.size), U2_eig)), ) diff --git a/src/struphy/eigenvalue_solvers/mhd_operators.py b/src/struphy/eigenvalue_solvers/mhd_operators.py index b2cb669ae..c0ffd0658 100644 --- a/src/struphy/eigenvalue_solvers/mhd_operators.py +++ b/src/struphy/eigenvalue_solvers/mhd_operators.py @@ -7,7 +7,7 @@ import struphy.eigenvalue_solvers.legacy.mass_matrices_3d_pre as mass_3d_pre from struphy.eigenvalue_solvers.mhd_operators_core import MHDOperatorsCore -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class MHDOperators: @@ -402,7 +402,7 @@ def __EF(self, u): out1 = self.int_N3.dot(self.dofs_EF[0].dot(u1).T).T + self.int_N3.dot(self.dofs_EF[1].dot(u3).T).T out3 = self.his_N3.dot(self.dofs_EF[2].dot(u1).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) elif self.core.basis_u == 2: u1, u3 = self.core.space.reshape_pol_2(u) @@ -410,7 +410,7 @@ def __EF(self, u): out1 = self.int_D3.dot(self.dofs_EF[0].dot(u1).T).T + self.int_N3.dot(self.dofs_EF[1].dot(u3).T).T out3 = self.his_D3.dot(self.dofs_EF[2].dot(u1).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) else: out = self.dofs_EF.dot(u) @@ -434,7 +434,7 @@ def __EF_transposed(self, e): ) out3 = self.int_N3.T.dot(self.dofs_EF[1].T.dot(e1).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) elif self.core.basis_u == 2: out1 = ( @@ -442,7 +442,7 @@ def __EF_transposed(self, e): ) out3 = self.int_N3.T.dot(self.dofs_EF[1].T.dot(e1).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) else: out = self.dofs_EF.T.dot(e) @@ -462,7 +462,7 @@ def __MF(self, u): out1 = self.his_N3.dot(self.dofs_MF[0].dot(u1).T).T out3 = self.int_N3.dot(self.dofs_MF[1].dot(u3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) elif self.core.basis_u == 2: u1, u3 = self.core.space.reshape_pol_2(u) @@ -470,7 +470,7 @@ def __MF(self, u): out1 = self.his_D3.dot(self.dofs_MF[0].dot(u1).T).T out3 = self.int_N3.dot(self.dofs_MF[1].dot(u3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) else: out = self.dofs_MF.dot(u) @@ -492,13 +492,13 @@ def __MF_transposed(self, f): out1 = self.his_N3.T.dot(self.dofs_MF[0].T.dot(f1).T).T out3 = self.int_N3.T.dot(self.dofs_MF[1].T.dot(f3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) elif self.core.basis_u == 2: out1 = self.his_D3.T.dot(self.dofs_MF[0].T.dot(f1).T).T out3 = self.int_N3.T.dot(self.dofs_MF[1].T.dot(f3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) else: out = self.dofs_MF.T.dot(f) @@ -518,7 +518,7 @@ def __PF(self, u): out1 = self.his_N3.dot(self.dofs_PF[0].dot(u1).T).T out3 = self.int_N3.dot(self.dofs_PF[1].dot(u3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) elif self.core.basis_u == 2: u1, u3 = self.core.space.reshape_pol_2(u) @@ -526,7 +526,7 @@ def __PF(self, u): out1 = self.his_D3.dot(self.dofs_PF[0].dot(u1).T).T out3 = self.int_N3.dot(self.dofs_PF[1].dot(u3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) else: out = self.dofs_PF.dot(u) @@ -548,13 +548,13 @@ def __PF_transposed(self, f): out1 = self.his_N3.T.dot(self.dofs_PF[0].T.dot(f1).T).T out3 = self.int_N3.T.dot(self.dofs_PF[1].T.dot(f3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) elif self.core.basis_u == 2: out1 = self.his_D3.T.dot(self.dofs_PF[0].T.dot(f1).T).T out3 = self.int_N3.T.dot(self.dofs_PF[1].T.dot(f3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) else: out = self.dofs_PF.T.dot(f) @@ -574,7 +574,7 @@ def __JF(self, u): out1 = self.his_N3.dot(self.dofs_JF[0].dot(u1).T).T out3 = self.int_N3.dot(self.dofs_JF[1].dot(u3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) elif self.core.basis_u == 2: u1, u3 = self.core.space.reshape_pol_2(u) @@ -582,7 +582,7 @@ def __JF(self, u): out1 = self.his_D3.dot(self.dofs_JF[0].dot(u1).T).T out3 = self.int_N3.dot(self.dofs_JF[1].dot(u3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) else: out = self.dofs_JF.dot(u) @@ -604,13 +604,13 @@ def __JF_transposed(self, f): out1 = self.his_N3.T.dot(self.dofs_JF[0].T.dot(f1).T).T out3 = self.int_N3.T.dot(self.dofs_JF[1].T.dot(f3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) elif self.core.basis_u == 2: out1 = self.his_D3.T.dot(self.dofs_JF[0].T.dot(f1).T).T out3 = self.int_N3.T.dot(self.dofs_JF[1].T.dot(f3).T).T - out = np.concatenate((out1.flatten(), out3.flatten())) + out = xp.concatenate((out1.flatten(), out3.flatten())) else: out = self.dofs_JF.T.dot(f) @@ -678,7 +678,7 @@ def __MJ(self, b): if self.MJ_as_tensor: if self.core.basis_u == 0: - out = np.zeros(self.core.space.Ev_0.shape[0], dtype=float) + out = xp.zeros(self.core.space.Ev_0.shape[0], dtype=float) elif self.core.basis_u == 2: out = self.core.space.apply_M2_ten( b, [[self.MJ_mat[0], self.core.space.M1_tor], [self.MJ_mat[1], self.core.space.M0_tor]] @@ -686,7 +686,7 @@ def __MJ(self, b): else: if self.core.basis_u == 0: - out = np.zeros(self.core.space.Ev_0.shape[0], dtype=float) + out = xp.zeros(self.core.space.Ev_0.shape[0], dtype=float) elif self.core.basis_u == 2: out = self.MJ_mat.dot(b) @@ -929,7 +929,7 @@ def guess_S2(self, u, b, kind): u_guess = u + self.dt_2 / 6 * (k1_u + 2 * k2_u + 2 * k3_u + k4_u) else: - u_guess = np.copy(u) + u_guess = xp.copy(u) return u_guess diff --git a/src/struphy/eigenvalue_solvers/mhd_operators_core.py b/src/struphy/eigenvalue_solvers/mhd_operators_core.py index 528ec2b78..45712afae 100644 --- a/src/struphy/eigenvalue_solvers/mhd_operators_core.py +++ b/src/struphy/eigenvalue_solvers/mhd_operators_core.py @@ -8,7 +8,7 @@ import struphy.eigenvalue_solvers.kernels_projectors_global_mhd as ker import struphy.eigenvalue_solvers.mass_matrices_2d as mass_2d import struphy.eigenvalue_solvers.mass_matrices_3d as mass_3d -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class MHDOperatorsCore: @@ -58,11 +58,11 @@ def __init__(self, space, equilibrium, basis_u): self.subs_cum = [space.projectors.subs_cum for space in self.space.spaces] # get 1D indices of non-vanishing values of expressions dofs_0(N), dofs_0(D), dofs_1(N) and dofs_1(D) - self.dofs_0_N_i = [list(np.nonzero(space.projectors.I.toarray())) for space in self.space.spaces] - self.dofs_1_D_i = [list(np.nonzero(space.projectors.H.toarray())) for space in self.space.spaces] + self.dofs_0_N_i = [list(xp.nonzero(space.projectors.I.toarray())) for space in self.space.spaces] + self.dofs_1_D_i = [list(xp.nonzero(space.projectors.H.toarray())) for space in self.space.spaces] - self.dofs_0_D_i = [list(np.nonzero(space.projectors.ID.toarray())) for space in self.space.spaces] - self.dofs_1_N_i = [list(np.nonzero(space.projectors.HN.toarray())) for space in self.space.spaces] + self.dofs_0_D_i = [list(xp.nonzero(space.projectors.ID.toarray())) for space in self.space.spaces] + self.dofs_1_N_i = [list(xp.nonzero(space.projectors.HN.toarray())) for space in self.space.spaces] for i in range(self.space.dim): for j in range(2): @@ -116,9 +116,9 @@ def get_blocks_EF(self, pol=True): B2_3_pts = B2_3_pts.reshape(self.nhis[0], self.nq[0], self.nint[1]) # assemble sparse matrix - val = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) ker.rhs11_2d( self.dofs_1_N_i[0][0], @@ -130,8 +130,8 @@ def get_blocks_EF(self, pol=True): self.wts[0], self.basis_his_N[0], self.basis_int_N[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_3_pts, val, row, @@ -150,9 +150,9 @@ def get_blocks_EF(self, pol=True): B2_2_pts = B2_2_pts.reshape(self.nhis[0], self.nq[0], self.nint[1]) # assemble sparse matrix - val = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) ker.rhs11_2d( self.dofs_1_N_i[0][0], @@ -164,8 +164,8 @@ def get_blocks_EF(self, pol=True): self.wts[0], self.basis_his_N[0], self.basis_int_N[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_2_pts, val, row, @@ -184,9 +184,9 @@ def get_blocks_EF(self, pol=True): B2_3_pts = B2_3_pts.reshape(self.nint[0], self.nhis[1], self.nq[1]) # assemble sparse matrix - val = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) ker.rhs12_2d( self.dofs_0_N_i[0][0], @@ -198,8 +198,8 @@ def get_blocks_EF(self, pol=True): self.wts[1], self.basis_int_N[0], self.basis_his_N[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_3_pts, val, row, @@ -218,9 +218,9 @@ def get_blocks_EF(self, pol=True): B2_1_pts = B2_1_pts.reshape(self.nint[0], self.nhis[1], self.nq[1]) # assemble sparse matrix - val = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) ker.rhs12_2d( self.dofs_0_N_i[0][0], @@ -232,8 +232,8 @@ def get_blocks_EF(self, pol=True): self.wts[1], self.basis_int_N[0], self.basis_his_N[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_1_pts, val, row, @@ -251,9 +251,9 @@ def get_blocks_EF(self, pol=True): B2_2_pts = self.equilibrium.b2_2(self.eta_int[0], self.eta_int[1], 0.0) # assemble sparse matrix - val = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) ker.rhs0_2d( self.dofs_0_N_i[0][0], @@ -279,9 +279,9 @@ def get_blocks_EF(self, pol=True): B2_1_pts = self.equilibrium.b2_1(self.eta_int[0], self.eta_int[1], 0.0) # assemble sparse matrix - val = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) ker.rhs0_2d( self.dofs_0_N_i[0][0], @@ -309,13 +309,13 @@ def get_blocks_EF(self, pol=True): B2_3_pts = B2_3_pts.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) @@ -332,8 +332,8 @@ def get_blocks_EF(self, pol=True): self.basis_his_N[0], self.basis_int_N[1], self.basis_int_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_3_pts, val, row, @@ -350,13 +350,13 @@ def get_blocks_EF(self, pol=True): B2_2_pts = B2_2_pts.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) @@ -373,8 +373,8 @@ def get_blocks_EF(self, pol=True): self.basis_his_N[0], self.basis_int_N[1], self.basis_int_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_2_pts, val, row, @@ -391,13 +391,13 @@ def get_blocks_EF(self, pol=True): B2_3_pts = B2_3_pts.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) @@ -414,8 +414,8 @@ def get_blocks_EF(self, pol=True): self.basis_int_N[0], self.basis_his_N[1], self.basis_int_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_3_pts, val, row, @@ -432,13 +432,13 @@ def get_blocks_EF(self, pol=True): B2_1_pts = B2_1_pts.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) @@ -455,8 +455,8 @@ def get_blocks_EF(self, pol=True): self.basis_int_N[0], self.basis_his_N[1], self.basis_int_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_1_pts, val, row, @@ -473,13 +473,13 @@ def get_blocks_EF(self, pol=True): B2_2_pts = B2_2_pts.reshape(self.nint[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int ) @@ -496,8 +496,8 @@ def get_blocks_EF(self, pol=True): self.basis_int_N[0], self.basis_int_N[1], self.basis_his_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_2_pts, val, row, @@ -514,13 +514,13 @@ def get_blocks_EF(self, pol=True): B2_1_pts = B2_1_pts.reshape(self.nint[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int ) @@ -537,8 +537,8 @@ def get_blocks_EF(self, pol=True): self.basis_int_N[0], self.basis_int_N[1], self.basis_his_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_1_pts, val, row, @@ -561,9 +561,9 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nint[1]) # assemble sparse matrix - val = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) ker.rhs11_2d( self.dofs_1_D_i[0][0], @@ -575,8 +575,8 @@ def get_blocks_EF(self, pol=True): self.wts[0], self.basis_his_D[0], self.basis_int_N[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_3_pts / det_dF, val, row, @@ -599,9 +599,9 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nint[1]) # assemble sparse matrix - val = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=float) - row = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=int) - col = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=int) ker.rhs11_2d( self.dofs_1_D_i[0][0], @@ -613,8 +613,8 @@ def get_blocks_EF(self, pol=True): self.wts[0], self.basis_his_D[0], self.basis_int_D[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_2_pts / det_dF, val, row, @@ -637,9 +637,9 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nint[0], self.nhis[1], self.nq[1]) # assemble sparse matrix - val = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) ker.rhs12_2d( self.dofs_0_N_i[0][0], @@ -651,8 +651,8 @@ def get_blocks_EF(self, pol=True): self.wts[1], self.basis_int_N[0], self.basis_his_D[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_3_pts / det_dF, val, row, @@ -675,9 +675,9 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nint[0], self.nhis[1], self.nq[1]) # assemble sparse matrix - val = np.empty(self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) ker.rhs12_2d( self.dofs_0_D_i[0][0], @@ -689,8 +689,8 @@ def get_blocks_EF(self, pol=True): self.wts[1], self.basis_int_D[0], self.basis_his_D[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_1_pts / det_dF, val, row, @@ -711,9 +711,9 @@ def get_blocks_EF(self, pol=True): det_dF = abs(self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_int[1], 0.0)) # assemble sparse matrix - val = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size, dtype=int) ker.rhs0_2d( self.dofs_0_N_i[0][0], @@ -742,9 +742,9 @@ def get_blocks_EF(self, pol=True): det_dF = abs(self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_int[1], 0.0)) # assemble sparse matrix - val = np.empty(self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) ker.rhs0_2d( self.dofs_0_D_i[0][0], @@ -778,13 +778,13 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=int ) @@ -801,8 +801,8 @@ def get_blocks_EF(self, pol=True): self.basis_his_D[0], self.basis_int_N[1], self.basis_int_D[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_3_pts / det_dF, val, row, @@ -825,13 +825,13 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) @@ -848,8 +848,8 @@ def get_blocks_EF(self, pol=True): self.basis_his_D[0], self.basis_int_D[1], self.basis_int_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_2_pts / det_dF, val, row, @@ -872,13 +872,13 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=int ) @@ -895,8 +895,8 @@ def get_blocks_EF(self, pol=True): self.basis_int_N[0], self.basis_his_D[1], self.basis_int_D[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_3_pts / det_dF, val, row, @@ -919,13 +919,13 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) @@ -942,8 +942,8 @@ def get_blocks_EF(self, pol=True): self.basis_int_D[0], self.basis_his_D[1], self.basis_int_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_1_pts / det_dF, val, row, @@ -966,13 +966,13 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nint[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) @@ -989,8 +989,8 @@ def get_blocks_EF(self, pol=True): self.basis_int_N[0], self.basis_int_D[1], self.basis_his_D[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), -B2_2_pts / det_dF, val, row, @@ -1013,13 +1013,13 @@ def get_blocks_EF(self, pol=True): det_dF = det_dF.reshape(self.nint[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) @@ -1036,8 +1036,8 @@ def get_blocks_EF(self, pol=True): self.basis_int_D[0], self.basis_int_N[1], self.basis_his_D[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), B2_1_pts / det_dF, val, row, @@ -1093,9 +1093,9 @@ def get_blocks_FL(self, which, pol=True): EQ = EQ.reshape(self.nint[0], self.nhis[1], self.nq[1]) # assemble sparse matrix - val = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) ker.rhs12_2d( self.dofs_0_N_i[0][0], @@ -1107,8 +1107,8 @@ def get_blocks_FL(self, which, pol=True): self.wts[1], self.basis_int_N[0], self.basis_his_N[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ, val, row, @@ -1133,9 +1133,9 @@ def get_blocks_FL(self, which, pol=True): EQ = EQ.reshape(self.nhis[0], self.nq[0], self.nint[1]) # assemble sparse matrix - val = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) ker.rhs11_2d( self.dofs_1_N_i[0][0], @@ -1147,8 +1147,8 @@ def get_blocks_FL(self, which, pol=True): self.wts[0], self.basis_his_N[0], self.basis_int_N[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ, val, row, @@ -1173,9 +1173,9 @@ def get_blocks_FL(self, which, pol=True): EQ = EQ.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1]) # assemble sparse matrix - val = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size, dtype=int) ker.rhs2_2d( self.dofs_1_N_i[0][0], @@ -1190,8 +1190,8 @@ def get_blocks_FL(self, which, pol=True): self.wts[1], self.basis_his_N[0], self.basis_his_N[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ, val, row, @@ -1219,13 +1219,13 @@ def get_blocks_FL(self, which, pol=True): EQ = EQ.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nhis[2], self.nq[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int ) @@ -1245,8 +1245,8 @@ def get_blocks_FL(self, which, pol=True): self.basis_int_N[0], self.basis_his_N[1], self.basis_his_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ, val, row, @@ -1271,13 +1271,13 @@ def get_blocks_FL(self, which, pol=True): EQ = EQ.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int ) @@ -1297,8 +1297,8 @@ def get_blocks_FL(self, which, pol=True): self.basis_his_N[0], self.basis_int_N[1], self.basis_his_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ, val, row, @@ -1323,13 +1323,13 @@ def get_blocks_FL(self, which, pol=True): EQ = EQ.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) @@ -1349,8 +1349,8 @@ def get_blocks_FL(self, which, pol=True): self.basis_his_N[0], self.basis_his_N[1], self.basis_int_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ, val, row, @@ -1377,9 +1377,9 @@ def get_blocks_FL(self, which, pol=True): det_dF = det_dF.reshape(self.nint[0], self.nhis[1], self.nq[1]) # assemble sparse matrix - val = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) - row = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) - col = np.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) ker.rhs12_2d( self.dofs_0_N_i[0][0], @@ -1391,8 +1391,8 @@ def get_blocks_FL(self, which, pol=True): self.wts[1], self.basis_int_N[0], self.basis_his_D[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ / det_dF, val, row, @@ -1419,9 +1419,9 @@ def get_blocks_FL(self, which, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nint[1]) # assemble sparse matrix - val = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) - row = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) - col = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size, dtype=int) ker.rhs11_2d( self.dofs_1_D_i[0][0], @@ -1433,8 +1433,8 @@ def get_blocks_FL(self, which, pol=True): self.wts[0], self.basis_his_D[0], self.basis_int_N[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ / det_dF, val, row, @@ -1463,9 +1463,9 @@ def get_blocks_FL(self, which, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1]) # assemble sparse matrix - val = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) - row = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) - col = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) ker.rhs2_2d( self.dofs_1_D_i[0][0], @@ -1480,8 +1480,8 @@ def get_blocks_FL(self, which, pol=True): self.wts[1], self.basis_his_D[0], self.basis_his_D[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ / det_dF, val, row, @@ -1513,13 +1513,13 @@ def get_blocks_FL(self, which, pol=True): det_dF = det_dF.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nhis[2], self.nq[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) @@ -1539,8 +1539,8 @@ def get_blocks_FL(self, which, pol=True): self.basis_int_N[0], self.basis_his_D[1], self.basis_his_D[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ / det_dF, val, row, @@ -1569,13 +1569,13 @@ def get_blocks_FL(self, which, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) @@ -1595,8 +1595,8 @@ def get_blocks_FL(self, which, pol=True): self.basis_his_D[0], self.basis_int_N[1], self.basis_his_D[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ / det_dF, val, row, @@ -1625,13 +1625,13 @@ def get_blocks_FL(self, which, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int ) @@ -1651,8 +1651,8 @@ def get_blocks_FL(self, which, pol=True): self.basis_his_D[0], self.basis_his_D[1], self.basis_int_N[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), EQ / det_dF, val, row, @@ -1696,9 +1696,9 @@ def get_blocks_PR(self, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1]) # assemble sparse matrix - val = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) - row = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) - col = np.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + val = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=float) + row = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) + col = xp.empty(self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size, dtype=int) ker.rhs2_2d( self.dofs_1_D_i[0][0], @@ -1713,8 +1713,8 @@ def get_blocks_PR(self, pol=True): self.wts[1], self.basis_his_D[0], self.basis_his_D[1], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), P3_pts / det_dF, val, row, @@ -1745,13 +1745,13 @@ def get_blocks_PR(self, pol=True): det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1], self.nhis[2], self.nq[2]) # assemble sparse matrix - val = np.empty( + val = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float ) - row = np.empty( + row = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) - col = np.empty( + col = xp.empty( self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int ) @@ -1774,8 +1774,8 @@ def get_blocks_PR(self, pol=True): self.basis_his_D[0], self.basis_his_D[1], self.basis_his_D[2], - np.array(self.space.NbaseN), - np.array(self.space.NbaseD), + xp.array(self.space.NbaseN), + xp.array(self.space.NbaseD), P3_pts / det_dF, val, row, diff --git a/src/struphy/eigenvalue_solvers/projectors_global.py b/src/struphy/eigenvalue_solvers/projectors_global.py index 208ff701a..7c66c3397 100644 --- a/src/struphy/eigenvalue_solvers/projectors_global.py +++ b/src/struphy/eigenvalue_solvers/projectors_global.py @@ -10,7 +10,7 @@ import struphy.bsplines.bsplines as bsp from struphy.linear_algebra.linalg_kron import kron_lusolve_2d, kron_lusolve_3d, kron_matvec_2d, kron_matvec_3d -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ======================= 1d ==================================== @@ -156,15 +156,15 @@ def __init__(self, spline_space, n_quad=6): self.n_quad = n_quad # Gauss - Legendre quadrature points and weights in (-1, 1) - self.pts_loc = np.polynomial.legendre.leggauss(self.n_quad)[0] - self.wts_loc = np.polynomial.legendre.leggauss(self.n_quad)[1] + self.pts_loc = xp.polynomial.legendre.leggauss(self.n_quad)[0] + self.wts_loc = xp.polynomial.legendre.leggauss(self.n_quad)[1] # set interpolation points (Greville points) self.x_int = spline_space.greville.copy() # set number of sub-intervals per integration interval between Greville points and integration boundaries - self.subs = np.ones(spline_space.NbaseD, dtype=int) - self.x_his = np.array([self.x_int[0]]) + self.subs = xp.ones(spline_space.NbaseD, dtype=int) + self.x_his = xp.array([self.x_int[0]]) for i in range(spline_space.NbaseD): for br in spline_space.el_b: @@ -181,16 +181,16 @@ def __init__(self, spline_space, n_quad=6): # compute subs and x_his if (br > xl + 1e-10) and (br < xr - 1e-10): self.subs[i] += 1 - self.x_his = np.append(self.x_his, br) + self.x_his = xp.append(self.x_his, br) elif br >= xr - 1e-10: - self.x_his = np.append(self.x_his, xr) + self.x_his = xp.append(self.x_his, xr) break if spline_space.spl_kind == True and spline_space.p % 2 == 0: - self.x_his = np.append(self.x_his, spline_space.el_b[-1] + self.x_his[0]) + self.x_his = xp.append(self.x_his, spline_space.el_b[-1] + self.x_his[0]) # cumulative number of sub-intervals for conversion local interval --> global interval - self.subs_cum = np.append(0, np.cumsum(self.subs - 1)[:-1]) + self.subs_cum = xp.append(0, xp.cumsum(self.subs - 1)[:-1]) # quadrature points and weights self.pts, self.wts = bsp.quadrature_grid(self.x_his, self.pts_loc, self.wts_loc) @@ -200,31 +200,31 @@ def __init__(self, spline_space, n_quad=6): self.x_hisG = self.x_int if spline_space.spl_kind == True: if spline_space.p % 2 == 0: - self.x_hisG = np.append(self.x_hisG, spline_space.el_b[-1] + self.x_hisG[0]) + self.x_hisG = xp.append(self.x_hisG, spline_space.el_b[-1] + self.x_hisG[0]) else: - self.x_hisG = np.append(self.x_hisG, spline_space.el_b[-1]) + self.x_hisG = xp.append(self.x_hisG, spline_space.el_b[-1]) self.ptsG, self.wtsG = bsp.quadrature_grid(self.x_hisG, self.pts_loc, self.wts_loc) self.ptsG = self.ptsG % spline_space.el_b[-1] # Knot span indices at interpolation points in format (greville, 0) - self.span_x_int_N = np.zeros(self.x_int[:, None].shape, dtype=int) - self.span_x_int_D = np.zeros(self.x_int[:, None].shape, dtype=int) + self.span_x_int_N = xp.zeros(self.x_int[:, None].shape, dtype=int) + self.span_x_int_D = xp.zeros(self.x_int[:, None].shape, dtype=int) for i in range(self.x_int.shape[0]): self.span_x_int_N[i, 0] = bsp.find_span(self.space.T, self.space.p, self.x_int[i]) self.span_x_int_D[i, 0] = bsp.find_span(self.space.t, self.space.p - 1, self.x_int[i]) # Knot span indices at quadrature points between x_int in format (i, iq) - self.span_ptsG_N = np.zeros(self.ptsG.shape, dtype=int) - self.span_ptsG_D = np.zeros(self.ptsG.shape, dtype=int) + self.span_ptsG_N = xp.zeros(self.ptsG.shape, dtype=int) + self.span_ptsG_D = xp.zeros(self.ptsG.shape, dtype=int) for i in range(self.ptsG.shape[0]): for iq in range(self.ptsG.shape[1]): self.span_ptsG_N[i, iq] = bsp.find_span(self.space.T, self.space.p, self.ptsG[i, iq]) self.span_ptsG_D[i, iq] = bsp.find_span(self.space.t, self.space.p - 1, self.ptsG[i, iq]) # Values of p + 1 non-zero basis functions at Greville points in format (greville, 0, basis function) - self.basis_x_int_N = np.zeros((*self.x_int[:, None].shape, self.space.p + 1), dtype=float) - self.basis_x_int_D = np.zeros((*self.x_int[:, None].shape, self.space.p), dtype=float) + self.basis_x_int_N = xp.zeros((*self.x_int[:, None].shape, self.space.p + 1), dtype=float) + self.basis_x_int_D = xp.zeros((*self.x_int[:, None].shape, self.space.p), dtype=float) N_temp = bsp.basis_ders_on_quad_grid(self.space.T, self.space.p, self.x_int[:, None], 0, normalize=False) D_temp = bsp.basis_ders_on_quad_grid(self.space.t, self.space.p - 1, self.x_int[:, None], 0, normalize=True) @@ -236,8 +236,8 @@ def __init__(self, spline_space, n_quad=6): self.basis_x_int_D[i, 0, b] = D_temp[i, b, 0, 0] # Values of p + 1 non-zero basis functions at quadrature points points between x_int in format (i, iq, basis function) - self.basis_ptsG_N = np.zeros((*self.ptsG.shape, self.space.p + 1), dtype=float) - self.basis_ptsG_D = np.zeros((*self.ptsG.shape, self.space.p), dtype=float) + self.basis_ptsG_N = xp.zeros((*self.ptsG.shape, self.space.p + 1), dtype=float) + self.basis_ptsG_D = xp.zeros((*self.ptsG.shape, self.space.p), dtype=float) N_temp = bsp.basis_ders_on_quad_grid(self.space.T, self.space.p, self.ptsG, 0, normalize=False) D_temp = bsp.basis_ders_on_quad_grid(self.space.t, self.space.p - 1, self.ptsG, 0, normalize=True) @@ -250,7 +250,7 @@ def __init__(self, spline_space, n_quad=6): self.basis_ptsG_D[i, iq, b] = D_temp[i, b, 0, iq] # quadrature matrix for performing integrations as matrix-vector products - self.Q = np.zeros((spline_space.NbaseD, self.wts.shape[0] * self.n_quad), dtype=float) + self.Q = xp.zeros((spline_space.NbaseD, self.wts.shape[0] * self.n_quad), dtype=float) for i in range(spline_space.NbaseD): for j in range(self.subs[i]): @@ -260,7 +260,7 @@ def __init__(self, spline_space, n_quad=6): self.Q = spa.csr_matrix(self.Q) # quadrature matrix for performing integrations as matrix-vector products, ignoring subs (less accurate integration for even degree) - self.QG = np.zeros((spline_space.NbaseD, self.wtsG.shape[0] * self.n_quad), dtype=float) + self.QG = xp.zeros((spline_space.NbaseD, self.wtsG.shape[0] * self.n_quad), dtype=float) for i in range(spline_space.NbaseD): self.QG[i, self.n_quad * i : self.n_quad * (i + 1)] = self.wtsG[i] @@ -399,17 +399,17 @@ def dofs_1d_bases_products(self, space): dofs_1_i(D_j*D_k). """ - dofs_0_NN = np.empty((space.NbaseN, space.NbaseN, space.NbaseN), dtype=float) - dofs_0_DN = np.empty((space.NbaseN, space.NbaseD, space.NbaseN), dtype=float) - dofs_0_DD = np.empty((space.NbaseN, space.NbaseD, space.NbaseD), dtype=float) + dofs_0_NN = xp.empty((space.NbaseN, space.NbaseN, space.NbaseN), dtype=float) + dofs_0_DN = xp.empty((space.NbaseN, space.NbaseD, space.NbaseN), dtype=float) + dofs_0_DD = xp.empty((space.NbaseN, space.NbaseD, space.NbaseD), dtype=float) - dofs_1_NN = np.empty((space.NbaseD, space.NbaseN, space.NbaseN), dtype=float) - dofs_1_DN = np.empty((space.NbaseD, space.NbaseD, space.NbaseN), dtype=float) - dofs_1_DD = np.empty((space.NbaseD, space.NbaseD, space.NbaseD), dtype=float) + dofs_1_NN = xp.empty((space.NbaseD, space.NbaseN, space.NbaseN), dtype=float) + dofs_1_DN = xp.empty((space.NbaseD, space.NbaseD, space.NbaseN), dtype=float) + dofs_1_DD = xp.empty((space.NbaseD, space.NbaseD, space.NbaseD), dtype=float) # ========= dofs_0_NN and dofs_1_NN ============== - cj = np.zeros(space.NbaseN, dtype=float) - ck = np.zeros(space.NbaseN, dtype=float) + cj = xp.zeros(space.NbaseN, dtype=float) + ck = xp.zeros(space.NbaseN, dtype=float) for j in range(space.NbaseN): for k in range(space.NbaseN): @@ -426,8 +426,8 @@ def N_jN_k(eta): dofs_1_NN[:, j, k] = self.dofs_1(N_jN_k) # ========= dofs_0_DN and dofs_1_DN ============== - cj = np.zeros(space.NbaseD, dtype=float) - ck = np.zeros(space.NbaseN, dtype=float) + cj = xp.zeros(space.NbaseD, dtype=float) + ck = xp.zeros(space.NbaseN, dtype=float) for j in range(space.NbaseD): for k in range(space.NbaseN): @@ -444,8 +444,8 @@ def D_jN_k(eta): dofs_1_DN[:, j, k] = self.dofs_1(D_jN_k) # ========= dofs_0_DD and dofs_1_DD ============= - cj = np.zeros(space.NbaseD, dtype=float) - ck = np.zeros(space.NbaseD, dtype=float) + cj = xp.zeros(space.NbaseD, dtype=float) + ck = xp.zeros(space.NbaseD, dtype=float) for j in range(space.NbaseD): for k in range(space.NbaseD): @@ -461,109 +461,109 @@ def D_jD_k(eta): dofs_0_DD[:, j, k] = self.dofs_0(D_jD_k) dofs_1_DD[:, j, k] = self.dofs_1(D_jD_k) - dofs_0_ND = np.transpose(dofs_0_DN, (0, 2, 1)) - dofs_1_ND = np.transpose(dofs_1_DN, (0, 2, 1)) + dofs_0_ND = xp.transpose(dofs_0_DN, (0, 2, 1)) + dofs_1_ND = xp.transpose(dofs_1_DN, (0, 2, 1)) # find non-zero entries - dofs_0_NN_indices = np.nonzero(dofs_0_NN) - dofs_0_DN_indices = np.nonzero(dofs_0_DN) - dofs_0_ND_indices = np.nonzero(dofs_0_ND) - dofs_0_DD_indices = np.nonzero(dofs_0_DD) - - dofs_1_NN_indices = np.nonzero(dofs_1_NN) - dofs_1_DN_indices = np.nonzero(dofs_1_DN) - dofs_1_ND_indices = np.nonzero(dofs_1_ND) - dofs_1_DD_indices = np.nonzero(dofs_1_DD) - - dofs_0_NN_i_red = np.empty(dofs_0_NN_indices[0].size, dtype=int) - dofs_0_DN_i_red = np.empty(dofs_0_DN_indices[0].size, dtype=int) - dofs_0_ND_i_red = np.empty(dofs_0_ND_indices[0].size, dtype=int) - dofs_0_DD_i_red = np.empty(dofs_0_DD_indices[0].size, dtype=int) - - dofs_1_NN_i_red = np.empty(dofs_1_NN_indices[0].size, dtype=int) - dofs_1_DN_i_red = np.empty(dofs_1_DN_indices[0].size, dtype=int) - dofs_1_ND_i_red = np.empty(dofs_1_ND_indices[0].size, dtype=int) - dofs_1_DD_i_red = np.empty(dofs_1_DD_indices[0].size, dtype=int) + dofs_0_NN_indices = xp.nonzero(dofs_0_NN) + dofs_0_DN_indices = xp.nonzero(dofs_0_DN) + dofs_0_ND_indices = xp.nonzero(dofs_0_ND) + dofs_0_DD_indices = xp.nonzero(dofs_0_DD) + + dofs_1_NN_indices = xp.nonzero(dofs_1_NN) + dofs_1_DN_indices = xp.nonzero(dofs_1_DN) + dofs_1_ND_indices = xp.nonzero(dofs_1_ND) + dofs_1_DD_indices = xp.nonzero(dofs_1_DD) + + dofs_0_NN_i_red = xp.empty(dofs_0_NN_indices[0].size, dtype=int) + dofs_0_DN_i_red = xp.empty(dofs_0_DN_indices[0].size, dtype=int) + dofs_0_ND_i_red = xp.empty(dofs_0_ND_indices[0].size, dtype=int) + dofs_0_DD_i_red = xp.empty(dofs_0_DD_indices[0].size, dtype=int) + + dofs_1_NN_i_red = xp.empty(dofs_1_NN_indices[0].size, dtype=int) + dofs_1_DN_i_red = xp.empty(dofs_1_DN_indices[0].size, dtype=int) + dofs_1_ND_i_red = xp.empty(dofs_1_ND_indices[0].size, dtype=int) + dofs_1_DD_i_red = xp.empty(dofs_1_DD_indices[0].size, dtype=int) # ================================ nv = space.NbaseN * dofs_0_NN_indices[1] + dofs_0_NN_indices[2] - un = np.unique(nv) + un = xp.unique(nv) for i in range(dofs_0_NN_indices[0].size): - dofs_0_NN_i_red[i] = np.nonzero(un == nv[i])[0] + dofs_0_NN_i_red[i] = xp.nonzero(un == nv[i])[0] # ================================ nv = space.NbaseN * dofs_0_DN_indices[1] + dofs_0_DN_indices[2] - un = np.unique(nv) + un = xp.unique(nv) for i in range(dofs_0_DN_indices[0].size): - dofs_0_DN_i_red[i] = np.nonzero(un == nv[i])[0] + dofs_0_DN_i_red[i] = xp.nonzero(un == nv[i])[0] # ================================ nv = space.NbaseD * dofs_0_ND_indices[1] + dofs_0_ND_indices[2] - un = np.unique(nv) + un = xp.unique(nv) for i in range(dofs_0_ND_indices[0].size): - dofs_0_ND_i_red[i] = np.nonzero(un == nv[i])[0] + dofs_0_ND_i_red[i] = xp.nonzero(un == nv[i])[0] # ================================ nv = space.NbaseD * dofs_0_DD_indices[1] + dofs_0_DD_indices[2] - un = np.unique(nv) + un = xp.unique(nv) for i in range(dofs_0_DD_indices[0].size): - dofs_0_DD_i_red[i] = np.nonzero(un == nv[i])[0] + dofs_0_DD_i_red[i] = xp.nonzero(un == nv[i])[0] # ================================ nv = space.NbaseN * dofs_1_NN_indices[1] + dofs_1_NN_indices[2] - un = np.unique(nv) + un = xp.unique(nv) for i in range(dofs_1_NN_indices[0].size): - dofs_1_NN_i_red[i] = np.nonzero(un == nv[i])[0] + dofs_1_NN_i_red[i] = xp.nonzero(un == nv[i])[0] # ================================ nv = space.NbaseN * dofs_1_DN_indices[1] + dofs_1_DN_indices[2] - un = np.unique(nv) + un = xp.unique(nv) for i in range(dofs_1_DN_indices[0].size): - dofs_1_DN_i_red[i] = np.nonzero(un == nv[i])[0] + dofs_1_DN_i_red[i] = xp.nonzero(un == nv[i])[0] # ================================ nv = space.NbaseD * dofs_1_ND_indices[1] + dofs_1_ND_indices[2] - un = np.unique(nv) + un = xp.unique(nv) for i in range(dofs_1_ND_indices[0].size): - dofs_1_ND_i_red[i] = np.nonzero(un == nv[i])[0] + dofs_1_ND_i_red[i] = xp.nonzero(un == nv[i])[0] # ================================ nv = space.NbaseD * dofs_1_DD_indices[1] + dofs_1_DD_indices[2] - un = np.unique(nv) + un = xp.unique(nv) for i in range(dofs_1_DD_indices[0].size): - dofs_1_DD_i_red[i] = np.nonzero(un == nv[i])[0] + dofs_1_DD_i_red[i] = xp.nonzero(un == nv[i])[0] - dofs_0_NN_indices = np.vstack( + dofs_0_NN_indices = xp.vstack( (dofs_0_NN_indices[0], dofs_0_NN_indices[1], dofs_0_NN_indices[2], dofs_0_NN_i_red) ) - dofs_0_DN_indices = np.vstack( + dofs_0_DN_indices = xp.vstack( (dofs_0_DN_indices[0], dofs_0_DN_indices[1], dofs_0_DN_indices[2], dofs_0_DN_i_red) ) - dofs_0_ND_indices = np.vstack( + dofs_0_ND_indices = xp.vstack( (dofs_0_ND_indices[0], dofs_0_ND_indices[1], dofs_0_ND_indices[2], dofs_0_ND_i_red) ) - dofs_0_DD_indices = np.vstack( + dofs_0_DD_indices = xp.vstack( (dofs_0_DD_indices[0], dofs_0_DD_indices[1], dofs_0_DD_indices[2], dofs_0_DD_i_red) ) - dofs_1_NN_indices = np.vstack( + dofs_1_NN_indices = xp.vstack( (dofs_1_NN_indices[0], dofs_1_NN_indices[1], dofs_1_NN_indices[2], dofs_1_NN_i_red) ) - dofs_1_DN_indices = np.vstack( + dofs_1_DN_indices = xp.vstack( (dofs_1_DN_indices[0], dofs_1_DN_indices[1], dofs_1_DN_indices[2], dofs_1_DN_i_red) ) - dofs_1_ND_indices = np.vstack( + dofs_1_ND_indices = xp.vstack( (dofs_1_ND_indices[0], dofs_1_ND_indices[1], dofs_1_ND_indices[2], dofs_1_ND_i_red) ) - dofs_1_DD_indices = np.vstack( + dofs_1_DD_indices = xp.vstack( (dofs_1_DD_indices[0], dofs_1_DD_indices[1], dofs_1_DD_indices[2], dofs_1_DD_i_red) ) @@ -642,8 +642,8 @@ def eval_for_PI(self, comp, fun): pts_PI = self.pts_PI[comp] - pts1, pts2 = np.meshgrid(pts_PI[0], pts_PI[1], indexing="ij") - # pts1, pts2 = np.meshgrid(pts_PI[0], pts_PI[1], indexing='ij', sparse=True) # numpy >1.7 + pts1, pts2 = xp.meshgrid(pts_PI[0], pts_PI[1], indexing="ij") + # pts1, pts2 = xp.meshgrid(pts_PI[0], pts_PI[1], indexing='ij', sparse=True) # numpy >1.7 return fun(pts1, pts2) @@ -906,8 +906,8 @@ def eval_for_PI(self, comp, fun): pts_PI = self.pts_PI[comp] - pts1, pts2, pts3 = np.meshgrid(pts_PI[0], pts_PI[1], pts_PI[2], indexing="ij") - # pts1, pts2, pts3 = np.meshgrid(pts_PI[0], pts_PI[1], pts_PI[2], indexing='ij', sparse=True) # numpy >1.7 + pts1, pts2, pts3 = xp.meshgrid(pts_PI[0], pts_PI[1], pts_PI[2], indexing="ij") + # pts1, pts2, pts3 = xp.meshgrid(pts_PI[0], pts_PI[1], pts_PI[2], indexing='ij', sparse=True) # numpy >1.7 return fun(pts1, pts2, pts3) @@ -939,25 +939,25 @@ def eval_for_PI(self, comp, fun): # rhs = mat_f # # elif comp=='11': - # rhs = np.empty( (self.d1, self.n2, self.n3) ) + # rhs = xp.empty( (self.d1, self.n2, self.n3) ) # # ker_glob.kernel_int_3d_eta1(self.subs1, self.subs_cum1, self.wts1, # mat_f.reshape(self.ne1, self.nq1, self.n2, self.n3), rhs # ) # elif comp=='12': - # rhs = np.empty( (self.n1, self.d2, self.n3) ) + # rhs = xp.empty( (self.n1, self.d2, self.n3) ) # # ker_glob.kernel_int_3d_eta2(self.subs2, self.subs_cum2, self.wts2, # mat_f.reshape(self.n1, self.ne2, self.nq2, self.n3), rhs # ) # elif comp=='13': - # rhs = np.empty( (self.n1, self.n2, self.d3) ) + # rhs = xp.empty( (self.n1, self.n2, self.d3) ) # # ker_glob.kernel_int_3d_eta3(self.subs3, self.subs_cum3, self.wts3, # mat_f.reshape(self.n1, self.n2, self.ne3, self.nq3), rhs # ) # elif comp=='21': - # rhs = np.empty( (self.n1, self.d2, self.d3) ) + # rhs = xp.empty( (self.n1, self.d2, self.d3) ) # # ker_glob.kernel_int_3d_eta2_eta3(self.subs2, self.subs3, # self.subs_cum2, self.subs_cum3, @@ -965,7 +965,7 @@ def eval_for_PI(self, comp, fun): # mat_f.reshape(self.n1, self.ne2, self.nq2, self.ne3, self.nq3), rhs # ) # elif comp=='22': - # rhs = np.empty( (self.d1, self.n2, self.d3) ) + # rhs = xp.empty( (self.d1, self.n2, self.d3) ) # # ker_glob.kernel_int_3d_eta1_eta3(self.subs1, self.subs3, # self.subs_cum1, self.subs_cum3, @@ -973,7 +973,7 @@ def eval_for_PI(self, comp, fun): # mat_f.reshape(self.ne1, self.nq1, self.n2, self.ne3, self.nq3), rhs # ) # elif comp=='23': - # rhs = np.empty( (self.d1, self.d2, self.n3) ) + # rhs = xp.empty( (self.d1, self.d2, self.n3) ) # # ker_glob.kernel_int_3d_eta1_eta2(self.subs1, self.subs2, # self.subs_cum1, self.subs_cum2, @@ -981,7 +981,7 @@ def eval_for_PI(self, comp, fun): # mat_f.reshape(self.ne1, self.nq1, self.ne2, self.nq2, self.n3), rhs # ) # elif comp=='3': - # rhs = np.empty( (self.d1, self.d2, self.d3) ) + # rhs = xp.empty( (self.d1, self.d2, self.d3) ) # # ker_glob.kernel_int_3d_eta1_eta2_eta3(self.subs1, self.subs2, self.subs3, # self.subs_cum1, self.subs_cum2, self.subs_cum3, @@ -1025,7 +1025,7 @@ def eval_for_PI(self, comp, fun): # # elif comp=='11': # assert mat_dofs.shape == (self.d1, self.n2, self.n3) - # rhs = np.empty( (self.ne1, self.nq1, self.n2, self.n3) ) + # rhs = xp.empty( (self.ne1, self.nq1, self.n2, self.n3) ) # # ker_glob.kernel_int_3d_eta1_transpose(self.subs1, self.subs_cum1, self.wts1, # mat_dofs, rhs) @@ -1034,7 +1034,7 @@ def eval_for_PI(self, comp, fun): # # elif comp=='12': # assert mat_dofs.shape == (self.n1, self.d2, self.n3) - # rhs = np.empty( (self.n1, self.ne2, self.nq2, self.n3) ) + # rhs = xp.empty( (self.n1, self.ne2, self.nq2, self.n3) ) # # ker_glob.kernel_int_3d_eta2_transpose(self.subs2, self.subs_cum2, self.wts2, # mat_dofs, rhs) @@ -1043,7 +1043,7 @@ def eval_for_PI(self, comp, fun): # # elif comp=='13': # assert mat_dofs.shape == (self.n1, self.n2, self.d3) - # rhs = np.empty( (self.n1, self.n2, self.ne3, self.nq3) ) + # rhs = xp.empty( (self.n1, self.n2, self.ne3, self.nq3) ) # # ker_glob.kernel_int_3d_eta3_transpose(self.subs3, self.subs_cum3, self.wts3, # mat_dofs, rhs) @@ -1052,7 +1052,7 @@ def eval_for_PI(self, comp, fun): # # elif comp=='21': # assert mat_dofs.shape == (self.n1, self.d2, self.d3) - # rhs = np.empty( (self.n1, self.ne2, self.nq2, self.ne3, self.nq3) ) + # rhs = xp.empty( (self.n1, self.ne2, self.nq2, self.ne3, self.nq3) ) # # ker_glob.kernel_int_3d_eta2_eta3_transpose(self.subs2, self.subs3, # self.subs_cum2, self.subs_cum3, @@ -1062,7 +1062,7 @@ def eval_for_PI(self, comp, fun): # # elif comp=='22': # assert mat_dofs.shape == (self.d1, self.n2, self.d3) - # rhs = np.empty( (self.ne1, self.nq1, self.n2, self.ne3, self.nq3) ) + # rhs = xp.empty( (self.ne1, self.nq1, self.n2, self.ne3, self.nq3) ) # # ker_glob.kernel_int_3d_eta1_eta3_transpose(self.subs1, self.subs3, # self.subs_cum1, self.subs_cum3, @@ -1072,7 +1072,7 @@ def eval_for_PI(self, comp, fun): # # elif comp=='23': # assert mat_dofs.shape == (self.d1, self.d2, self.n3) - # rhs = np.empty( (self.ne1, self.nq1, self.ne2, self.nq2, self.n3) ) + # rhs = xp.empty( (self.ne1, self.nq1, self.ne2, self.nq2, self.n3) ) # # ker_glob.kernel_int_3d_eta1_eta2_transpose(self.subs1, self.subs2, # self.subs_cum1, self.subs_cum2, @@ -1082,7 +1082,7 @@ def eval_for_PI(self, comp, fun): # # elif comp=='3': # assert mat_dofs.shape == (self.d1, self.d2, self.d3) - # rhs = np.empty( (self.ne1, self.nq1, self.ne2, self.nq2, self.ne3, self.nq3) ) + # rhs = xp.empty( (self.ne1, self.nq1, self.ne2, self.nq2, self.ne3, self.nq3) ) # # ker_glob.kernel_int_3d_eta1_eta2_eta3_transpose(self.subs1, self.subs2, self.subs3, # self.subs_cum1, self.subs_cum2, self.subs_cum3, @@ -1595,26 +1595,26 @@ def __init__(self, tensor_space): else: if tensor_space.n_tor == 0: - x_i3 = np.array([0.0]) - x_q3 = np.array([0.0]) - x_q3G = np.array([0.0]) + x_i3 = xp.array([0.0]) + x_q3 = xp.array([0.0]) + x_q3G = xp.array([0.0]) else: if tensor_space.basis_tor == "r": if tensor_space.n_tor > 0: - x_i3 = np.array([1.0, 0.25 / tensor_space.n_tor]) - x_q3 = np.array([1.0, 0.25 / tensor_space.n_tor]) - x_q3G = np.array([1.0, 0.25 / tensor_space.n_tor]) + x_i3 = xp.array([1.0, 0.25 / tensor_space.n_tor]) + x_q3 = xp.array([1.0, 0.25 / tensor_space.n_tor]) + x_q3G = xp.array([1.0, 0.25 / tensor_space.n_tor]) else: - x_i3 = np.array([1.0, 0.75 / (-tensor_space.n_tor)]) - x_q3 = np.array([1.0, 0.75 / (-tensor_space.n_tor)]) - x_q3G = np.array([1.0, 0.75 / (-tensor_space.n_tor)]) + x_i3 = xp.array([1.0, 0.75 / (-tensor_space.n_tor)]) + x_q3 = xp.array([1.0, 0.75 / (-tensor_space.n_tor)]) + x_q3G = xp.array([1.0, 0.75 / (-tensor_space.n_tor)]) else: - x_i3 = np.array([0.0]) - x_q3 = np.array([0.0]) - x_q3G = np.array([0.0]) + x_i3 = xp.array([0.0]) + x_q3 = xp.array([0.0]) + x_q3G = xp.array([0.0]) self.Q3 = spa.identity(tensor_space.NbaseN[2], format="csr") self.Q3G = spa.identity(tensor_space.NbaseN[2], format="csr") @@ -1756,11 +1756,11 @@ def eval_for_PI(self, comp, fun, eval_kind, with_subs=True): pts_PI = self.getpts_for_PI(comp, with_subs) # array of evaluated function - mat_f = np.empty((pts_PI[0].size, pts_PI[1].size, pts_PI[2].size), dtype=float) + mat_f = xp.empty((pts_PI[0].size, pts_PI[1].size, pts_PI[2].size), dtype=float) # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": - pts1, pts2, pts3 = np.meshgrid(pts_PI[0], pts_PI[1], pts_PI[2], indexing="ij") + pts1, pts2, pts3 = xp.meshgrid(pts_PI[0], pts_PI[1], pts_PI[2], indexing="ij") mat_f[:, :, :] = fun(pts1, pts2, pts3) # tensor-product evaluation is done by input function @@ -1783,13 +1783,13 @@ def eval_for_PI(self, comp, fun, eval_kind, with_subs=True): # n2 = self.pts_PI_0[1].size # # # apply (I0_22) to each column - # self.S0 = np.zeros(((n1 - 2)*n2, 3), dtype=float) + # self.S0 = xp.zeros(((n1 - 2)*n2, 3), dtype=float) # # for i in range(3): # self.S0[:, i] = kron_lusolve_2d(self.I0_22_LUs, self.I0_21[:, i].toarray().reshape(n1 - 2, n2)).flatten() # # # 3 x 3 matrix - # self.S0 = np.linalg.inv(self.I0_11.toarray() - self.I0_12.toarray().dot(self.S0)) + # self.S0 = xp.linalg.inv(self.I0_11.toarray() - self.I0_12.toarray().dot(self.S0)) # # # ====================================== @@ -1814,7 +1814,7 @@ def eval_for_PI(self, comp, fun, eval_kind, with_subs=True): # # solve for tensor-product coefficients # out2 = out2 - kron_lusolve_2d(self.I0_22_LUs, self.I0_21.dot(self.S0.dot(rhs1)).reshape(n1 - 2, n2)) + kron_lusolve_2d(self.I0_22_LUs, self.I0_21.dot(self.S0.dot(self.I0_12.dot(out2.flatten()))).reshape(n1 - 2, n2)) # - # return np.concatenate((out1, out2.flatten())) + # return xp.concatenate((out1, out2.flatten())) # ====================================== @@ -1857,7 +1857,7 @@ def solve_V1(self, dofs_1, include_bc): coeffs1 = self.I0_tor_LU.solve(self.I1_pol_0_LU.solve(dofs_11).T).T coeffs2 = self.H0_tor_LU.solve(self.I0_pol_0_LU.solve(dofs_12).T).T - return np.concatenate((coeffs1.flatten(), coeffs2.flatten())) + return xp.concatenate((coeffs1.flatten(), coeffs2.flatten())) # ====================================== def solve_V2(self, dofs_2, include_bc): @@ -1885,7 +1885,7 @@ def solve_V2(self, dofs_2, include_bc): coeffs1 = self.H0_tor_LU.solve(self.I2_pol_0_LU.solve(dofs_21).T).T coeffs2 = self.I0_tor_LU.solve(self.I3_pol_0_LU.solve(dofs_22).T).T - return np.concatenate((coeffs1.flatten(), coeffs2.flatten())) + return xp.concatenate((coeffs1.flatten(), coeffs2.flatten())) # ====================================== def solve_V3(self, dofs_3, include_bc): @@ -1947,7 +1947,7 @@ def apply_IinvT_V1(self, rhs, include_bc=False): rhs1 = self.I1_pol_0_T_LU.solve(self.I0_tor_T_LU.solve(rhs1.T).T) rhs2 = self.I0_pol_0_T_LU.solve(self.H0_tor_T_LU.solve(rhs2.T).T) - return np.concatenate((rhs1.flatten(), rhs2.flatten())) + return xp.concatenate((rhs1.flatten(), rhs2.flatten())) # ====================================== def apply_IinvT_V2(self, rhs, include_bc=False): @@ -1977,7 +1977,7 @@ def apply_IinvT_V2(self, rhs, include_bc=False): rhs1 = self.I2_pol_0_T_LU.solve(self.H0_tor_T_LU.solve(rhs1.T).T) rhs2 = self.I3_pol_0_T_LU.solve(self.I0_tor_T_LU.solve(rhs2.T).T) - return np.concatenate((rhs1.flatten(), rhs2.flatten())) + return xp.concatenate((rhs1.flatten(), rhs2.flatten())) # ====================================== def apply_IinvT_V3(self, rhs, include_bc=False): @@ -2042,9 +2042,9 @@ def dofs_1(self, fun, include_bc=True, eval_kind="meshgrid", with_subs=True): # apply extraction operator for dofs if include_bc: - dofs = self.P1.dot(np.concatenate((dofs_1.flatten(), dofs_2.flatten(), dofs_3.flatten()))) + dofs = self.P1.dot(xp.concatenate((dofs_1.flatten(), dofs_2.flatten(), dofs_3.flatten()))) else: - dofs = self.P1_0.dot(np.concatenate((dofs_1.flatten(), dofs_2.flatten(), dofs_3.flatten()))) + dofs = self.P1_0.dot(xp.concatenate((dofs_1.flatten(), dofs_2.flatten(), dofs_3.flatten()))) return dofs @@ -2075,9 +2075,9 @@ def dofs_2(self, fun, include_bc=True, eval_kind="meshgrid", with_subs=True): # apply extraction operator for dofs if include_bc: - dofs = self.P2.dot(np.concatenate((dofs_1.flatten(), dofs_2.flatten(), dofs_3.flatten()))) + dofs = self.P2.dot(xp.concatenate((dofs_1.flatten(), dofs_2.flatten(), dofs_3.flatten()))) else: - dofs = self.P2_0.dot(np.concatenate((dofs_1.flatten(), dofs_2.flatten(), dofs_3.flatten()))) + dofs = self.P2_0.dot(xp.concatenate((dofs_1.flatten(), dofs_2.flatten(), dofs_3.flatten()))) return dofs @@ -2122,18 +2122,18 @@ def pi_3(self, fun, include_bc=True, eval_kind="meshgrid", with_subs=True): def assemble_approx_inv(self, tol): if self.approx_Ik_0_inv == False or (self.approx_Ik_0_inv == True and self.approx_Ik_0_tol != tol): # poloidal plane - I0_pol_0_inv_approx = np.linalg.inv(self.I0_pol_0.toarray()) - I1_pol_0_inv_approx = np.linalg.inv(self.I1_pol_0.toarray()) - I2_pol_0_inv_approx = np.linalg.inv(self.I2_pol_0.toarray()) - I3_pol_0_inv_approx = np.linalg.inv(self.I3_pol_0.toarray()) - I0_pol_inv_approx = np.linalg.inv(self.I0_pol.toarray()) + I0_pol_0_inv_approx = xp.linalg.inv(self.I0_pol_0.toarray()) + I1_pol_0_inv_approx = xp.linalg.inv(self.I1_pol_0.toarray()) + I2_pol_0_inv_approx = xp.linalg.inv(self.I2_pol_0.toarray()) + I3_pol_0_inv_approx = xp.linalg.inv(self.I3_pol_0.toarray()) + I0_pol_inv_approx = xp.linalg.inv(self.I0_pol.toarray()) if tol > 1e-14: - I0_pol_0_inv_approx[np.abs(I0_pol_0_inv_approx) < tol] = 0.0 - I1_pol_0_inv_approx[np.abs(I1_pol_0_inv_approx) < tol] = 0.0 - I2_pol_0_inv_approx[np.abs(I2_pol_0_inv_approx) < tol] = 0.0 - I3_pol_0_inv_approx[np.abs(I3_pol_0_inv_approx) < tol] = 0.0 - I0_pol_inv_approx[np.abs(I0_pol_inv_approx) < tol] = 0.0 + I0_pol_0_inv_approx[xp.abs(I0_pol_0_inv_approx) < tol] = 0.0 + I1_pol_0_inv_approx[xp.abs(I1_pol_0_inv_approx) < tol] = 0.0 + I2_pol_0_inv_approx[xp.abs(I2_pol_0_inv_approx) < tol] = 0.0 + I3_pol_0_inv_approx[xp.abs(I3_pol_0_inv_approx) < tol] = 0.0 + I0_pol_inv_approx[xp.abs(I0_pol_inv_approx) < tol] = 0.0 I0_pol_0_inv_approx = spa.csr_matrix(I0_pol_0_inv_approx) I1_pol_0_inv_approx = spa.csr_matrix(I1_pol_0_inv_approx) @@ -2142,12 +2142,12 @@ def assemble_approx_inv(self, tol): I0_pol_inv_approx = spa.csr_matrix(I0_pol_inv_approx) # toroidal direction - I_inv_tor_approx = np.linalg.inv(self.I_tor.toarray()) - H_inv_tor_approx = np.linalg.inv(self.H_tor.toarray()) + I_inv_tor_approx = xp.linalg.inv(self.I_tor.toarray()) + H_inv_tor_approx = xp.linalg.inv(self.H_tor.toarray()) if tol > 1e-14: - I_inv_tor_approx[np.abs(I_inv_tor_approx) < tol] = 0.0 - H_inv_tor_approx[np.abs(H_inv_tor_approx) < tol] = 0.0 + I_inv_tor_approx[xp.abs(I_inv_tor_approx) < tol] = 0.0 + H_inv_tor_approx[xp.abs(H_inv_tor_approx) < tol] = 0.0 I_inv_tor_approx = spa.csr_matrix(I_inv_tor_approx) H_inv_tor_approx = spa.csr_matrix(H_inv_tor_approx) diff --git a/src/struphy/eigenvalue_solvers/spline_space.py b/src/struphy/eigenvalue_solvers/spline_space.py index cffc1902d..b71bc867b 100644 --- a/src/struphy/eigenvalue_solvers/spline_space.py +++ b/src/struphy/eigenvalue_solvers/spline_space.py @@ -9,7 +9,7 @@ import matplotlib import scipy.sparse as spa -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp matplotlib.rcParams.update({"font.size": 16}) import matplotlib.pyplot as plt @@ -50,19 +50,19 @@ class Spline_space_1d: Attributes ---------- - el_b : np.array + el_b : xp.array Element boundaries, equally spaced. delta : float Uniform grid spacing - T : np.array + T : xp.array Knot vector of 0-space. - t : np.arrray + t : xp.arrray Knot vector of 1-space. - greville : np.array + greville : xp.array Greville points. NbaseN : int @@ -71,22 +71,22 @@ class Spline_space_1d: NbaseD : int Dimension of 1-space. - indN : np.array + indN : xp.array Global indices of non-vanishing B-splines in each element in format (element, local basis function) - indD : np.array + indD : xp.array Global indices of non-vanishing M-splines in each element in format (element, local basis function) - pts : np.array + pts : xp.array Global GL quadrature points in format (element, local point). - wts : np.array + wts : xp.array Global GL quadrature weights in format (element, local point). - basisN : np.array + basisN : xp.array N-basis functions evaluated at quadrature points in format (element, local basis function, derivative, local point) - basisD : np.array + basisD : xp.array D-basis functions evaluated at quadrature points in format (element, local basis function, derivative, local point) E0 : csr_matrix @@ -140,7 +140,7 @@ def __init__(self, Nel, p, spl_kind, n_quad=6, bc=["f", "f"]): else: self.bc = bc - self.el_b = np.linspace(0.0, 1.0, Nel + 1) # element boundaries + self.el_b = xp.linspace(0.0, 1.0, Nel + 1) # element boundaries self.delta = 1 / self.Nel # element length self.T = bsp.make_knots(self.el_b, self.p, self.spl_kind) # spline knot vector for B-splines (N) @@ -152,13 +152,13 @@ def __init__(self, Nel, p, spl_kind, n_quad=6, bc=["f", "f"]): self.NbaseD = self.NbaseN - 1 + self.spl_kind # total number of M-splines (D) # global indices of non-vanishing splines in each element in format (Nel, p + 1) - self.indN = (np.indices((self.Nel, self.p + 1 - 0))[1] + np.arange(self.Nel)[:, None]) % self.NbaseN - self.indD = (np.indices((self.Nel, self.p + 1 - 1))[1] + np.arange(self.Nel)[:, None]) % self.NbaseD + self.indN = (xp.indices((self.Nel, self.p + 1 - 0))[1] + xp.arange(self.Nel)[:, None]) % self.NbaseN + self.indD = (xp.indices((self.Nel, self.p + 1 - 1))[1] + xp.arange(self.Nel)[:, None]) % self.NbaseD self.n_quad = n_quad # number of Gauss-Legendre points per grid cell (defined by break points) - self.pts_loc = np.polynomial.legendre.leggauss(self.n_quad)[0] # Gauss-Legendre points (GLQP) in (-1, 1) - self.wts_loc = np.polynomial.legendre.leggauss(self.n_quad)[1] # Gauss-Legendre weights (GLQW) in (-1, 1) + self.pts_loc = xp.polynomial.legendre.leggauss(self.n_quad)[0] # Gauss-Legendre points (GLQP) in (-1, 1) + self.wts_loc = xp.polynomial.legendre.leggauss(self.n_quad)[1] # Gauss-Legendre weights (GLQW) in (-1, 1) # global GLQP in format (element, local point) and total number of GLQP self.pts = bsp.quadrature_grid(self.el_b, self.pts_loc, self.wts_loc)[0] @@ -178,8 +178,8 @@ def __init__(self, Nel, p, spl_kind, n_quad=6, bc=["f", "f"]): d1 = self.NbaseD # boundary operators - self.B0 = np.identity(n1, dtype=float) - self.B1 = np.identity(d1, dtype=float) + self.B0 = xp.identity(n1, dtype=float) + self.B1 = xp.identity(d1, dtype=float) # extraction operators without boundary conditions self.E0 = spa.csr_matrix(self.B0.copy()) @@ -268,16 +268,16 @@ def evaluate_N(self, eta, coeff, kind=0): coeff = self.E0_0.T.dot(coeff) if isinstance(eta, float): - pts = np.array([eta]) - elif isinstance(eta, np.ndarray): + pts = xp.array([eta]) + elif isinstance(eta, xp.ndarray): pts = eta.flatten() - values = np.empty(pts.size, dtype=float) + values = xp.empty(pts.size, dtype=float) eva_1d.evaluate_vector(self.T, self.p, self.indN, coeff, pts, values, kind) if isinstance(eta, float): values = values[0] - elif isinstance(eta, np.ndarray): + elif isinstance(eta, xp.ndarray): values = values.reshape(eta.shape) return values @@ -304,16 +304,16 @@ def evaluate_D(self, eta, coeff): assert coeff.size == self.E1.shape[0] if isinstance(eta, float): - pts = np.array([eta]) - elif isinstance(eta, np.ndarray): + pts = xp.array([eta]) + elif isinstance(eta, xp.ndarray): pts = eta.flatten() - values = np.empty(pts.size, dtype=float) + values = xp.empty(pts.size, dtype=float) eva_1d.evaluate_vector(self.t, self.p - 1, self.indD, coeff, pts, values, 1) if isinstance(eta, float): values = values[0] - elif isinstance(eta, np.ndarray): + elif isinstance(eta, xp.ndarray): values = values.reshape(eta.shape) return values @@ -332,12 +332,12 @@ def plot_splines(self, n_pts=500, which="N"): which basis to plot. 'N', 'D' or 'dN' (optional, default='N') """ - etaplot = np.linspace(0.0, 1.0, n_pts) + etaplot = xp.linspace(0.0, 1.0, n_pts) degree = self.p if which == "N": - coeff = np.zeros(self.NbaseN, dtype=float) + coeff = xp.zeros(self.NbaseN, dtype=float) for i in range(self.NbaseN): coeff[:] = 0.0 @@ -345,7 +345,7 @@ def plot_splines(self, n_pts=500, which="N"): plt.plot(etaplot, self.evaluate_N(etaplot, coeff), label=str(i)) elif which == "D": - coeff = np.zeros(self.NbaseD, dtype=float) + coeff = xp.zeros(self.NbaseD, dtype=float) for i in range(self.NbaseD): coeff[:] = 0.0 @@ -355,7 +355,7 @@ def plot_splines(self, n_pts=500, which="N"): degree = self.p - 1 elif which == "dN": - coeff = np.zeros(self.NbaseN, dtype=float) + coeff = xp.zeros(self.NbaseN, dtype=float) for i in range(self.NbaseN): coeff[:] = 0.0 @@ -370,8 +370,8 @@ def plot_splines(self, n_pts=500, which="N"): else: bcs = "clamped" - (greville,) = plt.plot(self.greville, np.zeros(self.greville.shape), "ro", label="greville") - (breaks,) = plt.plot(self.el_b, np.zeros(self.el_b.shape), "k+", label="breaks") + (greville,) = plt.plot(self.greville, xp.zeros(self.greville.shape), "ro", label="greville") + (breaks,) = plt.plot(self.el_b, xp.zeros(self.el_b.shape), "k+", label="breaks") plt.title(which + f"$^{degree}$-splines, " + bcs + f", Nel={self.Nel}") plt.legend(handles=[greville, breaks]) @@ -555,8 +555,8 @@ def __init__(self, spline_spaces, ck=-1, cx=None, cy=None, n_tor=0, basis_tor="r self.M1_tor = spa.identity(1, format="csr") else: - self.M0_tor = spa.csr_matrix(np.identity(2) / 2) - self.M1_tor = spa.csr_matrix(np.identity(2) / 2) + self.M0_tor = spa.csr_matrix(xp.identity(2) / 2) + self.M1_tor = spa.csr_matrix(xp.identity(2) / 2) else: self.M0_tor = mass_1d.get_M(self.spaces[2], 0, 0) @@ -786,7 +786,7 @@ def apply_M1_ten(self, x, mats): out1 = mats[0][1].dot(mats[0][0].dot(x1).T).T out2 = mats[1][1].dot(mats[1][0].dot(x2).T).T - return np.concatenate((out1.flatten(), out2.flatten())) + return xp.concatenate((out1.flatten(), out2.flatten())) def apply_M2_ten(self, x, mats): """ @@ -798,7 +798,7 @@ def apply_M2_ten(self, x, mats): out1 = mats[0][1].dot(mats[0][0].dot(x1).T).T out2 = mats[1][1].dot(mats[1][0].dot(x2).T).T - return np.concatenate((out1.flatten(), out2.flatten())) + return xp.concatenate((out1.flatten(), out2.flatten())) def apply_M3_ten(self, x, mats): """ @@ -821,7 +821,7 @@ def apply_Mv_ten(self, x, mats): out1 = mats[0][1].dot(mats[0][0].dot(x1).T).T out2 = mats[1][1].dot(mats[1][0].dot(x2).T).T - return np.concatenate((out1.flatten(), out2.flatten())) + return xp.concatenate((out1.flatten(), out2.flatten())) def apply_M0_0_ten(self, x, mats): """ @@ -848,7 +848,7 @@ def apply_M1_0_ten(self, x, mats): mats[1][1].dot(self.B1_tor.T.dot(self.B0_pol.dot(mats[1][0].dot(self.B0_pol.T.dot(x2))).T)) ).T - return np.concatenate((out1.flatten(), out2.flatten())) + return xp.concatenate((out1.flatten(), out2.flatten())) def apply_M2_0_ten(self, x, mats): """ @@ -864,7 +864,7 @@ def apply_M2_0_ten(self, x, mats): mats[1][1].dot(self.B0_tor.T.dot(self.B3_pol.dot(mats[1][0].dot(self.B3_pol.T.dot(x2))).T)) ).T - return np.concatenate((out1.flatten(), out2.flatten())) + return xp.concatenate((out1.flatten(), out2.flatten())) def apply_M3_0_ten(self, x, mats): """ @@ -887,7 +887,7 @@ def apply_Mv_0_ten(self, x, mats): out1 = mats[0][1].dot(self.Bv_pol.dot(mats[0][0].dot(self.Bv_pol.T.dot(x1))).T).T out2 = self.B0_tor.dot(mats[1][1].dot(self.B0_tor.T.dot(mats[1][0].dot(x2).T))).T - return np.concatenate((out1.flatten(), out2.flatten())) + return xp.concatenate((out1.flatten(), out2.flatten())) def __assemble_M0(self, domain, as_tensor=False): """ @@ -1229,7 +1229,7 @@ def extract_1(self, coeff): else: coeff1 = self.E1_0.T.dot(coeff) - coeff1_1, coeff1_2, coeff1_3 = np.split(coeff1, [self.Ntot_1form_cum[0], self.Ntot_1form_cum[1]]) + coeff1_1, coeff1_2, coeff1_3 = xp.split(coeff1, [self.Ntot_1form_cum[0], self.Ntot_1form_cum[1]]) coeff1_1 = coeff1_1.reshape(self.Nbase_1form[0]) coeff1_2 = coeff1_2.reshape(self.Nbase_1form[1]) @@ -1260,7 +1260,7 @@ def extract_2(self, coeff): else: coeff2 = self.E2_0.T.dot(coeff) - coeff2_1, coeff2_2, coeff2_3 = np.split(coeff2, [self.Ntot_2form_cum[0], self.Ntot_2form_cum[1]]) + coeff2_1, coeff2_2, coeff2_3 = xp.split(coeff2, [self.Ntot_2form_cum[0], self.Ntot_2form_cum[1]]) coeff2_1 = coeff2_1.reshape(self.Nbase_2form[0]) coeff2_2 = coeff2_2.reshape(self.Nbase_2form[1]) @@ -1305,7 +1305,7 @@ def extract_v(self, coeff): else: coeffv = self.Ev_0.T.dot(coeff) - coeffv_1, coeffv_2, coeffv_3 = np.split(coeffv, [self.Ntot_0form, 2 * self.Ntot_0form]) + coeffv_1, coeffv_2, coeffv_3 = xp.split(coeffv, [self.Ntot_0form, 2 * self.Ntot_0form]) coeffv_1 = coeffv_1.reshape(self.Nbase_0form) coeffv_2 = coeffv_2.reshape(self.Nbase_0form) @@ -1358,15 +1358,15 @@ def evaluate_NN(self, eta1, eta2, eta3, coeff, which="V0", part="r"): assert coeff.shape[:2] == (self.NbaseN[0], self.NbaseN[1]) # get real and imaginary part - coeff_r = np.real(coeff) - coeff_i = np.imag(coeff) + coeff_r = xp.real(coeff) + coeff_i = xp.imag(coeff) # ------ evaluate FEM field at given points -------- - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor-product evaluation if eta1.ndim == 1: - values_r_1 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) - values_i_1 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_r_1 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_i_1 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) eva_2d.evaluate_tensor_product_2d( self.T[0], @@ -1396,8 +1396,8 @@ def evaluate_NN(self, eta1, eta2, eta3, coeff, which="V0", part="r"): ) if self.n_tor != 0 and self.basis_tor == "r": - values_r_2 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) - values_i_2 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_r_2 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_i_2 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) eva_2d.evaluate_tensor_product_2d( self.T[0], @@ -1428,8 +1428,8 @@ def evaluate_NN(self, eta1, eta2, eta3, coeff, which="V0", part="r"): # matrix evaluation else: - values_r_1 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) - values_i_1 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_r_1 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_i_1 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) eva_2d.evaluate_matrix_2d( self.T[0], @@ -1459,8 +1459,8 @@ def evaluate_NN(self, eta1, eta2, eta3, coeff, which="V0", part="r"): ) if self.n_tor != 0 and self.basis_tor == "r": - values_r_2 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) - values_i_2 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_r_2 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_i_2 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) eva_2d.evaluate_matrix_2d( self.T[0], @@ -1491,15 +1491,15 @@ def evaluate_NN(self, eta1, eta2, eta3, coeff, which="V0", part="r"): # multiply with Fourier basis in third direction if self.n_tor == 0: - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.ones(eta3.shape, dtype=float) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.ones(eta3.shape, dtype=float) else: if self.basis_tor == "r": - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.cos(2 * np.pi * self.n_tor * eta3) - out += (values_r_2 + 1j * values_i_2)[:, :, None] * np.sin(2 * np.pi * self.n_tor * eta3) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.cos(2 * xp.pi * self.n_tor * eta3) + out += (values_r_2 + 1j * values_i_2)[:, :, None] * xp.sin(2 * xp.pi * self.n_tor * eta3) else: - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.exp(1j * 2 * np.pi * self.n_tor * eta3) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.exp(1j * 2 * xp.pi * self.n_tor * eta3) # --------- evaluate FEM field at given point ------- else: @@ -1540,17 +1540,17 @@ def evaluate_NN(self, eta1, eta2, eta3, coeff, which="V0", part="r"): else: if self.basis_tor == "r": - out = (real_1 + 1j * imag_1) * np.cos(2 * np.pi * self.n_tor * eta3) - out += (real_2 + 1j * imag_2) * np.sin(2 * np.pi * self.n_tor * eta3) + out = (real_1 + 1j * imag_1) * xp.cos(2 * xp.pi * self.n_tor * eta3) + out += (real_2 + 1j * imag_2) * xp.sin(2 * xp.pi * self.n_tor * eta3) else: - out = (real_1 + 1j * imag_1) * np.exp(1j * 2 * np.pi * self.n_tor * eta3) + out = (real_1 + 1j * imag_1) * xp.exp(1j * 2 * xp.pi * self.n_tor * eta3) # return real or imaginary part if part == "r": - out = np.real(out) + out = xp.real(out) else: - out = np.imag(out) + out = xp.imag(out) return out @@ -1599,15 +1599,15 @@ def evaluate_DN(self, eta1, eta2, eta3, coeff, which="V1", part="r"): assert coeff.shape[:2] == (self.NbaseD[0], self.NbaseN[1]) # get real and imaginary part - coeff_r = np.real(coeff) - coeff_i = np.imag(coeff) + coeff_r = xp.real(coeff) + coeff_i = xp.imag(coeff) # ------ evaluate FEM field at given points -------- - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor-product evaluation if eta1.ndim == 1: - values_r_1 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) - values_i_1 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_r_1 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_i_1 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) eva_2d.evaluate_tensor_product_2d( self.t[0], @@ -1637,8 +1637,8 @@ def evaluate_DN(self, eta1, eta2, eta3, coeff, which="V1", part="r"): ) if self.n_tor != 0 and self.basis_tor == "r": - values_r_2 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) - values_i_2 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_r_2 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_i_2 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) eva_2d.evaluate_tensor_product_2d( self.t[0], @@ -1669,8 +1669,8 @@ def evaluate_DN(self, eta1, eta2, eta3, coeff, which="V1", part="r"): # matrix evaluation else: - values_r_1 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) - values_i_1 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_r_1 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_i_1 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) eva_2d.evaluate_matrix_2d( self.t[0], @@ -1700,8 +1700,8 @@ def evaluate_DN(self, eta1, eta2, eta3, coeff, which="V1", part="r"): ) if self.n_tor != 0 and self.basis_tor == "r": - values_r_2 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) - values_i_2 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_r_2 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_i_2 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) eva_2d.evaluate_matrix_2d( self.t[0], @@ -1732,15 +1732,15 @@ def evaluate_DN(self, eta1, eta2, eta3, coeff, which="V1", part="r"): # multiply with Fourier basis in third direction if self.n_tor == 0: - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.ones(eta3.shape, dtype=float) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.ones(eta3.shape, dtype=float) else: if self.basis_tor == "r": - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.cos(2 * np.pi * self.n_tor * eta3) - out += (values_r_2 + 1j * values_i_2)[:, :, None] * np.sin(2 * np.pi * self.n_tor * eta3) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.cos(2 * xp.pi * self.n_tor * eta3) + out += (values_r_2 + 1j * values_i_2)[:, :, None] * xp.sin(2 * xp.pi * self.n_tor * eta3) else: - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.exp(1j * 2 * np.pi * self.n_tor * eta3) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.exp(1j * 2 * xp.pi * self.n_tor * eta3) # --------- evaluate FEM field at given point ------- else: @@ -1797,17 +1797,17 @@ def evaluate_DN(self, eta1, eta2, eta3, coeff, which="V1", part="r"): else: if self.basis_tor == "r": - out = (real_1 + 1j * imag_1) * np.cos(2 * np.pi * self.n_tor * eta3) - out += (real_2 + 1j * imag_2) * np.sin(2 * np.pi * self.n_tor * eta3) + out = (real_1 + 1j * imag_1) * xp.cos(2 * xp.pi * self.n_tor * eta3) + out += (real_2 + 1j * imag_2) * xp.sin(2 * xp.pi * self.n_tor * eta3) else: - out = (real_1 + 1j * imag_1) * np.exp(1j * 2 * np.pi * self.n_tor * eta3) + out = (real_1 + 1j * imag_1) * xp.exp(1j * 2 * xp.pi * self.n_tor * eta3) # return real or imaginary part if part == "r": - out = np.real(out) + out = xp.real(out) else: - out = np.imag(out) + out = xp.imag(out) return out @@ -1856,15 +1856,15 @@ def evaluate_ND(self, eta1, eta2, eta3, coeff, which="V2", part="r"): assert coeff.shape[:2] == (self.NbaseN[0], self.NbaseD[1]) # get real and imaginary part - coeff_r = np.real(coeff) - coeff_i = np.imag(coeff) + coeff_r = xp.real(coeff) + coeff_i = xp.imag(coeff) # ------ evaluate FEM field at given points -------- - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor-product evaluation if eta1.ndim == 1: - values_r_1 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) - values_i_1 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_r_1 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_i_1 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) eva_2d.evaluate_tensor_product_2d( self.T[0], @@ -1894,8 +1894,8 @@ def evaluate_ND(self, eta1, eta2, eta3, coeff, which="V2", part="r"): ) if self.n_tor != 0 and self.basis_tor == "r": - values_r_2 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) - values_i_2 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_r_2 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_i_2 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) eva_2d.evaluate_tensor_product_2d( self.T[0], @@ -1926,8 +1926,8 @@ def evaluate_ND(self, eta1, eta2, eta3, coeff, which="V2", part="r"): # matrix evaluation else: - values_r_1 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) - values_i_1 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_r_1 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_i_1 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) eva_2d.evaluate_matrix_2d( self.T[0], @@ -1957,8 +1957,8 @@ def evaluate_ND(self, eta1, eta2, eta3, coeff, which="V2", part="r"): ) if self.n_tor != 0 and self.basis_tor == "r": - values_r_2 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) - values_i_2 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_r_2 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_i_2 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) eva_2d.evaluate_matrix_2d( self.T[0], @@ -1989,15 +1989,15 @@ def evaluate_ND(self, eta1, eta2, eta3, coeff, which="V2", part="r"): # multiply with Fourier basis in third direction if self.n_tor == 0: - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.ones(eta3.shape, dtype=float) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.ones(eta3.shape, dtype=float) else: if self.basis_tor == "r": - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.cos(2 * np.pi * self.n_tor * eta3) - out += (values_r_2 + 1j * values_i_2)[:, :, None] * np.sin(2 * np.pi * self.n_tor * eta3) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.cos(2 * xp.pi * self.n_tor * eta3) + out += (values_r_2 + 1j * values_i_2)[:, :, None] * xp.sin(2 * xp.pi * self.n_tor * eta3) else: - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.exp(1j * 2 * np.pi * self.n_tor * eta3) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.exp(1j * 2 * xp.pi * self.n_tor * eta3) # --------- evaluate FEM field at given point ------- else: @@ -2054,17 +2054,17 @@ def evaluate_ND(self, eta1, eta2, eta3, coeff, which="V2", part="r"): else: if self.basis_tor == "r": - out = (real_1 + 1j * imag_1) * np.cos(2 * np.pi * self.n_tor * eta3) - out += (real_2 + 1j * imag_2) * np.sin(2 * np.pi * self.n_tor * eta3) + out = (real_1 + 1j * imag_1) * xp.cos(2 * xp.pi * self.n_tor * eta3) + out += (real_2 + 1j * imag_2) * xp.sin(2 * xp.pi * self.n_tor * eta3) else: - out = (real_1 + 1j * imag_1) * np.exp(1j * 2 * np.pi * self.n_tor * eta3) + out = (real_1 + 1j * imag_1) * xp.exp(1j * 2 * xp.pi * self.n_tor * eta3) # return real or imaginary part if part == "r": - out = np.real(out) + out = xp.real(out) else: - out = np.imag(out) + out = xp.imag(out) return out @@ -2116,15 +2116,15 @@ def evaluate_DD(self, eta1, eta2, eta3, coeff, which="V3", part="r"): assert coeff.shape[:2] == (self.NbaseD[0], self.NbaseD[1]) # get real and imaginary part - coeff_r = np.real(coeff) - coeff_i = np.imag(coeff) + coeff_r = xp.real(coeff) + coeff_i = xp.imag(coeff) # ------ evaluate FEM field at given points -------- - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor-product evaluation if eta1.ndim == 1: - values_r_1 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) - values_i_1 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_r_1 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_i_1 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) eva_2d.evaluate_tensor_product_2d( self.t[0], @@ -2154,8 +2154,8 @@ def evaluate_DD(self, eta1, eta2, eta3, coeff, which="V3", part="r"): ) if self.n_tor != 0 and self.basis_tor == "r": - values_r_2 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) - values_i_2 = np.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_r_2 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) + values_i_2 = xp.empty((eta1.shape[0], eta2.shape[0]), dtype=float) eva_2d.evaluate_tensor_product_2d( self.t[0], @@ -2186,8 +2186,8 @@ def evaluate_DD(self, eta1, eta2, eta3, coeff, which="V3", part="r"): # matrix evaluation else: - values_r_1 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) - values_i_1 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_r_1 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_i_1 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) eva_2d.evaluate_matrix_2d( self.t[0], @@ -2217,8 +2217,8 @@ def evaluate_DD(self, eta1, eta2, eta3, coeff, which="V3", part="r"): ) if self.n_tor != 0 and self.basis_tor == "r": - values_r_2 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) - values_i_2 = np.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_r_2 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) + values_i_2 = xp.empty((eta1.shape[0], eta2.shape[1]), dtype=float) eva_2d.evaluate_matrix_2d( self.t[0], @@ -2249,15 +2249,15 @@ def evaluate_DD(self, eta1, eta2, eta3, coeff, which="V3", part="r"): # multiply with Fourier basis in third direction if self.n_tor == 0: - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.ones(eta3.shape, dtype=float) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.ones(eta3.shape, dtype=float) else: if self.basis_tor == "r": - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.cos(2 * np.pi * self.n_tor * eta3) - out += (values_r_2 + 1j * values_i_2)[:, :, None] * np.sin(2 * np.pi * self.n_tor * eta3) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.cos(2 * xp.pi * self.n_tor * eta3) + out += (values_r_2 + 1j * values_i_2)[:, :, None] * xp.sin(2 * xp.pi * self.n_tor * eta3) else: - out = (values_r_1 + 1j * values_i_1)[:, :, None] * np.exp(1j * 2 * np.pi * self.n_tor * eta3) + out = (values_r_1 + 1j * values_i_1)[:, :, None] * xp.exp(1j * 2 * xp.pi * self.n_tor * eta3) # --------- evaluate FEM field at given point ------- else: @@ -2314,17 +2314,17 @@ def evaluate_DD(self, eta1, eta2, eta3, coeff, which="V3", part="r"): else: if self.basis_tor == "r": - out = (real_1 + 1j * imag_1) * np.cos(2 * np.pi * self.n_tor * eta3) - out += (real_2 + 1j * imag_2) * np.sin(2 * np.pi * self.n_tor * eta3) + out = (real_1 + 1j * imag_1) * xp.cos(2 * xp.pi * self.n_tor * eta3) + out += (real_2 + 1j * imag_2) * xp.sin(2 * xp.pi * self.n_tor * eta3) else: - out = (real_1 + 1j * imag_1) * np.exp(1j * 2 * np.pi * self.n_tor * eta3) + out = (real_1 + 1j * imag_1) * xp.exp(1j * 2 * xp.pi * self.n_tor * eta3) # return real or imaginary part if part == "r": - out = np.real(out) + out = xp.real(out) else: - out = np.imag(out) + out = xp.imag(out) return out @@ -2335,13 +2335,13 @@ def evaluate_NNN(self, eta1, eta2, eta3, coeff): Parameters ---------- - eta1 : double or np.ndarray + eta1 : double or xp.ndarray 1st component of logical evaluation point - eta2 : double or np.ndarray + eta2 : double or xp.ndarray 2nd component of logical evaluation point - eta3 : double or np.ndarray + eta3 : double or xp.ndarray 3rd component of logical evaluation point coeff : array_like @@ -2356,10 +2356,10 @@ def evaluate_NNN(self, eta1, eta2, eta3, coeff): if coeff.ndim == 1: coeff = self.extract_0(coeff) - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor-product evaluation if eta1.ndim == 1: - values = np.empty((eta1.size, eta2.size, eta3.size), dtype=float) + values = xp.empty((eta1.size, eta2.size, eta3.size), dtype=float) eva_3d.evaluate_tensor_product( self.T[0], self.T[1], @@ -2380,7 +2380,7 @@ def evaluate_NNN(self, eta1, eta2, eta3, coeff): # matrix evaluation else: - values = np.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) + values = xp.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) # `eta1` is a sparse meshgrid. if max(eta1.shape) == eta1.size: eva_3d.evaluate_sparse( @@ -2450,13 +2450,13 @@ def evaluate_DNN(self, eta1, eta2, eta3, coeff): Parameters ---------- - eta1 : double or np.ndarray + eta1 : double or xp.ndarray 1st component of logical evaluation point - eta2 : double or np.ndarray + eta2 : double or xp.ndarray 2nd component of logical evaluation point - eta3 : double or np.ndarray + eta3 : double or xp.ndarray 3rd component of logical evaluation point coeff : array_like @@ -2471,10 +2471,10 @@ def evaluate_DNN(self, eta1, eta2, eta3, coeff): if coeff.ndim == 1: coeff = self.extract_1(coeff)[0] - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor product evaluation if eta1.ndim == 1: - values = np.empty((eta1.size, eta2.size, eta3.size), dtype=float) + values = xp.empty((eta1.size, eta2.size, eta3.size), dtype=float) eva_3d.evaluate_tensor_product( self.t[0], self.T[1], @@ -2495,7 +2495,7 @@ def evaluate_DNN(self, eta1, eta2, eta3, coeff): # matrix evaluation else: - values = np.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) + values = xp.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) # `eta1` is a sparse meshgrid. if max(eta1.shape) == eta1.size: eva_3d.evaluate_sparse( @@ -2564,13 +2564,13 @@ def evaluate_NDN(self, eta1, eta2, eta3, coeff): Parameters ---------- - eta1 : double or np.ndarray + eta1 : double or xp.ndarray 1st component of logical evaluation point - eta2 : double or np.ndarray + eta2 : double or xp.ndarray 2nd component of logical evaluation point - eta3 : double or np.ndarray + eta3 : double or xp.ndarray 3rd component of logical evaluation point coeff : array_like @@ -2585,10 +2585,10 @@ def evaluate_NDN(self, eta1, eta2, eta3, coeff): if coeff.ndim == 1: coeff = self.extract_1(coeff)[1] - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor product evaluation if eta1.ndim == 1: - values = np.empty((eta1.size, eta2.size, eta3.size), dtype=float) + values = xp.empty((eta1.size, eta2.size, eta3.size), dtype=float) eva_3d.evaluate_tensor_product( self.T[0], self.t[1], @@ -2609,7 +2609,7 @@ def evaluate_NDN(self, eta1, eta2, eta3, coeff): # matrix evaluation else: - values = np.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) + values = xp.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) # `eta1` is a sparse meshgrid. if max(eta1.shape) == eta1.size: eva_3d.evaluate_sparse( @@ -2678,13 +2678,13 @@ def evaluate_NND(self, eta1, eta2, eta3, coeff): Parameters ---------- - eta1 : double or np.ndarray + eta1 : double or xp.ndarray 1st component of logical evaluation point - eta2 : double or np.ndarray + eta2 : double or xp.ndarray 2nd component of logical evaluation point - eta3 : double or np.ndarray + eta3 : double or xp.ndarray 3rd component of logical evaluation point coeff : array_like @@ -2699,10 +2699,10 @@ def evaluate_NND(self, eta1, eta2, eta3, coeff): if coeff.ndim == 1: coeff = self.extract_1(coeff)[2] - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor product evaluation if eta1.ndim == 1: - values = np.empty((eta1.size, eta2.size, eta3.size), dtype=float) + values = xp.empty((eta1.size, eta2.size, eta3.size), dtype=float) eva_3d.evaluate_tensor_product( self.T[0], self.T[1], @@ -2723,7 +2723,7 @@ def evaluate_NND(self, eta1, eta2, eta3, coeff): # matrix evaluation else: - values = np.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) + values = xp.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) # `eta1` is a sparse meshgrid. if max(eta1.shape) == eta1.size: eva_3d.evaluate_sparse( @@ -2792,13 +2792,13 @@ def evaluate_NDD(self, eta1, eta2, eta3, coeff): Parameters ---------- - eta1 : double or np.ndarray + eta1 : double or xp.ndarray 1st component of logical evaluation point - eta2 : double or np.ndarray + eta2 : double or xp.ndarray 2nd component of logical evaluation point - eta3 : double or np.ndarray + eta3 : double or xp.ndarray 3rd component of logical evaluation point coeff : array_like @@ -2813,10 +2813,10 @@ def evaluate_NDD(self, eta1, eta2, eta3, coeff): if coeff.ndim == 1: coeff = self.extract_2(coeff)[0] - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor product evaluation if eta1.ndim == 1: - values = np.empty((eta1.size, eta2.size, eta3.size), dtype=float) + values = xp.empty((eta1.size, eta2.size, eta3.size), dtype=float) eva_3d.evaluate_tensor_product( self.T[0], self.t[1], @@ -2837,7 +2837,7 @@ def evaluate_NDD(self, eta1, eta2, eta3, coeff): # matrix evaluation else: - values = np.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) + values = xp.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) # `eta1` is a sparse meshgrid. if max(eta1.shape) == eta1.size: eva_3d.evaluate_sparse( @@ -2906,13 +2906,13 @@ def evaluate_DND(self, eta1, eta2, eta3, coeff): Parameters ---------- - eta1 : double or np.ndarray + eta1 : double or xp.ndarray 1st component of logical evaluation point - eta2 : double or np.ndarray + eta2 : double or xp.ndarray 2nd component of logical evaluation point - eta3 : double or np.ndarray + eta3 : double or xp.ndarray 3rd component of logical evaluation point coeff : array_like @@ -2927,10 +2927,10 @@ def evaluate_DND(self, eta1, eta2, eta3, coeff): if coeff.ndim == 1: coeff = self.extract_2(coeff)[1] - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor product evaluation if eta1.ndim == 1: - values = np.empty((eta1.size, eta2.size, eta3.size), dtype=float) + values = xp.empty((eta1.size, eta2.size, eta3.size), dtype=float) eva_3d.evaluate_tensor_product( self.t[0], self.T[1], @@ -2951,7 +2951,7 @@ def evaluate_DND(self, eta1, eta2, eta3, coeff): # matrix evaluation else: - values = np.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) + values = xp.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) # `eta1` is a sparse meshgrid. if max(eta1.shape) == eta1.size: eva_3d.evaluate_sparse( @@ -3020,13 +3020,13 @@ def evaluate_DDN(self, eta1, eta2, eta3, coeff): Parameters ---------- - eta1 : double or np.ndarray + eta1 : double or xp.ndarray 1st component of logical evaluation point - eta2 : double or np.ndarray + eta2 : double or xp.ndarray 2nd component of logical evaluation point - eta3 : double or np.ndarray + eta3 : double or xp.ndarray 3rd component of logical evaluation point coeff : array_like @@ -3041,10 +3041,10 @@ def evaluate_DDN(self, eta1, eta2, eta3, coeff): if coeff.ndim == 1: coeff = self.extract_2(coeff)[2] - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor product evaluation if eta1.ndim == 1: - values = np.empty((eta1.size, eta2.size, eta3.size), dtype=float) + values = xp.empty((eta1.size, eta2.size, eta3.size), dtype=float) eva_3d.evaluate_tensor_product( self.t[0], self.t[1], @@ -3065,7 +3065,7 @@ def evaluate_DDN(self, eta1, eta2, eta3, coeff): # matrix evaluation else: - values = np.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) + values = xp.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) # `eta1` is a sparse meshgrid. if max(eta1.shape) == eta1.size: eva_3d.evaluate_sparse( @@ -3134,13 +3134,13 @@ def evaluate_DDD(self, eta1, eta2, eta3, coeff): Parameters ---------- - eta1 : double or np.ndarray + eta1 : double or xp.ndarray 1st component of logical evaluation point - eta2 : double or np.ndarray + eta2 : double or xp.ndarray 2nd component of logical evaluation point - eta3 : double or np.ndarray + eta3 : double or xp.ndarray 3rd component of logical evaluation point coeff : array_like @@ -3155,10 +3155,10 @@ def evaluate_DDD(self, eta1, eta2, eta3, coeff): if coeff.ndim == 1: coeff = self.extract_3(coeff) - if isinstance(eta1, np.ndarray): + if isinstance(eta1, xp.ndarray): # tensor product evaluation if eta1.ndim == 1: - values = np.empty((eta1.size, eta2.size, eta3.size), dtype=float) + values = xp.empty((eta1.size, eta2.size, eta3.size), dtype=float) eva_3d.evaluate_tensor_product( self.t[0], self.t[1], @@ -3179,7 +3179,7 @@ def evaluate_DDD(self, eta1, eta2, eta3, coeff): # matrix evaluation else: - values = np.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) + values = xp.empty((eta1.shape[0], eta2.shape[1], eta3.shape[2]), dtype=float) # `eta1` is a sparse meshgrid. if max(eta1.shape) == eta1.size: eva_3d.evaluate_sparse( diff --git a/src/struphy/examples/_draw_parallel.py b/src/struphy/examples/_draw_parallel.py index b2f3cef83..62fae8cd2 100644 --- a/src/struphy/examples/_draw_parallel.py +++ b/src/struphy/examples/_draw_parallel.py @@ -3,7 +3,7 @@ from struphy.feec.psydac_derham import Derham from struphy.geometry import domains from struphy.pic.particles import Particles6D -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def main(): @@ -69,18 +69,18 @@ def main(): ) # are all markers in the correct domain? - conds = np.logical_and( + conds = xp.logical_and( particles.markers[:, :3] > derham.domain_array[rank, 0::3], particles.markers[:, :3] < derham.domain_array[rank, 1::3], ) holes = particles.markers[:, 0] == -1.0 - stay = np.all(conds, axis=1) + stay = xp.all(conds, axis=1) - error_mks = particles.markers[np.logical_and(~stay, ~holes)] + error_mks = particles.markers[xp.logical_and(~stay, ~holes)] print( - f"rank {rank} | markers not on correct process: {np.nonzero(np.logical_and(~stay, ~holes))} \ + f"rank {rank} | markers not on correct process: {xp.nonzero(xp.logical_and(~stay, ~holes))} \ \n corresponding positions:\n {error_mks[:, :3]}" ) diff --git a/src/struphy/examples/restelli2018/callables.py b/src/struphy/examples/restelli2018/callables.py index 4a4b3d5c5..54404ab27 100644 --- a/src/struphy/examples/restelli2018/callables.py +++ b/src/struphy/examples/restelli2018/callables.py @@ -1,6 +1,6 @@ "Analytical callables needed for the simulation of the Two-Fluid Quasi-Neutral Model by Restelli." -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class RestelliForcingTerm: @@ -74,9 +74,9 @@ def __init__(self, nu=1.0, R0=2.0, a=1.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0, self._eps_norm = eps def __call__(self, x, y, z): - R = np.sqrt(x**2 + y**2) - R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(-y, x) + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) force_Z = self._nu * ( self._alpha * (self._R0 - 4 * R) / (self._a * self._R0 * R) - self._beta * self._Bp * self._R0**2 / (self._B0 * self._a * R**3) @@ -197,31 +197,31 @@ def __init__( def __call__(self, x, y, z): A = self._alpha / (self._a * self._R0) C = self._beta * self._Bp * self._R0 / (self._B0 * self._a) - R = np.sqrt(x**2 + y**2) - R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(-y, x) + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) if self._species == "Ions": """Forceterm for ions on the right hand side.""" if self._dimension == "2D": fx = ( - -2.0 * np.pi * np.sin(2 * np.pi * x) - + np.cos(2 * np.pi * x) * np.cos(2 * np.pi * y) * self._B0 / self._eps_norm - - self._nu * 8.0 * np.pi**2 * np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y) + -2.0 * xp.pi * xp.sin(2 * xp.pi * x) + + xp.cos(2 * xp.pi * x) * xp.cos(2 * xp.pi * y) * self._B0 / self._eps_norm + - self._nu * 8.0 * xp.pi**2 * xp.sin(2 * xp.pi * x) * xp.sin(2 * xp.pi * y) ) fy = ( - 2.0 * np.pi * np.cos(2 * np.pi * y) - - np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y) * self._B0 / self._eps_norm - - self._nu * 8.0 * np.pi**2 * np.cos(2 * np.pi * x) * np.cos(2 * np.pi * y) + 2.0 * xp.pi * xp.cos(2 * xp.pi * y) + - xp.sin(2 * xp.pi * x) * xp.sin(2 * xp.pi * y) * self._B0 / self._eps_norm + - self._nu * 8.0 * xp.pi**2 * xp.cos(2 * xp.pi * x) * xp.cos(2 * xp.pi * y) ) fz = 0.0 * x elif self._dimension == "1D": fx = ( - 2.0 * np.pi * np.cos(2 * np.pi * x) - + self._nu * 4.0 * np.pi**2 * np.sin(2 * np.pi * x) - + (np.sin(2 * np.pi * x) + 1.0) / self._dt + 2.0 * xp.pi * xp.cos(2 * xp.pi * x) + + self._nu * 4.0 * xp.pi**2 * xp.sin(2 * xp.pi * x) + + (xp.sin(2 * xp.pi * x) + 1.0) / self._dt ) - fy = (np.sin(2 * np.pi * x) + 1.0) * self._B0 / self._eps_norm + fy = (xp.sin(2 * xp.pi * x) + 1.0) * self._B0 / self._eps_norm fz = 0.0 * x elif self._dimension == "Tokamak": @@ -234,8 +234,8 @@ def __call__(self, x, y, z): fZ = self._alpha * self._B0 * z / self._a + A * self._R0 / R * ((R - self._R0) * self._B0) fphi = A * self._R0 * self._Bp / (self._a * R**2) * ((R - self._R0) ** 2 + z**2) - fx = np.cos(phi) * fR - R * np.sin(phi) * fphi - fy = -np.sin(phi) * fR - R * np.cos(phi) * fphi + fx = xp.cos(phi) * fR - R * xp.sin(phi) * fphi + fy = -xp.sin(phi) * fR - R * xp.cos(phi) * fphi fz = fZ if self._comp == "0": @@ -251,26 +251,26 @@ def __call__(self, x, y, z): """Forceterm for electrons on the right hand side.""" if self._dimension == "2D": fx = ( - 2.0 * np.pi * np.sin(2 * np.pi * x) - - np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y) * self._B0 / self._eps_norm - - self._nu_e * 32.0 * np.pi**2 * np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y) - - self._stab_sigma * (-np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y)) + 2.0 * xp.pi * xp.sin(2 * xp.pi * x) + - xp.cos(4 * xp.pi * x) * xp.cos(4 * xp.pi * y) * self._B0 / self._eps_norm + - self._nu_e * 32.0 * xp.pi**2 * xp.sin(4 * xp.pi * x) * xp.sin(4 * xp.pi * y) + - self._stab_sigma * (-xp.sin(4 * xp.pi * x) * xp.sin(4 * xp.pi * y)) ) fy = ( - -2.0 * np.pi * np.cos(2 * np.pi * y) - + np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y) * self._B0 / self._eps_norm - - self._nu_e * 32.0 * np.pi**2 * np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y) - - self._stab_sigma * (-np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y)) + -2.0 * xp.pi * xp.cos(2 * xp.pi * y) + + xp.sin(4 * xp.pi * x) * xp.sin(4 * xp.pi * y) * self._B0 / self._eps_norm + - self._nu_e * 32.0 * xp.pi**2 * xp.cos(4 * xp.pi * x) * xp.cos(4 * xp.pi * y) + - self._stab_sigma * (-xp.cos(4 * xp.pi * x) * xp.cos(4 * xp.pi * y)) ) fz = 0.0 * x elif self._dimension == "1D": fx = ( - -2.0 * np.pi * np.cos(2 * np.pi * x) - + self._nu_e * 4.0 * np.pi**2 * np.sin(2 * np.pi * x) - - self._stab_sigma * np.sin(2 * np.pi * x) + -2.0 * xp.pi * xp.cos(2 * xp.pi * x) + + self._nu_e * 4.0 * xp.pi**2 * xp.sin(2 * xp.pi * x) + - self._stab_sigma * xp.sin(2 * xp.pi * x) ) - fy = -np.sin(2 * np.pi * x) * self._B0 / self._eps_norm + fy = -xp.sin(2 * xp.pi * x) * self._B0 / self._eps_norm fz = 0.0 * x elif self._dimension == "Tokamak": @@ -283,8 +283,8 @@ def __call__(self, x, y, z): fZ = -self._alpha * self._B0 * z / self._a - A * self._R0 / R * ((R - self._R0) * self._B0) fphi = -A * self._R0 * self._Bp / (self._a * R**2) * ((R - self._R0) ** 2 + z**2) - fx = np.cos(phi) * fR - R * np.sin(phi) * fphi - fy = -np.sin(phi) * fR - R * np.cos(phi) * fphi + fx = xp.cos(phi) * fR - R * xp.sin(phi) * fphi + fy = -xp.sin(phi) * fR - R * xp.cos(phi) * fphi fz = fZ if self._comp == "0": diff --git a/src/struphy/feec/basis_projection_ops.py b/src/struphy/feec/basis_projection_ops.py index 8d4f24d54..4a8b804b0 100644 --- a/src/struphy/feec/basis_projection_ops.py +++ b/src/struphy/feec/basis_projection_ops.py @@ -14,7 +14,7 @@ from struphy.feec.utilities import RotationMatrix from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.linear_operators import PolarExtractionOperator -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel @@ -51,7 +51,7 @@ def __init__(self, derham, domain, verbose=True, **weights): self._rank = derham.comm.Get_rank() if derham.comm is not None else 0 - if np.any([p == 1 and Nel > 1 for p, Nel in zip(derham.p, derham.Nel)]): + if xp.any([p == 1 and Nel > 1 for p, Nel in zip(derham.p, derham.Nel)]): if MPI.COMM_WORLD.Get_rank() == 0: print( f'\nWARNING: Class "BasisProjectionOperators" called with p={derham.p} (interpolation of piece-wise constants should be avoided).', @@ -1051,11 +1051,11 @@ def __init__( if isinstance(V, TensorFemSpace): self._Vspaces = [V.coeff_space] self._V1ds = [V.spaces] - self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) + self._VNbasis = xp.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) else: self._Vspaces = V.coeff_space self._V1ds = [comp.spaces for comp in V.spaces] - self._VNbasis = np.array( + self._VNbasis = xp.array( [ [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], [ @@ -1283,7 +1283,7 @@ def assemble(self, verbose=False): self._pds, self._periodic, self._p, - np.array([col0, col1, col2]), + xp.array([col0, col1, col2]), self._VNbasis, self._mat._data, coeff, @@ -1358,7 +1358,7 @@ def assemble(self, verbose=False): self._pds, self._periodic, self._p, - np.array( + xp.array( [ col0, col1, @@ -1438,7 +1438,7 @@ def assemble(self, verbose=False): self._pds[h], self._periodic, self._p, - np.array( + xp.array( [ col0, col1, @@ -1539,7 +1539,7 @@ def assemble(self, verbose=False): self._pds[h], self._periodic, self._p, - np.array( + xp.array( [ col0, col1, @@ -1613,7 +1613,7 @@ class BasisProjectionOperator(LinOpWithTransp): Finite element spline space (domain, input space). weights : list - Weight function(s) (callables or np.ndarrays) in a 2d list of shape corresponding to number of components of domain/codomain. + Weight function(s) (callables or xp.ndarrays) in a 2d list of shape corresponding to number of components of domain/codomain. V_extraction_op : PolarExtractionOperator | IdentityOperator Extraction operator to polar sub-space of V. @@ -1889,7 +1889,7 @@ def update_weights(self, weights): Parameters ---------- weights : list - Weight function(s) (callables or np.ndarrays) in a 2d list of shape corresponding to number of components of domain/codomain. + Weight function(s) (callables or xp.ndarrays) in a 2d list of shape corresponding to number of components of domain/codomain. """ self._weights = weights @@ -1945,13 +1945,13 @@ def assemble(self, weights=None, verbose=False): # input vector space (domain), column of block for j, (Vspace, V1d, loc_weight) in enumerate(zip(_Vspaces, _V1ds, weight_line)): - _starts_in = np.array(Vspace.starts) - _ends_in = np.array(Vspace.ends) - _pads_in = np.array(Vspace.pads) + _starts_in = xp.array(Vspace.starts) + _ends_in = xp.array(Vspace.ends) + _pads_in = xp.array(Vspace.pads) - _starts_out = np.array(Wspace.starts) - _ends_out = np.array(Wspace.ends) - _pads_out = np.array(Wspace.pads) + _starts_out = xp.array(Wspace.starts) + _ends_out = xp.array(Wspace.ends) + _pads_out = xp.array(Wspace.pads) # use cached information if asked if self._use_cache: @@ -1998,21 +1998,21 @@ def assemble(self, weights=None, verbose=False): # Evaluate weight function at quadrature points # evaluate weight at quadrature points if callable(loc_weight): - PTS = np.meshgrid(*_ptsG, indexing="ij") + PTS = xp.meshgrid(*_ptsG, indexing="ij") mat_w = loc_weight(*PTS).copy() - elif isinstance(loc_weight, np.ndarray): + elif isinstance(loc_weight, xp.ndarray): assert loc_weight.shape == (len(_ptsG[0]), len(_ptsG[1]), len(_ptsG[2])) mat_w = loc_weight elif loc_weight is not None: raise TypeError( - "weights must be np.ndarray, callable or None", + "weights must be xp.ndarray, callable or None", ) # Call the kernel if weight function is not zero or in the scalar case # to avoid calling _block of a StencilMatrix in the else - not_weight_zero = np.array( - int(loc_weight is not None and np.any(np.abs(mat_w) > 1e-14)), + not_weight_zero = xp.array( + int(loc_weight is not None and xp.any(xp.abs(mat_w) > 1e-14)), ) if self._mpi_comm is not None: diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 329ad1e9e..61a08e90e 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -10,7 +10,7 @@ from struphy.feec.utilities import apply_essential_bc_to_array from struphy.polar.basic import PolarDerhamSpace -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class LinOpWithTransp(LinearOperator): @@ -66,14 +66,14 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): if is_sparse == False: if out is None: # We declare the matrix form of our linear operator - out = np.zeros([self.codomain.dimension, self.domain.dimension], dtype=self.dtype) + out = xp.zeros([self.codomain.dimension, self.domain.dimension], dtype=self.dtype) else: - assert isinstance(out, np.ndarray) + assert isinstance(out, xp.ndarray) assert out.shape[0] == self.codomain.dimension assert out.shape[1] == self.domain.dimension # We use this matrix to store the partial results that we shall combine into the final matrix with a reduction at the end - result = np.zeros((self.codomain.dimension, self.domain.dimension), dtype=self.dtype) + result = xp.zeros((self.codomain.dimension, self.domain.dimension), dtype=self.dtype) else: if out is not None: raise Exception("If is_sparse is True then out must be set to None.") @@ -97,10 +97,10 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): ndim = [sp.ndim for sp in self.domain.spaces] # First each rank is going to need to know the starts and ends of all other ranks - startsarr = np.array([starts[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) + startsarr = xp.array([starts[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) # Create an array to store gathered data from all ranks - allstarts = np.empty(size * len(startsarr), dtype=int) + allstarts = xp.empty(size * len(startsarr), dtype=int) # Use Allgather to gather 'starts' from all ranks into 'allstarts' if comm is None or isinstance(comm, MockComm): @@ -111,9 +111,9 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): # Reshape 'allstarts' to have 9 columns and 'size' rows allstarts = allstarts.reshape((size, len(startsarr))) - endsarr = np.array([ends[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) + endsarr = xp.array([ends[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) # Create an array to store gathered data from all ranks - allends = np.empty(size * len(endsarr), dtype=int) + allends = xp.empty(size * len(endsarr), dtype=int) # Use Allgather to gather 'ends' from all ranks into 'allends' if comm is None or isinstance(comm, MockComm): @@ -148,13 +148,13 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): self.dot(v, out=tmp2) # Compute to which column this iteration belongs col = spoint - col += np.ravel_multi_index(i, npts[h]) + col += xp.ravel_multi_index(i, npts[h]) if is_sparse == False: result[:, col] = tmp2.toarray() else: aux = tmp2.toarray() # We now need to now which entries on tmp2 are non-zero and store then in our data list - for l in np.where(aux != 0)[0]: + for l in xp.where(aux != 0)[0]: data.append(aux[l]) colarr.append(col) row.append(l) @@ -179,9 +179,9 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): ndim = self.domain.ndim # First each rank is going to need to know the starts and ends of all other ranks - startsarr = np.array([starts[j] for j in range(ndim)], dtype=int) + startsarr = xp.array([starts[j] for j in range(ndim)], dtype=int) # Create an array to store gathered data from all ranks - allstarts = np.empty(size * len(startsarr), dtype=int) + allstarts = xp.empty(size * len(startsarr), dtype=int) # Use Allgather to gather 'starts' from all ranks into 'allstarts' if comm is None or isinstance(comm, MockComm): @@ -192,9 +192,9 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): # Reshape 'allstarts' to have 3 columns and 'size' rows allstarts = allstarts.reshape((size, len(startsarr))) - endsarr = np.array([ends[j] for j in range(ndim)], dtype=int) + endsarr = xp.array([ends[j] for j in range(ndim)], dtype=int) # Create an array to store gathered data from all ranks - allends = np.empty(size * len(endsarr), dtype=int) + allends = xp.empty(size * len(endsarr), dtype=int) # Use Allgather to gather 'ends' from all ranks into 'allends' if comm is None or isinstance(comm, MockComm): @@ -219,13 +219,13 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): # Compute dot product with the linear operator. self.dot(v, out=tmp2) # Compute to which column this iteration belongs - col = np.ravel_multi_index(i, npts) + col = xp.ravel_multi_index(i, npts) if is_sparse == False: result[:, col] = tmp2.toarray() else: aux = tmp2.toarray() # We now need to now which entries on tmp2 are non-zero and store then in our data list - for l in np.where(aux != 0)[0]: + for l in xp.where(aux != 0)[0]: data.append(aux[l]) colarr.append(col) row.append(l) diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 7a2d4cec2..1facbc2a0 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -16,7 +16,7 @@ from struphy.feec.utilities import RotationMatrix from struphy.geometry.base import Domain from struphy.polar.linear_operators import PolarExtractionOperator -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel @@ -777,7 +777,7 @@ def create_weighted_mass( 1. ``str`` : for square block matrices (V=W), a symmetry can be set in order to accelerate the assembly process. Possible strings are ``symm`` (symmetric), ``asym`` (anti-symmetric) and ``diag`` (diagonal). 2. ``None`` : all blocks are allocated, disregarding zero-blocks or any symmetry. 3. ``1D list`` : 1d list consisting of either a) strings or b) matrices (3x3 callables or 3x3 list) and can be mixed. Predefined names are ``G``, ``Ginv``, ``DFinv``, ``sqrt_g``. Access them using strings in the 1d list: ``weights=['']``. Possible choices for key-value pairs in **weights** are, at the moment: ``eq_mhd``: :class:`~struphy.fields_background.base.MHDequilibrium`. To access them, use for ```` the string ``eq_``, where ```` can be found in the just mentioned base classes for MHD equilibria. By default, all scalars are multiplied. For division of scalars use ``1/``. - 4. ``2D list`` : 2d list with the same number of rows/columns as the number of components of the domain/codomain spaces. The entries can be either a) callables or b) np.ndarrays representing the weights at the quadrature points. If an entry is zero or ``None``, the corresponding block is set to ``None`` to accelerate the dot product. + 4. ``2D list`` : 2d list with the same number of rows/columns as the number of components of the domain/codomain spaces. The entries can be either a) callables or b) xp.ndarrays representing the weights at the quadrature points. If an entry is zero or ``None``, the corresponding block is set to ``None`` to accelerate the dot product. assemble: bool Whether to assemble the weighted mass matrix, i.e. computes the integrals with @@ -986,11 +986,11 @@ def _get_range_rank(self, func): else: dummy_eta = (0.0, 0.0, 0.0) val = func(*dummy_eta) - assert isinstance(val, np.ndarray) + assert isinstance(val, xp.ndarray) out = len(val.shape) - 3 else: if isinstance(func, list): - if isinstance(func[0], np.ndarray): + if isinstance(func[0], xp.ndarray): out = 2 else: out = len(func) - 1 @@ -1052,7 +1052,7 @@ def __init__(self, derham, mass_ops, domain): ) grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) - self._f_values = np.zeros(grid_shape, dtype=float) + self._f_values = xp.zeros(grid_shape, dtype=float) metric = domain.metric(*integration_grid) self._mass_metric_term = deepcopy(metric) @@ -2083,7 +2083,7 @@ class WeightedMassOperator(LinOpWithTransp): 1. ``None`` : all blocks are allocated, disregarding zero-blocks or any symmetry. 2. ``str`` : for square block matrices (V=W), a symmetry can be set in order to accelerate the assembly process. Possible strings are ``symm`` (symmetric), ``asym`` (anti-symmetric) and ``diag`` (diagonal). - 3. ``list`` : 2d list with the same number of rows/columns as the number of components of the domain/codomain spaces. The entries can be either a) callables or b) np.ndarrays representing the weights at the quadrature points. If an entry is zero or ``None``, the corresponding block is set to ``None`` to accelerate the dot product. + 3. ``list`` : 2d list with the same number of rows/columns as the number of components of the domain/codomain spaces. The entries can be either a) callables or b) xp.ndarrays representing the weights at the quadrature points. If an entry is zero or ``None``, the corresponding block is set to ``None`` to accelerate the dot product. transposed : bool Whether to assemble the transposed operator. @@ -2364,16 +2364,16 @@ def __init__( ] if callable(weights_info[a][b]): - PTS = np.meshgrid(*pts, indexing="ij") + PTS = xp.meshgrid(*pts, indexing="ij") mat_w = weights_info[a][b](*PTS).copy() - elif isinstance(weights_info[a][b], np.ndarray): + elif isinstance(weights_info[a][b], xp.ndarray): mat_w = weights_info[a][b] assert mat_w.shape == tuple( [pt.size for pt in pts], ) - if np.any(np.abs(mat_w) > 1e-14): + if xp.any(xp.abs(mat_w) > 1e-14): if self._matrix_free: blocks[-1] += [ StencilMatrixFreeMassOperator( @@ -2689,7 +2689,7 @@ def assemble(self, weights=None, clear=True, verbose=True): Parameters ---------- weights : list | NoneType - Weight function(s) (callables or np.ndarrays) in a 2d list of shape corresponding to + Weight function(s) (callables or xp.ndarrays) in a 2d list of shape corresponding to number of components of domain/codomain. If ``weights=None``, the weight is taken from the given weights in the instanziation of the object, else it will be overriden. @@ -2712,7 +2712,7 @@ def assemble(self, weights=None, clear=True, verbose=True): if weight is not None: assert callable(weight) or isinstance( weight, - np.ndarray, + xp.ndarray, ) self._mat[a, b].weights = weight @@ -2822,13 +2822,13 @@ def assemble(self, weights=None, clear=True, verbose=True): # evaluate weight at quadrature points if callable(loc_weight): - PTS = np.meshgrid(*pts, indexing="ij") + PTS = xp.meshgrid(*pts, indexing="ij") mat_w = loc_weight(*PTS).copy() - elif isinstance(loc_weight, np.ndarray): + elif isinstance(loc_weight, xp.ndarray): mat_w = loc_weight elif loc_weight is not None: raise TypeError( - "weights must be callable or np.ndarray or None but is {}".format( + "weights must be callable or xp.ndarray or None but is {}".format( type(self._weights[a][b]), ), ) @@ -2836,8 +2836,8 @@ def assemble(self, weights=None, clear=True, verbose=True): if loc_weight is not None: assert mat_w.shape == tuple([pt.size for pt in pts]) - not_weight_zero = np.array( - int(loc_weight is not None and np.any(np.abs(mat_w) > 1e-14)), + not_weight_zero = xp.array( + int(loc_weight is not None and xp.any(xp.abs(mat_w) > 1e-14)), ) if self._mpi_comm is not None: self._mpi_comm.Allreduce( @@ -2861,7 +2861,7 @@ def assemble(self, weights=None, clear=True, verbose=True): mat = self._mat if loc_weight is None: # in case it's none we still need to have zeros weights to call the kernel - mat_w = np.zeros( + mat_w = xp.zeros( tuple([pt.size for pt in pts]), ) else: @@ -2997,12 +2997,12 @@ def eval_quad(W, coeffs, out=None): coeffs : StencilVector | BlockVector The coefficient vector corresponding to the FEM field. Ghost regions must be up-to-date! - out : np.ndarray | list/tuple of np.ndarrays, optional + out : xp.ndarray | list/tuple of xp.ndarrays, optional If given, the result will be written into these arrays in-place. Number of outs must be compatible with number of components of FEM field. Returns ------- - out : np.ndarray | list/tuple of np.ndarrays + out : xp.ndarray | list/tuple of xp.ndarrays The values of the FEM field at the quadrature points. """ @@ -3021,7 +3021,7 @@ def eval_quad(W, coeffs, out=None): out = () if isinstance(W, TensorFemSpace): out += ( - np.zeros( + xp.zeros( [ q_grid[nquad].points.size for q_grid, nquad in zip(self.derham.get_quad_grids(W, nquads=self.nquads), self.nquads) @@ -3032,7 +3032,7 @@ def eval_quad(W, coeffs, out=None): else: for space in W.spaces: out += ( - np.zeros( + xp.zeros( [ q_grid[nquad].points.size for q_grid, nquad in zip( @@ -3045,7 +3045,7 @@ def eval_quad(W, coeffs, out=None): else: if isinstance(W, TensorFemSpace): - assert isinstance(out, np.ndarray) + assert isinstance(out, xp.ndarray) out = (out,) else: assert isinstance(out, (list, tuple)) @@ -3161,7 +3161,7 @@ def __init__(self, derham, V, W, weights=None, nquads=None): ) shape = tuple(e - s + 1 for s, e in zip(V.coeff_space.starts, V.coeff_space.ends)) - self._diag_tmp = np.zeros((shape)) + self._diag_tmp = xp.zeros((shape)) # knot span indices of elements of local domain self._codomain_spans = [ @@ -3288,16 +3288,16 @@ def dot(self, v, out=None): # evaluate weight at quadrature points if callable(self._weights): - PTS = np.meshgrid(*self._pts, indexing="ij") + PTS = xp.meshgrid(*self._pts, indexing="ij") mat_w = self._weights(*PTS).copy() - elif isinstance(self._weights, np.ndarray): + elif isinstance(self._weights, xp.ndarray): mat_w = self._weights if self._weights is not None: assert mat_w.shape == tuple([pt.size for pt in self._pts]) # call kernel (if mat_w is not zero) by calling the appropriate kernel (1d, 2d or 3d) - if np.any(np.abs(mat_w) > 1e-14): + if xp.any(xp.abs(mat_w) > 1e-14): self._dot_kernel( *self._codomain_spans, *self._domain_spans, @@ -3357,9 +3357,9 @@ def diagonal(self, inverse=False, sqrt=False, out=None): # evaluate weight at quadrature points if callable(self._weights): - PTS = np.meshgrid(*self._pts, indexing="ij") + PTS = xp.meshgrid(*self._pts, indexing="ij") mat_w = self._weights(*PTS).copy() - elif isinstance(self._weights, np.ndarray): + elif isinstance(self._weights, xp.ndarray): mat_w = self._weights diag = self._diag_tmp @@ -3379,12 +3379,12 @@ def diagonal(self, inverse=False, sqrt=False, out=None): # Calculate entries of StencilDiagonalMatrix if sqrt: - diag = np.sqrt(diag) + diag = xp.sqrt(diag) if inverse: - data = np.divide(1, diag, out=data) + data = xp.divide(1, diag, out=data) elif out: - np.copyto(data, diag) + xp.copyto(data, diag) else: data = diag.copy() diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index b3a8744eb..783eb062d 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -11,7 +11,7 @@ from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.mass import WeightedMassOperator -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class MassMatrixPreconditioner(LinearOperator): @@ -94,12 +94,12 @@ def fun(e): s = e.shape[0] newshape = tuple([1 if i != d else s for i in range(n_dims)]) f = e.reshape(newshape) - return np.atleast_1d( + return xp.atleast_1d( loc_weights( - *[np.array(np.full_like(f, 0.5)) if i != d else np.array(f) for i in range(n_dims)], + *[xp.array(xp.full_like(f, 0.5)) if i != d else xp.array(f) for i in range(n_dims)], ).squeeze(), ) - elif isinstance(loc_weights, np.ndarray): + elif isinstance(loc_weights, xp.ndarray): s = loc_weights.shape if d == 0: fun = loc_weights[:, s[1] // 2, s[2] // 2] @@ -108,14 +108,14 @@ def fun(e): elif d == 2: fun = loc_weights[s[0] // 2, s[1] // 2, :] elif loc_weights is None: - fun = lambda e: np.ones(e.size, dtype=float) + fun = lambda e: xp.ones(e.size, dtype=float) else: raise TypeError( - "weights needs to be callable, np.ndarray or None but is{}".format(type(loc_weights)), + "weights needs to be callable, xp.ndarray or None but is{}".format(type(loc_weights)), ) fun = [[fun]] else: - fun = [[lambda e: np.ones(e.size, dtype=float)]] + fun = [[lambda e: xp.ones(e.size, dtype=float)]] # get 1D FEM space (serial, not distributed) and quadrature order femspace_1d = femspaces[c].spaces[d] @@ -207,7 +207,7 @@ def fun(e): M_local = StencilMatrix(V_local, V_local) - row_indices, col_indices = np.nonzero(M_arr) + row_indices, col_indices = xp.nonzero(M_arr) for row_i, col_i in zip(row_indices, col_indices): # only consider row indices on process @@ -220,7 +220,7 @@ def fun(e): ] = M_arr[row_i, col_i] # check if stencil matrix was built correctly - assert np.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) + assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) matrixcells += [M_local.copy()] # ======================================================================================================= @@ -487,7 +487,7 @@ def __init__(self, mass_operator, apply_bc=True): # loop over spatial directions for d in range(n_dims): - fun = [[lambda e: np.ones(e.size, dtype=float)]] + fun = [[lambda e: xp.ones(e.size, dtype=float)]] # get 1D FEM space (serial, not distributed) and quadrature order femspace_1d = femspaces[c].spaces[d] @@ -579,7 +579,7 @@ def __init__(self, mass_operator, apply_bc=True): M_local = StencilMatrix(V_local, V_local) - row_indices, col_indices = np.nonzero(M_arr) + row_indices, col_indices = xp.nonzero(M_arr) for row_i, col_i in zip(row_indices, col_indices): # only consider row indices on process @@ -592,7 +592,7 @@ def __init__(self, mass_operator, apply_bc=True): ] = M_arr[row_i, col_i] # check if stencil matrix was built correctly - assert np.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) + assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) matrixcells += [M_local.copy()] # ======================================================================================================= @@ -676,7 +676,7 @@ def __init__(self, mass_operator, apply_bc=True): # Need to assemble the logical mass matrix to extract the coefficients fun = [ - [lambda e1, e2, e3: np.ones_like(e1, dtype=float) if i == j else None for j in range(3)] for i in range(3) + [lambda e1, e2, e3: xp.ones_like(e1, dtype=float) if i == j else None for j in range(3)] for i in range(3) ] log_M = WeightedMassOperator( self._mass_operator.derham, @@ -864,15 +864,15 @@ class FFTSolver(BandedSolver): Parameters ---------- - circmat : np.ndarray + circmat : xp.ndarray Generic circulant matrix. """ def __init__(self, circmat): - assert isinstance(circmat, np.ndarray) + assert isinstance(circmat, xp.ndarray) assert is_circulant(circmat) - self._space = np.ndarray + self._space = xp.ndarray self._column = circmat[:, 0] # -------------------------------------- @@ -889,13 +889,13 @@ def solve(self, rhs, out=None, transposed=False): Parameters ---------- - rhs : np.ndarray + rhs : xp.ndarray The right-hand sides to solve for. The vectors are assumed to be given in C-contiguous order, i.e. if multiple right-hand sides are given, then rhs is a two-dimensional array with the 0-th index denoting the number of the right-hand side, and the 1-st index denoting the element inside a right-hand side. - out : np.ndarray, optional + out : xp.ndarray, optional Output vector. If given, it has to have the same shape and datatype as rhs. transposed : bool @@ -913,7 +913,7 @@ def solve(self, rhs, out=None, transposed=False): try: out[:] = solve_circulant(self._column, rhs.T).T - except np.linalg.LinAlgError: + except xp.linalg.LinAlgError: eps = 1e-4 print(f"Stabilizing singular preconditioning FFTSolver with {eps = }:") self._column[0] *= 1.0 + eps @@ -937,13 +937,13 @@ def is_circulant(mat): Whether the matrix is circulant (=True) or not (=False). """ - assert isinstance(mat, np.ndarray) + assert isinstance(mat, xp.ndarray) assert len(mat.shape) == 2 assert mat.shape[0] == mat.shape[1] if mat.shape[0] > 1: for i in range(mat.shape[0] - 1): - circulant = np.allclose(mat[i, :], np.roll(mat[i + 1, :], -1)) + circulant = xp.allclose(mat[i, :], xp.roll(mat[i + 1, :], -1)) if not circulant: return circulant else: diff --git a/src/struphy/feec/projectors.py b/src/struphy/feec/projectors.py index 1e9421c7e..5a43cfc74 100644 --- a/src/struphy/feec/projectors.py +++ b/src/struphy/feec/projectors.py @@ -37,7 +37,7 @@ from struphy.kernel_arguments.local_projectors_args_kernels import LocalProjectorsArguments from struphy.polar.basic import PolarVector from struphy.polar.linear_operators import PolarExtractionOperator -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class CommutingProjector: @@ -576,24 +576,24 @@ class CommutingProjectorLocal: fem_space : FemSpace FEEC space into which the functions shall be projected. - pts : list of np.array + pts : list of xp.array 3-list (or nested 3-list[3-list] for BlockVectors) of 2D arrays with the quasi-interpolation points (or Gauss-Legendre quadrature points for histopolation). In format [spatial direction](B-spline index, point) for StencilVector spaces or [vector component][spatial direction](B-spline index, point) for BlockVector spaces. - wts : list of np.array + wts : list of xp.array 3D (4D for BlockVectors) list of 2D array with the Gauss-Legendre quadrature weights (full of ones for interpolation). In format [spatial direction](B-spline index, point) for StencilVector spaces or [vector component][spatial direction](B-spline index, point) for BlockVector spaces. - wij : list of np.array + wij : list of xp.array List of 2D arrays for the coefficients :math:`\omega_j^i` obtained by inverting the local collocation matrix. Use for obtaining the FE coefficients of a function via interpolation. In format [spatial direction](B-spline index, point). - whij : list of np.array + whij : list of xp.array List of 2D arrays for the coefficients :math:`\hat{\omega}_j^i` obtained from the :math:`\omega_j^i`. Use for obtaining the FE coefficients of a function via histopolation. In format [spatial direction](D-spline index, point). @@ -639,22 +639,22 @@ def __init__( # FE space of zero forms. That means that we have B-splines in all three spatial directions. Bspaces_1d = [fem_space_B.spaces] - self._B_nbasis = np.array([space.nbasis for space in Bspaces_1d[0]]) + self._B_nbasis = xp.array([space.nbasis for space in Bspaces_1d[0]]) # Degree of the B-spline space, not to be confused with the degrees given by fem_space.spaces.degree since depending on the situation it will give the D-spline degree instead - self._p = np.zeros(3, dtype=int) + self._p = xp.zeros(3, dtype=int) for i, space in enumerate(fem_space_B.spaces): self._p[i] = space.degree # FE space of three forms. That means that we have D-splines in all three spatial directions. Dspaces_1d = [fem_space_D.spaces] - D_nbasis = np.array([space.nbasis for space in Dspaces_1d[0]]) + D_nbasis = xp.array([space.nbasis for space in Dspaces_1d[0]]) self._periodic = [] for space in fem_space.spaces: self._periodic.append(space.periodic) - self._periodic = np.array(self._periodic) + self._periodic = xp.array(self._periodic) if isinstance(fem_space, TensorFemSpace): # The comm, rank and size are only necessary for debugging. In particular, for printing stuff @@ -667,21 +667,21 @@ def __init__( self._size = self._comm.Get_size() # We get the start and endpoint for each sublist in out - self._starts = np.array(self.coeff_space.starts) - self._ends = np.array(self.coeff_space.ends) + self._starts = xp.array(self.coeff_space.starts) + self._ends = xp.array(self.coeff_space.ends) # We compute the number of FE coefficients the current MPI rank is responsible for - self._loc_num_coeff = np.array([self._ends[i] + 1 - self._starts[i] for i in range(3)], dtype=int) + self._loc_num_coeff = xp.array([self._ends[i] + 1 - self._starts[i] for i in range(3)], dtype=int) # We get the pads - self._pds = np.array(self.coeff_space.pads) + self._pds = xp.array(self.coeff_space.pads) # We get the number of spaces we have self._nsp = 1 self._localpts = [] self._index_translation = [] self._inv_index_translation = [] - self._original_pts_size = np.zeros((3), dtype=int) + self._original_pts_size = xp.zeros((3), dtype=int) elif isinstance(fem_space, VectorFemSpace): # The comm, rank and size are only necessary for debugging. In particular, for printing stuff @@ -694,17 +694,17 @@ def __init__( self._size = self._comm.Get_size() # we collect all starts and ends in two big lists - self._starts = np.array([vi.starts for vi in self.coeff_space.spaces]) - self._ends = np.array([vi.ends for vi in self.coeff_space.spaces]) + self._starts = xp.array([vi.starts for vi in self.coeff_space.spaces]) + self._ends = xp.array([vi.ends for vi in self.coeff_space.spaces]) # We compute the number of FE coefficients the current MPI rank is responsible for - self._loc_num_coeff = np.array( + self._loc_num_coeff = xp.array( [[self._ends[h][i] + 1 - self._starts[h][i] for i in range(3)] for h in range(3)], dtype=int, ) # We collect the pads - self._pds = np.array([vi.pads for vi in self.coeff_space.spaces]) + self._pds = xp.array([vi.pads for vi in self.coeff_space.spaces]) # We get the number of space we have self._nsp = len(self.coeff_space.spaces) @@ -720,7 +720,7 @@ def __init__( self._localpts = [[], [], []] # Here we will store the global number of points for each block entry and for each spatial direction. - self._original_pts_size = [np.zeros((3), dtype=int), np.zeros((3), dtype=int), np.zeros((3), dtype=int)] + self._original_pts_size = [xp.zeros((3), dtype=int), xp.zeros((3), dtype=int), xp.zeros((3), dtype=int)] # This will be a list of three elements (the first one for the first block element, the second one for the second block element, ...), each one being a list with three arrays, # each array will contain the B-spline indices of the corresponding spatial direction for which this MPI rank has to store at least one non-zero FE coefficient for the storage of the @@ -740,8 +740,8 @@ def __init__( self._are_zero_block_B_or_D_splines = [[], [], []] # self._Basis_function_indices_agreggated_B[i][j] = -1 if the jth B-spline is not necessary for any of the three block entries in the ith spatial direction, otherwise it is 0 - self._Basis_function_indices_agreggated_B = [-1 * np.ones(nbasis, dtype=int) for nbasis in self._B_nbasis] - self._Basis_function_indices_agreggated_D = [-1 * np.ones(nbasis, dtype=int) for nbasis in D_nbasis] + self._Basis_function_indices_agreggated_B = [-1 * xp.ones(nbasis, dtype=int) for nbasis in self._B_nbasis] + self._Basis_function_indices_agreggated_D = [-1 * xp.ones(nbasis, dtype=int) for nbasis in D_nbasis] # List that will contain the LocalProjectorsArguments for each value of h = 0,1,2. self._solve_args = [] @@ -753,20 +753,20 @@ def __init__( # List of list that tell us for each spatial direction whether we have Interpolation or Histopolation. IoH_for_indices = ["I", "I", "I"] # Same list as before but with bools instead of chars - self._IoH = np.array([False, False, False], dtype=bool) + self._IoH = xp.array([False, False, False], dtype=bool) # We make a list with the interpolation/histopolation weights we need for each block and each direction. self._geo_weights = [self._wij[0], self._wij[1], self._wij[2]] elif space_id == "L2": IoH_for_indices = ["H", "H", "H"] - self._IoH = np.array([True, True, True], dtype=bool) + self._IoH = xp.array([True, True, True], dtype=bool) self._geo_weights = [self._whij[0], self._whij[1], self._whij[2]] lenj1, lenj2, lenj3 = get_local_problem_size(self._periodic, self._p, self._IoH) lenj = [lenj1, lenj2, lenj3] - self._shift = np.array([0, 0, 0], dtype=int) + self._shift = xp.array([0, 0, 0], dtype=int) compute_shifts(self._IoH, self._p, self._B_nbasis, self._shift) split_points( @@ -788,7 +788,7 @@ def __init__( ) # We want to build the meshgrid for the evaluation of the degrees of freedom so it only contains the evaluation points that each specific MPI rank is actually going to use. - self._meshgrid = np.meshgrid( + self._meshgrid = xp.meshgrid( *[pt for pt in self._localpts], indexing="ij", ) @@ -927,18 +927,18 @@ def __init__( ) elif isinstance(fem_space, VectorFemSpace): - self._shift = [np.array([0, 0, 0], dtype=int) for _ in range(3)] + self._shift = [xp.array([0, 0, 0], dtype=int) for _ in range(3)] if space_id == "H1vec": # List of list that tell us for each block entry and for each spatial direction whether we have Interpolation or Histopolation. IoH_for_indices = [["I", "I", "I"], ["I", "I", "I"], ["I", "I", "I"]] # Same list as before but with bools instead of chars self._IoH = [ - np.array([False, False, False], dtype=bool), - np.array( + xp.array([False, False, False], dtype=bool), + xp.array( [False, False, False], dtype=bool, ), - np.array([False, False, False], dtype=bool), + xp.array([False, False, False], dtype=bool), ] # We make a list with the interpolation/histopolation weights we need for each block and each direction. self._geo_weights = [[self._wij[0], self._wij[1], self._wij[2]] for _ in range(3)] @@ -946,12 +946,12 @@ def __init__( elif space_id == "Hcurl": IoH_for_indices = [["H", "I", "I"], ["I", "H", "I"], ["I", "I", "H"]] self._IoH = [ - np.array([True, False, False], dtype=bool), - np.array( + xp.array([True, False, False], dtype=bool), + xp.array( [False, True, False], dtype=bool, ), - np.array([False, False, True], dtype=bool), + xp.array([False, False, True], dtype=bool), ] self._geo_weights = [ [self._whij[0], self._wij[1], self._wij[2]], @@ -966,12 +966,12 @@ def __init__( elif space_id == "Hdiv": IoH_for_indices = [["I", "H", "H"], ["H", "I", "H"], ["H", "H", "I"]] self._IoH = [ - np.array([False, True, True], dtype=bool), - np.array( + xp.array([False, True, True], dtype=bool), + xp.array( [True, False, True], dtype=bool, ), - np.array([True, True, False], dtype=bool), + xp.array([True, True, False], dtype=bool), ] self._geo_weights = [ [self._wij[0], self._whij[1], self._whij[2]], @@ -1010,7 +1010,7 @@ def __init__( # meshgrid for h component self._meshgrid.append( - np.meshgrid( + xp.meshgrid( *[pt for pt in self._localpts[h]], indexing="ij", ), @@ -1328,9 +1328,9 @@ def solve_weighted(self, rhs, out=None): if isinstance(self._fem_space, TensorFemSpace): if out is None: - out = np.zeros((self._loc_num_coeff[0], self._loc_num_coeff[1], self._loc_num_coeff[2]), dtype=float) + out = xp.zeros((self._loc_num_coeff[0], self._loc_num_coeff[1], self._loc_num_coeff[2]), dtype=float) else: - assert np.shape(out) == (self._loc_num_coeff[0], self._loc_num_coeff[1], self._loc_num_coeff[2]) + assert xp.shape(out) == (self._loc_num_coeff[0], self._loc_num_coeff[1], self._loc_num_coeff[2]) solve_local_main_loop_weighted( self._solve_args, @@ -1352,7 +1352,7 @@ def solve_weighted(self, rhs, out=None): out = [] for h in range(3): out.append( - np.zeros( + xp.zeros( ( self._loc_num_coeff[h][0], self._loc_num_coeff[h][1], @@ -1365,7 +1365,7 @@ def solve_weighted(self, rhs, out=None): else: assert len(out) == 3 for h in range(3): - assert np.shape(out[h]) == ( + assert xp.shape(out[h]) == ( self._loc_num_coeff[h][0], self._loc_num_coeff[h][1], self._loc_num_coeff[h][2], @@ -1375,7 +1375,7 @@ def solve_weighted(self, rhs, out=None): # the out block for which do_nothing tell us before hand they shall be zero. for h in range(3): if self._do_nothing[h] == 1: - out[h] = np.zeros( + out[h] = xp.zeros( ( self._loc_num_coeff[h][0], self._loc_num_coeff[h][1], @@ -1430,7 +1430,7 @@ def get_dofs(self, fun, dofs=None): fh = fun[h](*self._meshgrid[h]) # Array into which we will write the Dofs. - f_eval_aux = np.zeros(tuple(np.shape(dim)[0] for dim in self._localpts[h])) + f_eval_aux = xp.zeros(tuple(xp.shape(dim)[0] for dim in self._localpts[h])) # For 1-forms if self._space_key == "1": @@ -1442,7 +1442,7 @@ def get_dofs(self, fun, dofs=None): f_eval.append(f_eval_aux) elif self._space_key == "3": - f_eval = np.zeros(tuple(np.shape(dim)[0] for dim in self._localpts)) + f_eval = xp.zeros(tuple(xp.shape(dim)[0] for dim in self._localpts)) # Evaluation of the function at all Gauss-Legendre quadrature points faux = fun(*self._meshgrid) get_dofs_local_3_form(self._solve_args, faux, f_eval) @@ -1483,7 +1483,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non elif self._space_key == "1" or self._space_key == "2": assert len(fun) == 3, f"List input only for vector-valued spaces of size 3, but {len(fun) = }." - self._do_nothing = np.zeros(3, dtype=int) + self._do_nothing = xp.zeros(3, dtype=int) f_eval = [] # If this is the first time this rank has to evaluate the weights degrees of freedom we declare the list where to store them. @@ -1496,7 +1496,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non pre_computed_dofs.append(fun[h](*self._meshgrid[h])) # Array into which we will write the Dofs. - f_eval_aux = np.zeros(tuple(np.shape(dim)[0] for dim in self._localpts[h])) + f_eval_aux = xp.zeros(tuple(xp.shape(dim)[0] for dim in self._localpts[h])) # We check if the current set of basis functions is not one of those we have to compute in the current MPI rank. if ( @@ -1541,7 +1541,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non f_eval.append(f_eval_aux) elif self._space_key == "3": - f_eval = np.zeros(tuple(np.shape(dim)[0] for dim in self._localpts)) + f_eval = xp.zeros(tuple(xp.shape(dim)[0] for dim in self._localpts)) # Evaluation of the function at all Gauss-Legendre quadrature points if first_go == True: pre_computed_dofs = [fun(*self._meshgrid)] @@ -1563,7 +1563,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non elif self._space_key == "v": assert len(fun) == 3, f"List input only for vector-valued spaces of size 3, but {len(fun) = }." - self._do_nothing = np.zeros(3, dtype=int) + self._do_nothing = xp.zeros(3, dtype=int) for h in range(3): # We check if the current set of basis functions is not one of those we have to compute in the current MPI rank. if ( @@ -1641,13 +1641,13 @@ def __call__( set to false it means we computed it once already and we can reuse the dofs evaluation of the weights instead of recomputing them. - pre_computed_dofs : list of np.arrays + pre_computed_dofs : list of xp.arrays If we have already computed the evaluation of the weights at the dofs we can pass the arrays with their values here, so we do not have to compute them again. Returns ------- - coeffs : psydac.linalg.basic.vector | np.array 3D + coeffs : psydac.linalg.basic.vector | xp.array 3D The FEM spline coefficients after projection. """ if weighted == False: @@ -1858,7 +1858,7 @@ def __init__(self, space_id, mass_ops, **params): self._quad_grid_pts = self.mass_ops.derham.quad_grid_pts[self.space_key] if space_id in ("H1", "L2"): - self._quad_grid_mesh = np.meshgrid( + self._quad_grid_mesh = xp.meshgrid( *[pt.flatten() for pt in self.quad_grid_pts], indexing="ij", ) @@ -1868,12 +1868,12 @@ def __init__(self, space_id, mass_ops, **params): self._tmp = [] # tmp for matrix-vector product of geom_weights with fun for pts in self.quad_grid_pts: self._quad_grid_mesh += [ - np.meshgrid( + xp.meshgrid( *[pt.flatten() for pt in pts], indexing="ij", ), ] - self._tmp += [np.zeros_like(self.quad_grid_mesh[-1][0])] + self._tmp += [xp.zeros_like(self.quad_grid_mesh[-1][0])] # geometric weights evaluated at quadrature grid self._geom_weights = [] # loop over rows (different meshes) @@ -1884,7 +1884,7 @@ def __init__(self, space_id, mass_ops, **params): if weight is not None: self._geom_weights[-1] += [weight(*mesh)] else: - self._geom_weights[-1] += [np.zeros_like(mesh[0])] + self._geom_weights[-1] += [xp.zeros_like(mesh[0])] # other quad grid info if isinstance(self.space, TensorFemSpace): @@ -2010,7 +2010,7 @@ def get_dofs(self, fun, dofs=None, apply_bc=False, clear=True): Parameters ---------- fun : callable | list - Weight function(s) (callables or np.ndarrays) in a 1d list of shape corresponding to number of components. + Weight function(s) (callables or xp.ndarrays) in a 1d list of shape corresponding to number of components. dofs : StencilVector | BlockVector, optional The vector for the output. @@ -2025,7 +2025,7 @@ def get_dofs(self, fun, dofs=None, apply_bc=False, clear=True): # evaluate fun at quad_grid or check array size if callable(fun): fun_weights = fun(*self._quad_grid_mesh) - elif isinstance(fun, np.ndarray): + elif isinstance(fun, xp.ndarray): assert fun.shape == self._quad_grid_mesh[0].shape, ( f"Expected shape {self._quad_grid_mesh[0].shape}, got {fun.shape = } instead." ) @@ -2045,7 +2045,7 @@ def get_dofs(self, fun, dofs=None, apply_bc=False, clear=True): for f in fun: if callable(f): fun_weights[-1] += [f(*mesh)] - elif isinstance(f, np.ndarray): + elif isinstance(f, xp.ndarray): assert f.shape == mesh[0].shape, f"Expected shape {mesh[0].shape}, got {f.shape = } instead." fun_weights[-1] += [f] else: @@ -2062,7 +2062,7 @@ def get_dofs(self, fun, dofs=None, apply_bc=False, clear=True): # compute matrix data for kernel, i.e. fun * geom_weight tot_weights = [] - if isinstance(fun_weights, np.ndarray): + if isinstance(fun_weights, xp.ndarray): tot_weights += [fun_weights * self.geom_weights] else: # loop over rows (differnt meshes) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index bc7583ab8..fbbddae28 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -33,7 +33,7 @@ from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Derham: @@ -117,7 +117,7 @@ def __init__( if dirichlet_bc is not None: assert len(dirichlet_bc) == 3 # make sure that boundary conditions are compatible with spline space - assert np.all([bc == (False, False) for i, bc in enumerate(dirichlet_bc) if spl_kind[i]]) + assert xp.all([bc == (False, False) for i, bc in enumerate(dirichlet_bc) if spl_kind[i]]) self._dirichlet_bc = dirichlet_bc @@ -300,7 +300,7 @@ def __init__( fag.basis, ] - self._spline_types_pyccel[sp_form][-1] = np.array( + self._spline_types_pyccel[sp_form][-1] = xp.array( self._spline_types_pyccel[sp_form][-1], ) # In this case we are working with a scalar valued space @@ -352,7 +352,7 @@ def __init__( self._quad_grid_spans[sp_form] += [fag.spans] self._quad_grid_bases[sp_form] += [fag.basis] - self._spline_types_pyccel[sp_form] = np.array( + self._spline_types_pyccel[sp_form] = xp.array( self._spline_types_pyccel[sp_form], ) else: @@ -364,8 +364,8 @@ def __init__( # index arrays self._indN = [ ( - np.indices((space.ncells, space.degree + 1))[1] - + np.arange( + xp.indices((space.ncells, space.degree + 1))[1] + + xp.arange( space.ncells, )[:, None] ) @@ -374,8 +374,8 @@ def __init__( ] self._indD = [ ( - np.indices((space.ncells, space.degree + 1))[1] - + np.arange( + xp.indices((space.ncells, space.degree + 1))[1] + + xp.arange( space.ncells, )[:, None] ) @@ -525,11 +525,11 @@ def __init__( # collect arguments for kernels self._args_derham = DerhamArguments( - np.array(self.p), + xp.array(self.p), self.Vh_fem["0"].knots[0], self.Vh_fem["0"].knots[1], self.Vh_fem["0"].knots[2], - np.array(self.Vh["0"].starts), + xp.array(self.Vh["0"].starts), ) @property @@ -1062,7 +1062,7 @@ def _discretize_space( ) # Create uniform grid - grids = [np.linspace(xmin, xmax, num=ne + 1) for xmin, xmax, ne in zip(min_coords, max_coords, ncells)] + grids = [xp.linspace(xmin, xmax, num=ne + 1) for xmin, xmax, ne in zip(min_coords, max_coords, ncells)] # Create 1D finite element spaces and precompute quadrature data spaces_1d = [ @@ -1107,7 +1107,7 @@ def _get_domain_array(self): Returns ------- - dom_arr : np.ndarray + dom_arr : xp.ndarray A 2d array of shape (#MPI processes, 9). The row index denotes the process rank. The columns are for n=0,1,2: - arr[i, 3*n + 0] holds the LEFT domain boundary of process i in direction eta_(n+1). - arr[i, 3*n + 1] holds the RIGHT domain boundary of process i in direction eta_(n+1). @@ -1121,10 +1121,10 @@ def _get_domain_array(self): nproc = 1 # send buffer - dom_arr_loc = np.zeros(9, dtype=float) + dom_arr_loc = xp.zeros(9, dtype=float) # main array (receive buffers) - dom_arr = np.zeros(nproc * 9, dtype=float) + dom_arr = xp.zeros(nproc * 9, dtype=float) # Get global starts and ends of domain decomposition gl_s = self.domain_decomposition.starts @@ -1155,7 +1155,7 @@ def _get_index_array(self, decomposition): Returns ------- - ind_arr : np.ndarray + ind_arr : xp.ndarray A 2d array of shape (#MPI processes, 6). The row index denotes the process rank. The columns are for n=0,1,2: - arr[i, 2*n + 0] holds the global start index process i in direction eta_(n+1). - arr[i, 2*n + 1] holds the global end index of process i in direction eta_(n+1). @@ -1168,10 +1168,10 @@ def _get_index_array(self, decomposition): nproc = 1 # send buffer - ind_arr_loc = np.zeros(6, dtype=int) + ind_arr_loc = xp.zeros(6, dtype=int) # main array (receive buffers) - ind_arr = np.zeros(nproc * 6, dtype=int) + ind_arr = xp.zeros(nproc * 6, dtype=int) # Get global starts and ends of cart OR domain decomposition gl_s = decomposition.starts @@ -1214,13 +1214,13 @@ def _get_neighbours(self): Returns ------- - neighbours : np.ndarray + neighbours : xp.ndarray A 3d array of shape (3,3,3). The i-th axis is the direction eta_(i+1). Neighbours along the faces have index with two 1s, neighbours along the edges only have one 1, neighbours along the edges have no 1 in the index. """ - neighs = np.empty((3, 3, 3), dtype=int) + neighs = xp.empty((3, 3, 3), dtype=int) for i in range(3): for j in range(3): @@ -1265,8 +1265,8 @@ def _get_neighbour_one_component(self, comp): if comp == [1, 1, 1]: return neigh_id - comp = np.array(comp) - kinds = np.array(kinds) + comp = xp.array(comp) + kinds = xp.array(kinds) # if only one process: check if comp is neighbour in non-peridic directions, if this is not the case then return the rank as neighbour id if size == 1: @@ -1301,15 +1301,15 @@ def _get_neighbour_one_component(self, comp): "Wrong value for component; must be 0 or 1 or 2 !", ) - neigh_inds = np.array(neigh_inds) + neigh_inds = xp.array(neigh_inds) # only use indices where information is present to find the neighbours rank - inds = np.where(neigh_inds != None) + inds = xp.where(neigh_inds != None) # find ranks (row index of domain_array) which agree in start/end indices - index_temp = np.squeeze(self.index_array[:, inds]) - unique_ranks = np.where( - np.equal(index_temp, neigh_inds[inds]).all(1), + index_temp = xp.squeeze(self.index_array[:, inds]) + unique_ranks = xp.where( + xp.equal(index_temp, neigh_inds[inds]).all(1), )[0] # if any row satisfies condition, return its index (=rank of neighbour) @@ -1329,7 +1329,7 @@ def _get_span_and_basis_for_eval_mpi(self, etas, Nspace, end): Parameters ---------- - etas : np.array + etas : xp.array 1d array of evaluation points (ascending). Nspace : SplineSpace @@ -1340,13 +1340,13 @@ def _get_span_and_basis_for_eval_mpi(self, etas, Nspace, end): Returns ------- - spans : np.array + spans : xp.array 1d array of knot span indices. - bn : np.array + bn : xp.array 2d array of pn + 1 values of N-splines indexed by (eta, spline value). - bd : np.array + bd : xp.array 2d array of pn values of D-splines indexed by (eta, spline value). """ @@ -1356,11 +1356,11 @@ def _get_span_and_basis_for_eval_mpi(self, etas, Nspace, end): Tn = Nspace.knots pn = Nspace.degree - spans = np.zeros(etas.size, dtype=int) - bns = np.zeros((etas.size, pn + 1), dtype=float) - bds = np.zeros((etas.size, pn), dtype=float) - bn = np.zeros(pn + 1, dtype=float) - bd = np.zeros(pn, dtype=float) + spans = xp.zeros(etas.size, dtype=int) + bns = xp.zeros((etas.size, pn + 1), dtype=float) + bds = xp.zeros((etas.size, pn), dtype=float) + bn = xp.zeros(pn + 1, dtype=float) + bd = xp.zeros(pn, dtype=float) for n in range(etas.size): # avoid 1. --> 0. for clamped interpolation @@ -1538,7 +1538,7 @@ def vector(self, value): """In-place setter for Stencil-/Block-/PolarVector.""" if isinstance(self._vector, StencilVector): - assert isinstance(value, (StencilVector, np.ndarray)) + assert isinstance(value, (StencilVector, xp.ndarray)) s1, s2, s3 = self.starts e1, e2, e3 = self.ends @@ -1561,10 +1561,10 @@ def vector(self, value): self._vector.set_vector(value) else: if isinstance(self._vector.tp, StencilVector): - assert isinstance(value[0], np.ndarray) + assert isinstance(value[0], xp.ndarray) assert isinstance( value[1], - (StencilVector, np.ndarray), + (StencilVector, xp.ndarray), ) self._vector.pol[0][:] = value[0][:] @@ -1577,10 +1577,10 @@ def vector(self, value): ] else: for n in range(3): - assert isinstance(value[n][0], np.ndarray) + assert isinstance(value[n][0], xp.ndarray) assert isinstance( value[n][1], - (StencilVector, np.ndarray), + (StencilVector, xp.ndarray), ) self._vector.pol[n][:] = value[n][0][:] @@ -1884,7 +1884,7 @@ def eval_tp_fixed_loc(self, spans, bases, out=None): assert [span.size for span in spans] == [base.shape[0] for base in bases] if out is None: - out = np.empty([span.size for span in spans], dtype=float) + out = xp.empty([span.size for span in spans], dtype=float) else: assert out.shape == tuple([span.size for span in spans]) @@ -1893,8 +1893,8 @@ def eval_tp_fixed_loc(self, spans, bases, out=None): *bases, vec._data, self.derham.spline_types_pyccel[self.space_key], - np.array(self.derham.p), - np.array(self.starts), + xp.array(self.derham.p), + xp.array(self.starts), out, ) @@ -1908,7 +1908,7 @@ def eval_tp_fixed_loc(self, spans, bases, out=None): assert [span.size for span in spans] == [base.shape[0] for base in bases[i]] if out_is_none: - out += np.empty( + out += xp.empty( [span.size for span in spans], dtype=float, ) @@ -1922,10 +1922,10 @@ def eval_tp_fixed_loc(self, spans, bases, out=None): *bases[i], vec[i]._data, self.derham.spline_types_pyccel[self.space_key][i], - np.array( + xp.array( self.derham.p, ), - np.array( + xp.array( self.starts[i], ), out[i], @@ -1992,14 +1992,14 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): # prepare arrays for AllReduce if tmp is None: - tmp = np.zeros( + tmp = xp.zeros( tmp_shape, dtype=float, ) else: - assert isinstance(tmp, np.ndarray) + assert isinstance(tmp, xp.ndarray) assert tmp.shape == tmp_shape - assert tmp.dtype.type is np.float64 + assert tmp.dtype.type is xp.float64 tmp[:] = 0.0 # scalar-valued field @@ -2014,11 +2014,11 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): E3, self._vector_stencil._data, kind, - np.array(self.derham.p), + xp.array(self.derham.p), T1, T2, T3, - np.array(self.starts), + xp.array(self.starts), tmp, ) elif marker_evaluation: @@ -2027,11 +2027,11 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): markers, self._vector_stencil._data, kind, - np.array(self.derham.p), + xp.array(self.derham.p), T1, T2, T3, - np.array(self.starts), + xp.array(self.starts), tmp, ) else: @@ -2042,11 +2042,11 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): E3, self._vector_stencil._data, kind, - np.array(self.derham.p), + xp.array(self.derham.p), T1, T2, T3, - np.array(self.starts), + xp.array(self.starts), tmp, ) @@ -2066,7 +2066,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): out += tmp if squeeze_out: - out = np.squeeze(out) + out = xp.squeeze(out) if out.ndim == 0: out = out.item() @@ -2085,11 +2085,11 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): E3, self._vector_stencil[n]._data, kind, - np.array(self.derham.p), + xp.array(self.derham.p), T1, T2, T3, - np.array(self.starts[n]), + xp.array(self.starts[n]), tmp, ) elif marker_evaluation: @@ -2098,11 +2098,11 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): markers, self._vector_stencil[n]._data, kind, - np.array(self.derham.p), + xp.array(self.derham.p), T1, T2, T3, - np.array(self.starts[n]), + xp.array(self.starts[n]), tmp, ) else: @@ -2113,11 +2113,11 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): E3, self._vector_stencil[n]._data, kind, - np.array(self.derham.p), + xp.array(self.derham.p), T1, T2, T3, - np.array(self.starts[n]), + xp.array(self.starts[n]), tmp, ) @@ -2139,7 +2139,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): tmp[:] = 0.0 if squeeze_out: - out[-1] = np.squeeze(out[-1]) + out[-1] = xp.squeeze(out[-1]) if out[-1].ndim == 0: out[-1] = out[-1].item() @@ -2172,11 +2172,11 @@ def _flag_pts_not_on_proc(self, *etas): markers = etas[0] # check which particles are on the current process domain - is_on_proc_domain = np.logical_and( + is_on_proc_domain = xp.logical_and( markers[:, :3] >= dom_arr[rank, 0::3], markers[:, :3] <= dom_arr[rank, 1::3], ) - on_proc = np.all(is_on_proc_domain, axis=1) + on_proc = xp.all(is_on_proc_domain, axis=1) markers[~on_proc, :] = -1.0 @@ -2202,15 +2202,15 @@ def _flag_pts_not_on_proc(self, *etas): E3[E3 == dom_arr[rank, 7]] += 1e-8 # True for eval points on current process - E1_on_proc = np.logical_and( + E1_on_proc = xp.logical_and( E1 >= dom_arr[rank, 0], E1 <= dom_arr[rank, 1], ) - E2_on_proc = np.logical_and( + E2_on_proc = xp.logical_and( E2 >= dom_arr[rank, 3], E2 <= dom_arr[rank, 4], ) - E3_on_proc = np.logical_and( + E3_on_proc = xp.logical_and( E3 >= dom_arr[rank, 6], E3 <= dom_arr[rank, 7], ) @@ -2371,7 +2371,7 @@ def _tmp_noise_for_mpi(self, *shapes, direction="e3", amp=0.0001, seed=None): Returns ------- - _amps : np.array + _amps : xp.array The noisy FE coefficients in the desired direction (1d, 2d or 3d array).""" if self.derham.comm is not None: @@ -2386,40 +2386,40 @@ def _tmp_noise_for_mpi(self, *shapes, direction="e3", amp=0.0001, seed=None): domain_array = self.derham.domain_array if seed is not None: - np.random.seed(seed) + xp.random.seed(seed) # temporary - _amps = np.zeros(shapes) + _amps = xp.zeros(shapes) # no process has been drawn for yet - already_drawn = np.zeros(nprocs) == 1.0 + already_drawn = xp.zeros(nprocs) == 1.0 # 1d mid point arrays in each direction mid_points = [] for npr in nprocs: delta = 1.0 / npr - mid_points_i = np.zeros(npr) + mid_points_i = xp.zeros(npr) for n in range(npr): mid_points_i[n] = delta * (n + 1 / 2) mid_points += [mid_points_i] if direction == "e1": - tmp_arrays = np.zeros(nprocs[0]).tolist() + tmp_arrays = xp.zeros(nprocs[0]).tolist() elif direction == "e2": - tmp_arrays = np.zeros(nprocs[1]).tolist() + tmp_arrays = xp.zeros(nprocs[1]).tolist() elif direction == "e3": - tmp_arrays = np.zeros(nprocs[2]).tolist() + tmp_arrays = xp.zeros(nprocs[2]).tolist() elif direction == "e1e2": - tmp_arrays = np.zeros((nprocs[0], nprocs[1])).tolist() + tmp_arrays = xp.zeros((nprocs[0], nprocs[1])).tolist() Warning, f"2d noise in the directions {direction} is not correctly initilaized for MPI !!" elif direction == "e1e3": - tmp_arrays = np.zeros((nprocs[0], nprocs[2])).tolist() + tmp_arrays = xp.zeros((nprocs[0], nprocs[2])).tolist() Warning, f"2d noise in the directions {direction} is not correctly initilaized for MPI !!" elif direction == "e2e3": - tmp_arrays = np.zeros((nprocs[1], nprocs[2])).tolist() + tmp_arrays = xp.zeros((nprocs[1], nprocs[2])).tolist() Warning, f"2d noise in the directions {direction} is not correctly initilaized for MPI !!" elif direction == "e1e2e3": - tmp_arrays = np.zeros((nprocs[0], nprocs[1], nprocs[2])).tolist() + tmp_arrays = xp.zeros((nprocs[0], nprocs[1], nprocs[2])).tolist() Warning, f"3d noise in the directions {direction} is not correctly initilaized for MPI !!" else: raise ValueError("Invalid direction for tmp_arrays.") @@ -2428,7 +2428,7 @@ def _tmp_noise_for_mpi(self, *shapes, direction="e3", amp=0.0001, seed=None): inds_current = [] for n in range(3): mid_pt_current = (domain_array[rank, 3 * n] + domain_array[rank, 3 * n + 1]) / 2.0 - inds_current += [np.argmin(np.abs(mid_points[n] - mid_pt_current))] + inds_current += [xp.argmin(xp.abs(mid_points[n] - mid_pt_current))] # loop over processes for i in range(comm_size): @@ -2436,7 +2436,7 @@ def _tmp_noise_for_mpi(self, *shapes, direction="e3", amp=0.0001, seed=None): inds = [] for n in range(3): mid_pt = (domain_array[i, 3 * n] + domain_array[i, 3 * n + 1]) / 2.0 - inds += [np.argmin(np.abs(mid_points[n] - mid_pt))] + inds += [xp.argmin(xp.abs(mid_points[n] - mid_pt))] if already_drawn[inds[0], inds[1], inds[2]]: if direction == "e1": @@ -2458,7 +2458,7 @@ def _tmp_noise_for_mpi(self, *shapes, direction="e3", amp=0.0001, seed=None): if direction == "e1": tmp_arrays[inds[0]] = ( ( - np.random.rand( + xp.random.rand( *shapes, ) - 0.5 @@ -2471,7 +2471,7 @@ def _tmp_noise_for_mpi(self, *shapes, direction="e3", amp=0.0001, seed=None): elif direction == "e2": tmp_arrays[inds[1]] = ( ( - np.random.rand( + xp.random.rand( *shapes, ) - 0.5 @@ -2484,7 +2484,7 @@ def _tmp_noise_for_mpi(self, *shapes, direction="e3", amp=0.0001, seed=None): elif direction == "e3": tmp_arrays[inds[2]] = ( ( - np.random.rand( + xp.random.rand( *shapes, ) - 0.5 @@ -2495,23 +2495,23 @@ def _tmp_noise_for_mpi(self, *shapes, direction="e3", amp=0.0001, seed=None): already_drawn[:, :, inds[2]] = True _amps[:] = tmp_arrays[inds[2]] elif direction == "e1e2": - tmp_arrays[inds[0]][inds[1]] = (np.random.rand(*shapes) - 0.5) * 2.0 * amp + tmp_arrays[inds[0]][inds[1]] = (xp.random.rand(*shapes) - 0.5) * 2.0 * amp already_drawn[inds[0], inds[1], :] = True _amps[:] = tmp_arrays[inds[0]][inds[1]] elif direction == "e1e3": - tmp_arrays[inds[0]][inds[2]] = (np.random.rand(*shapes) - 0.5) * 2.0 * amp + tmp_arrays[inds[0]][inds[2]] = (xp.random.rand(*shapes) - 0.5) * 2.0 * amp already_drawn[inds[0], :, inds[2]] = True _amps[:] = tmp_arrays[inds[0]][inds[2]] elif direction == "e2e3": - tmp_arrays[inds[1]][inds[2]] = (np.random.rand(*shapes) - 0.5) * 2.0 * amp + tmp_arrays[inds[1]][inds[2]] = (xp.random.rand(*shapes) - 0.5) * 2.0 * amp already_drawn[:, inds[1], inds[2]] = True _amps[:] = tmp_arrays[inds[1]][inds[2]] elif direction == "e1e2e3": - tmp_arrays[inds[0]][inds[1]][inds[2]] = (np.random.rand(*shapes) - 0.5) * 2.0 * amp + tmp_arrays[inds[0]][inds[1]][inds[2]] = (xp.random.rand(*shapes) - 0.5) * 2.0 * amp already_drawn[inds[0], inds[1], inds[2]] = True _amps[:] = tmp_arrays[inds[0]][inds[1]][inds[2]] - if np.all(np.array([ind_c == ind for ind_c, ind in zip(inds_current, inds)])): + if xp.all(xp.array([ind_c == ind for ind_c, ind in zip(inds_current, inds)])): return _amps @@ -2762,16 +2762,16 @@ def get_pts_and_wts(space_1d, start, end, n_quad=None, polar_shift=False): histopol_loc = space_1d.histopolation_grid[start : end + 2].copy() # make sure that greville points used for interpolation are in [0, 1] - assert np.all(np.logical_and(greville_loc >= 0.0, greville_loc <= 1.0)) + assert xp.all(xp.logical_and(greville_loc >= 0.0, greville_loc <= 1.0)) # interpolation if space_1d.basis == "B": x_grid = greville_loc pts = greville_loc[:, None] - wts = np.ones(pts.shape, dtype=float) + wts = xp.ones(pts.shape, dtype=float) # sub-interval index is always 0 for interpolation. - subs = np.zeros(pts.shape[0], dtype=int) + subs = xp.zeros(pts.shape[0], dtype=int) # !! shift away first interpolation point in eta_1 direction for polar domains !! if pts[0] == 0.0 and polar_shift: @@ -2785,27 +2785,27 @@ def get_pts_and_wts(space_1d, start, end, n_quad=None, polar_shift=False): union_breaks = space_1d.breaks[:-1] # Make union of Greville and break points - tmp = set(np.round(space_1d.histopolation_grid, decimals=14)).union( - np.round(union_breaks, decimals=14), + tmp = set(xp.round(space_1d.histopolation_grid, decimals=14)).union( + xp.round(union_breaks, decimals=14), ) tmp = list(tmp) tmp.sort() - tmp_a = np.array(tmp) + tmp_a = xp.array(tmp) x_grid = tmp_a[ - np.logical_and( + xp.logical_and( tmp_a - >= np.min( + >= xp.min( histopol_loc, ) - 1e-14, - tmp_a <= np.max(histopol_loc) + 1e-14, + tmp_a <= xp.max(histopol_loc) + 1e-14, ) ] # determine subinterval index (= 0 or 1): - subs = np.zeros(x_grid[:-1].size, dtype=int) + subs = xp.zeros(x_grid[:-1].size, dtype=int) for n, x_h in enumerate(x_grid[:-1]): add = 1 for x_g in histopol_loc: @@ -2818,7 +2818,7 @@ def get_pts_and_wts(space_1d, start, end, n_quad=None, polar_shift=False): # products of basis functions are integrated exactly n_quad = space_1d.degree + 1 - pts_loc, wts_loc = np.polynomial.legendre.leggauss(n_quad) + pts_loc, wts_loc = xp.polynomial.legendre.leggauss(n_quad) x, wts = bsp.quadrature_grid(x_grid, pts_loc, wts_loc) @@ -2881,12 +2881,12 @@ def get_pts_and_wts_quasi( # interpolation if space_1d.basis == "B": if p == 1 and h != 1.0: - x_grid = np.linspace(-(p - 1) * h, 1.0 - h + (h / 2.0), (N + p - 1) * 2) + x_grid = xp.linspace(-(p - 1) * h, 1.0 - h + (h / 2.0), (N + p - 1) * 2) else: - x_grid = np.linspace(-(p - 1) * h, 1.0 - h, (N + p - 1) * 2 - 1) + x_grid = xp.linspace(-(p - 1) * h, 1.0 - h, (N + p - 1) * 2 - 1) pts = x_grid[:, None] % 1.0 - wts = np.ones(pts.shape, dtype=float) + wts = xp.ones(pts.shape, dtype=float) # !! shift away first interpolation point in eta_1 direction for polar domains !! if pts[0] == 0.0 and polar_shift: @@ -2897,16 +2897,16 @@ def get_pts_and_wts_quasi( # The computation of histopolation points breaks in case we have Nel=1 and periodic boundary conditions since we end up with only one x_grid point. # We need to build the histopolation points by hand in this scenario. if p == 0 and h == 1.0: - x_grid = np.array([0.0, 0.5, 1.0]) + x_grid = xp.array([0.0, 0.5, 1.0]) elif p == 0 and h != 1.0: - x_grid = np.linspace(-p * h, 1.0 - h + (h / 2.0), (N + p) * 2) + x_grid = xp.linspace(-p * h, 1.0 - h + (h / 2.0), (N + p) * 2) else: - x_grid = np.linspace(-p * h, 1.0 - h, (N + p) * 2 - 1) + x_grid = xp.linspace(-p * h, 1.0 - h, (N + p) * 2 - 1) n_quad = p + 1 # Gauss - Legendre quadrature points and weights # products of basis functions are integrated exactly - pts_loc, wts_loc = np.polynomial.legendre.leggauss(n_quad) + pts_loc, wts_loc = xp.polynomial.legendre.leggauss(n_quad) x, wts = bsp.quadrature_grid(x_grid, pts_loc, wts_loc) pts = x % 1.0 @@ -2920,26 +2920,26 @@ def get_pts_and_wts_quasi( N_b = N + p # Filling the quasi-interpolation points for i=0 and i=1 (since they are equal) - x_grid = np.linspace(0.0, knots[p + 1], p + 1) - x_aux = np.linspace(0.0, knots[p + 1], p + 1) - x_grid = np.append(x_grid, x_aux) + x_grid = xp.linspace(0.0, knots[p + 1], p + 1) + x_aux = xp.linspace(0.0, knots[p + 1], p + 1) + x_grid = xp.append(x_grid, x_aux) # Now we append those for 1 V2):") res_PSY = OPS_PSY.Q1.dot(x1_st) - res_STR = OPS_STR.Q1_dot(np.concatenate((x1[0].flatten(), x1[1].flatten(), x1[2].flatten()))) + res_STR = OPS_STR.Q1_dot(xp.concatenate((x1[0].flatten(), x1[1].flatten(), x1[2].flatten()))) res_STR_0, res_STR_1, res_STR_2 = SPACES.extract_2(res_STR) MPI_COMM.Barrier() @@ -284,7 +284,7 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): Q1T = OPS_PSY.Q1.transpose() res_PSY = Q1T.dot(x2_st) - res_STR = OPS_STR.transpose_Q1_dot(np.concatenate((x2[0].flatten(), x2[1].flatten(), x2[2].flatten()))) + res_STR = OPS_STR.transpose_Q1_dot(xp.concatenate((x2[0].flatten(), x2[1].flatten(), x2[2].flatten()))) res_STR_0, res_STR_1, res_STR_2 = SPACES.extract_1(res_STR) MPI_COMM.Barrier() @@ -310,7 +310,7 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): print("\nW1 (V1 --> V1, Identity operator in this case):") res_PSY = OPS_PSY.W1.dot(x1_st) - res_STR = OPS_STR.W1_dot(np.concatenate((x1[0].flatten(), x1[1].flatten(), x1[2].flatten()))) + res_STR = OPS_STR.W1_dot(xp.concatenate((x1[0].flatten(), x1[1].flatten(), x1[2].flatten()))) res_STR_0, res_STR_1, res_STR_2 = SPACES.extract_1(res_STR) MPI_COMM.barrier() @@ -333,7 +333,7 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): W1T = OPS_PSY.W1.transpose() res_PSY = W1T.dot(x1_st) - res_STR = OPS_STR.transpose_W1_dot(np.concatenate((x1[0].flatten(), x1[1].flatten(), x1[2].flatten()))) + res_STR = OPS_STR.transpose_W1_dot(xp.concatenate((x1[0].flatten(), x1[1].flatten(), x1[2].flatten()))) res_STR_0, res_STR_1, res_STR_2 = SPACES.extract_1(res_STR) MPI_COMM.barrier() @@ -359,7 +359,7 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): print("\nQ2 (V2 --> V2, Identity operator in this case):") res_PSY = OPS_PSY.Q2.dot(x2_st) - res_STR = OPS_STR.Q2_dot(np.concatenate((x2[0].flatten(), x2[1].flatten(), x2[2].flatten()))) + res_STR = OPS_STR.Q2_dot(xp.concatenate((x2[0].flatten(), x2[1].flatten(), x2[2].flatten()))) res_STR_0, res_STR_1, res_STR_2 = SPACES.extract_2(res_STR) MPI_COMM.Barrier() @@ -382,7 +382,7 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): Q2T = OPS_PSY.Q2.transpose() res_PSY = Q2T.dot(x2_st) - res_STR = OPS_STR.transpose_Q2_dot(np.concatenate((x2[0].flatten(), x2[1].flatten(), x2[2].flatten()))) + res_STR = OPS_STR.transpose_Q2_dot(xp.concatenate((x2[0].flatten(), x2[1].flatten(), x2[2].flatten()))) res_STR_0, res_STR_1, res_STR_2 = SPACES.extract_2(res_STR) MPI_COMM.Barrier() @@ -408,7 +408,7 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): print("\nX1 (V1 --> V0 x V0 x V0):") res_PSY = OPS_PSY.X1.dot(x1_st) - res_STR = OPS_STR.X1_dot(np.concatenate((x1[0].flatten(), x1[1].flatten(), x1[2].flatten()))) + res_STR = OPS_STR.X1_dot(xp.concatenate((x1[0].flatten(), x1[1].flatten(), x1[2].flatten()))) res_STR_0 = SPACES.extract_0(res_STR[0]) res_STR_1 = SPACES.extract_0(res_STR[1]) res_STR_2 = SPACES.extract_0(res_STR[2]) @@ -474,7 +474,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal from struphy.fields_background.equils import ScrewPinch from struphy.geometry import domains from struphy.polar.basic import PolarVector - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -584,11 +584,11 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal x2_pol_psy.tp = x2_psy x3_pol_psy.tp = x3_psy - np.random.seed(1607) - x0_pol_psy.pol = [np.random.rand(x0_pol_psy.pol[0].shape[0], x0_pol_psy.pol[0].shape[1])] - x1_pol_psy.pol = [np.random.rand(x1_pol_psy.pol[n].shape[0], x1_pol_psy.pol[n].shape[1]) for n in range(3)] - x2_pol_psy.pol = [np.random.rand(x2_pol_psy.pol[n].shape[0], x2_pol_psy.pol[n].shape[1]) for n in range(3)] - x3_pol_psy.pol = [np.random.rand(x3_pol_psy.pol[0].shape[0], x3_pol_psy.pol[0].shape[1])] + xp.random.seed(1607) + x0_pol_psy.pol = [xp.random.rand(x0_pol_psy.pol[0].shape[0], x0_pol_psy.pol[0].shape[1])] + x1_pol_psy.pol = [xp.random.rand(x1_pol_psy.pol[n].shape[0], x1_pol_psy.pol[n].shape[1]) for n in range(3)] + x2_pol_psy.pol = [xp.random.rand(x2_pol_psy.pol[n].shape[0], x2_pol_psy.pol[n].shape[1]) for n in range(3)] + x3_pol_psy.pol = [xp.random.rand(x3_pol_psy.pol[0].shape[0], x3_pol_psy.pol[0].shape[1])] # apply boundary conditions to legacy vectors for right shape x0_pol_str = space.B0.dot(x0_pol_psy.toarray(True)) @@ -614,7 +614,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal r_str = mhd_ops_str.PR(x3_pol_str) print(f"Rank {mpi_rank} | Asserting MHD operator K3.") - np.allclose(space.B3.T.dot(r_str), r_psy.toarray(True)) + xp.allclose(space.B3.T.dot(r_str), r_psy.toarray(True)) print(f"Rank {mpi_rank} | Assertion passed.") mpi_comm.Barrier() @@ -627,7 +627,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal r_str = mhd_ops_str.PR.T(x3_pol_str) print(f"Rank {mpi_rank} | Asserting transpose MHD operator K3.T.") - np.allclose(space.B3.T.dot(r_str), r_psy.toarray(True)) + xp.allclose(space.B3.T.dot(r_str), r_psy.toarray(True)) print(f"Rank {mpi_rank} | Assertion passed.") # ===== operator Q2 (V2 --> V2) ============ @@ -644,7 +644,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal r_str = mhd_ops_str.MF(x2_pol_str) print(f"Rank {mpi_rank} | Asserting MHD operator Q2.") - np.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) + xp.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) print(f"Rank {mpi_rank} | Assertion passed.") mpi_comm.Barrier() @@ -657,7 +657,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal r_str = mhd_ops_str.MF.T(x2_pol_str) print(f"Rank {mpi_rank} | Asserting transposed MHD operator Q2.T.") - np.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) + xp.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) print(f"Rank {mpi_rank} | Assertion passed.") # ===== operator T2 (V2 --> V1) ============ @@ -674,7 +674,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal r_str = mhd_ops_str.EF(x2_pol_str) print(f"Rank {mpi_rank} | Asserting MHD operator T2.") - np.allclose(space.B1.T.dot(r_str), r_psy.toarray(True)) + xp.allclose(space.B1.T.dot(r_str), r_psy.toarray(True)) print(f"Rank {mpi_rank} | Assertion passed.") mpi_comm.Barrier() @@ -687,7 +687,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal r_str = mhd_ops_str.EF.T(x1_pol_str) print(f"Rank {mpi_rank} | Asserting transposed MHD operator T2.T.") - np.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) + xp.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) print(f"Rank {mpi_rank} | Assertion passed.") # ===== operator S2 (V2 --> V2) ============ @@ -704,7 +704,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal r_str = mhd_ops_str.PF(x2_pol_str) print(f"Rank {mpi_rank} | Asserting MHD operator S2.") - np.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) + xp.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) print(f"Rank {mpi_rank} | Assertion passed.") mpi_comm.Barrier() @@ -717,7 +717,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal r_str = mhd_ops_str.PF.T(x2_pol_str) print(f"Rank {mpi_rank} | Asserting transposed MHD operator S2.T.") - np.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) + xp.allclose(space.B2.T.dot(r_str), r_psy.toarray(True)) print(f"Rank {mpi_rank} | Assertion passed.") @@ -726,7 +726,7 @@ def assert_ops(mpi_rank, res_PSY, res_STR, verbose=False, MPI_COMM=None): TODO """ - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp if verbose: if MPI_COMM is not None: @@ -789,8 +789,8 @@ def assert_ops(mpi_rank, res_PSY, res_STR, verbose=False, MPI_COMM=None): print( f"Rank {mpi_rank} | Maximum absolute diference (result):\n", - np.max( - np.abs( + xp.max( + xp.abs( res_PSY[ res_PSY.starts[0] : res_PSY.ends[0] + 1, res_PSY.starts[1] : res_PSY.ends[1] + 1, @@ -809,7 +809,7 @@ def assert_ops(mpi_rank, res_PSY, res_STR, verbose=False, MPI_COMM=None): MPI_COMM.Barrier() # Compare results. (Works only for Nel=[N, N, N] so far! TODO: Find this bug!) - assert np.allclose( + assert xp.allclose( res_PSY[ res_PSY.starts[0] : res_PSY.ends[0] + 1, res_PSY.starts[1] : res_PSY.ends[1] + 1, diff --git a/src/struphy/feec/tests/test_derham.py b/src/struphy/feec/tests/test_derham.py index e5cf181c9..fff60bbfc 100644 --- a/src/struphy/feec/tests/test_derham.py +++ b/src/struphy/feec/tests/test_derham.py @@ -14,7 +14,7 @@ def test_psydac_derham(Nel, p, spl_kind): from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import compare_arrays - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -47,11 +47,11 @@ def test_psydac_derham(Nel, p, spl_kind): N3_tot = DR_STR.Ntot_3form # Random vectors for testing - np.random.seed(1981) - x0 = np.random.rand(N0_tot) - x1 = np.random.rand(np.sum(N1_tot)) - x2 = np.random.rand(np.sum(N2_tot)) - x3 = np.random.rand(N3_tot) + xp.random.seed(1981) + x0 = xp.random.rand(N0_tot) + x1 = xp.random.rand(xp.sum(N1_tot)) + x2 = xp.random.rand(xp.sum(N2_tot)) + x3 = xp.random.rand(N3_tot) ############################ ### TEST STENCIL VECTORS ### @@ -174,7 +174,7 @@ def test_psydac_derham(Nel, p, spl_kind): zero2_STR = curl_STR.dot(d1_STR) zero2_PSY = derham.curl.dot(d1_PSY) - assert np.allclose(zero2_STR, np.zeros_like(zero2_STR)) + assert xp.allclose(zero2_STR, xp.zeros_like(zero2_STR)) if rank == 0: print("\nCompare curl of grad:") compare_arrays(zero2_PSY, DR_STR.extract_2(zero2_STR), rank) @@ -183,7 +183,7 @@ def test_psydac_derham(Nel, p, spl_kind): zero3_STR = div_STR.dot(d2_STR) zero3_PSY = derham.div.dot(d2_PSY) - assert np.allclose(zero3_STR, np.zeros_like(zero3_STR)) + assert xp.allclose(zero3_STR, xp.zeros_like(zero3_STR)) if rank == 0: print("\nCompare div of curl:") compare_arrays(zero3_PSY, DR_STR.extract_3(zero3_STR), rank) @@ -201,7 +201,7 @@ def test_psydac_derham(Nel, p, spl_kind): # compare projectors def f(eta1, eta2, eta3): - return np.sin(4 * np.pi * eta1) * np.cos(2 * np.pi * eta2) + np.exp(np.cos(2 * np.pi * eta3)) + return xp.sin(4 * xp.pi * eta1) * xp.cos(2 * xp.pi * eta2) + xp.exp(xp.cos(2 * xp.pi * eta3)) fh0_STR = PI("0", f) fh0_PSY = derham.P["0"](f) diff --git a/src/struphy/feec/tests/test_eval_field.py b/src/struphy/feec/tests/test_eval_field.py index c2c803a1d..8213853b1 100644 --- a/src/struphy/feec/tests/test_eval_field.py +++ b/src/struphy/feec/tests/test_eval_field.py @@ -2,7 +2,7 @@ from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @@ -54,33 +54,33 @@ def test_eval_field(Nel, p, spl_kind): uv.initialize_coeffs(perturbations=[pert_uv_1, pert_uv_2, pert_uv_3]) # evaluation points for meshgrid - eta1 = np.linspace(0, 1, 11) - eta2 = np.linspace(0, 1, 14) - eta3 = np.linspace(0, 1, 18) + eta1 = xp.linspace(0, 1, 11) + eta2 = xp.linspace(0, 1, 14) + eta3 = xp.linspace(0, 1, 18) # evaluation points for markers Np = 33 - markers = np.random.rand(Np, 3) - markers_1 = np.zeros((eta1.size, 3)) + markers = xp.random.rand(Np, 3) + markers_1 = xp.zeros((eta1.size, 3)) markers_1[:, 0] = eta1 - markers_2 = np.zeros((eta2.size, 3)) + markers_2 = xp.zeros((eta2.size, 3)) markers_2[:, 1] = eta2 - markers_3 = np.zeros((eta3.size, 3)) + markers_3 = xp.zeros((eta3.size, 3)) markers_3[:, 2] = eta3 # arrays for legacy evaluation arr1, arr2, arr3, is_sparse_meshgrid = Domain.prepare_eval_pts(eta1, eta2, eta3) - tmp = np.zeros_like(arr1) + tmp = xp.zeros_like(arr1) ###### # V0 # ###### # create legacy arrays with same coeffs - coeffs_loc = np.reshape(p0.vector.toarray(), p0.nbasis) + coeffs_loc = xp.reshape(p0.vector.toarray(), p0.nbasis) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(p0.vector, coeffs, rank) @@ -102,12 +102,12 @@ def test_eval_field(Nel, p, spl_kind): tmp, 0, ) - val_legacy = np.squeeze(tmp.copy()) + val_legacy = xp.squeeze(tmp.copy()) tmp[:] = 0 # distributed evaluation and comparison val = p0(eta1, eta2, eta3, squeeze_out=True) - assert np.allclose(val, val_legacy) + assert xp.allclose(val, val_legacy) # marker evaluation m_vals = p0(markers) @@ -120,19 +120,19 @@ def test_eval_field(Nel, p, spl_kind): m_vals_ref_2 = p0(0.0, eta2, 0.0, squeeze_out=True) m_vals_ref_3 = p0(0.0, 0.0, eta3, squeeze_out=True) - assert np.allclose(m_vals_1, m_vals_ref_1) - assert np.allclose(m_vals_2, m_vals_ref_2) - assert np.allclose(m_vals_3, m_vals_ref_3) + assert xp.allclose(m_vals_1, m_vals_ref_1) + assert xp.allclose(m_vals_2, m_vals_ref_2) + assert xp.allclose(m_vals_3, m_vals_ref_3) ###### # V1 # ###### # create legacy arrays with same coeffs - coeffs_loc = np.reshape(E1.vector[0].toarray(), E1.nbasis[0]) + coeffs_loc = xp.reshape(E1.vector[0].toarray(), E1.nbasis[0]) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(E1.vector[0], coeffs, rank) @@ -154,15 +154,15 @@ def test_eval_field(Nel, p, spl_kind): tmp, 11, ) - val_legacy_1 = np.squeeze(tmp.copy()) + val_legacy_1 = xp.squeeze(tmp.copy()) tmp[:] = 0 # create legacy arrays with same coeffs - coeffs_loc = np.reshape(E1.vector[1].toarray(), E1.nbasis[1]) + coeffs_loc = xp.reshape(E1.vector[1].toarray(), E1.nbasis[1]) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(E1.vector[1], coeffs, rank) @@ -184,15 +184,15 @@ def test_eval_field(Nel, p, spl_kind): tmp, 12, ) - val_legacy_2 = np.squeeze(tmp.copy()) + val_legacy_2 = xp.squeeze(tmp.copy()) tmp[:] = 0 # create legacy arrays with same coeffs - coeffs_loc = np.reshape(E1.vector[2].toarray(), E1.nbasis[2]) + coeffs_loc = xp.reshape(E1.vector[2].toarray(), E1.nbasis[2]) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(E1.vector[2], coeffs, rank) @@ -214,14 +214,14 @@ def test_eval_field(Nel, p, spl_kind): tmp, 13, ) - val_legacy_3 = np.squeeze(tmp.copy()) + val_legacy_3 = xp.squeeze(tmp.copy()) tmp[:] = 0 # distributed evaluation and comparison val1, val2, val3 = E1(eta1, eta2, eta3, squeeze_out=True) - assert np.allclose(val1, val_legacy_1) - assert np.allclose(val2, val_legacy_2) - assert np.allclose(val3, val_legacy_3) + assert xp.allclose(val1, val_legacy_1) + assert xp.allclose(val2, val_legacy_2) + assert xp.allclose(val3, val_legacy_3) # marker evaluation m_vals = E1(markers) @@ -234,25 +234,25 @@ def test_eval_field(Nel, p, spl_kind): m_vals_ref_2 = E1(0.0, eta2, 0.0, squeeze_out=True) m_vals_ref_3 = E1(0.0, 0.0, eta3, squeeze_out=True) - assert np.all( - [np.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)] + assert xp.all( + [xp.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)] ) - assert np.all( - [np.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)] + assert xp.all( + [xp.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)] ) - assert np.all( - [np.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] + assert xp.all( + [xp.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] ) ###### # V2 # ###### # create legacy arrays with same coeffs - coeffs_loc = np.reshape(B2.vector[0].toarray(), B2.nbasis[0]) + coeffs_loc = xp.reshape(B2.vector[0].toarray(), B2.nbasis[0]) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(B2.vector[0], coeffs, rank) @@ -274,15 +274,15 @@ def test_eval_field(Nel, p, spl_kind): tmp, 21, ) - val_legacy_1 = np.squeeze(tmp.copy()) + val_legacy_1 = xp.squeeze(tmp.copy()) tmp[:] = 0 # create legacy arrays with same coeffs - coeffs_loc = np.reshape(B2.vector[1].toarray(), B2.nbasis[1]) + coeffs_loc = xp.reshape(B2.vector[1].toarray(), B2.nbasis[1]) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(B2.vector[1], coeffs, rank) @@ -304,15 +304,15 @@ def test_eval_field(Nel, p, spl_kind): tmp, 22, ) - val_legacy_2 = np.squeeze(tmp.copy()) + val_legacy_2 = xp.squeeze(tmp.copy()) tmp[:] = 0 # create legacy arrays with same coeffs - coeffs_loc = np.reshape(B2.vector[2].toarray(), B2.nbasis[2]) + coeffs_loc = xp.reshape(B2.vector[2].toarray(), B2.nbasis[2]) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(B2.vector[2], coeffs, rank) @@ -334,14 +334,14 @@ def test_eval_field(Nel, p, spl_kind): tmp, 23, ) - val_legacy_3 = np.squeeze(tmp.copy()) + val_legacy_3 = xp.squeeze(tmp.copy()) tmp[:] = 0 # distributed evaluation and comparison val1, val2, val3 = B2(eta1, eta2, eta3, squeeze_out=True) - assert np.allclose(val1, val_legacy_1) - assert np.allclose(val2, val_legacy_2) - assert np.allclose(val3, val_legacy_3) + assert xp.allclose(val1, val_legacy_1) + assert xp.allclose(val2, val_legacy_2) + assert xp.allclose(val3, val_legacy_3) # marker evaluation m_vals = B2(markers) @@ -354,25 +354,25 @@ def test_eval_field(Nel, p, spl_kind): m_vals_ref_2 = B2(0.0, eta2, 0.0, squeeze_out=True) m_vals_ref_3 = B2(0.0, 0.0, eta3, squeeze_out=True) - assert np.all( - [np.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)] + assert xp.all( + [xp.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)] ) - assert np.all( - [np.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)] + assert xp.all( + [xp.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)] ) - assert np.all( - [np.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] + assert xp.all( + [xp.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] ) ###### # V3 # ###### # create legacy arrays with same coeffs - coeffs_loc = np.reshape(n3.vector.toarray(), n3.nbasis) + coeffs_loc = xp.reshape(n3.vector.toarray(), n3.nbasis) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(n3.vector, coeffs, rank) @@ -394,12 +394,12 @@ def test_eval_field(Nel, p, spl_kind): tmp, 3, ) - val_legacy = np.squeeze(tmp.copy()) + val_legacy = xp.squeeze(tmp.copy()) tmp[:] = 0 # distributed evaluation and comparison val = n3(eta1, eta2, eta3, squeeze_out=True) - assert np.allclose(val, val_legacy) + assert xp.allclose(val, val_legacy) # marker evaluation m_vals = n3(markers) @@ -412,19 +412,19 @@ def test_eval_field(Nel, p, spl_kind): m_vals_ref_2 = n3(0.0, eta2, 0.0, squeeze_out=True) m_vals_ref_3 = n3(0.0, 0.0, eta3, squeeze_out=True) - assert np.allclose(m_vals_1, m_vals_ref_1) - assert np.allclose(m_vals_2, m_vals_ref_2) - assert np.allclose(m_vals_3, m_vals_ref_3) + assert xp.allclose(m_vals_1, m_vals_ref_1) + assert xp.allclose(m_vals_2, m_vals_ref_2) + assert xp.allclose(m_vals_3, m_vals_ref_3) ######### # V0vec # ######### # create legacy arrays with same coeffs - coeffs_loc = np.reshape(uv.vector[0].toarray(), uv.nbasis[0]) + coeffs_loc = xp.reshape(uv.vector[0].toarray(), uv.nbasis[0]) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(uv.vector[0], coeffs, rank) @@ -446,15 +446,15 @@ def test_eval_field(Nel, p, spl_kind): tmp, 0, ) - val_legacy_1 = np.squeeze(tmp.copy()) + val_legacy_1 = xp.squeeze(tmp.copy()) tmp[:] = 0 # create legacy arrays with same coeffs - coeffs_loc = np.reshape(uv.vector[1].toarray(), uv.nbasis[1]) + coeffs_loc = xp.reshape(uv.vector[1].toarray(), uv.nbasis[1]) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(uv.vector[1], coeffs, rank) @@ -476,15 +476,15 @@ def test_eval_field(Nel, p, spl_kind): tmp, 0, ) - val_legacy_2 = np.squeeze(tmp.copy()) + val_legacy_2 = xp.squeeze(tmp.copy()) tmp[:] = 0 # create legacy arrays with same coeffs - coeffs_loc = np.reshape(uv.vector[2].toarray(), uv.nbasis[2]) + coeffs_loc = xp.reshape(uv.vector[2].toarray(), uv.nbasis[2]) if isinstance(comm, MockComm): coeffs = coeffs_loc else: - coeffs = np.zeros_like(coeffs_loc) + coeffs = xp.zeros_like(coeffs_loc) comm.Allreduce(coeffs_loc, coeffs, op=MPI.SUM) compare_arrays(uv.vector[2], coeffs, rank) @@ -506,14 +506,14 @@ def test_eval_field(Nel, p, spl_kind): tmp, 0, ) - val_legacy_3 = np.squeeze(tmp.copy()) + val_legacy_3 = xp.squeeze(tmp.copy()) tmp[:] = 0 # distributed evaluation and comparison val1, val2, val3 = uv(eta1, eta2, eta3, squeeze_out=True) - assert np.allclose(val1, val_legacy_1) - assert np.allclose(val2, val_legacy_2) - assert np.allclose(val3, val_legacy_3) + assert xp.allclose(val1, val_legacy_1) + assert xp.allclose(val2, val_legacy_2) + assert xp.allclose(val3, val_legacy_3) # marker evaluation m_vals = uv(markers) @@ -526,14 +526,14 @@ def test_eval_field(Nel, p, spl_kind): m_vals_ref_2 = uv(0.0, eta2, 0.0, squeeze_out=True) m_vals_ref_3 = uv(0.0, 0.0, eta3, squeeze_out=True) - assert np.all( - [np.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)] + assert xp.all( + [xp.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)] ) - assert np.all( - [np.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)] + assert xp.all( + [xp.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)] ) - assert np.all( - [np.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] + assert xp.all( + [xp.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] ) print("\nAll assertions passed.") diff --git a/src/struphy/feec/tests/test_field_init.py b/src/struphy/feec/tests/test_field_init.py index 9c075782a..248c32d58 100644 --- a/src/struphy/feec/tests/test_field_init.py +++ b/src/struphy/feec/tests/test_field_init.py @@ -13,7 +13,7 @@ def test_bckgr_init_const(Nel, p, spl_kind, spaces, vec_comps): from struphy.feec.psydac_derham import Derham from struphy.io.options import FieldsBackground - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -22,14 +22,14 @@ def test_bckgr_init_const(Nel, p, spl_kind, spaces, vec_comps): derham = Derham(Nel, p, spl_kind, comm=comm) # evaluation grids for comparisons - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) - meshgrids = np.meshgrid(e1, e2, e3, indexing="ij") + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) + meshgrids = xp.meshgrid(e1, e2, e3, indexing="ij") # test values - np.random.seed(1234) - val = np.random.rand() + xp.random.seed(1234) + val = xp.random.rand() if val > 0.5: val = int(val * 10) @@ -40,20 +40,20 @@ def test_bckgr_init_const(Nel, p, spl_kind, spaces, vec_comps): background = FieldsBackground(type="LogicalConst", values=(val,)) field.initialize_coeffs(backgrounds=background) print( - f"\n{rank = }, {space = }, after init:\n {np.max(np.abs(field(*meshgrids) - val)) = }", + f"\n{rank = }, {space = }, after init:\n {xp.max(xp.abs(field(*meshgrids) - val)) = }", ) # print(f'{field(*meshgrids) = }') - assert np.allclose(field(*meshgrids), val) + assert xp.allclose(field(*meshgrids), val) else: background = FieldsBackground(type="LogicalConst", values=(val, None, val)) field.initialize_coeffs(backgrounds=background) for j, val in enumerate(background.values): if val is not None: print( - f"\n{rank = }, {space = }, after init:\n {j = }, {np.max(np.abs(field(*meshgrids)[j] - val)) = }", + f"\n{rank = }, {space = }, after init:\n {j = }, {xp.max(xp.abs(field(*meshgrids)[j] - val)) = }", ) # print(f'{field(*meshgrids)[i] = }') - assert np.allclose(field(*meshgrids)[j], val) + assert xp.allclose(field(*meshgrids)[j], val) @pytest.mark.parametrize("Nel", [[18, 24, 12]]) @@ -72,7 +72,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB from struphy.geometry import domains from struphy.io.options import FieldsBackground - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -88,10 +88,10 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show bckgr_4 = FieldsBackground(type="FluidEquilibrium", variable="uv") # evaluation grids for comparisons - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) - meshgrids = np.meshgrid(e1, e2, e3, indexing="ij") + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) + meshgrids = xp.meshgrid(e1, e2, e3, indexing="ij") # test for key, val in inspect.getmembers(equils): @@ -132,8 +132,8 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show elif "ShearedSlab" in key: mhd_equil.domain = domains.Cuboid( r1=mhd_equil.params["a"], - r2=mhd_equil.params["a"] * 2 * np.pi, - r3=mhd_equil.params["R0"] * 2 * np.pi, + r2=mhd_equil.params["a"] * 2 * xp.pi, + r3=mhd_equil.params["R0"] * 2 * xp.pi, ) elif "ShearFluid" in key: mhd_equil.domain = domains.Cuboid( @@ -145,7 +145,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show mhd_equil.domain = domains.HollowCylinder( a1=1e-3, a2=mhd_equil.params["a"], - Lz=mhd_equil.params["R0"] * 2 * np.pi, + Lz=mhd_equil.params["R0"] * 2 * xp.pi, ) else: try: @@ -186,57 +186,57 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show # scalar spaces print( - f"{np.max(np.abs(field_3(*meshgrids) - mhd_equil.p3(*meshgrids))) / np.max(np.abs(mhd_equil.p3(*meshgrids)))}" + f"{xp.max(xp.abs(field_3(*meshgrids) - mhd_equil.p3(*meshgrids))) / xp.max(xp.abs(mhd_equil.p3(*meshgrids)))}" ) assert ( - np.max( - np.abs(field_3(*meshgrids) - mhd_equil.p3(*meshgrids)), + xp.max( + xp.abs(field_3(*meshgrids) - mhd_equil.p3(*meshgrids)), ) - / np.max(np.abs(mhd_equil.p3(*meshgrids))) + / xp.max(xp.abs(mhd_equil.p3(*meshgrids))) < 0.54 ) if isinstance(mhd_equil, FluidEquilibriumWithB): print( - f"{np.max(np.abs(field_0(*meshgrids) - mhd_equil.absB0(*meshgrids))) / np.max(np.abs(mhd_equil.absB0(*meshgrids)))}" + f"{xp.max(xp.abs(field_0(*meshgrids) - mhd_equil.absB0(*meshgrids))) / xp.max(xp.abs(mhd_equil.absB0(*meshgrids)))}" ) assert ( - np.max( - np.abs(field_0(*meshgrids) - mhd_equil.absB0(*meshgrids)), + xp.max( + xp.abs(field_0(*meshgrids) - mhd_equil.absB0(*meshgrids)), ) - / np.max(np.abs(mhd_equil.absB0(*meshgrids))) + / xp.max(xp.abs(mhd_equil.absB0(*meshgrids))) < 0.057 ) print("Scalar asserts passed.") # vector-valued spaces ref = mhd_equil.u1(*meshgrids) - if np.max(np.abs(ref[0])) < 1e-11: + if xp.max(xp.abs(ref[0])) < 1e-11: denom = 1.0 else: - denom = np.max(np.abs(ref[0])) + denom = xp.max(xp.abs(ref[0])) print( - f"{np.max(np.abs(field_1(*meshgrids)[0] - ref[0])) / denom = }", + f"{xp.max(xp.abs(field_1(*meshgrids)[0] - ref[0])) / denom = }", ) - assert np.max(np.abs(field_1(*meshgrids)[0] - ref[0])) / denom < 0.28 - if np.max(np.abs(ref[1])) < 1e-11: + assert xp.max(xp.abs(field_1(*meshgrids)[0] - ref[0])) / denom < 0.28 + if xp.max(xp.abs(ref[1])) < 1e-11: denom = 1.0 else: - denom = np.max(np.abs(ref[1])) + denom = xp.max(xp.abs(ref[1])) print( - f"{np.max(np.abs(field_1(*meshgrids)[1] - ref[1])) / denom = }", + f"{xp.max(xp.abs(field_1(*meshgrids)[1] - ref[1])) / denom = }", ) - assert np.max(np.abs(field_1(*meshgrids)[1] - ref[1])) / denom < 0.33 - if np.max(np.abs(ref[2])) < 1e-11: + assert xp.max(xp.abs(field_1(*meshgrids)[1] - ref[1])) / denom < 0.33 + if xp.max(xp.abs(ref[2])) < 1e-11: denom = 1.0 else: - denom = np.max(np.abs(ref[2])) + denom = xp.max(xp.abs(ref[2])) print( - f"{np.max(np.abs(field_1(*meshgrids)[2] - ref[2])) / denom = }", + f"{xp.max(xp.abs(field_1(*meshgrids)[2] - ref[2])) / denom = }", ) assert ( - np.max( - np.abs( + xp.max( + xp.abs( field_1(*meshgrids)[2] - ref[2], ), ) @@ -246,75 +246,75 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show print("u1 asserts passed.") ref = mhd_equil.u2(*meshgrids) - if np.max(np.abs(ref[0])) < 1e-11: + if xp.max(xp.abs(ref[0])) < 1e-11: denom = 1.0 else: - denom = np.max(np.abs(ref[0])) + denom = xp.max(xp.abs(ref[0])) print( - f"{np.max(np.abs(field_2(*meshgrids)[0] - ref[0])) / denom = }", + f"{xp.max(xp.abs(field_2(*meshgrids)[0] - ref[0])) / denom = }", ) - assert np.max(np.abs(field_2(*meshgrids)[0] - ref[0])) / denom < 0.86 - if np.max(np.abs(ref[1])) < 1e-11: + assert xp.max(xp.abs(field_2(*meshgrids)[0] - ref[0])) / denom < 0.86 + if xp.max(xp.abs(ref[1])) < 1e-11: denom = 1.0 else: - denom = np.max(np.abs(ref[1])) + denom = xp.max(xp.abs(ref[1])) print( - f"{np.max(np.abs(field_2(*meshgrids)[1] - ref[1])) / denom = }", + f"{xp.max(xp.abs(field_2(*meshgrids)[1] - ref[1])) / denom = }", ) assert ( - np.max( - np.abs( + xp.max( + xp.abs( field_2(*meshgrids)[1] - ref[1], ), ) / denom < 0.4 ) - if np.max(np.abs(ref[2])) < 1e-11: + if xp.max(xp.abs(ref[2])) < 1e-11: denom = 1.0 else: - denom = np.max(np.abs(ref[2])) + denom = xp.max(xp.abs(ref[2])) print( - f"{np.max(np.abs(field_2(*meshgrids)[2] - ref[2])) / denom = }", + f"{xp.max(xp.abs(field_2(*meshgrids)[2] - ref[2])) / denom = }", ) - assert np.max(np.abs(field_2(*meshgrids)[2] - ref[2])) / denom < 0.21 + assert xp.max(xp.abs(field_2(*meshgrids)[2] - ref[2])) / denom < 0.21 print("u2 asserts passed.") ref = mhd_equil.uv(*meshgrids) - if np.max(np.abs(ref[0])) < 1e-11: + if xp.max(xp.abs(ref[0])) < 1e-11: denom = 1.0 else: - denom = np.max(np.abs(ref[0])) + denom = xp.max(xp.abs(ref[0])) print( - f"{np.max(np.abs(field_4(*meshgrids)[0] - ref[0])) / denom = }", + f"{xp.max(xp.abs(field_4(*meshgrids)[0] - ref[0])) / denom = }", ) - assert np.max(np.abs(field_4(*meshgrids)[0] - ref[0])) / denom < 0.6 - if np.max(np.abs(ref[1])) < 1e-11: + assert xp.max(xp.abs(field_4(*meshgrids)[0] - ref[0])) / denom < 0.6 + if xp.max(xp.abs(ref[1])) < 1e-11: denom = 1.0 else: - denom = np.max(np.abs(ref[1])) + denom = xp.max(xp.abs(ref[1])) print( - f"{np.max(np.abs(field_4(*meshgrids)[1] - ref[1])) / denom = }", + f"{xp.max(xp.abs(field_4(*meshgrids)[1] - ref[1])) / denom = }", ) assert ( - np.max( - np.abs( + xp.max( + xp.abs( field_4(*meshgrids)[1] - ref[1], ), ) / denom < 0.2 ) - if np.max(np.abs(ref[2])) < 1e-11: + if xp.max(xp.abs(ref[2])) < 1e-11: denom = 1.0 else: - denom = np.max(np.abs(ref[2])) + denom = xp.max(xp.abs(ref[2])) print( - f"{np.max(np.abs(field_4(*meshgrids)[2] - ref[2])) / denom = }", + f"{xp.max(xp.abs(field_4(*meshgrids)[2] - ref[2])) / denom = }", ) assert ( - np.max( - np.abs( + xp.max( + xp.abs( field_4(*meshgrids)[2] - ref[2], ), ) @@ -355,7 +355,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show absB0_h = mhd_equil.domain.push(field_0, *meshgrids) absB0 = mhd_equil.domain.push(mhd_equil.absB0, *meshgrids) - levels = np.linspace(np.min(absB0) - 1e-10, np.max(absB0), 20) + levels = xp.linspace(xp.min(absB0) - 1e-10, xp.max(absB0), 20) plt.figure(f"0/3-forms top, {mhd_equil = }") plt.subplot(2, 3, 1) @@ -493,7 +493,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show p3_h = mhd_equil.domain.push(field_3, *meshgrids) p3 = mhd_equil.domain.push(mhd_equil.p3, *meshgrids) - levels = np.linspace(np.min(p3) - 1e-10, np.max(p3), 20) + levels = xp.linspace(xp.min(p3) - 1e-10, xp.max(p3), 20) plt.figure(f"0/3-forms top, {mhd_equil = }") plt.subplot(2, 3, 2) @@ -640,7 +640,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show ) for i, (bh, b) in enumerate(zip(b1h, b1)): - levels = np.linspace(np.min(b) - 1e-10, np.max(b), 20) + levels = xp.linspace(xp.min(b) - 1e-10, xp.max(b), 20) plt.figure(f"1-forms top, {mhd_equil = }") plt.subplot(2, 3, 1 + i) @@ -789,7 +789,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show ) for i, (bh, b) in enumerate(zip(b2h, b2)): - levels = np.linspace(np.min(b) - 1e-10, np.max(b), 20) + levels = xp.linspace(xp.min(b) - 1e-10, xp.max(b), 20) plt.figure(f"2-forms top, {mhd_equil = }") plt.subplot(2, 3, 1 + i) @@ -938,7 +938,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show ) for i, (bh, b) in enumerate(zip(bvh, bv)): - levels = np.linspace(np.min(b) - 1e-10, np.max(b), 20) + levels = xp.linspace(xp.min(b) - 1e-10, xp.max(b), 20) plt.figure(f"vector-fields top, {mhd_equil = }") plt.subplot(2, 3, 1 + i) @@ -1089,7 +1089,7 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): from struphy.feec.psydac_derham import Derham from struphy.initial.perturbations import ModesCos, ModesSin from struphy.io.options import FieldsBackground - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -1167,10 +1167,10 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): field_2 = derham.create_spline_function("name_2", "Hdiv", backgrounds=bckgr_2, perturbations=[f_cos_22]) # evaluation grids for comparisons - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) - meshgrids = np.meshgrid(e1, e2, e3, indexing="ij") + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) + meshgrids = xp.meshgrid(e1, e2, e3, indexing="ij") fun_0 = avg_0 + f_sin_0(*meshgrids) + f_cos_0(*meshgrids) @@ -1189,24 +1189,24 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): f1_h = field_1(*meshgrids) f2_h = field_2(*meshgrids) - print(f"{np.max(np.abs(fun_0 - f0_h)) = }") - print(f"{np.max(np.abs(fun_1[0] - f1_h[0])) = }") - print(f"{np.max(np.abs(fun_1[1] - f1_h[1])) = }") - print(f"{np.max(np.abs(fun_1[2] - f1_h[2])) = }") - print(f"{np.max(np.abs(fun_2[0] - f2_h[0])) = }") - print(f"{np.max(np.abs(fun_2[1] - f2_h[1])) = }") - print(f"{np.max(np.abs(fun_2[2] - f2_h[2])) = }") - - assert np.max(np.abs(fun_0 - f0_h)) < 3e-5 - assert np.max(np.abs(fun_1[0] - f1_h[0])) < 3e-5 - assert np.max(np.abs(fun_1[1] - f1_h[1])) < 3e-5 - assert np.max(np.abs(fun_1[2] - f1_h[2])) < 3e-5 - assert np.max(np.abs(fun_2[0] - f2_h[0])) < 3e-5 - assert np.max(np.abs(fun_2[1] - f2_h[1])) < 3e-5 - assert np.max(np.abs(fun_2[2] - f2_h[2])) < 3e-5 + print(f"{xp.max(xp.abs(fun_0 - f0_h)) = }") + print(f"{xp.max(xp.abs(fun_1[0] - f1_h[0])) = }") + print(f"{xp.max(xp.abs(fun_1[1] - f1_h[1])) = }") + print(f"{xp.max(xp.abs(fun_1[2] - f1_h[2])) = }") + print(f"{xp.max(xp.abs(fun_2[0] - f2_h[0])) = }") + print(f"{xp.max(xp.abs(fun_2[1] - f2_h[1])) = }") + print(f"{xp.max(xp.abs(fun_2[2] - f2_h[2])) = }") + + assert xp.max(xp.abs(fun_0 - f0_h)) < 3e-5 + assert xp.max(xp.abs(fun_1[0] - f1_h[0])) < 3e-5 + assert xp.max(xp.abs(fun_1[1] - f1_h[1])) < 3e-5 + assert xp.max(xp.abs(fun_1[2] - f1_h[2])) < 3e-5 + assert xp.max(xp.abs(fun_2[0] - f2_h[0])) < 3e-5 + assert xp.max(xp.abs(fun_2[1] - f2_h[1])) < 3e-5 + assert xp.max(xp.abs(fun_2[2] - f2_h[2])) < 3e-5 if show_plot and rank == 0: - levels = np.linspace(np.min(fun_0) - 1e-10, np.max(fun_0), 40) + levels = xp.linspace(xp.min(fun_0) - 1e-10, xp.max(fun_0), 40) plt.figure("0-form", figsize=(10, 16)) plt.subplot(2, 1, 1) @@ -1239,7 +1239,7 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): plt.figure("1-form", figsize=(30, 16)) for i, (f_h, fun) in enumerate(zip(f1_h, fun_1)): - levels = np.linspace(np.min(fun) - 1e-10, np.max(fun), 40) + levels = xp.linspace(xp.min(fun) - 1e-10, xp.max(fun), 40) plt.subplot(2, 3, 1 + i) plt.contourf( @@ -1271,7 +1271,7 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): plt.figure("2-form", figsize=(30, 16)) for i, (f_h, fun) in enumerate(zip(f2_h, fun_2)): - levels = np.linspace(np.min(fun) - 1e-10, np.max(fun), 40) + levels = xp.linspace(xp.min(fun) - 1e-10, xp.max(fun), 40) plt.subplot(2, 3, 1 + i) plt.contourf( @@ -1317,7 +1317,7 @@ def test_noise_init(Nel, p, spl_kind, space, direction): from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import compare_arrays from struphy.initial.perturbations import Noise - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/feec/tests/test_l2_projectors.py b/src/struphy/feec/tests/test_l2_projectors.py index 999215c36..500190075 100644 --- a/src/struphy/feec/tests/test_l2_projectors.py +++ b/src/struphy/feec/tests/test_l2_projectors.py @@ -8,7 +8,7 @@ from struphy.feec.projectors import L2Projector from struphy.feec.psydac_derham import Derham from struphy.geometry import domains -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize("Nel", [[16, 32, 1]]) @@ -28,7 +28,7 @@ def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_desc, do_plo derham = Derham(Nel, p, spl_kind, comm=comm) # constant function - f = lambda e1, e2, e3: np.sin(np.pi * e1) * np.cos(2 * np.pi * e2) + f = lambda e1, e2, e3: xp.sin(xp.pi * e1) * xp.cos(2 * xp.pi * e2) # create domain object dom_types = [] @@ -39,11 +39,11 @@ def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_desc, do_plo dom_classes += [val] # evaluation points - e1 = np.linspace(0.0, 1.0, 30) - e2 = np.linspace(0.0, 1.0, 40) + e1 = xp.linspace(0.0, 1.0, 30) + e2 = xp.linspace(0.0, 1.0, 40) e3 = 0.0 - ee1, ee2, ee3 = np.meshgrid(e1, e2, e3, indexing="ij") + ee1, ee2, ee3 = xp.meshgrid(e1, e2, e3, indexing="ij") for dom_type, dom_class in zip(dom_types, dom_classes): print("#" * 80) @@ -76,12 +76,12 @@ def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_desc, do_plo if array_input: pts_q = derham.quad_grid_pts[sp_key] if sp_id in ("H1", "L2"): - ee = np.meshgrid(*[pt.flatten() for pt in pts_q], indexing="ij") + ee = xp.meshgrid(*[pt.flatten() for pt in pts_q], indexing="ij") f_array = f(*ee) else: f_array = [] for pts in pts_q: - ee = np.meshgrid(*[pt.flatten() for pt in pts], indexing="ij") + ee = xp.meshgrid(*[pt.flatten() for pt in pts], indexing="ij") f_array += [f(*ee)] f_args = f_array else: @@ -91,27 +91,27 @@ def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_desc, do_plo veco = P_L2(f_args, out=out) assert veco is out - assert np.all(vec.toarray() == veco.toarray()) + assert xp.all(vec.toarray() == veco.toarray()) field.vector = vec field_vals = field(e1, e2, e3) if sp_id in ("H1", "L2"): - err = np.max(np.abs(f_analytic(ee1, ee2, ee3) - field_vals)) + err = xp.max(xp.abs(f_analytic(ee1, ee2, ee3) - field_vals)) f_plot = field_vals else: - err = [np.max(np.abs(exact(ee1, ee2, ee3) - field_v)) for exact, field_v in zip(f_analytic, field_vals)] + err = [xp.max(xp.abs(exact(ee1, ee2, ee3) - field_v)) for exact, field_v in zip(f_analytic, field_vals)] f_plot = field_vals[0] - print(f"{sp_id = }, {np.max(err) = }") + print(f"{sp_id = }, {xp.max(err) = }") if sp_id in ("H1", "H1vec"): - assert np.max(err) < 0.004 + assert xp.max(err) < 0.004 else: - assert np.max(err) < 0.12 + assert xp.max(err) < 0.12 if do_plot and rank == 0: plt.figure(f"{dom_type}, {sp_id}") - plt.contourf(e1, e2, np.squeeze(f_plot[:, :, 0].T)) + plt.contourf(e1, e2, xp.squeeze(f_plot[:, :, 0].T)) plt.show() @@ -134,7 +134,7 @@ def test_l2_projectors_convergence(direction, pi, spl_kindi, do_plot=False): for n, Neli in enumerate(Nels): # test function def fun(eta): - return np.cos(4 * np.pi * eta) + return xp.cos(4 * xp.pi * eta) # create derham object, test functions and evaluation points e1 = 0.0 @@ -144,7 +144,7 @@ def fun(eta): Nel = [Neli, 1, 1] p = [pi, 1, 1] spl_kind = [spl_kindi, True, True] - e1 = np.linspace(0.0, 1.0, 100) + e1 = xp.linspace(0.0, 1.0, 100) e = e1 c = 0 @@ -154,7 +154,7 @@ def f(x, y, z): Nel = [1, Neli, 1] p = [1, pi, 1] spl_kind = [True, spl_kindi, True] - e2 = np.linspace(0.0, 1.0, 100) + e2 = xp.linspace(0.0, 1.0, 100) e = e2 c = 1 @@ -164,7 +164,7 @@ def f(x, y, z): Nel = [1, 1, Neli] p = [1, 1, pi] spl_kind = [True, True, spl_kindi] - e3 = np.linspace(0.0, 1.0, 100) + e3 = xp.linspace(0.0, 1.0, 100) e = e3 c = 2 @@ -199,19 +199,19 @@ def f(x, y, z): vec = P_L2(f_analytic) veco = P_L2(f_analytic, out=out) assert veco is out - assert np.all(vec.toarray() == veco.toarray()) + assert xp.all(vec.toarray() == veco.toarray()) field.vector = vec field_vals = field(e1, e2, e3, squeeze_out=True) if sp_id in ("H1", "L2"): - err = np.max(np.abs(f_analytic(e1, e2, e3) - field_vals)) + err = xp.max(xp.abs(f_analytic(e1, e2, e3) - field_vals)) f_plot = field_vals else: - err = [np.max(np.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(f_analytic, field_vals)] + err = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(f_analytic, field_vals)] f_plot = field_vals[0] - errors[sp_id] += [np.max(err)] + errors[sp_id] += [xp.max(err)] if do_plot: plt.figure(sp_id + ", L2-proj. convergence") @@ -232,7 +232,7 @@ def f(x, y, z): line_for_rate_p1 = [Ne ** (-rate_p1) * errors[sp_id][0] / Nels[0] ** (-rate_p1) for Ne in Nels] line_for_rate_p0 = [Ne ** (-rate_p0) * errors[sp_id][0] / Nels[0] ** (-rate_p0) for Ne in Nels] - m, _ = np.polyfit(np.log(Nels), np.log(errors[sp_id]), deg=1) + m, _ = xp.polyfit(xp.log(Nels), xp.log(errors[sp_id]), deg=1) print(f"{sp_id = }, fitted convergence rate = {-m}, degree = {pi}") if sp_id in ("H1", "H1vec"): assert -m > (pi + 1 - 0.05) diff --git a/src/struphy/feec/tests/test_local_projectors.py b/src/struphy/feec/tests/test_local_projectors.py index 21e2392a4..a7549600c 100644 --- a/src/struphy/feec/tests/test_local_projectors.py +++ b/src/struphy/feec/tests/test_local_projectors.py @@ -12,7 +12,7 @@ from struphy.feec.local_projectors_kernels import fill_matrix_column from struphy.feec.psydac_derham import Derham from struphy.feec.utilities_local_projectors import get_one_spline, get_span_and_basis, get_values_and_indices_splines -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def get_span_and_basis(pts, space): @@ -20,7 +20,7 @@ def get_span_and_basis(pts, space): Parameters ---------- - pts : np.array + pts : xp.array 2d array of points (ii, iq) = (interval, quadrature point). space : SplineSpace @@ -28,10 +28,10 @@ def get_span_and_basis(pts, space): Returns ------- - span : np.array + span : xp.array 2d array indexed by (n, nq), where n is the interval and nq is the quadrature point in the interval. - basis : np.array + basis : xp.array 3d array of values of basis functions indexed by (n, nq, basis function). """ @@ -41,8 +41,8 @@ def get_span_and_basis(pts, space): T = space.knots p = space.degree - span = np.zeros(pts.shape, dtype=int) - basis = np.zeros((*pts.shape, p + 1), dtype=float) + span = xp.zeros(pts.shape, dtype=int) + basis = xp.zeros((*pts.shape, p + 1), dtype=float) for n in range(pts.shape[0]): for nq in range(pts.shape[1]): @@ -79,15 +79,15 @@ def test_local_projectors_compare_global(Nel, p, spl_kind): # constant function def f(e1, e2, e3): - return np.sin(2.0 * np.pi * e1) * np.cos(4.0 * np.pi * e2) * np.sin(6.0 * np.pi * e3) + return xp.sin(2.0 * xp.pi * e1) * xp.cos(4.0 * xp.pi * e2) * xp.sin(6.0 * xp.pi * e3) - # f = lambda e1, e2, e3: np.sin(2.0*np.pi*e1) * np.cos(4.0*np.pi*e2) + # f = lambda e1, e2, e3: xp.sin(2.0*xp.pi*e1) * xp.cos(4.0*xp.pi*e2) # evaluation points - e1 = np.linspace(0.0, 1.0, 10) - e2 = np.linspace(0.0, 1.0, 9) - e3 = np.linspace(0.0, 1.0, 8) + e1 = xp.linspace(0.0, 1.0, 10) + e2 = xp.linspace(0.0, 1.0, 9) + e3 = xp.linspace(0.0, 1.0, 8) - ee1, ee2, ee3 = np.meshgrid(e1, e2, e3, indexing="ij") + ee1, ee2, ee3 = xp.meshgrid(e1, e2, e3, indexing="ij") # loop over spaces for sp_id, sp_key in derham.space_to_form.items(): @@ -126,29 +126,29 @@ def f(e1, e2, e3): fieldg_vals = fieldg(e1, e2, e3) if sp_id in ("H1", "L2"): - err = np.max(np.abs(f_analytic(ee1, ee2, ee3) - field_vals)) + err = xp.max(xp.abs(f_analytic(ee1, ee2, ee3) - field_vals)) # Error comparing the global and local projectors - errg = np.max(np.abs(fieldg_vals - field_vals)) + errg = xp.max(xp.abs(fieldg_vals - field_vals)) else: - err = np.zeros(3) - err[0] = np.max(np.abs(f(ee1, ee2, ee3) - field_vals[0])) - err[1] = np.max(np.abs(f(ee1, ee2, ee3) - field_vals[1])) - err[2] = np.max(np.abs(f(ee1, ee2, ee3) - field_vals[2])) + err = xp.zeros(3) + err[0] = xp.max(xp.abs(f(ee1, ee2, ee3) - field_vals[0])) + err[1] = xp.max(xp.abs(f(ee1, ee2, ee3) - field_vals[1])) + err[2] = xp.max(xp.abs(f(ee1, ee2, ee3) - field_vals[2])) # Error comparing the global and local projectors - errg = np.zeros(3) - errg[0] = np.max(np.abs(fieldg_vals[0] - field_vals[0])) - errg[1] = np.max(np.abs(fieldg_vals[1] - field_vals[1])) - errg[2] = np.max(np.abs(fieldg_vals[2] - field_vals[2])) + errg = xp.zeros(3) + errg[0] = xp.max(xp.abs(fieldg_vals[0] - field_vals[0])) + errg[1] = xp.max(xp.abs(fieldg_vals[1] - field_vals[1])) + errg[2] = xp.max(xp.abs(fieldg_vals[2] - field_vals[2])) - print(f"{sp_id = }, {np.max(err) = }, {np.max(errg) = },{exectime = }") + print(f"{sp_id = }, {xp.max(err) = }, {xp.max(errg) = },{exectime = }") if sp_id in ("H1", "H1vec"): - assert np.max(err) < 0.011 - assert np.max(errg) < 0.011 + assert xp.max(err) < 0.011 + assert xp.max(errg) < 0.011 else: - assert np.max(err) < 0.1 - assert np.max(errg) < 0.1 + assert xp.max(err) < 0.1 + assert xp.max(errg) < 0.1 @pytest.mark.parametrize("direction", [0, 1, 2]) @@ -173,7 +173,7 @@ def test_local_projectors_convergence(direction, pi, spl_kindi, do_plot=False): for n, Neli in enumerate(Nels): # test function def fun(eta): - return np.cos(4 * np.pi * eta) + return xp.cos(4 * xp.pi * eta) # create derham object, test functions and evaluation points e1 = 0.0 @@ -183,7 +183,7 @@ def fun(eta): Nel = [Neli, 1, 1] p = [pi, 1, 1] spl_kind = [spl_kindi, True, True] - e1 = np.linspace(0.0, 1.0, 100) + e1 = xp.linspace(0.0, 1.0, 100) e = e1 c = 0 @@ -193,7 +193,7 @@ def f(x, y, z): Nel = [1, Neli, 1] p = [1, pi, 1] spl_kind = [True, spl_kindi, True] - e2 = np.linspace(0.0, 1.0, 100) + e2 = xp.linspace(0.0, 1.0, 100) e = e2 c = 1 @@ -203,7 +203,7 @@ def f(x, y, z): Nel = [1, 1, Neli] p = [1, 1, pi] spl_kind = [True, True, spl_kindi] - e3 = np.linspace(0.0, 1.0, 100) + e3 = xp.linspace(0.0, 1.0, 100) e = e3 c = 2 @@ -232,13 +232,13 @@ def f(x, y, z): field_vals = field(e1, e2, e3, squeeze_out=True) if sp_id in ("H1", "L2"): - err = np.max(np.abs(f_analytic(e1, e2, e3) - field_vals)) + err = xp.max(xp.abs(f_analytic(e1, e2, e3) - field_vals)) f_plot = field_vals else: - err = [np.max(np.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(f_analytic, field_vals)] + err = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(f_analytic, field_vals)] f_plot = field_vals[0] - errors[sp_id] += [np.max(err)] + errors[sp_id] += [xp.max(err)] if do_plot: plt.figure(sp_id + ", Local-proj. convergence") @@ -257,20 +257,20 @@ def f(x, y, z): line_for_rate_p1 = [Ne ** (-rate_p1) * errors[sp_id][0] / Nels[0] ** (-rate_p1) for Ne in Nels] line_for_rate_p0 = [Ne ** (-rate_p0) * errors[sp_id][0] / Nels[0] ** (-rate_p0) for Ne in Nels] - m, _ = np.polyfit(np.log(Nels), np.log(errors[sp_id]), deg=1) + m, _ = xp.polyfit(xp.log(Nels), xp.log(errors[sp_id]), deg=1) if sp_id in ("H1", "H1vec"): # Sometimes for very large number of elements the convergance rate falls of a bit since the error is already so small floating point impressions become relevant # for those cases is better to compute the convergance rate using only the information of Nel with smaller number if -m <= (pi + 1 - 0.1): - m = -np.log2(errors[sp_id][1] / errors[sp_id][2]) + m = -xp.log2(errors[sp_id][1] / errors[sp_id][2]) print(f"{sp_id = }, fitted convergence rate = {-m}, degree = {pi}") assert -m > (pi + 1 - 0.1) else: # Sometimes for very large number of elements the convergance rate falls of a bit since the error is already so small floating point impressions become relevant # for those cases is better to compute the convergance rate using only the information of Nel with smaller number if -m <= (pi - 0.1): - m = -np.log2(errors[sp_id][1] / errors[sp_id][2]) + m = -xp.log2(errors[sp_id][1] / errors[sp_id][2]) print(f"{sp_id = }, fitted convergence rate = {-m}, degree = {pi}") assert -m > (pi - 0.1) @@ -315,12 +315,12 @@ def aux_test_replication_of_basis(Nel, plist, spl_kind): def make_basis_fun(i): def fun(etas, eta2, eta3): if isinstance(etas, float) or isinstance(etas, int): - etas = np.array([etas]) - out = np.zeros_like(etas) + etas = xp.array([etas]) + out = xp.zeros_like(etas) for j, eta in enumerate(etas): span = find_span(T, p, eta) - inds = np.arange(span - p, span + 1) % N - pos = np.argwhere(inds == i) + inds = xp.arange(span - p, span + 1) % N + pos = xp.argwhere(inds == i) # print(f'{pos = }') if pos.size > 0: pos = pos[0, 0] @@ -335,18 +335,18 @@ def fun(etas, eta2, eta3): fun = make_basis_fun(j) lambdas = P_Loc(fun).toarray() - etas = np.linspace(0.0, 1.0, 100) - fun_h = np.zeros(100) + etas = xp.linspace(0.0, 1.0, 100) + fun_h = xp.zeros(100) for k, eta in enumerate(etas): span = find_span(T, p, eta) - ind1 = np.arange(span - p, span + 1) % N + ind1 = xp.arange(span - p, span + 1) % N basis = basis_funs(T, p, eta, span, normalize=normalize) fun_h[k] = evaluation_kernel_1d(p, basis, ind1, lambdas) - if np.max(np.abs(fun(etas, 0.0, 0.0) - fun_h)) >= 10.0**-10: - print(np.max(np.abs(fun(etas, 0.0, 0.0) - fun_h))) - assert np.max(np.abs(fun(etas, 0.0, 0.0) - fun_h)) < 10.0**-10 - # print(f'{j = }, max error: {np.max(np.abs(fun(etas,0.0,0.0) - fun_h))}') + if xp.max(xp.abs(fun(etas, 0.0, 0.0) - fun_h)) >= 10.0**-10: + print(xp.max(xp.abs(fun(etas, 0.0, 0.0) - fun_h))) + assert xp.max(xp.abs(fun(etas, 0.0, 0.0) - fun_h)) < 10.0**-10 + # print(f'{j = }, max error: {xp.max(xp.abs(fun(etas,0.0,0.0) - fun_h))}') # For D-splines @@ -421,7 +421,7 @@ def test_basis_projection_operator_local(Nel, plist, spl_kind, out_sp_key, in_sp # Helper function to handle reshaping and getting spans and basis def process_eta(eta, w1d): if isinstance(eta, (float, int)): - eta = np.array([eta]) + eta = xp.array([eta]) if len(eta.shape) == 1: eta = eta.reshape((eta.shape[0], 1)) spans, values = get_span_and_basis(eta, w1d) @@ -434,7 +434,7 @@ def fun(eta1, eta2, eta3): eta = eta_map[dim_idx] w1d = W1ds[0][dim_idx] if is_B else V1ds[0][dim_idx] - out = np.zeros_like(eta) + out = xp.zeros_like(eta) for j1 in range(eta.shape[0]): for j2 in range(eta.shape[1]): for j3 in range(eta.shape[2]): @@ -477,21 +477,21 @@ def fun(eta1, eta2, eta3): if out_sp_key == "0" or out_sp_key == "3": npts_out = derham.Vh[out_sp_key].npts - starts = np.array(out.starts, dtype=int) - ends = np.array(out.ends, dtype=int) - pds = np.array(out.pads, dtype=int) + starts = xp.array(out.starts, dtype=int) + ends = xp.array(out.ends, dtype=int) + pds = xp.array(out.pads, dtype=int) VFEM1ds = [VFEM.spaces] - nbasis_out = np.array([VFEM1ds[0][0].nbasis, VFEM1ds[0][1].nbasis, VFEM1ds[0][2].nbasis]) + nbasis_out = xp.array([VFEM1ds[0][0].nbasis, VFEM1ds[0][1].nbasis, VFEM1ds[0][2].nbasis]) else: - npts_out = np.array([sp.npts for sp in P_Loc.coeff_space.spaces]) - pds = np.array([vi.pads for vi in P_Loc.coeff_space.spaces]) - starts = np.array([vi.starts for vi in P_Loc.coeff_space.spaces]) - ends = np.array([vi.ends for vi in P_Loc.coeff_space.spaces]) - starts = np.array(starts, dtype=int) - ends = np.array(ends, dtype=int) - pds = np.array(pds, dtype=int) + npts_out = xp.array([sp.npts for sp in P_Loc.coeff_space.spaces]) + pds = xp.array([vi.pads for vi in P_Loc.coeff_space.spaces]) + starts = xp.array([vi.starts for vi in P_Loc.coeff_space.spaces]) + ends = xp.array([vi.ends for vi in P_Loc.coeff_space.spaces]) + starts = xp.array(starts, dtype=int) + ends = xp.array(ends, dtype=int) + pds = xp.array(pds, dtype=int) VFEM1ds = [comp.spaces for comp in VFEM.spaces] - nbasis_out = np.array( + nbasis_out = xp.array( [ [VFEM1ds[0][0].nbasis, VFEM1ds[0][1].nbasis, VFEM1ds[0][2].nbasis], [ @@ -506,7 +506,7 @@ def fun(eta1, eta2, eta3): if in_sp_key == "0" or in_sp_key == "3": npts_in = derham.Vh[in_sp_key].npts else: - npts_in = np.array([sp.npts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) + npts_in = xp.array([sp.npts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) def define_basis(in_sp_key): def wrapper(dim, index, h=None): @@ -556,13 +556,13 @@ def basis3(i3, h=None): input[random_i0, random_i1, random_i2] = 1.0 input.update_ghost_regions() else: - npts_in = np.array([sp.npts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) + npts_in = xp.array([sp.npts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) random_h = random.randrange(0, 3) random_i0 = random.randrange(0, npts_in[random_h][0]) random_i1 = random.randrange(0, npts_in[random_h][1]) random_i2 = random.randrange(0, npts_in[random_h][2]) - starts_in = np.array([sp.starts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) - ends_in = np.array([sp.ends for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) + starts_in = xp.array([sp.starts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) + ends_in = xp.array([sp.ends for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) if starts_in[random_h][0] <= random_i0 and random_i0 <= ends_in[random_h][0]: input[random_h][random_i0, random_i1, random_i2] = 1.0 input.update_ghost_regions() @@ -570,9 +570,9 @@ def basis3(i3, h=None): # We define the matrix if out_sp_key == "0" or out_sp_key == "3": if in_sp_key == "0" or in_sp_key == "3": - matrix = np.zeros((npts_out[0] * npts_out[1] * npts_out[2], npts_in[0] * npts_in[1] * npts_in[2])) + matrix = xp.zeros((npts_out[0] * npts_out[1] * npts_out[2], npts_in[0] * npts_in[1] * npts_in[2])) else: - matrix = np.zeros( + matrix = xp.zeros( ( npts_out[0] * npts_out[1] * npts_out[2], npts_in[0][0] * npts_in[0][1] * npts_in[0][2] @@ -583,61 +583,61 @@ def basis3(i3, h=None): else: if in_sp_key == "0" or in_sp_key == "3": - matrix0 = np.zeros((npts_out[0][0] * npts_out[0][1] * npts_out[0][2], npts_in[0] * npts_in[1] * npts_in[2])) - matrix1 = np.zeros((npts_out[1][0] * npts_out[1][1] * npts_out[1][2], npts_in[0] * npts_in[1] * npts_in[2])) - matrix2 = np.zeros((npts_out[2][0] * npts_out[2][1] * npts_out[2][2], npts_in[0] * npts_in[1] * npts_in[2])) + matrix0 = xp.zeros((npts_out[0][0] * npts_out[0][1] * npts_out[0][2], npts_in[0] * npts_in[1] * npts_in[2])) + matrix1 = xp.zeros((npts_out[1][0] * npts_out[1][1] * npts_out[1][2], npts_in[0] * npts_in[1] * npts_in[2])) + matrix2 = xp.zeros((npts_out[2][0] * npts_out[2][1] * npts_out[2][2], npts_in[0] * npts_in[1] * npts_in[2])) else: - matrix00 = np.zeros( + matrix00 = xp.zeros( ( npts_out[0][0] * npts_out[0][1] * npts_out[0][2], npts_in[0][0] * npts_in[0][1] * npts_in[0][2], ) ) - matrix10 = np.zeros( + matrix10 = xp.zeros( ( npts_out[1][0] * npts_out[1][1] * npts_out[1][2], npts_in[0][0] * npts_in[0][1] * npts_in[0][2], ) ) - matrix20 = np.zeros( + matrix20 = xp.zeros( ( npts_out[2][0] * npts_out[2][1] * npts_out[2][2], npts_in[0][0] * npts_in[0][1] * npts_in[0][2], ) ) - matrix01 = np.zeros( + matrix01 = xp.zeros( ( npts_out[0][0] * npts_out[0][1] * npts_out[0][2], npts_in[1][0] * npts_in[1][1] * npts_in[1][2], ) ) - matrix11 = np.zeros( + matrix11 = xp.zeros( ( npts_out[1][0] * npts_out[1][1] * npts_out[1][2], npts_in[1][0] * npts_in[1][1] * npts_in[1][2], ) ) - matrix21 = np.zeros( + matrix21 = xp.zeros( ( npts_out[2][0] * npts_out[2][1] * npts_out[2][2], npts_in[1][0] * npts_in[1][1] * npts_in[1][2], ) ) - matrix02 = np.zeros( + matrix02 = xp.zeros( ( npts_out[0][0] * npts_out[0][1] * npts_out[0][2], npts_in[2][0] * npts_in[2][1] * npts_in[2][2], ) ) - matrix12 = np.zeros( + matrix12 = xp.zeros( ( npts_out[1][0] * npts_out[1][1] * npts_out[1][2], npts_in[2][0] * npts_in[2][1] * npts_in[2][2], ) ) - matrix22 = np.zeros( + matrix22 = xp.zeros( ( npts_out[2][0] * npts_out[2][1] * npts_out[2][2], npts_in[2][0] * npts_in[2][1] * npts_in[2][2], @@ -647,7 +647,7 @@ def basis3(i3, h=None): # We build the BasisProjectionOperator by hand if out_sp_key == "0" or out_sp_key == "3": if in_sp_key == "0" or in_sp_key == "3": - # def f_analytic(e1,e2,e3): return (np.sin(2.0*np.pi*e1)+np.cos(4.0*np.pi*e2))*basis1(random_i0)(e1,e2,e3)*basis2(random_i1)(e1,e2,e3)*basis3(random_i2)(e1,e2,e3) + # def f_analytic(e1,e2,e3): return (xp.sin(2.0*xp.pi*e1)+xp.cos(4.0*xp.pi*e2))*basis1(random_i0)(e1,e2,e3)*basis2(random_i1)(e1,e2,e3)*basis3(random_i2)(e1,e2,e3) # out = P_Loc(f_analytic) counter = 0 @@ -657,7 +657,7 @@ def basis3(i3, h=None): def f_analytic(e1, e2, e3): return ( - (np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e2)) + (xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e2)) * basis1(col0)(e1, e2, e3) * basis2(col1)(e1, e2, e3) * basis3(col2)(e1, e2, e3) @@ -677,7 +677,7 @@ def f_analytic(e1, e2, e3): def f_analytic(e1, e2, e3): return ( - (np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e2)) + (xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e2)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -696,7 +696,7 @@ def f_analytic(e1, e2, e3): def f_analytic1(e1, e2, e3): return ( - (np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e2)) + (xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e2)) * basis1(col0)(e1, e2, e3) * basis2(col1)(e1, e2, e3) * basis3(col2)(e1, e2, e3) @@ -704,7 +704,7 @@ def f_analytic1(e1, e2, e3): def f_analytic2(e1, e2, e3): return ( - (np.cos(2.0 * np.pi * e2) + np.cos(6.0 * np.pi * e3)) + (xp.cos(2.0 * xp.pi * e2) + xp.cos(6.0 * xp.pi * e3)) * basis1(col0)(e1, e2, e3) * basis2(col1)(e1, e2, e3) * basis3(col2)(e1, e2, e3) @@ -712,7 +712,7 @@ def f_analytic2(e1, e2, e3): def f_analytic3(e1, e2, e3): return ( - (np.sin(6.0 * np.pi * e1) + np.sin(4.0 * np.pi * e3)) + (xp.sin(6.0 * xp.pi * e1) + xp.sin(4.0 * xp.pi * e3)) * basis1(col0)(e1, e2, e3) * basis2(col1)(e1, e2, e3) * basis3(col2)(e1, e2, e3) @@ -724,7 +724,7 @@ def f_analytic3(e1, e2, e3): fill_matrix_column(starts[2], ends[2], pds[2], counter, nbasis_out[2], matrix2, out[2]._data) counter += 1 - matrix = np.vstack((matrix0, matrix1, matrix2)) + matrix = xp.vstack((matrix0, matrix1, matrix2)) else: for h in range(3): @@ -736,7 +736,7 @@ def f_analytic3(e1, e2, e3): def f_analytic0(e1, e2, e3): return ( - (np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e2)) + (xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e2)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -744,7 +744,7 @@ def f_analytic0(e1, e2, e3): def f_analytic1(e1, e2, e3): return ( - (np.sin(10.0 * np.pi * e1) + np.cos(41.0 * np.pi * e2)) + (xp.sin(10.0 * xp.pi * e1) + xp.cos(41.0 * xp.pi * e2)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -752,7 +752,7 @@ def f_analytic1(e1, e2, e3): def f_analytic2(e1, e2, e3): return ( - (np.sin(25.0 * np.pi * e1) + np.cos(49.0 * np.pi * e2)) + (xp.sin(25.0 * xp.pi * e1) + xp.cos(49.0 * xp.pi * e2)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -762,7 +762,7 @@ def f_analytic2(e1, e2, e3): def f_analytic0(e1, e2, e3): return ( - (np.cos(2.0 * np.pi * e2) + np.cos(6.0 * np.pi * e3)) + (xp.cos(2.0 * xp.pi * e2) + xp.cos(6.0 * xp.pi * e3)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -770,7 +770,7 @@ def f_analytic0(e1, e2, e3): def f_analytic1(e1, e2, e3): return ( - (np.cos(12.0 * np.pi * e2) + np.cos(62.0 * np.pi * e3)) + (xp.cos(12.0 * xp.pi * e2) + xp.cos(62.0 * xp.pi * e3)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -778,7 +778,7 @@ def f_analytic1(e1, e2, e3): def f_analytic2(e1, e2, e3): return ( - (np.cos(25.0 * np.pi * e2) + np.cos(68.0 * np.pi * e3)) + (xp.cos(25.0 * xp.pi * e2) + xp.cos(68.0 * xp.pi * e3)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -787,7 +787,7 @@ def f_analytic2(e1, e2, e3): def f_analytic0(e1, e2, e3): return ( - (np.sin(6.0 * np.pi * e1) + np.sin(4.0 * np.pi * e3)) + (xp.sin(6.0 * xp.pi * e1) + xp.sin(4.0 * xp.pi * e3)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -795,7 +795,7 @@ def f_analytic0(e1, e2, e3): def f_analytic1(e1, e2, e3): return ( - (np.sin(16.0 * np.pi * e1) + np.sin(43.0 * np.pi * e3)) + (xp.sin(16.0 * xp.pi * e1) + xp.sin(43.0 * xp.pi * e3)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -803,7 +803,7 @@ def f_analytic1(e1, e2, e3): def f_analytic2(e1, e2, e3): return ( - (np.sin(65.0 * np.pi * e1) + np.sin(47.0 * np.pi * e3)) + (xp.sin(65.0 * xp.pi * e1) + xp.sin(47.0 * xp.pi * e3)) * basis1(col0, h)(e1, e2, e3) * basis2(col1, h)(e1, e2, e3) * basis3(col2, h)(e1, e2, e3) @@ -898,23 +898,23 @@ def f_analytic2(e1, e2, e3): ) counter += 1 - matrix0 = np.hstack((matrix00, matrix01, matrix02)) - matrix1 = np.hstack((matrix10, matrix11, matrix12)) - matrix2 = np.hstack((matrix20, matrix21, matrix22)) - matrix = np.vstack((matrix0, matrix1, matrix2)) + matrix0 = xp.hstack((matrix00, matrix01, matrix02)) + matrix1 = xp.hstack((matrix10, matrix11, matrix12)) + matrix2 = xp.hstack((matrix20, matrix21, matrix22)) + matrix = xp.vstack((matrix0, matrix1, matrix2)) # Now we build the same matrix using the BasisProjectionOperatorLocal if out_sp_key == "0" or out_sp_key == "3": if in_sp_key == "0" or in_sp_key == "3": def f_analytic(e1, e2, e3): - return np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e2) + return xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e2) matrix_new = BasisProjectionOperatorLocal(P_Loc, derham.Vh_fem[in_sp_key], [[f_analytic]], transposed=False) else: def f_analytic(e1, e2, e3): - return np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e2) + return xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e2) matrix_new = BasisProjectionOperatorLocal( P_Loc, @@ -929,13 +929,13 @@ def f_analytic(e1, e2, e3): if in_sp_key == "0" or in_sp_key == "3": def f_analytic1(e1, e2, e3): - return np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e2) + return xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e2) def f_analytic2(e1, e2, e3): - return np.cos(2.0 * np.pi * e2) + np.cos(6.0 * np.pi * e3) + return xp.cos(2.0 * xp.pi * e2) + xp.cos(6.0 * xp.pi * e3) def f_analytic3(e1, e2, e3): - return np.sin(6.0 * np.pi * e1) + np.sin(4.0 * np.pi * e3) + return xp.sin(6.0 * xp.pi * e1) + xp.sin(4.0 * xp.pi * e3) matrix_new = BasisProjectionOperatorLocal( P_Loc, @@ -952,31 +952,31 @@ def f_analytic3(e1, e2, e3): else: def f_analytic00(e1, e2, e3): - return np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e2) + return xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e2) def f_analytic01(e1, e2, e3): - return np.cos(2.0 * np.pi * e2) + np.cos(6.0 * np.pi * e3) + return xp.cos(2.0 * xp.pi * e2) + xp.cos(6.0 * xp.pi * e3) def f_analytic02(e1, e2, e3): - return np.sin(6.0 * np.pi * e1) + np.sin(4.0 * np.pi * e3) + return xp.sin(6.0 * xp.pi * e1) + xp.sin(4.0 * xp.pi * e3) def f_analytic10(e1, e2, e3): - return np.sin(10.0 * np.pi * e1) + np.cos(41.0 * np.pi * e2) + return xp.sin(10.0 * xp.pi * e1) + xp.cos(41.0 * xp.pi * e2) def f_analytic11(e1, e2, e3): - return np.cos(12.0 * np.pi * e2) + np.cos(62.0 * np.pi * e3) + return xp.cos(12.0 * xp.pi * e2) + xp.cos(62.0 * xp.pi * e3) def f_analytic12(e1, e2, e3): - return np.sin(16.0 * np.pi * e1) + np.sin(43.0 * np.pi * e3) + return xp.sin(16.0 * xp.pi * e1) + xp.sin(43.0 * xp.pi * e3) def f_analytic20(e1, e2, e3): - return np.sin(25.0 * np.pi * e1) + np.cos(49.0 * np.pi * e2) + return xp.sin(25.0 * xp.pi * e1) + xp.cos(49.0 * xp.pi * e2) def f_analytic21(e1, e2, e3): - return np.cos(25.0 * np.pi * e2) + np.cos(68.0 * np.pi * e3) + return xp.cos(25.0 * xp.pi * e2) + xp.cos(68.0 * xp.pi * e3) def f_analytic22(e1, e2, e3): - return np.sin(65.0 * np.pi * e1) + np.sin(47.0 * np.pi * e3) + return xp.sin(65.0 * xp.pi * e1) + xp.sin(47.0 * xp.pi * e3) matrix_new = BasisProjectionOperatorLocal( P_Loc, @@ -993,7 +993,7 @@ def f_analytic22(e1, e2, e3): transposed=False, ) - compare_arrays(matrix_new.dot(v), np.matmul(matrix, varr), rank) + compare_arrays(matrix_new.dot(v), xp.matmul(matrix, varr), rank) print("BasisProjectionOperatorLocal test passed.") @@ -1029,7 +1029,7 @@ def test_basis_projection_operator_local_new(Nel, plist, spl_kind, out_sp_key, i # Helper function to handle reshaping and getting spans and basis def process_eta(eta, w1d): if isinstance(eta, (float, int)): - eta = np.array([eta]) + eta = xp.array([eta]) if len(eta.shape) == 1: eta = eta.reshape((eta.shape[0], 1)) spans, values = get_span_and_basis(eta, w1d) @@ -1042,7 +1042,7 @@ def fun(eta1, eta2, eta3): eta = eta_map[dim_idx] w1d = W1ds[0][dim_idx] if is_B else V1ds[0][dim_idx] - out = np.zeros_like(eta) + out = xp.zeros_like(eta) for j1 in range(eta.shape[0]): for j2 in range(eta.shape[1]): for j3 in range(eta.shape[2]): @@ -1119,22 +1119,22 @@ def basis3(i3, h=None): input[random_i0, random_i1, random_i2] = 1.0 input.update_ghost_regions() else: - npts_in = np.array([sp.npts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) + npts_in = xp.array([sp.npts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) random_h = random.randrange(0, 3) random_i0 = random.randrange(0, npts_in[random_h][0]) random_i1 = random.randrange(0, npts_in[random_h][1]) random_i2 = random.randrange(0, npts_in[random_h][2]) - starts = np.array([sp.starts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) - ends = np.array([sp.ends for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) + starts = xp.array([sp.starts for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) + ends = xp.array([sp.ends for sp in derham.Vh_fem[in_sp_key].coeff_space.spaces]) if starts[random_h][0] <= random_i0 and random_i0 <= ends[random_h][0]: input[random_h][random_i0, random_i1, random_i2] = 1.0 input.update_ghost_regions() - etas1 = np.linspace(0.0, 1.0, 1000) - etas2 = np.array([0.5]) + etas1 = xp.linspace(0.0, 1.0, 1000) + etas2 = xp.array([0.5]) - etas3 = np.array([0.5]) - meshgrid = np.meshgrid(*[etas1, etas2, etas3], indexing="ij") + etas3 = xp.array([0.5]) + meshgrid = xp.meshgrid(*[etas1, etas2, etas3], indexing="ij") # Now we build the same matrix using the BasisProjectionOperatorLocal and BasisProjectionOperator @@ -1142,7 +1142,7 @@ def basis3(i3, h=None): if in_sp_key == "0" or in_sp_key == "3": def f_analytic(e1, e2, e3): - return np.sin(2.0 * np.pi * e1) + np.sin(4.0 * np.pi * e1) + return xp.sin(2.0 * xp.pi * e1) + xp.sin(4.0 * xp.pi * e1) matrix_new = BasisProjectionOperatorLocal(P_Loc, derham.Vh_fem[in_sp_key], [[f_analytic]], transposed=False) matrix_global = BasisProjectionOperator(P, derham.Vh_fem[in_sp_key], [[f_analytic]], transposed=False) @@ -1156,7 +1156,7 @@ def f_analytic(e1, e2, e3): else: def f_analytic(e1, e2, e3): - return np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e1) + return xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e1) matrix_new = BasisProjectionOperatorLocal( P_Loc, @@ -1186,13 +1186,13 @@ def f_analytic(e1, e2, e3): if in_sp_key == "0" or in_sp_key == "3": def f_analytic1(e1, e2, e3): - return np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e1) + return xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e1) def f_analytic2(e1, e2, e3): - return np.cos(2.0 * np.pi * e1) + np.cos(6.0 * np.pi * e1) + return xp.cos(2.0 * xp.pi * e1) + xp.cos(6.0 * xp.pi * e1) def f_analytic3(e1, e2, e3): - return np.sin(6.0 * np.pi * e1) + np.sin(4.0 * np.pi * e1) + return xp.sin(6.0 * xp.pi * e1) + xp.sin(4.0 * xp.pi * e1) matrix_new = BasisProjectionOperatorLocal( P_Loc, @@ -1219,7 +1219,7 @@ def f_analytic3(e1, e2, e3): transposed=False, ) - analytic_vals = np.array( + analytic_vals = xp.array( [ f_analytic1(*meshgrid) * basis1(random_i0)(*meshgrid) @@ -1238,31 +1238,31 @@ def f_analytic3(e1, e2, e3): else: def f_analytic00(e1, e2, e3): - return np.sin(2.0 * np.pi * e1) + np.cos(4.0 * np.pi * e1) + return xp.sin(2.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e1) def f_analytic01(e1, e2, e3): - return np.cos(2.0 * np.pi * e1) + np.cos(6.0 * np.pi * e1) + return xp.cos(2.0 * xp.pi * e1) + xp.cos(6.0 * xp.pi * e1) def f_analytic02(e1, e2, e3): - return np.sin(6.0 * np.pi * e1) + np.sin(4.0 * np.pi * e1) + return xp.sin(6.0 * xp.pi * e1) + xp.sin(4.0 * xp.pi * e1) def f_analytic10(e1, e2, e3): - return np.sin(3.0 * np.pi * e1) + np.cos(4.0 * np.pi * e1) + return xp.sin(3.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e1) def f_analytic11(e1, e2, e3): - return np.cos(2.0 * np.pi * e1) + np.cos(3.0 * np.pi * e1) + return xp.cos(2.0 * xp.pi * e1) + xp.cos(3.0 * xp.pi * e1) def f_analytic12(e1, e2, e3): - return np.sin(5.0 * np.pi * e1) + np.sin(3.0 * np.pi * e1) + return xp.sin(5.0 * xp.pi * e1) + xp.sin(3.0 * xp.pi * e1) def f_analytic20(e1, e2, e3): - return np.sin(5.0 * np.pi * e1) + np.cos(4.0 * np.pi * e1) + return xp.sin(5.0 * xp.pi * e1) + xp.cos(4.0 * xp.pi * e1) def f_analytic21(e1, e2, e3): - return np.cos(5.0 * np.pi * e1) + np.cos(6.0 * np.pi * e1) + return xp.cos(5.0 * xp.pi * e1) + xp.cos(6.0 * xp.pi * e1) def f_analytic22(e1, e2, e3): - return np.sin(5.0 * np.pi * e1) + np.sin(4.0 * np.pi * e1) + return xp.sin(5.0 * xp.pi * e1) + xp.sin(4.0 * xp.pi * e1) matrix_new = BasisProjectionOperatorLocal( P_Loc, @@ -1300,7 +1300,7 @@ def f_analytic22(e1, e2, e3): } # Use the map to get analytic values - analytic_vals = np.array( + analytic_vals = xp.array( [ f_analytic_map[dim][random_h](*meshgrid) * basis1(random_i0, random_h)(*meshgrid) @@ -1330,14 +1330,14 @@ def f_analytic22(e1, e2, e3): fieldglo = derham.create_spline_function("fh", out_sp_id) fieldglo.vector = FE_glo - errorloc = np.abs(fieldloc(*meshgrid) - analytic_vals) - errorglo = np.abs(fieldglo(*meshgrid) - analytic_vals) + errorloc = xp.abs(fieldloc(*meshgrid) - analytic_vals) + errorglo = xp.abs(fieldglo(*meshgrid) - analytic_vals) - meanlocal = np.mean(errorloc) - maxlocal = np.max(errorloc) + meanlocal = xp.mean(errorloc) + maxlocal = xp.max(errorloc) - meanglobal = np.mean(errorglo) - maxglobal = np.max(errorglo) + meanglobal = xp.mean(errorglo) + maxglobal = xp.max(errorglo) if isinstance(comm, MockComm): reducemeanlocal = meanlocal @@ -1424,7 +1424,7 @@ def aux_test_spline_evaluation(Nel, plist, spl_kind): # Helper function to handle reshaping and getting spans and basis def process_eta(eta, w1d): if isinstance(eta, (float, int)): - eta = np.array([eta]) + eta = xp.array([eta]) if len(eta.shape) == 1: eta = eta.reshape((eta.shape[0], 1)) spans, values = get_span_and_basis(eta, w1d) @@ -1437,7 +1437,7 @@ def fun(eta1, eta2, eta3): eta = eta_map[dim_idx] w1d = W1ds[0][dim_idx] if is_B else V1ds[0][dim_idx] - out = np.zeros_like(eta) + out = xp.zeros_like(eta) for j1 in range(eta.shape[0]): for j2 in range(eta.shape[1]): for j3 in range(eta.shape[2]): @@ -1471,10 +1471,10 @@ def fun(eta1, eta2, eta3): fieldD = derham.create_spline_function("fh", "L2") npts_in_D = derham.Vh["3"].npts - etas1 = np.linspace(0.0, 1.0, 20) - etas2 = np.linspace(0.0, 1.0, 20) - etas3 = np.linspace(0.0, 1.0, 20) - meshgrid = np.meshgrid(*[etas1, etas2, etas3], indexing="ij") + etas1 = xp.linspace(0.0, 1.0, 20) + etas2 = xp.linspace(0.0, 1.0, 20) + etas3 = xp.linspace(0.0, 1.0, 20) + meshgrid = xp.meshgrid(*[etas1, etas2, etas3], indexing="ij") maxerrorB = 0.0 @@ -1487,7 +1487,7 @@ def fun(eta1, eta2, eta3): fieldB.vector = inputB def error(e1, e2, e3): - return np.abs( + return xp.abs( fieldB(e1, e2, e3) - ( make_basis_fun(True, 0, col0)(e1, e2, e3) @@ -1496,7 +1496,7 @@ def error(e1, e2, e3): ), ) - auxerror = np.max(error(*meshgrid)) + auxerror = xp.max(error(*meshgrid)) if auxerror > maxerrorB: maxerrorB = auxerror @@ -1515,7 +1515,7 @@ def error(e1, e2, e3): fieldD.vector = inputD def error(e1, e2, e3): - return np.abs( + return xp.abs( fieldD(e1, e2, e3) - ( make_basis_fun(False, 0, col0)(e1, e2, e3) @@ -1524,7 +1524,7 @@ def error(e1, e2, e3): ), ) - auxerror = np.max(error(*meshgrid)) + auxerror = xp.max(error(*meshgrid)) if auxerror > maxerrorD: maxerrorD = auxerror diff --git a/src/struphy/feec/tests/test_lowdim_nel_is_1.py b/src/struphy/feec/tests/test_lowdim_nel_is_1.py index cdc7e0705..476c8f0c4 100644 --- a/src/struphy/feec/tests/test_lowdim_nel_is_1.py +++ b/src/struphy/feec/tests/test_lowdim_nel_is_1.py @@ -13,7 +13,7 @@ def test_lowdim_derham(Nel, p, spl_kind, do_plot=False): from psydac.linalg.stencil import StencilVector from struphy.feec.psydac_derham import Derham - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -74,17 +74,17 @@ def test_lowdim_derham(Nel, p, spl_kind, do_plot=False): ### TEST COMMUTING PROJECTORS ### ################################# def fun(eta): - return np.cos(2 * np.pi * eta) + return xp.cos(2 * xp.pi * eta) def dfun(eta): - return -2 * np.pi * np.sin(2 * np.pi * eta) + return -2 * xp.pi * xp.sin(2 * xp.pi * eta) # evaluation points and gradient e1 = 0.0 e2 = 0.0 e3 = 0.0 if Nel[0] > 1: - e1 = np.linspace(0.0, 1.0, 100) + e1 = xp.linspace(0.0, 1.0, 100) e = e1 c = 0 @@ -95,12 +95,12 @@ def dfx(x, y, z): return dfun(x) def dfy(x, y, z): - return np.zeros_like(x) + return xp.zeros_like(x) def dfz(x, y, z): - return np.zeros_like(x) + return xp.zeros_like(x) elif Nel[1] > 1: - e2 = np.linspace(0.0, 1.0, 100) + e2 = xp.linspace(0.0, 1.0, 100) e = e2 c = 1 @@ -108,15 +108,15 @@ def f(x, y, z): return fun(y) def dfx(x, y, z): - return np.zeros_like(y) + return xp.zeros_like(y) def dfy(x, y, z): return dfun(y) def dfz(x, y, z): - return np.zeros_like(y) + return xp.zeros_like(y) elif Nel[2] > 1: - e3 = np.linspace(0.0, 1.0, 100) + e3 = xp.linspace(0.0, 1.0, 100) e = e3 c = 2 @@ -124,10 +124,10 @@ def f(x, y, z): return fun(z) def dfx(x, y, z): - return np.zeros_like(z) + return xp.zeros_like(z) def dfy(x, y, z): - return np.zeros_like(z) + return xp.zeros_like(z) def dfz(x, y, z): return dfun(z) @@ -160,22 +160,22 @@ def div_f(x, y, z): field_f0_vals = field_f0(e1, e2, e3, squeeze_out=True) # a) projection error - err_f0 = np.max(np.abs(f(e1, e2, e3) - field_f0_vals)) + err_f0 = xp.max(xp.abs(f(e1, e2, e3) - field_f0_vals)) print(f"\n{err_f0 = }") assert err_f0 < 1e-2 # b) commuting property df0_h = derham.grad.dot(f0_h) - assert np.allclose(df0_h.toarray(), proj_of_grad_f.toarray()) + assert xp.allclose(df0_h.toarray(), proj_of_grad_f.toarray()) # c) derivative error field_df0 = derham.create_spline_function("df0", "Hcurl") field_df0.vector = df0_h field_df0_vals = field_df0(e1, e2, e3, squeeze_out=True) - err_df0 = [np.max(np.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(grad_f, field_df0_vals)] + err_df0 = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(grad_f, field_df0_vals)] print(f"{err_df0 = }") - assert np.max(err_df0) < 0.64 + assert xp.max(err_df0) < 0.64 # d) plotting plt.figure(figsize=(8, 12)) @@ -202,22 +202,22 @@ def div_f(x, y, z): field_f1_vals = field_f1(e1, e2, e3, squeeze_out=True) # a) projection error - err_f1 = [np.max(np.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip([f, f, f], field_f1_vals)] + err_f1 = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip([f, f, f], field_f1_vals)] print(f"{err_f1 = }") - assert np.max(err_f1) < 0.09 + assert xp.max(err_f1) < 0.09 # b) commuting property df1_h = derham.curl.dot(f1_h) - assert np.allclose(df1_h.toarray(), proj_of_curl_fff.toarray()) + assert xp.allclose(df1_h.toarray(), proj_of_curl_fff.toarray()) # c) derivative error field_df1 = derham.create_spline_function("df1", "Hdiv") field_df1.vector = df1_h field_df1_vals = field_df1(e1, e2, e3, squeeze_out=True) - err_df1 = [np.max(np.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(curl_f, field_df1_vals)] + err_df1 = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(curl_f, field_df1_vals)] print(f"{err_df1 = }") - assert np.max(err_df1) < 0.64 + assert xp.max(err_df1) < 0.64 # d) plotting plt.figure(figsize=(8, 12)) @@ -249,22 +249,22 @@ def div_f(x, y, z): field_f2_vals = field_f2(e1, e2, e3, squeeze_out=True) # a) projection error - err_f2 = [np.max(np.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip([f, f, f], field_f2_vals)] + err_f2 = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip([f, f, f], field_f2_vals)] print(f"{err_f2 = }") - assert np.max(err_f2) < 0.09 + assert xp.max(err_f2) < 0.09 # b) commuting property df2_h = derham.div.dot(f2_h) - assert np.allclose(df2_h.toarray(), proj_of_div_fff.toarray()) + assert xp.allclose(df2_h.toarray(), proj_of_div_fff.toarray()) # c) derivative error field_df2 = derham.create_spline_function("df2", "L2") field_df2.vector = df2_h field_df2_vals = field_df2(e1, e2, e3, squeeze_out=True) - err_df2 = np.max(np.abs(div_f(e1, e2, e3) - field_df2_vals)) + err_df2 = xp.max(xp.abs(div_f(e1, e2, e3) - field_df2_vals)) print(f"{err_df2 = }") - assert np.max(err_df2) < 0.64 + assert xp.max(err_df2) < 0.64 # d) plotting plt.figure(figsize=(8, 12)) @@ -291,7 +291,7 @@ def div_f(x, y, z): field_f3_vals = field_f3(e1, e2, e3, squeeze_out=True) # a) projection error - err_f3 = np.max(np.abs(f(e1, e2, e3) - field_f3_vals)) + err_f3 = xp.max(xp.abs(f(e1, e2, e3) - field_f3_vals)) print(f"{err_f3 = }") assert err_f3 < 0.09 diff --git a/src/struphy/feec/tests/test_mass_matrices.py b/src/struphy/feec/tests/test_mass_matrices.py index eb4d988e4..c106abc0b 100644 --- a/src/struphy/feec/tests/test_mass_matrices.py +++ b/src/struphy/feec/tests/test_mass_matrices.py @@ -21,7 +21,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): from struphy.feec.utilities import RotationMatrix, compare_arrays, create_equal_random_arrays from struphy.fields_background.equils import ScrewPinch, ShearedSlab from struphy.geometry import domains - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -48,7 +48,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): eq_mhd = ShearedSlab( **{ "a": (mapping[1]["r1"] - mapping[1]["l1"]), - "R0": (mapping[1]["r3"] - mapping[1]["l3"]) / (2 * np.pi), + "R0": (mapping[1]["r3"] - mapping[1]["l3"]) / (2 * xp.pi), "B0": 1.0, "q0": 1.05, "q1": 1.8, @@ -63,7 +63,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): eq_mhd = ShearedSlab( **{ "a": mapping[1]["Lx"], - "R0": mapping[1]["Lz"] / (2 * np.pi), + "R0": mapping[1]["Lz"] / (2 * xp.pi), "B0": 1.0, "q0": 1.05, "q1": 1.8, @@ -386,7 +386,7 @@ def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): from struphy.fields_background.equils import ScrewPinch from struphy.geometry import domains from struphy.polar.basic import PolarVector - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -499,11 +499,11 @@ def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): x2_pol_psy.tp = x2_psy x3_pol_psy.tp = x3_psy - np.random.seed(1607) - x0_pol_psy.pol = [np.random.rand(x0_pol_psy.pol[0].shape[0], x0_pol_psy.pol[0].shape[1])] - x1_pol_psy.pol = [np.random.rand(x1_pol_psy.pol[n].shape[0], x1_pol_psy.pol[n].shape[1]) for n in range(3)] - x2_pol_psy.pol = [np.random.rand(x2_pol_psy.pol[n].shape[0], x2_pol_psy.pol[n].shape[1]) for n in range(3)] - x3_pol_psy.pol = [np.random.rand(x3_pol_psy.pol[0].shape[0], x3_pol_psy.pol[0].shape[1])] + xp.random.seed(1607) + x0_pol_psy.pol = [xp.random.rand(x0_pol_psy.pol[0].shape[0], x0_pol_psy.pol[0].shape[1])] + x1_pol_psy.pol = [xp.random.rand(x1_pol_psy.pol[n].shape[0], x1_pol_psy.pol[n].shape[1]) for n in range(3)] + x2_pol_psy.pol = [xp.random.rand(x2_pol_psy.pol[n].shape[0], x2_pol_psy.pol[n].shape[1]) for n in range(3)] + x3_pol_psy.pol = [xp.random.rand(x3_pol_psy.pol[0].shape[0], x3_pol_psy.pol[0].shape[1])] # apply boundary conditions to old STRUPHY x0_pol_str = x0_pol_psy.toarray(True) @@ -533,12 +533,12 @@ def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): rn_pol_psy = mass_mats.M2n.dot(x2_pol_psy, apply_bc=True) rJ_pol_psy = mass_mats.M2J.dot(x2_pol_psy, apply_bc=True) - assert np.allclose(r0_pol_str, r0_pol_psy.toarray(True)) - assert np.allclose(r1_pol_str, r1_pol_psy.toarray(True)) - assert np.allclose(r2_pol_str, r2_pol_psy.toarray(True)) - assert np.allclose(r3_pol_str, r3_pol_psy.toarray(True)) - assert np.allclose(rn_pol_str, rn_pol_psy.toarray(True)) - assert np.allclose(rJ_pol_str, rJ_pol_psy.toarray(True)) + assert xp.allclose(r0_pol_str, r0_pol_psy.toarray(True)) + assert xp.allclose(r1_pol_str, r1_pol_psy.toarray(True)) + assert xp.allclose(r2_pol_str, r2_pol_psy.toarray(True)) + assert xp.allclose(r3_pol_str, r3_pol_psy.toarray(True)) + assert xp.allclose(rn_pol_str, rn_pol_psy.toarray(True)) + assert xp.allclose(rJ_pol_str, rJ_pol_psy.toarray(True)) # perfrom matrix-vector products (without boundary conditions) r0_pol_str = space.M0(x0_pol_str) @@ -551,12 +551,12 @@ def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): r2_pol_psy = mass_mats.M2.dot(x2_pol_psy, apply_bc=False) r3_pol_psy = mass_mats.M3.dot(x3_pol_psy, apply_bc=False) - assert np.allclose(r0_pol_str, r0_pol_psy.toarray(True)) - assert np.allclose(r1_pol_str, r1_pol_psy.toarray(True)) - assert np.allclose(r2_pol_str, r2_pol_psy.toarray(True)) - assert np.allclose(r3_pol_str, r3_pol_psy.toarray(True)) - assert np.allclose(rn_pol_str, rn_pol_psy.toarray(True)) - assert np.allclose(rJ_pol_str, rJ_pol_psy.toarray(True)) + assert xp.allclose(r0_pol_str, r0_pol_psy.toarray(True)) + assert xp.allclose(r1_pol_str, r1_pol_psy.toarray(True)) + assert xp.allclose(r2_pol_str, r2_pol_psy.toarray(True)) + assert xp.allclose(r3_pol_str, r3_pol_psy.toarray(True)) + assert xp.allclose(rn_pol_str, rn_pol_psy.toarray(True)) + assert xp.allclose(rJ_pol_str, rJ_pol_psy.toarray(True)) print(f"Rank {mpi_rank} | All tests passed!") @@ -584,7 +584,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots from struphy.feec.utilities import create_equal_random_arrays from struphy.fields_background.equils import ScrewPinch, ShearedSlab from struphy.geometry import domains - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -611,7 +611,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots eq_mhd = ShearedSlab( **{ "a": (mapping[1]["r1"] - mapping[1]["l1"]), - "R0": (mapping[1]["r3"] - mapping[1]["l3"]) / (2 * np.pi), + "R0": (mapping[1]["r3"] - mapping[1]["l3"]) / (2 * xp.pi), "B0": 1.0, "q0": 1.05, "q1": 1.8, @@ -626,7 +626,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots eq_mhd = ShearedSlab( **{ "a": mapping[1]["Lx"], - "R0": mapping[1]["Lz"] / (2 * np.pi), + "R0": mapping[1]["Lz"] / (2 * xp.pi), "B0": 1.0, "q0": 1.05, "q1": 1.8, @@ -751,27 +751,27 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots print("Done") # compare output arrays - assert np.allclose(r0.toarray(), r0_pre.toarray()) - assert np.allclose(r1.toarray(), r1_pre.toarray()) - assert np.allclose(r2.toarray(), r2_pre.toarray()) - assert np.allclose(r3.toarray(), r3_pre.toarray()) - assert np.allclose(rv.toarray(), rv_pre.toarray()) + assert xp.allclose(r0.toarray(), r0_pre.toarray()) + assert xp.allclose(r1.toarray(), r1_pre.toarray()) + assert xp.allclose(r2.toarray(), r2_pre.toarray()) + assert xp.allclose(r3.toarray(), r3_pre.toarray()) + assert xp.allclose(rv.toarray(), rv_pre.toarray()) - assert np.allclose(r1n.toarray(), r1n_pre.toarray()) - assert np.allclose(r2n.toarray(), r2n_pre.toarray()) - assert np.allclose(rvn.toarray(), rvn_pre.toarray()) + assert xp.allclose(r1n.toarray(), r1n_pre.toarray()) + assert xp.allclose(r2n.toarray(), r2n_pre.toarray()) + assert xp.allclose(rvn.toarray(), rvn_pre.toarray()) - assert np.allclose(r1Bninv.toarray(), r1Bninv_pre.toarray()) - assert np.allclose(r1Bninv.toarray(), r1Bninvold_pre.toarray()) - assert np.allclose(r1Bninvold.toarray(), r1Bninv_pre.toarray()) + assert xp.allclose(r1Bninv.toarray(), r1Bninv_pre.toarray()) + assert xp.allclose(r1Bninv.toarray(), r1Bninvold_pre.toarray()) + assert xp.allclose(r1Bninvold.toarray(), r1Bninv_pre.toarray()) # test if preconditioner satisfies PC * M = Identity if mapping[0] == "Cuboid" or mapping[0] == "HollowCylinder": - assert np.allclose(mass_mats.M0.dot(M0pre.solve(x0)).toarray(), derham.boundary_ops["0"].dot(x0).toarray()) - assert np.allclose(mass_mats.M1.dot(M1pre.solve(x1)).toarray(), derham.boundary_ops["1"].dot(x1).toarray()) - assert np.allclose(mass_mats.M2.dot(M2pre.solve(x2)).toarray(), derham.boundary_ops["2"].dot(x2).toarray()) - assert np.allclose(mass_mats.M3.dot(M3pre.solve(x3)).toarray(), derham.boundary_ops["3"].dot(x3).toarray()) - assert np.allclose(mass_mats.Mv.dot(Mvpre.solve(xv)).toarray(), derham.boundary_ops["v"].dot(xv).toarray()) + assert xp.allclose(mass_mats.M0.dot(M0pre.solve(x0)).toarray(), derham.boundary_ops["0"].dot(x0).toarray()) + assert xp.allclose(mass_mats.M1.dot(M1pre.solve(x1)).toarray(), derham.boundary_ops["1"].dot(x1).toarray()) + assert xp.allclose(mass_mats.M2.dot(M2pre.solve(x2)).toarray(), derham.boundary_ops["2"].dot(x2).toarray()) + assert xp.allclose(mass_mats.M3.dot(M3pre.solve(x3)).toarray(), derham.boundary_ops["3"].dot(x3).toarray()) + assert xp.allclose(mass_mats.Mv.dot(Mvpre.solve(xv)).toarray(), derham.boundary_ops["v"].dot(xv).toarray()) # test preconditioner in iterative solver M0inv = inverse(mass_mats.M0, "pcg", pc=M0pre, tol=1e-8, maxiter=1000) @@ -892,7 +892,7 @@ def test_mass_preconditioner_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show from struphy.fields_background.equils import ScrewPinch from struphy.geometry import domains from struphy.polar.basic import PolarVector - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -986,11 +986,11 @@ def test_mass_preconditioner_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show x2_pol.tp = x2 x3_pol.tp = x3 - np.random.seed(1607) - x0_pol.pol = [np.random.rand(x0_pol.pol[0].shape[0], x0_pol.pol[0].shape[1])] - x1_pol.pol = [np.random.rand(x1_pol.pol[n].shape[0], x1_pol.pol[n].shape[1]) for n in range(3)] - x2_pol.pol = [np.random.rand(x2_pol.pol[n].shape[0], x2_pol.pol[n].shape[1]) for n in range(3)] - x3_pol.pol = [np.random.rand(x3_pol.pol[0].shape[0], x3_pol.pol[0].shape[1])] + xp.random.seed(1607) + x0_pol.pol = [xp.random.rand(x0_pol.pol[0].shape[0], x0_pol.pol[0].shape[1])] + x1_pol.pol = [xp.random.rand(x1_pol.pol[n].shape[0], x1_pol.pol[n].shape[1]) for n in range(3)] + x2_pol.pol = [xp.random.rand(x2_pol.pol[n].shape[0], x2_pol.pol[n].shape[1]) for n in range(3)] + x3_pol.pol = [xp.random.rand(x3_pol.pol[0].shape[0], x3_pol.pol[0].shape[1])] # test preconditioner in iterative solver and compare to case without preconditioner M0inv = inverse(mass_mats.M0, "pcg", pc=M0pre, tol=1e-8, maxiter=500) diff --git a/src/struphy/feec/tests/test_toarray_struphy.py b/src/struphy/feec/tests/test_toarray_struphy.py index c1d03249d..db266d108 100644 --- a/src/struphy/feec/tests/test_toarray_struphy.py +++ b/src/struphy/feec/tests/test_toarray_struphy.py @@ -18,7 +18,7 @@ def test_toarray_struphy(Nel, p, spl_kind, mapping): from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import compare_arrays, create_equal_random_arrays from struphy.geometry import domains - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -70,30 +70,30 @@ def test_toarray_struphy(Nel, p, spl_kind, mapping): v3arr = v3arr[0].flatten() # not in-place - compare_arrays(M0.dot(v0), np.matmul(M0arr, v0arr), rank) - compare_arrays(M1.dot(v1), np.matmul(M1arr, v1arr), rank) - compare_arrays(M2.dot(v2), np.matmul(M2arr, v2arr), rank) - compare_arrays(M3.dot(v3), np.matmul(M3arr, v3arr), rank) + compare_arrays(M0.dot(v0), xp.matmul(M0arr, v0arr), rank) + compare_arrays(M1.dot(v1), xp.matmul(M1arr, v1arr), rank) + compare_arrays(M2.dot(v2), xp.matmul(M2arr, v2arr), rank) + compare_arrays(M3.dot(v3), xp.matmul(M3arr, v3arr), rank) # Now we test the in-place version - IM0 = np.zeros([M0.codomain.dimension, M0.domain.dimension], dtype=M0.dtype) - IM1 = np.zeros([M1.codomain.dimension, M1.domain.dimension], dtype=M1.dtype) - IM2 = np.zeros([M2.codomain.dimension, M2.domain.dimension], dtype=M2.dtype) - IM3 = np.zeros([M3.codomain.dimension, M3.domain.dimension], dtype=M3.dtype) + IM0 = xp.zeros([M0.codomain.dimension, M0.domain.dimension], dtype=M0.dtype) + IM1 = xp.zeros([M1.codomain.dimension, M1.domain.dimension], dtype=M1.dtype) + IM2 = xp.zeros([M2.codomain.dimension, M2.domain.dimension], dtype=M2.dtype) + IM3 = xp.zeros([M3.codomain.dimension, M3.domain.dimension], dtype=M3.dtype) M0.toarray_struphy(out=IM0) M1.toarray_struphy(out=IM1) M2.toarray_struphy(out=IM2) M3.toarray_struphy(out=IM3) - compare_arrays(M0.dot(v0), np.matmul(IM0, v0arr), rank) - compare_arrays(M1.dot(v1), np.matmul(IM1, v1arr), rank) - compare_arrays(M2.dot(v2), np.matmul(IM2, v2arr), rank) - compare_arrays(M3.dot(v3), np.matmul(IM3, v3arr), rank) + compare_arrays(M0.dot(v0), xp.matmul(IM0, v0arr), rank) + compare_arrays(M1.dot(v1), xp.matmul(IM1, v1arr), rank) + compare_arrays(M2.dot(v2), xp.matmul(IM2, v2arr), rank) + compare_arrays(M3.dot(v3), xp.matmul(IM3, v3arr), rank) print("test_toarray_struphy passed!") - # assert np.allclose(out1.toarray(), v1.toarray(), atol=1e-5) + # assert xp.allclose(out1.toarray(), v1.toarray(), atol=1e-5) if __name__ == "__main__": diff --git a/src/struphy/feec/tests/test_tosparse_struphy.py b/src/struphy/feec/tests/test_tosparse_struphy.py index e4dc75de9..d9b96aee0 100644 --- a/src/struphy/feec/tests/test_tosparse_struphy.py +++ b/src/struphy/feec/tests/test_tosparse_struphy.py @@ -21,7 +21,7 @@ def test_tosparse_struphy(Nel, p, spl_kind, mapping): from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import create_equal_random_arrays from struphy.geometry import domains - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -102,13 +102,13 @@ def test_tosparse_struphy(Nel, p, spl_kind, mapping): comm.Allreduce(v3_local, v3_global, op=MPI.SUM) # not in-place - assert np.allclose(v0_global, M0arr.dot(v0arr)) - assert np.allclose(v1_global, M1arr.dot(v1arr)) - assert np.allclose(v2_global, M2arr.dot(v2arr)) - assert np.allclose(v3_global, M3arr.dot(v3arr)) - assert np.allclose(v0_global, M0arrad.dot(v0arr)) - assert np.allclose(v1_global, M1arrad.dot(v1arr)) - assert np.allclose(v2_global, M2arrad.dot(v2arr)) + assert xp.allclose(v0_global, M0arr.dot(v0arr)) + assert xp.allclose(v1_global, M1arr.dot(v1arr)) + assert xp.allclose(v2_global, M2arr.dot(v2arr)) + assert xp.allclose(v3_global, M3arr.dot(v3arr)) + assert xp.allclose(v0_global, M0arrad.dot(v0arr)) + assert xp.allclose(v1_global, M1arrad.dot(v1arr)) + assert xp.allclose(v2_global, M2arrad.dot(v2arr)) print("test_tosparse_struphy passed!") diff --git a/src/struphy/feec/tests/xx_test_preconds.py b/src/struphy/feec/tests/xx_test_preconds.py index a5e14b8d7..1bd145fc6 100644 --- a/src/struphy/feec/tests/xx_test_preconds.py +++ b/src/struphy/feec/tests/xx_test_preconds.py @@ -21,7 +21,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, mapping): from struphy.feec.preconditioner import MassMatrixPreconditioner from struphy.feec.psydac_derham import Derham from struphy.geometry import domains - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp MPI_COMM = MPI.COMM_WORLD @@ -40,22 +40,22 @@ def test_mass_preconditioner(Nel, p, spl_kind, mapping): v = [] v += [StencilVector(derham.V0.coeff_space)] - v[-1]._data = np.random.rand(*v[-1]._data.shape) + v[-1]._data = xp.random.rand(*v[-1]._data.shape) v += [BlockVector(derham.V1.coeff_space)] for v1i in v[-1]: - v1i._data = np.random.rand(*v1i._data.shape) + v1i._data = xp.random.rand(*v1i._data.shape) v += [BlockVector(derham.V2.coeff_space)] for v1i in v[-1]: - v1i._data = np.random.rand(*v1i._data.shape) + v1i._data = xp.random.rand(*v1i._data.shape) v += [StencilVector(derham.V3.coeff_space)] - v[-1]._data = np.random.rand(*v[-1]._data.shape) + v[-1]._data = xp.random.rand(*v[-1]._data.shape) v += [BlockVector(derham.V0vec.coeff_space)] for v1i in v[-1]: - v1i._data = np.random.rand(*v1i._data.shape) + v1i._data = xp.random.rand(*v1i._data.shape) # assemble preconditioners M_pre = [] @@ -68,7 +68,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, mapping): n = "v" if domain.kind_map == 10 or domain.kind_map == 11: - assert np.allclose(M._mat.toarray(), M_p.matrix.toarray()) + assert xp.allclose(M._mat.toarray(), M_p.matrix.toarray()) print(f'Matrix assertion for space {n} case "Cuboid/HollowCylinder" passed.') inv_A = InverseLinearOperator(M, pc=M_p, tol=1e-8, maxiter=5000) diff --git a/src/struphy/feec/utilities.py b/src/struphy/feec/utilities.py index d70ad4b0e..e2952572f 100644 --- a/src/struphy/feec/utilities.py +++ b/src/struphy/feec/utilities.py @@ -8,7 +8,7 @@ import struphy.feec.utilities_kernels as kernels from struphy.feec import banded_to_stencil_kernels as bts from struphy.polar.basic import PolarVector -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class RotationMatrix: @@ -41,7 +41,7 @@ def __init__(self, *vec_fun): def __call__(self, e1, e2, e3): # array from 2d list gives 3x3 array is in the first two indices - tmp = np.array( + tmp = xp.array( [ [self._cross_mask[m][n] * fun(e1, e2, e3) for n, fun in enumerate(row)] for m, row in enumerate(self._funs) @@ -49,7 +49,7 @@ def __call__(self, e1, e2, e3): ) # numpy operates on the last two indices with @ - return np.transpose(tmp, axes=(2, 3, 4, 0, 1)) + return xp.transpose(tmp, axes=(2, 3, 4, 0, 1)) def create_equal_random_arrays(V, seed=123, flattened=False): @@ -77,7 +77,7 @@ def create_equal_random_arrays(V, seed=123, flattened=False): assert isinstance(V, (TensorFemSpace, VectorFemSpace)) - np.random.seed(seed) + xp.random.seed(seed) arr = [] @@ -93,7 +93,7 @@ def create_equal_random_arrays(V, seed=123, flattened=False): dims = V.coeff_space.npts - arr += [np.random.rand(*dims)] + arr += [xp.random.rand(*dims)] s = arr_psy.starts e = arr_psy.ends @@ -111,7 +111,7 @@ def create_equal_random_arrays(V, seed=123, flattened=False): for d, block in enumerate(arr_psy.blocks): dims = V.spaces[d].coeff_space.npts - arr += [np.random.rand(*dims)] + arr += [xp.random.rand(*dims)] s = block.starts e = block.ends @@ -121,7 +121,7 @@ def create_equal_random_arrays(V, seed=123, flattened=False): ] if flattened: - arr = np.concatenate( + arr = xp.concatenate( ( arr[0].flatten(), arr[1].flatten(), @@ -168,11 +168,11 @@ def compare_arrays(arr_psy, arr, rank, atol=1e-14, verbose=False): arr_psy.space.npts[2], )[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] - assert np.allclose(tmp1, tmp2, atol=atol) + assert xp.allclose(tmp1, tmp2, atol=atol) elif isinstance(arr_psy, BlockVector): if not (isinstance(arr, tuple) or isinstance(arr, list)): - arrs = np.split( + arrs = xp.split( arr, [ arr_psy.blocks[0].shape[0], @@ -197,7 +197,7 @@ def compare_arrays(arr_psy, arr, rank, atol=1e-14, verbose=False): s[2] : e[2] + 1, ] - assert np.allclose(tmp1, tmp2, atol=atol) + assert xp.allclose(tmp1, tmp2, atol=atol) elif isinstance(arr_psy, StencilMatrix): s = arr_psy.codomain.starts @@ -216,7 +216,7 @@ def compare_arrays(arr_psy, arr, rank, atol=1e-14, verbose=False): if tmp_arr.shape == tmp1.shape: tmp2 = tmp_arr else: - tmp2 = np.zeros( + tmp2 = xp.zeros( ( e[0] + 1 - s[0], e[1] + 1 - s[1], @@ -229,7 +229,7 @@ def compare_arrays(arr_psy, arr, rank, atol=1e-14, verbose=False): ) bts.band_to_stencil_3d(tmp_arr, tmp2) - assert np.allclose(tmp1, tmp2, atol=atol) + assert xp.allclose(tmp1, tmp2, atol=atol) elif isinstance(arr_psy, BlockLinearOperator): for row_psy, row in zip(arr_psy.blocks, arr): @@ -260,7 +260,7 @@ def compare_arrays(arr_psy, arr, rank, atol=1e-14, verbose=False): if tmp_mat.shape == tmp1.shape: tmp2 = tmp_mat else: - tmp2 = np.zeros( + tmp2 = xp.zeros( ( e[0] + 1 - s[0], e[1] + 1 - s[1], @@ -273,7 +273,7 @@ def compare_arrays(arr_psy, arr, rank, atol=1e-14, verbose=False): ) bts.band_to_stencil_3d(tmp_mat, tmp2) - assert np.allclose(tmp1, tmp2, atol=atol) + assert xp.allclose(tmp1, tmp2, atol=atol) else: raise AssertionError("Wrong input type.") diff --git a/src/struphy/feec/utilities_local_projectors.py b/src/struphy/feec/utilities_local_projectors.py index 5aa9a5b61..09fed1e54 100644 --- a/src/struphy/feec/utilities_local_projectors.py +++ b/src/struphy/feec/utilities_local_projectors.py @@ -1,5 +1,5 @@ from struphy.feec.local_projectors_kernels import are_quadrature_points_zero, get_rows, select_quasi_points -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def split_points( @@ -32,7 +32,7 @@ def split_points( shifts : 1d int array For each one of the three spatial directions it determines by which amount to shift the position index (pos) in case we have to loop over the evaluation points. - pts : list of np.array + pts : list of xp.array 3D list of 2D array with the quasi-interpolation points (or Gauss-Legendre quadrature points for histopolation). In format (ns, nb, np) = (spatial direction, B-spline index, point) for StencilVector spaces . @@ -49,7 +49,7 @@ def split_points( npts : list of ints Contains the number of B-splines for each one of the three spatial directions. - periodic : 1D bool np.array + periodic : 1D bool xp.array For each one of the three spatial directions contains the information of whether the B-splines are periodic or not. wij: 3d float array @@ -74,18 +74,18 @@ def split_points( """ # We iterate over the three spatial directions for n, pt in enumerate(pts): - original_pts_size[n] = np.shape(pt)[0] + original_pts_size[n] = xp.shape(pt)[0] # We initialize localpts with as many entries as the global pt, but with all entries being -1 # This function will change the values of the needed entries from -1 to the value of the point. if IoH[n] == "I": - localpts = np.full((np.shape(pt)[0]), fill_value=-1, dtype=float) + localpts = xp.full((xp.shape(pt)[0]), fill_value=-1, dtype=float) elif IoH[n] == "H": - localpts = np.full((np.shape(pt)), fill_value=-1, dtype=float) + localpts = xp.full((xp.shape(pt)), fill_value=-1, dtype=float) for i in range(starts[n], ends[n] + 1): startj1, endj1 = select_quasi_points(int(i), int(p[n]), int(npts[n]), bool(periodic[n])) for j1 in range(lenj[n]): - if startj1 + j1 < np.shape(pt)[0]: + if startj1 + j1 < xp.shape(pt)[0]: pos = startj1 + j1 else: pos = int(startj1 + j1 + shift[n]) @@ -97,42 +97,42 @@ def split_points( localpts[pos] = pt[pos] # We get the local points by grabing only the values different from -1. if IoH[n] == "I": - localpos = np.where(localpts != -1)[0] + localpos = xp.where(localpts != -1)[0] elif IoH[n] == "H": - localpos = np.where(localpts[:, 0] != -1)[0] + localpos = xp.where(localpts[:, 0] != -1)[0] localpts = localpts[localpos] - localptsout.append(np.array(localpts)) + localptsout.append(xp.array(localpts)) ## # We build the index_translation array that shall turn global indices into local indices ## - mini_indextranslation = np.full( - (np.shape(pt)[0]), + mini_indextranslation = xp.full( + (xp.shape(pt)[0]), fill_value=-1, dtype=int, ) for i, j in enumerate(localpos): mini_indextranslation[j] = i - index_translation.append(np.array(mini_indextranslation)) + index_translation.append(xp.array(mini_indextranslation)) ## # We build the inv_index_translation that shall turn local indices into global indices ## - inv_mini_indextranslation = np.full( - (np.shape(localptsout[-1])[0]), + inv_mini_indextranslation = xp.full( + (xp.shape(localptsout[-1])[0]), fill_value=-1, dtype=int, ) for i, j in enumerate(localpos): inv_mini_indextranslation[i] = j - inv_index_translation.append(np.array(inv_mini_indextranslation)) + inv_index_translation.append(xp.array(inv_mini_indextranslation)) def get_values_and_indices_splines(Nbasis, degree, periodic, spans, values): - """Given an array with the values of the splines evaluated at certain points this function returns a np.array that tell us the index of each spline. So we can know to which spline each + """Given an array with the values of the splines evaluated at certain points this function returns a xp.array that tell us the index of each spline. So we can know to which spline each value corresponds. It also modifies the evaluation values in the case we have one spline of degree one with periodic boundary conditions, so it is artificially equal to the identity. Parameters @@ -146,31 +146,31 @@ def get_values_and_indices_splines(Nbasis, degree, periodic, spans, values): periodic : bool Whether we have periodic boundary conditions or nor. - span : np.array + span : xp.array 2d array indexed by (n, nq), where n is the interval and nq is the quadrature point in the interval. - values : np.array + values : xp.array 3d array of values of basis functions indexed by (n, nq, basis function). Returns ------- - eval_indeces : np.array + eval_indeces : xp.array 3d array of basis functions indices, indexed by (n, nq, basis function). - values : np.array + values : xp.array 3d array of values of basis functions indexed by (n, nq, basis function). """ # In this case we want this spatial direction to be "neglected", that means we artificially set the values of the B-spline to 1 at all points. So it becomes the multiplicative identity. if Nbasis == 1 and degree == 1 and periodic: # Set all values to 1 for the identity case - values = np.ones((values.shape[0], values.shape[1], 1)) - eval_indeces = np.zeros_like(values, dtype=int) + values = xp.ones((values.shape[0], values.shape[1], 1)) + eval_indeces = xp.zeros_like(values, dtype=int) else: - eval_indeces = np.zeros_like(values, dtype=int) - for i in range(np.shape(spans)[0]): - for k in range(np.shape(spans)[1]): + eval_indeces = xp.zeros_like(values, dtype=int) + for i in range(xp.shape(spans)[0]): + for k in range(xp.shape(spans)[1]): for j in range(degree + 1): eval_indeces[i, k, j] = (spans[i][k] - degree + j) % Nbasis @@ -179,31 +179,31 @@ def get_values_and_indices_splines(Nbasis, degree, periodic, spans, values): def get_one_spline(a, values, eval_indeces): """Given the spline index, an array with the splines evaluated at the evaluation points and another array with the indices indicating to which spline each value corresponds, this function returns - a 1d np.array with the desired spline evaluated at all evaluation points. + a 1d xp.array with the desired spline evaluated at all evaluation points. Parameters ---------- a : int Spline index - values : np.array + values : xp.array 3d array of values of basis functions indexed by (n, nq, basis function). - eval_indeces : np.array + eval_indeces : xp.array 3d array of basis functions indices, indexed by (n, nq, basis function). Returns ------- - my_values : np.array + my_values : xp.array 1d array of values for the spline evaluated at all evaluation points. """ - my_values = np.zeros(np.shape(values)[0] * np.shape(values)[1]) - for i in range(np.shape(values)[0]): - for j in range(np.shape(values)[1]): - for k in range(np.shape(values)[2]): + my_values = xp.zeros(xp.shape(values)[0] * xp.shape(values)[1]) + for i in range(xp.shape(values)[0]): + for j in range(xp.shape(values)[1]): + for k in range(xp.shape(values)[2]): if eval_indeces[i, j, k] == a: - my_values[i * np.shape(values)[1] + j] = values[i, j, k] + my_values[i * xp.shape(values)[1] + j] = values[i, j, k] break return my_values @@ -213,7 +213,7 @@ def get_span_and_basis(pts, space): Parameters ---------- - pts : np.array + pts : xp.array 2d array of points (ii, iq) = (interval, quadrature point). space : SplineSpace @@ -221,10 +221,10 @@ def get_span_and_basis(pts, space): Returns ------- - span : np.array + span : xp.array 2d array indexed by (n, nq), where n is the interval and nq is the quadrature point in the interval. - basis : np.array + basis : xp.array 3d array of values of basis functions indexed by (n, nq, basis function). """ @@ -234,8 +234,8 @@ def get_span_and_basis(pts, space): T = space.knots p = space.degree - span = np.zeros(pts.shape, dtype=int) - basis = np.zeros((*pts.shape, p + 1), dtype=float) + span = xp.zeros(pts.shape, dtype=int) + basis = xp.zeros((*pts.shape, p + 1), dtype=float) for n in range(pts.shape[0]): for nq in range(pts.shape[1]): @@ -463,7 +463,7 @@ def get_non_zero_B_spline_indices(periodic, IoH, p, B_nbasis, starts, ends, Basi stuck, ) - Basis_functions_indices_B.append(np.array(aux_indices)) + Basis_functions_indices_B.append(xp.array(aux_indices)) def get_non_zero_D_spline_indices(periodic, IoH, p, D_nbasis, starts, ends, Basis_functions_indices_D): @@ -521,7 +521,7 @@ def get_non_zero_D_spline_indices(periodic, IoH, p, D_nbasis, starts, ends, Basi stuck, ) - Basis_functions_indices_D.append(np.array(aux_indices)) + Basis_functions_indices_D.append(xp.array(aux_indices)) def build_translation_list_for_non_zero_spline_indices( @@ -583,17 +583,17 @@ def build_translation_list_for_non_zero_spline_indices( """ translation_indices_B_or_D_splines = [ { - "B": np.full((B_nbasis[h]), fill_value=-1, dtype=int), - "D": np.full((D_nbasis[h]), fill_value=-1, dtype=int), + "B": xp.full((B_nbasis[h]), fill_value=-1, dtype=int), + "D": xp.full((D_nbasis[h]), fill_value=-1, dtype=int), } for h in range(3) ] for h in range(3): - translation_indices_B_or_D_splines[h]["B"][Basis_functions_indices_B[h]] = np.arange( + translation_indices_B_or_D_splines[h]["B"][Basis_functions_indices_B[h]] = xp.arange( len(Basis_functions_indices_B[h]) ) - translation_indices_B_or_D_splines[h]["D"][Basis_functions_indices_D[h]] = np.arange( + translation_indices_B_or_D_splines[h]["D"][Basis_functions_indices_D[h]] = xp.arange( len(Basis_functions_indices_D[h]) ) @@ -653,7 +653,7 @@ def evaluate_relevant_splines_at_relevant_points( for h in range(3): # Reshape localpts[h] if necessary localpts_reshaped = ( - localpts[h].reshape((np.shape(localpts[h])[0], 1)) if len(np.shape(localpts[h])) == 1 else localpts[h] + localpts[h].reshape((xp.shape(localpts[h])[0], 1)) if len(xp.shape(localpts[h])) == 1 else localpts[h] ) # Get spans and evaluation values for B-splines and D-splines @@ -759,7 +759,7 @@ def determine_non_zero_rows_for_each_spline( def process_splines(indices, nbasis, is_D, h): for i in indices[h]: - aux = np.zeros((ends[h] + 1 - starts[h]), dtype=int) + aux = xp.zeros((ends[h] + 1 - starts[h]), dtype=int) get_rows( int(i), int(starts[h]), @@ -773,8 +773,8 @@ def process_splines(indices, nbasis, is_D, h): ) rangestart, rangeend = transform_into_ranges(aux) key = "D" if is_D else "B" - rows_B_or_D_splines[h][key].append(np.array(rangestart, dtype=int)) - rowe_B_or_D_splines[h][key].append(np.array(rangeend, dtype=int)) + rows_B_or_D_splines[h][key].append(xp.array(rangestart, dtype=int)) + rowe_B_or_D_splines[h][key].append(xp.array(rangeend, dtype=int)) for h in range(3): process_splines(Basis_functions_indices_B, B_nbasis, False, h) @@ -890,7 +890,7 @@ def is_spline_zero_at_quadrature_points( for h in range(3): if necessary_direction[h]: for i in Basis_functions_indices_B[h]: - Auxiliar = np.ones((np.shape(localpts[h])[0]), dtype=int) + Auxiliar = xp.ones((xp.shape(localpts[h])[0]), dtype=int) are_quadrature_points_zero( Auxiliar, int( @@ -901,7 +901,7 @@ def is_spline_zero_at_quadrature_points( are_zero_B_or_D_splines[h]["B"].append(Auxiliar) for i in Basis_functions_indices_D[h]: - Auxiliar = np.ones((np.shape(localpts[h])[0]), dtype=int) + Auxiliar = xp.ones((xp.shape(localpts[h])[0]), dtype=int) are_quadrature_points_zero( Auxiliar, int( diff --git a/src/struphy/feec/variational_utilities.py b/src/struphy/feec/variational_utilities.py index 61773122e..b2549d42b 100644 --- a/src/struphy/feec/variational_utilities.py +++ b/src/struphy/feec/variational_utilities.py @@ -12,7 +12,7 @@ ) from struphy.feec.linear_operators import LinOpWithTransp from struphy.feec.psydac_derham import Derham -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class BracketOperator(LinOpWithTransp): @@ -192,10 +192,10 @@ def __init__( # Create tmps for later use in evaluating on the grid grid_shape = tuple([len(loc_grid) for loc_grid in interpolation_grid]) - self._vf_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._gvf1_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._gvf2_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._gvf3_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] + self._vf_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._gvf1_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._gvf2_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._gvf3_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] # gradient of the component of the vector field grad = derham.grad_bcfree @@ -378,13 +378,13 @@ def __init__(self, derham, transposed=False, weights=None): ) grid_shape = tuple([len(loc_grid) for loc_grid in hist_grid_0]) - self._f_0_values = np.zeros(grid_shape, dtype=float) + self._f_0_values = xp.zeros(grid_shape, dtype=float) grid_shape = tuple([len(loc_grid) for loc_grid in hist_grid_1]) - self._f_1_values = np.zeros(grid_shape, dtype=float) + self._f_1_values = xp.zeros(grid_shape, dtype=float) grid_shape = tuple([len(loc_grid) for loc_grid in hist_grid_2]) - self._f_2_values = np.zeros(grid_shape, dtype=float) + self._f_2_values = xp.zeros(grid_shape, dtype=float) @property def domain(self): @@ -533,7 +533,7 @@ def __init__(self, derham, transposed=False, weights=None): ) grid_shape = tuple([len(loc_grid) for loc_grid in hist_grid_0]) - self._bf0_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] + self._bf0_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] self.hist_grid_0_b = [ [self.hist_grid_0_bn[0], self.hist_grid_0_bd[1], self.hist_grid_0_bd[2]], [ @@ -544,7 +544,7 @@ def __init__(self, derham, transposed=False, weights=None): [self.hist_grid_0_bd[0], self.hist_grid_0_bd[1], self.hist_grid_0_bn[2]], ] grid_shape = tuple([len(loc_grid) for loc_grid in hist_grid_1]) - self._bf1_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] + self._bf1_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] self.hist_grid_1_b = [ [self.hist_grid_1_bn[0], self.hist_grid_1_bd[1], self.hist_grid_1_bd[2]], [ @@ -556,7 +556,7 @@ def __init__(self, derham, transposed=False, weights=None): ] grid_shape = tuple([len(loc_grid) for loc_grid in hist_grid_2]) - self._bf2_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] + self._bf2_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] self.hist_grid_2_b = [ [self.hist_grid_2_bn[0], self.hist_grid_2_bd[1], self.hist_grid_2_bd[2]], [ @@ -727,8 +727,8 @@ def __init__(self, derham, phys_domain, Uv, gamma, transposed=False, weights1=No self._proj_p_metric = deepcopy(metric) grid_shape = tuple([len(loc_grid) for loc_grid in int_grid]) - self._pf_values = np.zeros(grid_shape, dtype=float) - self._mapped_pf_values = np.zeros(grid_shape, dtype=float) + self._pf_values = xp.zeros(grid_shape, dtype=float) + self._mapped_pf_values = xp.zeros(grid_shape, dtype=float) # gradient of the component of the vector field @@ -749,13 +749,13 @@ def __init__(self, derham, phys_domain, Uv, gamma, transposed=False, weights1=No ) grid_shape = tuple([len(loc_grid) for loc_grid in hist_grid_20]) - self._pf_0_values = np.zeros(grid_shape, dtype=float) + self._pf_0_values = xp.zeros(grid_shape, dtype=float) grid_shape = tuple([len(loc_grid) for loc_grid in hist_grid_21]) - self._pf_1_values = np.zeros(grid_shape, dtype=float) + self._pf_1_values = xp.zeros(grid_shape, dtype=float) grid_shape = tuple([len(loc_grid) for loc_grid in hist_grid_22]) - self._pf_2_values = np.zeros(grid_shape, dtype=float) + self._pf_2_values = xp.zeros(grid_shape, dtype=float) @property def domain(self): @@ -877,21 +877,21 @@ def __init__(self, derham, gamma): self.rhof1 = self._derham.create_spline_function("rhof1", "L2") grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) - self._rhof_values = np.zeros(grid_shape, dtype=float) - self._rhof1_values = np.zeros(grid_shape, dtype=float) - self._sf_values = np.zeros(grid_shape, dtype=float) - self._sf1_values = np.zeros(grid_shape, dtype=float) - self._delta_values = np.zeros(grid_shape, dtype=float) - self._rhof_mid_values = np.zeros(grid_shape, dtype=float) - self._sf_mid_values = np.zeros(grid_shape, dtype=float) - self._eta_values = np.zeros(grid_shape, dtype=float) - self._en_values = np.zeros(grid_shape, dtype=float) - self._en1_values = np.zeros(grid_shape, dtype=float) - self._de_values = np.zeros(grid_shape, dtype=float) - self._d2e_values = np.zeros(grid_shape, dtype=float) - self._tmp_int_grid = np.zeros(grid_shape, dtype=float) - self._tmp_int_grid2 = np.zeros(grid_shape, dtype=float) - self._DG_values = np.zeros(grid_shape, dtype=float) + self._rhof_values = xp.zeros(grid_shape, dtype=float) + self._rhof1_values = xp.zeros(grid_shape, dtype=float) + self._sf_values = xp.zeros(grid_shape, dtype=float) + self._sf1_values = xp.zeros(grid_shape, dtype=float) + self._delta_values = xp.zeros(grid_shape, dtype=float) + self._rhof_mid_values = xp.zeros(grid_shape, dtype=float) + self._sf_mid_values = xp.zeros(grid_shape, dtype=float) + self._eta_values = xp.zeros(grid_shape, dtype=float) + self._en_values = xp.zeros(grid_shape, dtype=float) + self._en1_values = xp.zeros(grid_shape, dtype=float) + self._de_values = xp.zeros(grid_shape, dtype=float) + self._d2e_values = xp.zeros(grid_shape, dtype=float) + self._tmp_int_grid = xp.zeros(grid_shape, dtype=float) + self._tmp_int_grid2 = xp.zeros(grid_shape, dtype=float) + self._DG_values = xp.zeros(grid_shape, dtype=float) def ener(self, rho, s, out=None): r"""Themodynamical energy as a function of rho and s, usign the perfect gaz hypothesis. @@ -901,13 +901,13 @@ def ener(self, rho, s, out=None): """ gam = self._gamma if out is None: - out = np.power(rho, gam) * np.exp(s / rho) + out = xp.power(rho, gam) * xp.exp(s / rho) else: out *= 0.0 out += s out /= rho - np.exp(out, out=out) - np.power(rho, gam, out=self._tmp_int_grid) + xp.exp(out, out=out) + xp.power(rho, gam, out=self._tmp_int_grid) out *= self._tmp_int_grid return out @@ -919,17 +919,17 @@ def dener_drho(self, rho, s, out=None): """ gam = self._gamma if out is None: - out = (gam * np.power(rho, gam - 1) - s * np.power(rho, gam - 2)) * np.exp(s / rho) + out = (gam * xp.power(rho, gam - 1) - s * xp.power(rho, gam - 2)) * xp.exp(s / rho) else: out *= 0.0 out += s out /= rho - np.exp(out, out=out) + xp.exp(out, out=out) - np.power(rho, gam - 1, out=self._tmp_int_grid) + xp.power(rho, gam - 1, out=self._tmp_int_grid) self._tmp_int_grid *= gam - np.power(rho, gam - 2, out=self._tmp_int_grid2) + xp.power(rho, gam - 2, out=self._tmp_int_grid2) self._tmp_int_grid2 *= s self._tmp_int_grid -= self._tmp_int_grid2 @@ -944,13 +944,13 @@ def dener_ds(self, rho, s, out=None): """ gam = self._gamma if out is None: - out = np.power(rho, gam - 1) * np.exp(s / rho) + out = xp.power(rho, gam - 1) * xp.exp(s / rho) else: out *= 0.0 out += s out /= rho - np.exp(out, out=out) - np.power(rho, gam - 1, out=self._tmp_int_grid) + xp.exp(out, out=out) + xp.power(rho, gam - 1, out=self._tmp_int_grid) out *= self._tmp_int_grid return out @@ -963,25 +963,25 @@ def d2ener_drho2(self, rho, s, out=None): gam = self._gamma if out is None: out = ( - gam * (gam - 1) * np.power(rho, gam - 2) - - s * 2 * (gam - 1) * np.power(rho, gam - 3) - + s**2 * np.power(rho, gam - 4) - ) * np.exp(s / rho) + gam * (gam - 1) * xp.power(rho, gam - 2) + - s * 2 * (gam - 1) * xp.power(rho, gam - 3) + + s**2 * xp.power(rho, gam - 4) + ) * xp.exp(s / rho) else: out *= 0.0 out += s out /= rho - np.exp(out, out=out) + xp.exp(out, out=out) - np.power(rho, gam - 2, out=self._tmp_int_grid) + xp.power(rho, gam - 2, out=self._tmp_int_grid) self._tmp_int_grid *= gam * (gam - 1) - np.power(rho, gam - 3, out=self._tmp_int_grid2) + xp.power(rho, gam - 3, out=self._tmp_int_grid2) self._tmp_int_grid2 *= s self._tmp_int_grid2 *= 2 * (gam - 1) self._tmp_int_grid -= self._tmp_int_grid2 - np.power(rho, gam - 4, out=self._tmp_int_grid2) + xp.power(rho, gam - 4, out=self._tmp_int_grid2) self._tmp_int_grid2 *= s self._tmp_int_grid2 *= s self._tmp_int_grid += self._tmp_int_grid2 @@ -996,27 +996,27 @@ def d2ener_ds2(self, rho, s, out=None): """ gam = self._gamma if out is None: - out = np.power(rho, gam - 2) * np.exp(s / rho) + out = xp.power(rho, gam - 2) * xp.exp(s / rho) else: out *= 0.0 out += s out /= rho - np.exp(out, out=out) - np.power(rho, gam - 2, out=self._tmp_int_grid) + xp.exp(out, out=out) + xp.power(rho, gam - 2, out=self._tmp_int_grid) out *= self._tmp_int_grid return out def eta(self, delta_x, out=None): r"""Switch function :math:`\eta(\delta) = 1- \text{exp}((-\delta/10^{-5})^2)`.""" if out is None: - out = 1.0 - np.exp(-((delta_x / 1e-5) ** 2)) + out = 1.0 - xp.exp(-((delta_x / 1e-5) ** 2)) else: out *= 0.0 out += delta_x out /= 1e-5 out **= 2 out *= -1 - np.exp(out, out=out) + xp.exp(out, out=out) out *= -1 out += 1.0 return out @@ -1328,7 +1328,7 @@ def __init__(self, derham, mass_ops, domain): ) grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) - self._f_values = np.zeros(grid_shape, dtype=float) + self._f_values = xp.zeros(grid_shape, dtype=float) metric = domain.metric(*integration_grid) self._mass_metric_term = deepcopy(metric) @@ -1437,10 +1437,10 @@ def __init__(self, derham, domain, mass_ops): self.uf = derham.create_spline_function("uf", "H1vec") self.uf1 = derham.create_spline_function("uf1", "H1vec") - self._uf_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._uf1_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._Guf_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._tmp_int_grid = np.zeros(grid_shape, dtype=float) + self._uf_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._uf1_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._Guf_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._tmp_int_grid = xp.zeros(grid_shape, dtype=float) metric = domain.metric( *integration_grid, diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index a1056fd73..c6178d579 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -6,7 +6,7 @@ from pyevtk.hl import gridToVTK from struphy.geometry.base import Domain -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class FluidEquilibrium(metaclass=ABCMeta): @@ -140,7 +140,7 @@ def t3(self, *etas, squeeze_out=False): def vth0(self, *etas, squeeze_out=False): """0-form thermal velocity on logical cube [0, 1]^3.""" - return np.sqrt(self.t0(*etas, squeeze_out=squeeze_out)) + return xp.sqrt(self.t0(*etas, squeeze_out=squeeze_out)) def vth3(self, *etas, squeeze_out=False): """3-form thermal velocity on logical cube [0, 1]^3.""" @@ -156,7 +156,7 @@ def q0(self, *etas, squeeze_out=False): """0-form square root of the pressure on logical cube [0, 1]^3.""" # xyz = self.domain(*etas, squeeze_out=False) p = self.p0(*etas) - q = np.sqrt(p) + q = xp.sqrt(p) return self.domain.pull(q, *etas, kind="0", squeeze_out=squeeze_out) def q3(self, *etas, squeeze_out=False): @@ -176,7 +176,7 @@ def s0_monoatomic(self, *etas, squeeze_out=False): # xyz = self.domain(*etas, squeeze_out=False) p = self.p0(*etas) n = self.n0(*etas) - s = n * np.log(p / (2 / 3 * np.power(n, 5 / 3))) + s = n * xp.log(p / (2 / 3 * xp.power(n, 5 / 3))) return self.domain.pull(s, *etas, kind="0", squeeze_out=squeeze_out) def s3_monoatomic(self, *etas, squeeze_out=False): @@ -198,7 +198,7 @@ def s0_diatomic(self, *etas, squeeze_out=False): # xyz = self.domain(*etas, squeeze_out=False) p = self.p0(*etas) n = self.n0(*etas) - s = n * np.log(p / (2 / 5 * np.power(n, 7 / 5))) + s = n * xp.log(p / (2 / 5 * xp.power(n, 7 / 5))) return self.domain.pull(s, *etas, kind="0", squeeze_out=squeeze_out) def s3_diatomic(self, *etas, squeeze_out=False): @@ -395,7 +395,7 @@ def unit_b_cart(self, *etas, squeeze_out=False): """Unit vector Cartesian components of magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z).""" b, xyz = self.b_cart(*etas, squeeze_out=squeeze_out) absB = self.absB0(*etas, squeeze_out=squeeze_out) - out = np.array([b[0] / absB, b[1] / absB, b[2] / absB], dtype=float) + out = xp.array([b[0] / absB, b[1] / absB, b[2] / absB], dtype=float) return out, xyz def gradB1(self, *etas, squeeze_out=False): @@ -481,7 +481,7 @@ def av(self, *etas, squeeze_out=False): def absB0(self, *etas, squeeze_out=False): """0-form absolute value of magnetic field on logical cube [0, 1]^3.""" b, xyz = self.b_cart(*etas, squeeze_out=squeeze_out) - return np.sqrt(b[0] ** 2 + b[1] ** 2 + b[2] ** 2) + return xp.sqrt(b[0] ** 2 + b[1] ** 2 + b[2] ** 2) def absB3(self, *etas, squeeze_out=False): """3-form absolute value of magnetic field on logical cube [0, 1]^3.""" @@ -804,7 +804,7 @@ def curl_unit_b_cart(self, *etas, squeeze_out=False): j, xyz = self.j_cart(*etas, squeeze_out=squeeze_out) gradB, xyz = self.gradB_cart(*etas, squeeze_out=squeeze_out) absB = self.absB0(*etas, squeeze_out=squeeze_out) - out = np.array( + out = xp.array( [ j[0] / absB + (b[1] * gradB[2] - b[2] * gradB[1]) / absB**2, j[1] / absB + (b[2] * gradB[0] - b[0] * gradB[2]) / absB**2, @@ -908,9 +908,9 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): "HollowTorus", ) - e1 = np.linspace(0.0001, 1, n1) - e2 = np.linspace(0, 1, n2) - e3 = np.linspace(0, 1, n3) + e1 = xp.linspace(0.0001, 1, n1) + e2 = xp.linspace(0, 1, n2) + e3 = xp.linspace(0, 1, n3) if self.domain.__class__.__name__ in ("GVECunit", "DESCunit"): if n_planes > 1: @@ -935,7 +935,7 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): print("Computation of abs(B) done.") j_cart, xyz = self.j_cart(e1, e2, e3) print("Computation of current density done.") - absJ = np.sqrt(j_cart[0] ** 2 + j_cart[1] ** 2 + j_cart[2] ** 2) + absJ = xp.sqrt(j_cart[0] ** 2 + j_cart[1] ** 2 + j_cart[2] ** 2) _path = struphy.__path__[0] + "/fields_background/mhd_equil/gvec/output/" gridToVTK( @@ -958,14 +958,14 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): print(key, ": ", val) # poloidal plane grid - fig = plt.figure(figsize=(13, np.ceil(n_planes / 2) * 6.5)) + fig = plt.figure(figsize=(13, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): xp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = np.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" @@ -975,7 +975,7 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): l1 = "x" l2 = "y" - ax = fig.add_subplot(int(np.ceil(n_planes / 2)), 2, n + 1) + ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1) for i in range(pc1.shape[0]): for j in range(pc1.shape[1] - 1): if i < pc1.shape[0] - 1: @@ -1004,9 +1004,9 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): ) # top view - e1 = np.linspace(0, 1, n1) # radial coordinate in [0, 1] - e2 = np.linspace(0, 1, 3) # poloidal angle in [0, 1] - e3 = np.linspace(0, 1, n3) # toroidal angle in [0, 1] + e1 = xp.linspace(0, 1, n1) # radial coordinate in [0, 1] + e2 = xp.linspace(0, 1, 3) # poloidal angle in [0, 1] + e3 = xp.linspace(0, 1, n3) # toroidal angle in [0, 1] xt, yt, zt = self.domain(e1, e2, e3) @@ -1058,14 +1058,14 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): ax.set_title("Device top view") # Jacobian determinant - fig = plt.figure(figsize=(13, np.ceil(n_planes / 2) * 6.5)) + fig = plt.figure(figsize=(13, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): xp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = np.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" @@ -1077,7 +1077,7 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): detp = det_df[:, :, int(n * jump)].squeeze() - ax = fig.add_subplot(int(np.ceil(n_planes / 2)), 2, n + 1) + ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1) map = ax.contourf(pc1, pc2, detp, 30) ax.set_xlabel(l1) ax.set_ylabel(l2) @@ -1088,14 +1088,14 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): fig.colorbar(map, ax=ax, location="right") # pressure - fig = plt.figure(figsize=(15, np.ceil(n_planes / 2) * 6.5)) + fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): xp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = np.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" @@ -1107,7 +1107,7 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): pp = p[:, :, int(n * jump)].squeeze() - ax = fig.add_subplot(int(np.ceil(n_planes / 2)), 2, n + 1) + ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1) map = ax.contourf(pc1, pc2, pp, 30) ax.set_xlabel(l1) ax.set_ylabel(l2) @@ -1118,14 +1118,14 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): fig.colorbar(map, ax=ax, location="right") # density - fig = plt.figure(figsize=(15, np.ceil(n_planes / 2) * 6.5)) + fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): xp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = np.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" @@ -1137,7 +1137,7 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): nn = n_dens[:, :, int(n * jump)].squeeze() - ax = fig.add_subplot(int(np.ceil(n_planes / 2)), 2, n + 1) + ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1) map = ax.contourf(pc1, pc2, nn, 30) ax.set_xlabel(l1) ax.set_ylabel(l2) @@ -1148,14 +1148,14 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): fig.colorbar(map, ax=ax, location="right") # magnetic field strength - fig = plt.figure(figsize=(15, np.ceil(n_planes / 2) * 6.5)) + fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): xp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = np.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" @@ -1167,7 +1167,7 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): ab = absB[:, :, int(n * jump)].squeeze() - ax = fig.add_subplot(int(np.ceil(n_planes / 2)), 2, n + 1) + ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1) map = ax.contourf(pc1, pc2, ab, 30) ax.set_xlabel(l1) ax.set_ylabel(l2) @@ -1178,14 +1178,14 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): fig.colorbar(map, ax=ax, location="right") # current density - fig = plt.figure(figsize=(15, np.ceil(n_planes / 2) * 6.5)) + fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): xp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = np.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" @@ -1197,7 +1197,7 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): ab = absJ[:, :, int(n * jump)].squeeze() - ax = fig.add_subplot(int(np.ceil(n_planes / 2)), 2, n + 1) + ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1) map = ax.contourf(pc1, pc2, ab, 30) ax.set_xlabel(l1) ax.set_ylabel(l2) @@ -1307,8 +1307,8 @@ def b_xyz(self, x, y, z): BZ = self.psi(R, Z, dR=1) / R # push-forward to Cartesian components - Bx = BR * np.cos(Phi) - BP * np.sin(Phi) - By = BR * np.sin(Phi) + BP * np.cos(Phi) + Bx = BR * xp.cos(Phi) - BP * xp.sin(Phi) + By = BR * xp.sin(Phi) + BP * xp.cos(Phi) Bz = 1 * BZ return Bx, By, Bz @@ -1324,8 +1324,8 @@ def j_xyz(self, x, y, z): jZ = self.g_tor(R, Z, dR=1) / R # push-forward to Cartesian components - jx = jR * np.cos(Phi) - jP * np.sin(Phi) - jy = jR * np.sin(Phi) + jP * np.cos(Phi) + jx = jR * xp.cos(Phi) - jP * xp.sin(Phi) + jy = jR * xp.sin(Phi) + jP * xp.cos(Phi) jz = 1 * jZ return jx, jy, jz @@ -1335,7 +1335,7 @@ def gradB_xyz(self, x, y, z): R, Phi, Z = self.inverse_map(x, y, z) - RabsB = np.sqrt( + RabsB = xp.sqrt( self.psi(R, Z, dZ=1) ** 2 + self.g_tor(R, Z) ** 2 + self.psi(R, Z, dR=1) ** 2, ) @@ -1363,8 +1363,8 @@ def gradB_xyz(self, x, y, z): ) # push-forward to Cartesian components - gradBx = gradBR * np.cos(Phi) - gradBP * np.sin(Phi) - gradBy = gradBR * np.sin(Phi) + gradBP * np.cos(Phi) + gradBx = gradBR * xp.cos(Phi) - gradBP * xp.sin(Phi) + gradBy = gradBR * xp.sin(Phi) + gradBP * xp.cos(Phi) gradBz = 1 * gradBZ return gradBx, gradBy, gradBz @@ -1373,8 +1373,8 @@ def gradB_xyz(self, x, y, z): def inverse_map(x, y, z): """Inverse cylindrical mapping.""" - R = np.sqrt(x**2 + y**2) - P = np.arctan2(y, x) + R = xp.sqrt(x**2 + y**2) + P = xp.arctan2(y, x) Z = 1 * z return R, P, Z diff --git a/src/struphy/fields_background/coil_fields/base.py b/src/struphy/fields_background/coil_fields/base.py index 540268ef3..55b0720e8 100644 --- a/src/struphy/fields_background/coil_fields/base.py +++ b/src/struphy/fields_background/coil_fields/base.py @@ -3,7 +3,7 @@ from matplotlib import pyplot as plt from pyevtk.hl import gridToVTK -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class CoilMagneticField(metaclass=ABCMeta): @@ -32,7 +32,7 @@ def b_xyz(self, x, y, z): def absB0(self, *etas, squeeze_out=False): """0-form absolute value of equilibrium magnetic field on logical cube [0, 1]^3.""" b, xyz = self.b_cart(*etas, squeeze_out=squeeze_out) - return np.sqrt(b[0] ** 2 + b[1] ** 2 + b[2] ** 2) + return xp.sqrt(b[0] ** 2 + b[1] ** 2 + b[2] ** 2) def absB3(self, *etas, squeeze_out=False): """3-form absolute value of equilibrium magnetic field on logical cube [0, 1]^3.""" @@ -92,7 +92,7 @@ def unit_b_cart(self, *etas, squeeze_out=False): """Unit vector Cartesian components of equilibrium magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z).""" b, xyz = self.b_cart(*etas, squeeze_out=squeeze_out) absB = self.absB0(*etas, squeeze_out=squeeze_out) - out = np.array([b[0] / absB, b[1] / absB, b[2] / absB], dtype=float) + out = xp.array([b[0] / absB, b[1] / absB, b[2] / absB], dtype=float) return out, xyz diff --git a/src/struphy/fields_background/coil_fields/coil_fields.py b/src/struphy/fields_background/coil_fields/coil_fields.py index 75895c465..2629cc707 100644 --- a/src/struphy/fields_background/coil_fields/coil_fields.py +++ b/src/struphy/fields_background/coil_fields/coil_fields.py @@ -1,6 +1,6 @@ from struphy.feec.psydac_derham import Derham from struphy.fields_background.coil_fields.base import CoilMagneticField, load_csv_data -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class RatGUI(CoilMagneticField): @@ -79,8 +79,8 @@ def bfield_RZphi(self): def b_xyz(self, x, y, z): """Cartesian coil magnetic field in physical space. Must return the components as a tuple.""" # compute (R, Z, phi) corrdinates from (x, y, z), for example: - R = np.sqrt(x**2 + y**2) + R = xp.sqrt(x**2 + y**2) Z = z - phi = -np.arctan2(y / x) + phi = -xp.arctan2(y / x) return self.bfield_RZphi(R, Z, phi) diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index d2eb441f0..25d3d00ad 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -28,7 +28,7 @@ NumericalMHDequilibrium, ) from struphy.fields_background.mhd_equil.eqdsk import readeqdsk -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.utils import read_state, subp_run @@ -237,10 +237,10 @@ def q_x(self, x, der=0): else: if der == 0: - qout = self.params["q0"] + self.params["q1"] * np.sin(2.0 * np.pi * x / self.params["a"]) + qout = self.params["q0"] + self.params["q1"] * xp.sin(2.0 * xp.pi * x / self.params["a"]) else: qout = ( - 2.0 * np.pi / self.params["a"] * self.params["q1"] * np.cos(2.0 * np.pi * x / self.params["a"]) + 2.0 * xp.pi / self.params["a"] * self.params["q1"] * xp.cos(2.0 * xp.pi * x / self.params["a"]) ) return qout @@ -251,7 +251,7 @@ def p_x(self, x): eps = self.params["a"] / self.params["R0"] - if np.all(q >= 100.0): + if xp.all(q >= 100.0): pout = self.params["B0"] ** 2 * self.params["beta"] / 2.0 - 0 * x else: pout = self.params["B0"] ** 2 * self.params["beta"] / 2.0 * (1 + eps**2 / q**2) + self.params[ @@ -273,7 +273,7 @@ def plot_profiles(self, n_pts=501): import matplotlib.pyplot as plt - x = np.linspace(0.0, self.params["a"], n_pts) + x = xp.linspace(0.0, self.params["a"], n_pts) fig, ax = plt.subplots(1, 3) @@ -307,7 +307,7 @@ def b_xyz(self, x, y, z): q = self.q_x(x) eps = self.params["a"] / self.params["R0"] - if np.all(q >= 100.0): + if xp.all(q >= 100.0): by = 0 * x bz = self.params["B0"] - 0 * x else: @@ -324,7 +324,7 @@ def j_xyz(self, x, y, z): q = self.q_x(x) eps = self.params["a"] / self.params["R0"] - if np.all(q >= 100.0): + if xp.all(q >= 100.0): jz = 0 * x else: jz = -self.params["B0"] * eps * self.q_x(x, der=1) / q**2 @@ -353,13 +353,13 @@ def gradB_xyz(self, x, y, z): q = self.q_x(x) eps = self.params["a"] / self.params["R0"] - if np.all(q >= 100.0): + if xp.all(q >= 100.0): gradBx = 0 * x else: gradBx = ( -self.params["B0"] * eps**2 - / np.sqrt(1 + eps**2 / self.q_x(x) ** 2) + / xp.sqrt(1 + eps**2 / self.q_x(x) ** 2) * self.q_x(x, der=1) / self.q_x(x) ** 3 ) @@ -457,8 +457,8 @@ def __init__( def T_z(self, z): r"""Swap function T(z) = \tanh(z - z_1)/\delta) - \tanh(z - z_2)/\delta)""" Tout = ( - np.tanh((z - self.params["z1"]) / self.params["delta"]) - - np.tanh((z - self.params["z2"]) / self.params["delta"]) + xp.tanh((z - self.params["z1"]) / self.params["delta"]) + - xp.tanh((z - self.params["z2"]) / self.params["delta"]) ) / 2.0 return Tout @@ -480,7 +480,7 @@ def plot_profiles(self, n_pts=501): import matplotlib.pyplot as plt - z = np.linspace(0.0, self.params["c"], n_pts) + z = xp.linspace(0.0, self.params["c"], n_pts) fig, ax = plt.subplots(1, 3) @@ -647,8 +647,8 @@ def __init__( self.params = copy.deepcopy(locals()) # inverse cylindrical coordinate transformation (x, y, z) --> (r, theta, phi) - self.r = lambda x, y, z: np.sqrt(x**2 + y**2) - self.theta = lambda x, y, z: np.arctan2(y, x) + self.r = lambda x, y, z: xp.sqrt(x**2 + y**2) + self.theta = lambda x, y, z: xp.arctan2(y, x) self.z = lambda x, y, z: 1 * z # =============================================================== @@ -707,7 +707,7 @@ def plot_profiles(self, n_pts=501): import matplotlib.pyplot as plt - r = np.linspace(0.0, self.params["a"], n_pts) + r = xp.linspace(0.0, self.params["a"], n_pts) fig, ax = plt.subplots(1, 3) @@ -718,7 +718,7 @@ def plot_profiles(self, n_pts=501): ax[0].set_xlabel("r") ax[0].set_ylabel("q") - ax[0].plot(r, np.ones(r.size), "k--") + ax[0].plot(r, xp.ones(r.size), "k--") ax[1].plot(r, self.p_r(r)) ax[1].set_xlabel("r") @@ -743,13 +743,13 @@ def b_xyz(self, x, y, z): theta = self.theta(x, y, z) q = self.q_r(r) # azimuthal component - if np.all(q >= 100.0): + if xp.all(q >= 100.0): b_theta = 0 * r else: b_theta = self.params["B0"] * r / (self.params["R0"] * q) # cartesian x-component - bx = -b_theta * np.sin(theta) - by = b_theta * np.cos(theta) + bx = -b_theta * xp.sin(theta) + by = b_theta * xp.cos(theta) bz = self.params["B0"] - 0 * x return bx, by, bz @@ -763,7 +763,7 @@ def j_xyz(self, x, y, z): r = self.r(x, y, z) q = self.q_r(r) q_p = self.q_r(r, der=1) - if np.all(q >= 100.0): + if xp.all(q >= 100.0): jz = 0 * x else: jz = self.params["B0"] / (self.params["R0"] * q**2) * (2 * q - r * q_p) @@ -790,13 +790,13 @@ def gradB_xyz(self, x, y, z): r = self.r(x, y, z) theta = self.theta(x, y, z) q = self.q_r(r) - if np.all(q >= 100.0): + if xp.all(q >= 100.0): gradBr = 0 * x else: gradBr = ( self.params["B0"] / self.params["R0"] ** 2 - / np.sqrt( + / xp.sqrt( 1 + r**2 / self.q_r( @@ -807,8 +807,8 @@ def gradB_xyz(self, x, y, z): ) * (r / self.q_r(r) ** 2 - r**2 / self.q_r(r) ** 3 * self.q_r(r, der=1)) ) - gradBx = gradBr * np.cos(theta) - gradBy = gradBr * np.sin(theta) + gradBx = gradBr * xp.cos(theta) + gradBy = gradBr * xp.sin(theta) gradBz = 0 * x return gradBx, gradBy, gradBz @@ -947,10 +947,10 @@ def __init__( self.params = copy.deepcopy(locals()) # plasma boundary contour - ths = np.linspace(0.0, 2 * np.pi, 201) + ths = xp.linspace(0.0, 2 * xp.pi, 201) - self._rbs = self.params["R0"] * (1 + self.params["a"] / self.params["R0"] * np.cos(ths)) - self._zbs = self.params["a"] * np.sin(ths) + self._rbs = self.params["R0"] * (1 + self.params["a"] / self.params["R0"] * xp.cos(ths)) + self._zbs = self.params["a"] * xp.sin(ths) # set on-axis and boundary fluxes if self.params["q_kind"] == 0: @@ -961,12 +961,12 @@ def __init__( self._p_i = None else: - r_i = np.linspace(0.0, self.params["a"], self.params["psi_nel"] + 1) + r_i = xp.linspace(0.0, self.params["a"], self.params["psi_nel"] + 1) def dpsi_dr(r): - return self.params["B0"] * r / (self.q_r(r) * np.sqrt(1 - r**2 / self.params["R0"] ** 2)) + return self.params["B0"] * r / (self.q_r(r) * xp.sqrt(1 - r**2 / self.params["R0"] ** 2)) - psis = np.zeros_like(r_i) + psis = xp.zeros_like(r_i) for i, rr in enumerate(r_i): psis[i] = quad(dpsi_dr, 0.0, rr)[0] @@ -989,7 +989,7 @@ def dp_dr(r): * (2 * self.q_r(r) - r * self.q_r(r, der=1)) ) - ps = np.zeros_like(r_i) + ps = xp.zeros_like(r_i) for i, rr in enumerate(r_i): ps[i] = quad(dp_dr, 0.0, rr)[0] @@ -1044,7 +1044,7 @@ def psi_r(self, r, der=0): dq = q1 - q0 # geometric correction factor and its first derivative - gf_0 = np.sqrt(1 - (r / self.params["R0"]) ** 2) + gf_0 = xp.sqrt(1 - (r / self.params["R0"]) ** 2) gf_1 = -r / (self.params["R0"] ** 2 * gf_0) # safety factors @@ -1055,9 +1055,9 @@ def psi_r(self, r, der=0): q_bar_1 = q_1 * gf_0 + q_0 * gf_1 if der == 0: - out = -self.params["B0"] * self.params["a"] ** 2 / np.sqrt(dq * q0 * eps**2 + dq**2) - out *= np.arctanh( - np.sqrt((dq - dq * (r / self.params["R0"]) ** 2) / (q0 * eps**2 + dq)), + out = -self.params["B0"] * self.params["a"] ** 2 / xp.sqrt(dq * q0 * eps**2 + dq**2) + out *= xp.arctanh( + xp.sqrt((dq - dq * (r / self.params["R0"]) ** 2) / (q0 * eps**2 + dq)), ) elif der == 1: out = self.params["B0"] * r / q_bar_0 @@ -1127,10 +1127,10 @@ def q_r(self, r, der=0): r_flat = r.flatten() - r_zeros = np.where(r_flat == 0.0)[0] - r_nzero = np.where(r_flat != 0.0)[0] + r_zeros = xp.where(r_flat == 0.0)[0] + r_nzero = xp.where(r_flat != 0.0)[0] - qout = np.zeros(r_flat.size, dtype=float) + qout = xp.zeros(r_flat.size, dtype=float) if der == 0: if self.params["q0"] == self.params["q1"]: @@ -1223,7 +1223,7 @@ def plot_profiles(self, n_pts=501): import matplotlib.pyplot as plt - r = np.linspace(0.0, self.params["a"], n_pts) + r = xp.linspace(0.0, self.params["a"], n_pts) fig, ax = plt.subplots(2, 2) @@ -1257,7 +1257,7 @@ def plot_profiles(self, n_pts=501): def psi(self, R, Z, dR=0, dZ=0): """Poloidal flux function psi = psi(R, Z).""" - r = np.sqrt(Z**2 + (R - self.params["R0"]) ** 2) + r = xp.sqrt(Z**2 + (R - self.params["R0"]) ** 2) if dR == 0 and dZ == 0: out = self.psi_r(r, der=0) @@ -1305,7 +1305,7 @@ def g_tor(self, R, Z, dR=0, dZ=0): def p_xyz(self, x, y, z): """Pressure p = p(x, y, z).""" - r = np.sqrt((np.sqrt(x**2 + y**2) - self._params["R0"]) ** 2 + z**2) + r = xp.sqrt((xp.sqrt(x**2 + y**2) - self._params["R0"]) ** 2 + z**2) pp = self.p_r(r) @@ -1313,7 +1313,7 @@ def p_xyz(self, x, y, z): def n_xyz(self, x, y, z): """Number density n = n(x, y, z).""" - r = np.sqrt((np.sqrt(x**2 + y**2) - self._params["R0"]) ** 2 + z**2) + r = xp.sqrt((xp.sqrt(x**2 + y**2) - self._params["R0"]) ** 2 + z**2) nn = self.n_r(r) @@ -1441,10 +1441,10 @@ def __init__( self.params = copy.deepcopy(locals()) # plasma boundary contour - ths = np.linspace(0.0, 2 * np.pi, 201) + ths = xp.linspace(0.0, 2 * xp.pi, 201) - self._rbs = self.params["R0"] * (1 + self.params["a"] / self.params["R0"] * np.cos(ths)) - self._zbs = self.params["a"] * np.sin(ths) + self._rbs = self.params["R0"] * (1 + self.params["a"] / self.params["R0"] * xp.cos(ths)) + self._zbs = self.params["a"] * xp.sin(ths) # on-axis flux (arbitrary value) self._psi0 = -10.0 @@ -1465,12 +1465,12 @@ def dpsi_dr(psi, r, psi1): q = q0 + psi_norm * (q1 - q0 + (q1p - q1 + q0) * (1 - psi_s) * (psi_norm - 1) / (psi_norm - psi_s)) - out = B0 * r / (q * np.sqrt(1 - r**2 / R0**2)) + out = B0 * r / (q * xp.sqrt(1 - r**2 / R0**2)) return out # solve differential equation and fix boundary flux - r_i = np.linspace(0.0, self.params["a"], self.params["psi_nel"] + 1) + r_i = xp.linspace(0.0, self.params["a"], self.params["psi_nel"] + 1) def fun(psi1): out = odeint(dpsi_dr, self._psi0, r_i, args=(psi1,)).flatten() @@ -1557,13 +1557,13 @@ def p_psi(self, psi, der=0): psi_norm = (psi - self._psi0) / (self._psi1 - self._psi0) if der == 0: - out = self.params["beta"] * self.params["B0"] ** 2 / 2.0 * np.exp(-psi_norm / p1) + out = self.params["beta"] * self.params["B0"] ** 2 / 2.0 * xp.exp(-psi_norm / p1) else: out = ( -self.params["beta"] * self.params["B0"] ** 2 / 2.0 - * np.exp(-psi_norm / p1) + * xp.exp(-psi_norm / p1) / (p1 * (self._psi1 - self._psi0)) ) @@ -1592,8 +1592,8 @@ def plot_profiles(self, n_pts=501): import matplotlib.pyplot as plt - r = np.linspace(0.0, self.params["a"], n_pts) - psi = np.linspace(self._psi0, self._psi1, n_pts) + r = xp.linspace(0.0, self.params["a"], n_pts) + psi = xp.linspace(self._psi0, self._psi1, n_pts) fig, ax = plt.subplots(2, 2) @@ -1627,7 +1627,7 @@ def plot_profiles(self, n_pts=501): def psi(self, R, Z, dR=0, dZ=0): """Poloidal flux function psi = psi(R, Z).""" - r = np.sqrt(Z**2 + (R - self.params["R0"]) ** 2) + r = xp.sqrt(Z**2 + (R - self.params["R0"]) ** 2) if dR == 0 and dZ == 0: out = self.psi_r(r, der=0) @@ -1671,13 +1671,13 @@ def g_tor(self, R, Z, dR=0, dZ=0): def p_xyz(self, x, y, z): """Pressure p = p(x, y, z).""" - r = np.sqrt((np.sqrt(x**2 + y**2) - self._params["R0"]) ** 2 + z**2) + r = xp.sqrt((xp.sqrt(x**2 + y**2) - self._params["R0"]) ** 2 + z**2) return self.p_psi(self.psi_r(r)) def n_xyz(self, x, y, z): """Number density n = n(x, y, z).""" - r = np.sqrt((np.sqrt(x**2 + y**2) - self._params["R0"]) ** 2 + z**2) + r = xp.sqrt((xp.sqrt(x**2 + y**2) - self._params["R0"]) ** 2 + z**2) return self.n_psi(self.psi_r(r)) @@ -1815,8 +1815,8 @@ def __init__( self._r_range = [rleft, rleft + rdim] self._z_range = [zmid - zdim / 2, zmid + zdim / 2] - R = np.linspace(self._r_range[0], self._r_range[1], nR) - Z = np.linspace(self._z_range[0], self._z_range[1], nZ) + R = xp.linspace(self._r_range[0], self._r_range[1], nR) + Z = xp.linspace(self._z_range[0], self._z_range[1], nZ) smooth_steps = [ int(1 / (self.params["psi_resolution"][0] * 0.01)), @@ -1846,7 +1846,7 @@ def __init__( self._psi1 = psi_edge # interpolate toroidal field function, pressure profile and q-profile on unifrom flux grid from axis to boundary - flux_grid = np.linspace(self._psi0, self._psi1, g_profile.size) + flux_grid = xp.linspace(self._psi0, self._psi1, g_profile.size) smooth_step = int(1 / (self.params["flux_resolution"] * 0.01)) @@ -2018,7 +2018,7 @@ def g_tor(self, R, Z, dR=0, dZ=0): def p_xyz(self, x, y, z): """Pressure p = p(x, y, z) in units 1 Tesla^2/mu_0.""" - R = np.sqrt(x**2 + y**2) + R = xp.sqrt(x**2 + y**2) Z = 1 * z out = self.p_psi(self.psi(R, Z)) @@ -2031,7 +2031,7 @@ def p_xyz(self, x, y, z): def n_xyz(self, x, y, z): """Number density in physical space. Units from parameter file.""" - R = np.sqrt(x**2 + y**2) + R = xp.sqrt(x**2 + y**2) Z = 1 * z out = self.n_psi(self.psi(R, Z)) @@ -2211,9 +2211,9 @@ def bv(self, *etas, squeeze_out=False): bt += "_B" bz += "_B" self.state.compute(ev, bt, bz) - bv_2 = getattr(ev, bt).data / (2 * np.pi) - bv_3 = getattr(ev, bz).data / (2 * np.pi) * self._nfp - out = (np.zeros_like(bv_2), bv_2, bv_3) + bv_2 = getattr(ev, bt).data / (2 * xp.pi) + bv_3 = getattr(ev, bz).data / (2 * xp.pi) * self._nfp + out = (xp.zeros_like(bv_2), bv_2, bv_3) # apply struphy units for o in out: @@ -2231,8 +2231,8 @@ def jv(self, *etas, squeeze_out=False): self.state.compute(ev, jr, jt, jz) rmin = self._params["rmin"] jv_1 = ev.J_contra_r.data / (1.0 - rmin) - jv_2 = ev.J_contra_t.data / (2 * np.pi) - jv_3 = ev.J_contra_z.data / (2 * np.pi) * self._nfp + jv_2 = ev.J_contra_t.data / (2 * xp.pi) + jv_3 = ev.J_contra_z.data / (2 * xp.pi) * self._nfp if self.params["use_boozer"]: warnings.warn("GVEC current density in Boozer coords not yet implemented, set to zero.") # jr += "_B" @@ -2257,11 +2257,11 @@ def p0(self, *etas, squeeze_out=False): if not flat_eval: eta2 = etas[1] eta3 = etas[2] - if isinstance(eta2, np.ndarray): + if isinstance(eta2, xp.ndarray): if eta2.ndim == 3: eta2 = eta2[0, :, 0] eta3 = eta3[0, 0, :] - tmp, _1, _2 = np.meshgrid(ev.p.data, eta2, eta3, indexing="ij") + tmp, _1, _2 = xp.meshgrid(ev.p.data, eta2, eta3, indexing="ij") else: tmp = ev.p.data @@ -2321,7 +2321,7 @@ def _gvec_evaluations(self, *etas): etas = list(etas) for i, eta in enumerate(etas): if isinstance(eta, (float, int)): - etas[i] = np.array((eta,)) + etas[i] = xp.array((eta,)) assert etas[0].ndim == etas[1].ndim == etas[2].ndim if etas[0].ndim == 1: eta1 = etas[0] @@ -2338,8 +2338,8 @@ def _gvec_evaluations(self, *etas): # gvec coordinates rho = rmin + eta1 * (1.0 - rmin) - theta = 2 * np.pi * eta2 - zeta = 2 * np.pi * eta3 + theta = 2 * xp.pi * eta2 + zeta = 2 * xp.pi * eta3 # evaluate if self.params["use_boozer"]: @@ -2509,7 +2509,7 @@ def bv(self, *etas, squeeze_out=False): li = [] for gi, ei in zip(grid, etas): if gi.shape == ei.shape: - li += [np.allclose(gi, ei)] + li += [xp.allclose(gi, ei)] else: li += [False] if all(li): @@ -2561,9 +2561,9 @@ def _eval_bv(self, *etas, squeeze_out=False): if var == "B^rho": tmp /= 1.0 - self.rmin elif var == "B^theta": - tmp /= 2.0 * np.pi + tmp /= 2.0 * xp.pi elif var == "B^zeta": - tmp /= 2.0 * np.pi / nfp + tmp /= 2.0 * xp.pi / nfp # adjust for Struphy units tmp /= self.units["B"] / self.units["x"] out += [tmp] @@ -2582,7 +2582,7 @@ def jv(self, *etas, squeeze_out=False): li = [] for gi, ei in zip(grid, etas): if gi.shape == ei.shape: - li += [np.allclose(gi, ei)] + li += [xp.allclose(gi, ei)] else: li += [False] if all(li): @@ -2634,9 +2634,9 @@ def _eval_jv(self, *etas, squeeze_out=False): if var == "J^rho": tmp /= 1.0 - self.rmin elif var == "J^theta": - tmp /= 2.0 * np.pi + tmp /= 2.0 * xp.pi elif var == "J^zeta": - tmp /= 2.0 * np.pi / nfp + tmp /= 2.0 * xp.pi / nfp # adjust for Struphy units tmp /= self.units["j"] / self.units["x"] out += [tmp] @@ -2706,7 +2706,7 @@ def gradB1(self, *etas, squeeze_out=False): li = [] for gi, ei in zip(grid, etas): if gi.shape == ei.shape: - li += [np.allclose(gi, ei)] + li += [xp.allclose(gi, ei)] else: li += [False] if all(li): @@ -2757,9 +2757,9 @@ def _eval_gradB1(self, *etas, squeeze_out=False): if var == "|B|_r": tmp *= 1.0 - self.rmin elif var == "|B|_t": - tmp *= 2.0 * np.pi + tmp *= 2.0 * xp.pi elif var == "|B|_z": - tmp *= 2.0 * np.pi / nfp + tmp *= 2.0 * xp.pi / nfp # adjust for Struphy units tmp /= self.units["B"] out += [tmp] @@ -2769,9 +2769,9 @@ def _eval_gradB1(self, *etas, squeeze_out=False): def desc_eval( self, var: str, - e1: np.ndarray, - e2: np.ndarray, - e3: np.ndarray, + e1: xp.ndarray, + e2: xp.ndarray, + e3: xp.ndarray, flat_eval: bool = False, nfp: int = 1, verbose: bool = False, @@ -2785,7 +2785,7 @@ def desc_eval( Desc equilibrium quantitiy to evaluate, from `https://desc-docs.readthedocs.io/en/latest/variables.html#list-of-variables`_. - e1, e2, e3 : np.ndarray + e1, e2, e3 : xp.ndarray Input grids, either 1d or 3d. flat_eval : bool @@ -2804,21 +2804,21 @@ def desc_eval( warnings.filterwarnings("ignore") ttime = time() # Fix issue 353 with float dummy etas - e1 = np.array([e1]) if isinstance(e1, float) else e1 - e2 = np.array([e2]) if isinstance(e2, float) else e2 - e3 = np.array([e3]) if isinstance(e3, float) else e3 + e1 = xp.array([e1]) if isinstance(e1, float) else e1 + e2 = xp.array([e2]) if isinstance(e2, float) else e2 + e3 = xp.array([e3]) if isinstance(e3, float) else e3 # transform input grids if e1.ndim == 3: assert e1.shape == e2.shape == e3.shape rho = self.rmin + e1[:, 0, 0] * (1.0 - self.rmin) - theta = 2 * np.pi * e2[0, :, 0] - zeta = 2 * np.pi * e3[0, 0, :] / nfp + theta = 2 * xp.pi * e2[0, :, 0] + zeta = 2 * xp.pi * e3[0, 0, :] / nfp else: assert e1.ndim == e2.ndim == e3.ndim == 1 rho = self.rmin + e1 * (1.0 - self.rmin) - theta = 2 * np.pi * e2 - zeta = 2 * np.pi * e3 / nfp + theta = 2 * xp.pi * e2 + zeta = 2 * xp.pi * e3 / nfp # eval type if flat_eval: @@ -2827,13 +2827,13 @@ def desc_eval( t = theta z = zeta else: - r, t, z = np.meshgrid(rho, theta, zeta, indexing="ij") + r, t, z = xp.meshgrid(rho, theta, zeta, indexing="ij") r = r.flatten() t = t.flatten() z = z.flatten() - nodes = np.stack((r, t, z)).T - grid_3d = Grid(nodes, spacing=np.ones_like(nodes), jitable=False) + nodes = xp.stack((r, t, z)).T + grid_3d = Grid(nodes, spacing=xp.ones_like(nodes), jitable=False) # compute output corresponding to the generated desc grid node_values = self.eq.compute( @@ -2874,9 +2874,9 @@ def desc_eval( )[0, 0, :] # make sure the desc grid is correct - assert np.all(rho == rho1) - assert np.all(theta == theta1) - assert np.all(zeta == zeta1) + assert xp.all(rho == rho1) + assert xp.all(theta == theta1) + assert xp.all(zeta == zeta1) if verbose: # import sys @@ -2899,7 +2899,7 @@ def desc_eval( print(f"{zeta1 = }") # make c-contiguous - out = np.ascontiguousarray(out) + out = xp.ascontiguousarray(out) print(f"desc_eval for {var}: {time() - ttime} seconds") return out @@ -2947,12 +2947,12 @@ def n_xyz(self, x, y, z): elif self.params["density_profile"] == "affine": return self.params["n"] + self.params["n1"] * x elif self.params["density_profile"] == "gaussian_xy": - return self.params["n"] * np.exp(-(x**2 + y**2) / self.params["p0"]) + return self.params["n"] * xp.exp(-(x**2 + y**2) / self.params["p0"]) elif self.params["density_profile"] == "step_function_x": out = 1e-8 + 0 * x - # mask_x = np.logical_and(x < .6, x > .4) - # mask_y = np.logical_and(y < .6, y > .4) - # mask = np.logical_and(mask_x, mask_y) + # mask_x = xp.logical_and(x < .6, x > .4) + # mask_y = xp.logical_and(y < .6, y > .4) + # mask = xp.logical_and(mask_x, mask_y) mask = x < -2.0 out[mask] = self.params["n"] return out @@ -3278,7 +3278,7 @@ def plot_profiles(self, n_pts=501): import matplotlib.pyplot as plt - r = np.linspace(0.0, self.params["a"], n_pts) + r = xp.linspace(0.0, self.params["a"], n_pts) fig, ax = plt.subplots(1, 3) @@ -3289,7 +3289,7 @@ def plot_profiles(self, n_pts=501): ax[0].set_xlabel("r") ax[0].set_ylabel("q") - ax[0].plot(r, np.ones(r.size), "k--") + ax[0].plot(r, xp.ones(r.size), "k--") ax[1].plot(r, self.p_r(r)) ax[1].set_xlabel("r") @@ -3312,8 +3312,8 @@ def b_xyz(self, x, y, z): """Magnetic field.""" bz = 0 * x - by = np.tanh(z / self._params["delta"]) - bx = np.sqrt(1 - by**2) + by = xp.tanh(z / self._params["delta"]) + bx = xp.sqrt(1 - by**2) bxs = self._params["amp"] * bx bys = self._params["amp"] * by diff --git a/src/struphy/fields_background/tests/test_desc_equil.py b/src/struphy/fields_background/tests/test_desc_equil.py index c8f1dc6fe..e7ae0872c 100644 --- a/src/struphy/fields_background/tests/test_desc_equil.py +++ b/src/struphy/fields_background/tests/test_desc_equil.py @@ -3,7 +3,7 @@ import pytest from matplotlib import pyplot as plt -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp desc_spec = importlib.util.find_spec("desc") @@ -33,9 +33,9 @@ def test_desc_equil(do_plot=False): n2 = 9 n3 = 11 - e1 = np.linspace(0.0001, 1, n1) - e2 = np.linspace(0, 1, n2) - e3 = np.linspace(0, 1 - 1e-6, n3) + e1 = xp.linspace(0.0001, 1, n1) + e2 = xp.linspace(0, 1, n2) + e3 = xp.linspace(0, 1 - 1e-6, n3) # desc grid and evaluation vars = [ @@ -70,43 +70,43 @@ def test_desc_equil(do_plot=False): outs[nfp] = {} rho = rmin + e1 * (1.0 - rmin) - theta = 2 * np.pi * e2 - zeta = 2 * np.pi * e3 / nfp + theta = 2 * xp.pi * e2 + zeta = 2 * xp.pi * e3 / nfp - r, t, ze = np.meshgrid(rho, theta, zeta, indexing="ij") + r, t, ze = xp.meshgrid(rho, theta, zeta, indexing="ij") r = r.flatten() t = t.flatten() ze = ze.flatten() - nodes = np.stack((r, t, ze)).T - grid_3d = Grid(nodes, spacing=np.ones_like(nodes), jitable=False) + nodes = xp.stack((r, t, ze)).T + grid_3d = Grid(nodes, spacing=xp.ones_like(nodes), jitable=False) for var in vars: node_values = desc_eq.compute(var, grid=grid_3d, override_grid=False) if node_values[var].ndim == 1: out = node_values[var].reshape((rho.size, theta.size, zeta.size), order="C") - outs[nfp][var] = np.ascontiguousarray(out) + outs[nfp][var] = xp.ascontiguousarray(out) else: B = [] for i in range(3): Bcomp = node_values[var][:, i].reshape((rho.size, theta.size, zeta.size), order="C") - Bcomp = np.ascontiguousarray(Bcomp) + Bcomp = xp.ascontiguousarray(Bcomp) B += [Bcomp] outs[nfp][var + str(i + 1)] = Bcomp - outs[nfp][var] = np.sqrt(B[0] ** 2 + B[1] ** 2 + B[2] ** 2) + outs[nfp][var] = xp.sqrt(B[0] ** 2 + B[1] ** 2 + B[2] ** 2) - assert np.allclose(outs[nfp]["B1"], outs[nfp]["B_R"]) - assert np.allclose(outs[nfp]["B2"], outs[nfp]["B_phi"]) - assert np.allclose(outs[nfp]["B3"], outs[nfp]["B_Z"]) + assert xp.allclose(outs[nfp]["B1"], outs[nfp]["B_R"]) + assert xp.allclose(outs[nfp]["B2"], outs[nfp]["B_phi"]) + assert xp.allclose(outs[nfp]["B3"], outs[nfp]["B_Z"]) - assert np.allclose(outs[nfp]["J1"], outs[nfp]["J_R"]) - assert np.allclose(outs[nfp]["J2"], outs[nfp]["J_phi"]) - assert np.allclose(outs[nfp]["J3"], outs[nfp]["J_Z"]) + assert xp.allclose(outs[nfp]["J1"], outs[nfp]["J_R"]) + assert xp.allclose(outs[nfp]["J2"], outs[nfp]["J_phi"]) + assert xp.allclose(outs[nfp]["J3"], outs[nfp]["J_Z"]) - outs[nfp]["Bx"] = np.cos(outs[nfp]["phi"]) * outs[nfp]["B_R"] - np.sin(outs[nfp]["phi"]) * outs[nfp]["B_phi"] + outs[nfp]["Bx"] = xp.cos(outs[nfp]["phi"]) * outs[nfp]["B_R"] - xp.sin(outs[nfp]["phi"]) * outs[nfp]["B_phi"] - outs[nfp]["By"] = np.sin(outs[nfp]["phi"]) * outs[nfp]["B_R"] + np.cos(outs[nfp]["phi"]) * outs[nfp]["B_phi"] + outs[nfp]["By"] = xp.sin(outs[nfp]["phi"]) * outs[nfp]["B_R"] + xp.cos(outs[nfp]["phi"]) * outs[nfp]["B_phi"] outs[nfp]["Bz"] = outs[nfp]["B_Z"] @@ -123,32 +123,32 @@ def test_desc_equil(do_plot=False): outs_struphy[nfp]["Y"] = y outs_struphy[nfp]["Z"] = z - outs_struphy[nfp]["R"] = np.sqrt(x**2 + y**2) - tmp = np.arctan2(y, x) - tmp[tmp < -1e-6] += 2 * np.pi + outs_struphy[nfp]["R"] = xp.sqrt(x**2 + y**2) + tmp = xp.arctan2(y, x) + tmp[tmp < -1e-6] += 2 * xp.pi outs_struphy[nfp]["phi"] = tmp - outs_struphy[nfp]["sqrt(g)"] = s_eq.domain.jacobian_det(e1, e2, e3) / (4 * np.pi**2 / nfp) + outs_struphy[nfp]["sqrt(g)"] = s_eq.domain.jacobian_det(e1, e2, e3) / (4 * xp.pi**2 / nfp) outs_struphy[nfp]["p"] = s_eq.p0(e1, e2, e3) # include push forward to DESC logical coordinates bv = s_eq.bv(e1, e2, e3) outs_struphy[nfp]["B^rho"] = bv[0] * (1 - rmin) - outs_struphy[nfp]["B^theta"] = bv[1] * 2 * np.pi - outs_struphy[nfp]["B^zeta"] = bv[2] * 2 * np.pi / nfp + outs_struphy[nfp]["B^theta"] = bv[1] * 2 * xp.pi + outs_struphy[nfp]["B^zeta"] = bv[2] * 2 * xp.pi / nfp outs_struphy[nfp]["B"] = s_eq.absB0(e1, e2, e3) # include push forward to DESC logical coordinates jv = s_eq.jv(e1, e2, e3) outs_struphy[nfp]["J^rho"] = jv[0] * (1 - rmin) - outs_struphy[nfp]["J^theta"] = jv[1] * 2 * np.pi - outs_struphy[nfp]["J^zeta"] = jv[2] * 2 * np.pi / nfp + outs_struphy[nfp]["J^theta"] = jv[1] * 2 * xp.pi + outs_struphy[nfp]["J^zeta"] = jv[2] * 2 * xp.pi / nfp j1 = s_eq.j1(e1, e2, e3) - outs_struphy[nfp]["J"] = np.sqrt(jv[0] * j1[0] + jv[1] * j1[1] + jv[2] * j1[2]) + outs_struphy[nfp]["J"] = xp.sqrt(jv[0] * j1[0] + jv[1] * j1[1] + jv[2] * j1[2]) b_cart, xyz = s_eq.b_cart(e1, e2, e3) outs_struphy[nfp]["Bx"] = b_cart[0] @@ -158,8 +158,8 @@ def test_desc_equil(do_plot=False): # include push forward to DESC logical coordinates gradB1 = s_eq.gradB1(e1, e2, e3) outs_struphy[nfp]["|B|_r"] = gradB1[0] / (1 - rmin) - outs_struphy[nfp]["|B|_t"] = gradB1[1] / (2 * np.pi) - outs_struphy[nfp]["|B|_z"] = gradB1[2] / (2 * np.pi / nfp) + outs_struphy[nfp]["|B|_t"] = gradB1[1] / (2 * xp.pi) + outs_struphy[nfp]["|B|_z"] = gradB1[2] / (2 * xp.pi / nfp) # comparisons vars += ["Bx", "By", "Bz"] @@ -173,10 +173,10 @@ def test_desc_equil(do_plot=False): if var in ("B_R", "B_phi", "B_Z", "J_R", "J_phi", "J_Z"): continue else: - max_norm = np.max(np.abs(outs[nfp][var])) + max_norm = xp.max(xp.abs(outs[nfp][var])) if max_norm < 1e-16: max_norm = 1.0 - err = np.max(np.abs(outs[nfp][var] - outs_struphy[nfp][var])) / max_norm + err = xp.max(xp.abs(outs[nfp][var] - outs_struphy[nfp][var])) / max_norm assert err < err_lim print( @@ -186,7 +186,7 @@ def test_desc_equil(do_plot=False): if do_plot: fig = plt.figure(figsize=(12, 13)) - levels = np.linspace(np.min(outs[nfp][var]) - 1e-10, np.max(outs[nfp][var]), 20) + levels = xp.linspace(xp.min(outs[nfp][var]) - 1e-10, xp.max(outs[nfp][var]), 20) # poloidal plot R = outs[nfp]["R"][:, :, 0].squeeze() diff --git a/src/struphy/fields_background/tests/test_generic_equils.py b/src/struphy/fields_background/tests/test_generic_equils.py index 8c10eb80e..6ed2c6178 100644 --- a/src/struphy/fields_background/tests/test_generic_equils.py +++ b/src/struphy/fields_background/tests/test_generic_equils.py @@ -5,12 +5,12 @@ GenericCartesianFluidEquilibrium, GenericCartesianFluidEquilibriumWithB, ) -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def test_generic_equils(show=False): - fun_vec = lambda x, y, z: (np.cos(2 * np.pi * x), np.cos(2 * np.pi * y), z) - fun_n = lambda x, y, z: np.exp(-((x - 1) ** 2) - (y) ** 2) + fun_vec = lambda x, y, z: (xp.cos(2 * xp.pi * x), xp.cos(2 * xp.pi * y), z) + fun_n = lambda x, y, z: xp.exp(-((x - 1) ** 2) - (y) ** 2) fun_p = lambda x, y, z: x**2 gen_eq = GenericCartesianFluidEquilibrium( u_xyz=fun_vec, @@ -25,22 +25,22 @@ def test_generic_equils(show=False): gradB_xyz=fun_vec, ) - x = np.linspace(-3, 3, 32) - y = np.linspace(-4, 4, 32) + x = xp.linspace(-3, 3, 32) + y = xp.linspace(-4, 4, 32) z = 1.0 - xx, yy, zz = np.meshgrid(x, y, z) + xx, yy, zz = xp.meshgrid(x, y, z) # gen_eq - assert all([np.all(tmp == fun_i) for tmp, fun_i in zip(gen_eq.u_xyz(xx, yy, zz), fun_vec(xx, yy, zz))]) - assert np.all(gen_eq.p_xyz(xx, yy, zz) == fun_p(xx, yy, zz)) - assert np.all(gen_eq.n_xyz(xx, yy, zz) == fun_n(xx, yy, zz)) + assert all([xp.all(tmp == fun_i) for tmp, fun_i in zip(gen_eq.u_xyz(xx, yy, zz), fun_vec(xx, yy, zz))]) + assert xp.all(gen_eq.p_xyz(xx, yy, zz) == fun_p(xx, yy, zz)) + assert xp.all(gen_eq.n_xyz(xx, yy, zz) == fun_n(xx, yy, zz)) # gen_eq_B - assert all([np.all(tmp == fun_i) for tmp, fun_i in zip(gen_eq_B.u_xyz(xx, yy, zz), fun_vec(xx, yy, zz))]) - assert np.all(gen_eq_B.p_xyz(xx, yy, zz) == fun_p(xx, yy, zz)) - assert np.all(gen_eq_B.n_xyz(xx, yy, zz) == fun_n(xx, yy, zz)) - assert all([np.all(tmp == fun_i) for tmp, fun_i in zip(gen_eq_B.b_xyz(xx, yy, zz), fun_vec(xx, yy, zz))]) - assert all([np.all(tmp == fun_i) for tmp, fun_i in zip(gen_eq_B.gradB_xyz(xx, yy, zz), fun_vec(xx, yy, zz))]) + assert all([xp.all(tmp == fun_i) for tmp, fun_i in zip(gen_eq_B.u_xyz(xx, yy, zz), fun_vec(xx, yy, zz))]) + assert xp.all(gen_eq_B.p_xyz(xx, yy, zz) == fun_p(xx, yy, zz)) + assert xp.all(gen_eq_B.n_xyz(xx, yy, zz) == fun_n(xx, yy, zz)) + assert all([xp.all(tmp == fun_i) for tmp, fun_i in zip(gen_eq_B.b_xyz(xx, yy, zz), fun_vec(xx, yy, zz))]) + assert all([xp.all(tmp == fun_i) for tmp, fun_i in zip(gen_eq_B.gradB_xyz(xx, yy, zz), fun_vec(xx, yy, zz))]) if show: plt.figure(figsize=(12, 12)) diff --git a/src/struphy/fields_background/tests/test_mhd_equils.py b/src/struphy/fields_background/tests/test_mhd_equils.py index bd750e9d9..076891914 100644 --- a/src/struphy/fields_background/tests/test_mhd_equils.py +++ b/src/struphy/fields_background/tests/test_mhd_equils.py @@ -1,7 +1,7 @@ import pytest from struphy.fields_background import equils -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize( @@ -9,44 +9,44 @@ [ ("HomogenSlab", {}, "Cuboid", {}), ("HomogenSlab", {}, "Colella", {"alpha": 0.06}), - ("ShearedSlab", {"a": 0.75, "R0": 3.5}, "Cuboid", {"r1": 0.75, "r2": 2 * np.pi * 0.75, "r3": 2 * np.pi * 3.5}), + ("ShearedSlab", {"a": 0.75, "R0": 3.5}, "Cuboid", {"r1": 0.75, "r2": 2 * xp.pi * 0.75, "r3": 2 * xp.pi * 3.5}), ( "ShearedSlab", {"a": 0.75, "R0": 3.5, "q0": "inf", "q1": "inf"}, "Cuboid", - {"r1": 0.75, "r2": 2 * np.pi * 0.75, "r3": 2 * np.pi * 3.5}, + {"r1": 0.75, "r2": 2 * xp.pi * 0.75, "r3": 2 * xp.pi * 3.5}, ), ( "ShearedSlab", {"a": 0.55, "R0": 4.5}, "Orthogonal", - {"Lx": 0.55, "Ly": 2 * np.pi * 0.55, "Lz": 2 * np.pi * 4.5}, + {"Lx": 0.55, "Ly": 2 * xp.pi * 0.55, "Lz": 2 * xp.pi * 4.5}, ), - ("ScrewPinch", {"a": 0.45, "R0": 2.5}, "HollowCylinder", {"a1": 0.05, "a2": 0.45, "Lz": 2 * np.pi * 2.5}), - ("ScrewPinch", {"a": 1.45, "R0": 6.5}, "IGAPolarCylinder", {"a": 1.45, "Lz": 2 * np.pi * 6.5}), + ("ScrewPinch", {"a": 0.45, "R0": 2.5}, "HollowCylinder", {"a1": 0.05, "a2": 0.45, "Lz": 2 * xp.pi * 2.5}), + ("ScrewPinch", {"a": 1.45, "R0": 6.5}, "IGAPolarCylinder", {"a": 1.45, "Lz": 2 * xp.pi * 6.5}), ( "ScrewPinch", {"a": 0.45, "R0": 2.5, "q0": 1.5, "q1": 1.5}, "HollowCylinder", - {"a1": 0.05, "a2": 0.45, "Lz": 2 * np.pi * 2.5}, + {"a1": 0.05, "a2": 0.45, "Lz": 2 * xp.pi * 2.5}, ), ( "ScrewPinch", {"a": 1.45, "R0": 6.5, "q0": 1.5, "q1": 1.5}, "IGAPolarCylinder", - {"a": 1.45, "Lz": 2 * np.pi * 6.5}, + {"a": 1.45, "Lz": 2 * xp.pi * 6.5}, ), ( "ScrewPinch", {"a": 0.45, "R0": 2.5, "q0": "inf", "q1": "inf"}, "HollowCylinder", - {"a1": 0.05, "a2": 0.45, "Lz": 2 * np.pi * 2.5}, + {"a1": 0.05, "a2": 0.45, "Lz": 2 * xp.pi * 2.5}, ), ( "ScrewPinch", {"a": 1.45, "R0": 6.5, "q0": "inf", "q1": "inf"}, "IGAPolarCylinder", - {"a": 1.45, "Lz": 2 * np.pi * 6.5}, + {"a": 1.45, "Lz": 2 * xp.pi * 6.5}, ), ( "AdhocTorus", @@ -141,24 +141,24 @@ def test_equils(equil_domain_pair): from struphy.geometry import domains # logical evalution point - pt = (np.random.rand(), np.random.rand(), np.random.rand()) + pt = (xp.random.rand(), xp.random.rand(), xp.random.rand()) # logical arrays: - e1 = np.random.rand(4) - e2 = np.random.rand(5) - e3 = np.random.rand(6) + e1 = xp.random.rand(4) + e2 = xp.random.rand(5) + e3 = xp.random.rand(6) # 2d slices - mat_12_1, mat_12_2 = np.meshgrid(e1, e2, indexing="ij") - mat_13_1, mat_13_3 = np.meshgrid(e1, e3, indexing="ij") - mat_23_2, mat_23_3 = np.meshgrid(e2, e3, indexing="ij") + mat_12_1, mat_12_2 = xp.meshgrid(e1, e2, indexing="ij") + mat_13_1, mat_13_3 = xp.meshgrid(e1, e3, indexing="ij") + mat_23_2, mat_23_3 = xp.meshgrid(e2, e3, indexing="ij") # 3d - mat_123_1, mat_123_2, mat_123_3 = np.meshgrid(e1, e2, e3, indexing="ij") - mat_123_1_sp, mat_123_2_sp, mat_123_3_sp = np.meshgrid(e1, e2, e3, indexing="ij", sparse=True) + mat_123_1, mat_123_2, mat_123_3 = xp.meshgrid(e1, e2, e3, indexing="ij") + mat_123_1_sp, mat_123_2_sp, mat_123_3_sp = xp.meshgrid(e1, e2, e3, indexing="ij", sparse=True) # markers - markers = np.random.rand(33, 10) + markers = xp.random.rand(33, 10) # create MHD equilibrium eq_mhd = getattr(equils, equil_domain_pair[0])(**equil_domain_pair[1]) @@ -274,8 +274,8 @@ def test_equils(equil_domain_pair): # --------- eta1 evaluation --------- results = [] - e2_pt = np.random.rand() - e3_pt = np.random.rand() + e2_pt = xp.random.rand() + e3_pt = xp.random.rand() # scalar functions results.append(eq_mhd.absB0(e1, e2_pt, e3_pt, squeeze_out=True)) @@ -321,8 +321,8 @@ def test_equils(equil_domain_pair): # --------- eta2 evaluation --------- results = [] - e1_pt = np.random.rand() - e3_pt = np.random.rand() + e1_pt = xp.random.rand() + e3_pt = xp.random.rand() # scalar functions results.append(eq_mhd.absB0(e1_pt, e2, e3_pt, squeeze_out=True)) @@ -370,8 +370,8 @@ def test_equils(equil_domain_pair): # --------- eta3 evaluation --------- results = [] - e1_pt = np.random.rand() - e2_pt = np.random.rand() + e1_pt = xp.random.rand() + e2_pt = xp.random.rand() # scalar functions results.append(eq_mhd.absB0(e1_pt, e2_pt, e3, squeeze_out=True)) @@ -419,7 +419,7 @@ def test_equils(equil_domain_pair): # --------- eta1-eta2 evaluation --------- results = [] - e3_pt = np.random.rand() + e3_pt = xp.random.rand() # scalar functions results.append(eq_mhd.absB0(e1, e2, e3_pt, squeeze_out=True)) @@ -467,7 +467,7 @@ def test_equils(equil_domain_pair): # --------- eta1-eta3 evaluation --------- results = [] - e2_pt = np.random.rand() + e2_pt = xp.random.rand() # scalar functions results.append(eq_mhd.absB0(e1, e2_pt, e3, squeeze_out=True)) @@ -515,7 +515,7 @@ def test_equils(equil_domain_pair): # --------- eta2-eta3 evaluation --------- results = [] - e1_pt = np.random.rand() + e1_pt = xp.random.rand() # scalar functions results.append(eq_mhd.absB0(e1_pt, e2, e3, squeeze_out=True)) @@ -609,7 +609,7 @@ def test_equils(equil_domain_pair): # --------- 12 matrix evaluation --------- results = [] - e3_pt = np.random.rand() + e3_pt = xp.random.rand() # scalar functions results.append(eq_mhd.absB0(mat_12_1, mat_12_2, e3_pt, squeeze_out=True)) @@ -657,7 +657,7 @@ def test_equils(equil_domain_pair): # --------- 13 matrix evaluation --------- results = [] - e2_pt = np.random.rand() + e2_pt = xp.random.rand() # scalar functions results.append(eq_mhd.absB0(mat_13_1, e2_pt, mat_13_3, squeeze_out=True)) @@ -705,7 +705,7 @@ def test_equils(equil_domain_pair): # --------- 23 matrix evaluation --------- results = [] - e1_pt = np.random.rand() + e1_pt = xp.random.rand() # scalar functions results.append(eq_mhd.absB0(e1_pt, mat_23_2, mat_23_3, squeeze_out=True)) @@ -848,22 +848,22 @@ def assert_scalar(result, kind, *etas): markers = etas[0] n_p = markers.shape[0] - assert isinstance(result, np.ndarray) + assert isinstance(result, xp.ndarray) assert result.shape == (n_p,) for ip in range(n_p): assert isinstance(result[ip], float) - assert not np.isnan(result[ip]) + assert not xp.isnan(result[ip]) else: # point-wise if kind == "point": assert isinstance(result, float) - assert not np.isnan(result) + assert not xp.isnan(result) # slices else: - assert isinstance(result, np.ndarray) + assert isinstance(result, xp.ndarray) # eta1-array if kind == "e1": @@ -915,27 +915,27 @@ def assert_vector(result, kind, *etas): markers = etas[0] n_p = markers.shape[0] - assert isinstance(result, np.ndarray) + assert isinstance(result, xp.ndarray) assert result.shape == (3, n_p) for c in range(3): for ip in range(n_p): assert isinstance(result[c, ip], float) - assert not np.isnan(result[c, ip]) + assert not xp.isnan(result[c, ip]) else: # point-wise if kind == "point": - assert isinstance(result, np.ndarray) + assert isinstance(result, xp.ndarray) assert result.shape == (3,) for c in range(3): assert isinstance(result[c], float) - assert not np.isnan(result[c]) + assert not xp.isnan(result[c]) # slices else: - assert isinstance(result, np.ndarray) + assert isinstance(result, xp.ndarray) # eta1-array if kind == "e1": diff --git a/src/struphy/fields_background/tests/test_numerical_mhd_equil.py b/src/struphy/fields_background/tests/test_numerical_mhd_equil.py index 4d34e8352..1002c37c8 100644 --- a/src/struphy/fields_background/tests/test_numerical_mhd_equil.py +++ b/src/struphy/fields_background/tests/test_numerical_mhd_equil.py @@ -1,7 +1,7 @@ import pytest from struphy.fields_background.base import FluidEquilibrium, LogicalMHDequilibrium -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize( @@ -50,53 +50,53 @@ def test_transformations(mapping, mhd_equil): num_equil = NumEqTest(domain, proxy) # compare values: - eta1 = np.random.rand(4) - eta2 = np.random.rand(5) - eta3 = np.random.rand(6) + eta1 = xp.random.rand(4) + eta2 = xp.random.rand(5) + eta3 = xp.random.rand(6) - assert np.allclose(ana_equil.absB0(eta1, eta2, eta3), num_equil.absB0(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.absB0(eta1, eta2, eta3), num_equil.absB0(eta1, eta2, eta3)) - assert np.allclose(ana_equil.bv(eta1, eta2, eta3)[0], num_equil.bv(eta1, eta2, eta3)[0]) - assert np.allclose(ana_equil.bv(eta1, eta2, eta3)[1], num_equil.bv(eta1, eta2, eta3)[1]) - assert np.allclose(ana_equil.bv(eta1, eta2, eta3)[2], num_equil.bv(eta1, eta2, eta3)[2]) + assert xp.allclose(ana_equil.bv(eta1, eta2, eta3)[0], num_equil.bv(eta1, eta2, eta3)[0]) + assert xp.allclose(ana_equil.bv(eta1, eta2, eta3)[1], num_equil.bv(eta1, eta2, eta3)[1]) + assert xp.allclose(ana_equil.bv(eta1, eta2, eta3)[2], num_equil.bv(eta1, eta2, eta3)[2]) - assert np.allclose(ana_equil.b1_1(eta1, eta2, eta3), num_equil.b1_1(eta1, eta2, eta3)) - assert np.allclose(ana_equil.b1_2(eta1, eta2, eta3), num_equil.b1_2(eta1, eta2, eta3)) - assert np.allclose(ana_equil.b1_3(eta1, eta2, eta3), num_equil.b1_3(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.b1_1(eta1, eta2, eta3), num_equil.b1_1(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.b1_2(eta1, eta2, eta3), num_equil.b1_2(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.b1_3(eta1, eta2, eta3), num_equil.b1_3(eta1, eta2, eta3)) - assert np.allclose(ana_equil.b2_1(eta1, eta2, eta3), num_equil.b2_1(eta1, eta2, eta3)) - assert np.allclose(ana_equil.b2_2(eta1, eta2, eta3), num_equil.b2_2(eta1, eta2, eta3)) - assert np.allclose(ana_equil.b2_3(eta1, eta2, eta3), num_equil.b2_3(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.b2_1(eta1, eta2, eta3), num_equil.b2_1(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.b2_2(eta1, eta2, eta3), num_equil.b2_2(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.b2_3(eta1, eta2, eta3), num_equil.b2_3(eta1, eta2, eta3)) - assert np.allclose(ana_equil.unit_bv(eta1, eta2, eta3)[0], num_equil.unit_bv(eta1, eta2, eta3)[0]) - assert np.allclose(ana_equil.unit_bv(eta1, eta2, eta3)[1], num_equil.unit_bv(eta1, eta2, eta3)[1]) - assert np.allclose(ana_equil.unit_bv(eta1, eta2, eta3)[2], num_equil.unit_bv(eta1, eta2, eta3)[2]) + assert xp.allclose(ana_equil.unit_bv(eta1, eta2, eta3)[0], num_equil.unit_bv(eta1, eta2, eta3)[0]) + assert xp.allclose(ana_equil.unit_bv(eta1, eta2, eta3)[1], num_equil.unit_bv(eta1, eta2, eta3)[1]) + assert xp.allclose(ana_equil.unit_bv(eta1, eta2, eta3)[2], num_equil.unit_bv(eta1, eta2, eta3)[2]) - assert np.allclose(ana_equil.unit_b1_1(eta1, eta2, eta3), num_equil.unit_b1_1(eta1, eta2, eta3)) - assert np.allclose(ana_equil.unit_b1_2(eta1, eta2, eta3), num_equil.unit_b1_2(eta1, eta2, eta3)) - assert np.allclose(ana_equil.unit_b1_3(eta1, eta2, eta3), num_equil.unit_b1_3(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.unit_b1_1(eta1, eta2, eta3), num_equil.unit_b1_1(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.unit_b1_2(eta1, eta2, eta3), num_equil.unit_b1_2(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.unit_b1_3(eta1, eta2, eta3), num_equil.unit_b1_3(eta1, eta2, eta3)) - assert np.allclose(ana_equil.unit_b2_1(eta1, eta2, eta3), num_equil.unit_b2_1(eta1, eta2, eta3)) - assert np.allclose(ana_equil.unit_b2_2(eta1, eta2, eta3), num_equil.unit_b2_2(eta1, eta2, eta3)) - assert np.allclose(ana_equil.unit_b2_3(eta1, eta2, eta3), num_equil.unit_b2_3(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.unit_b2_1(eta1, eta2, eta3), num_equil.unit_b2_1(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.unit_b2_2(eta1, eta2, eta3), num_equil.unit_b2_2(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.unit_b2_3(eta1, eta2, eta3), num_equil.unit_b2_3(eta1, eta2, eta3)) - assert np.allclose(ana_equil.jv(eta1, eta2, eta3)[0], num_equil.jv(eta1, eta2, eta3)[0]) - assert np.allclose(ana_equil.jv(eta1, eta2, eta3)[1], num_equil.jv(eta1, eta2, eta3)[1]) - assert np.allclose(ana_equil.jv(eta1, eta2, eta3)[2], num_equil.jv(eta1, eta2, eta3)[2]) + assert xp.allclose(ana_equil.jv(eta1, eta2, eta3)[0], num_equil.jv(eta1, eta2, eta3)[0]) + assert xp.allclose(ana_equil.jv(eta1, eta2, eta3)[1], num_equil.jv(eta1, eta2, eta3)[1]) + assert xp.allclose(ana_equil.jv(eta1, eta2, eta3)[2], num_equil.jv(eta1, eta2, eta3)[2]) - assert np.allclose(ana_equil.j1_1(eta1, eta2, eta3), num_equil.j1_1(eta1, eta2, eta3)) - assert np.allclose(ana_equil.j1_2(eta1, eta2, eta3), num_equil.j1_2(eta1, eta2, eta3)) - assert np.allclose(ana_equil.j1_3(eta1, eta2, eta3), num_equil.j1_3(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.j1_1(eta1, eta2, eta3), num_equil.j1_1(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.j1_2(eta1, eta2, eta3), num_equil.j1_2(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.j1_3(eta1, eta2, eta3), num_equil.j1_3(eta1, eta2, eta3)) - assert np.allclose(ana_equil.j2_1(eta1, eta2, eta3), num_equil.j2_1(eta1, eta2, eta3)) - assert np.allclose(ana_equil.j2_2(eta1, eta2, eta3), num_equil.j2_2(eta1, eta2, eta3)) - assert np.allclose(ana_equil.j2_3(eta1, eta2, eta3), num_equil.j2_3(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.j2_1(eta1, eta2, eta3), num_equil.j2_1(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.j2_2(eta1, eta2, eta3), num_equil.j2_2(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.j2_3(eta1, eta2, eta3), num_equil.j2_3(eta1, eta2, eta3)) - assert np.allclose(ana_equil.p0(eta1, eta2, eta3), num_equil.p0(eta1, eta2, eta3)) - assert np.allclose(ana_equil.p3(eta1, eta2, eta3), num_equil.p3(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.p0(eta1, eta2, eta3), num_equil.p0(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.p3(eta1, eta2, eta3), num_equil.p3(eta1, eta2, eta3)) - assert np.allclose(ana_equil.n0(eta1, eta2, eta3), num_equil.n0(eta1, eta2, eta3)) - assert np.allclose(ana_equil.n3(eta1, eta2, eta3), num_equil.n3(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.n0(eta1, eta2, eta3), num_equil.n0(eta1, eta2, eta3)) + assert xp.allclose(ana_equil.n3(eta1, eta2, eta3), num_equil.n3(eta1, eta2, eta3)) class NumEqTest(LogicalMHDequilibrium): diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index 2a3141e96..9de4acede 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -11,7 +11,7 @@ from struphy.geometry import evaluation_kernels, transform_kernels from struphy.kernel_arguments.pusher_args_kernels import DomainArguments from struphy.linear_algebra import linalg_kron -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Domain(metaclass=ABCMeta): @@ -56,12 +56,12 @@ def __init__( self._NbaseN = [Nel + p - kind * p for Nel, p, kind in zip(Nel, p, spl_kind)] - el_b = [np.linspace(0.0, 1.0, Nel + 1) for Nel in Nel] + el_b = [xp.linspace(0.0, 1.0, Nel + 1) for Nel in Nel] self._T = [bsp.make_knots(el_b, p, kind) for el_b, p, kind in zip(el_b, p, spl_kind)] self._indN = [ - (np.indices((Nel, p + 1))[1] + np.arange(Nel)[:, None]) % NbaseN + (xp.indices((Nel, p + 1))[1] + xp.arange(Nel)[:, None]) % NbaseN for Nel, p, NbaseN in zip(Nel, p, self._NbaseN) ] @@ -71,15 +71,15 @@ def __init__( self._p = (*self._p, 0) self._NbaseN = self._NbaseN + [0] - self._T = self._T + [np.zeros((1,), dtype=float)] + self._T = self._T + [xp.zeros((1,), dtype=float)] - self._indN = self._indN + [np.zeros((1, 1), dtype=int)] + self._indN = self._indN + [xp.zeros((1, 1), dtype=int)] # create dummy attributes for analytical mappings if self.kind_map >= 10: - self._cx = np.zeros((1, 1, 1), dtype=float) - self._cy = np.zeros((1, 1, 1), dtype=float) - self._cz = np.zeros((1, 1, 1), dtype=float) + self._cx = xp.zeros((1, 1, 1), dtype=float) + self._cy = xp.zeros((1, 1, 1), dtype=float) + self._cz = xp.zeros((1, 1, 1), dtype=float) self._transformation_ids = { "pull": 0, @@ -120,7 +120,7 @@ def __init__( self._args_domain = DomainArguments( self.kind_map, self.params_numpy, - np.array(self.p), + xp.array(self.p), self.T[0], self.T[1], self.T[2], @@ -165,15 +165,15 @@ def params(self, new): self._params = new @property - def params_numpy(self) -> np.ndarray: + def params_numpy(self) -> xp.ndarray: """Mapping parameters as numpy array (can be empty).""" if not hasattr(self, "_params_numpy"): - self._params_numpy = np.array([0], dtype=float) + self._params_numpy = xp.array([0], dtype=float) return self._params_numpy @params_numpy.setter def params_numpy(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) assert new.ndim == 1 self._params_numpy = new @@ -768,7 +768,7 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs): markers = etas[0] # to keep C-ordering the (3, 3)-part is in the last indices - out = np.empty((markers.shape[0], 3, 3), dtype=float) + out = xp.empty((markers.shape[0], 3, 3), dtype=float) n_inside = evaluation_kernels.kernel_evaluate_pic( markers, @@ -780,24 +780,24 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs): ) # move the (3, 3)-part to front - out = np.transpose(out, axes=(1, 2, 0)) + out = xp.transpose(out, axes=(1, 2, 0)) # remove holes out = out[:, :, :n_inside] if transposed: - out = np.transpose(out, axes=(1, 0, 2)) + out = xp.transpose(out, axes=(1, 0, 2)) # change size of "out" depending on which metric coeff has been evaluated if which == 0 or which == -1: out = out[:, 0, :] if change_out_order: - out = np.transpose(out, axes=(1, 0)) + out = xp.transpose(out, axes=(1, 0)) elif which == 2: out = out[0, 0, :] else: if change_out_order: - out = np.transpose(out, axes=(2, 0, 1)) + out = xp.transpose(out, axes=(2, 0, 1)) # tensor-product/slice evaluation else: @@ -809,7 +809,7 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs): ) # to keep C-ordering the (3, 3)-part is in the last indices - out = np.empty( + out = xp.empty( (E1.shape[0], E2.shape[1], E3.shape[2], 3, 3), dtype=float, ) @@ -825,20 +825,20 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs): ) # move the (3, 3)-part to front - out = np.transpose(out, axes=(3, 4, 0, 1, 2)) + out = xp.transpose(out, axes=(3, 4, 0, 1, 2)) if transposed: - out = np.transpose(out, axes=(1, 0, 2, 3, 4)) + out = xp.transpose(out, axes=(1, 0, 2, 3, 4)) if which == 0: out = out[:, 0, :, :, :] if change_out_order: - out = np.transpose(out, axes=(1, 2, 3, 0)) + out = xp.transpose(out, axes=(1, 2, 3, 0)) elif which == 2: out = out[0, 0, :, :, :] else: if change_out_order: - out = np.transpose(out, axes=(2, 3, 4, 0, 1)) + out = xp.transpose(out, axes=(2, 3, 4, 0, 1)) # remove singleton dimensions for slice evaluation if squeeze_out: @@ -903,7 +903,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa assert len(etas) == 3 assert etas[0].shape == etas[1].shape == etas[2].shape assert etas[0].ndim == 1 - markers = np.stack(etas, axis=1) + markers = xp.stack(etas, axis=1) else: markers = etas[0] @@ -955,7 +955,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa A_has_holes = False # call evaluation kernel - out = np.empty((markers.shape[0], 3), dtype=float) + out = xp.empty((markers.shape[0], 3), dtype=float) # make sure we don't have stride = 0 A = A.copy() @@ -971,7 +971,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa ) # move the (3, 3)-part to front - out = np.transpose(out, axes=(1, 0)) + out = xp.transpose(out, axes=(1, 0)) # remove holes out = out[:, :n_inside] @@ -985,7 +985,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa out = out[0, :] else: if change_out_order: - out = np.transpose(out, axes=(1, 0)) + out = xp.transpose(out, axes=(1, 0)) # tensor-product/slice evaluation else: @@ -1012,7 +1012,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa A = Domain.prepare_arg(a, X[0], X[1], X[2], a_kwargs=a_kwargs) # call evaluation kernel - out = np.empty( + out = xp.empty( (E1.shape[0], E2.shape[1], E3.shape[2], 3), dtype=float, ) @@ -1029,14 +1029,14 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa ) # move the (3, 3)-part to front - out = np.transpose(out, axes=(3, 0, 1, 2)) + out = xp.transpose(out, axes=(3, 0, 1, 2)) # change output order if kind_int < 10: out = out[0, :, :, :] else: if change_out_order: - out = np.transpose(out, axes=(1, 2, 3, 0)) + out = xp.transpose(out, axes=(1, 2, 3, 0)) # remove singleton dimensions for slice evaluation if squeeze_out: @@ -1083,22 +1083,22 @@ def prepare_eval_pts(x, y, z, flat_eval=False): if flat_eval: # convert list type data to numpy array: if isinstance(x, list): - arg_x = np.array(x) - elif isinstance(x, np.ndarray): + arg_x = xp.array(x) + elif isinstance(x, xp.ndarray): arg_x = x else: raise ValueError("Input x must be a 1d list or numpy array") if isinstance(y, list): - arg_y = np.array(y) - elif isinstance(y, np.ndarray): + arg_y = xp.array(y) + elif isinstance(y, xp.ndarray): arg_y = y else: raise ValueError("Input y must be a 1d list or numpy array") if isinstance(z, list): - arg_z = np.array(z) - elif isinstance(z, np.ndarray): + arg_z = xp.array(z) + elif isinstance(z, xp.ndarray): arg_z = z else: raise ValueError("Input z must be a 1d list or numpy array") @@ -1117,56 +1117,56 @@ def prepare_eval_pts(x, y, z, flat_eval=False): else: # convert list type data to numpy array: if isinstance(x, float): - arg_x = np.array([x]) + arg_x = xp.array([x]) elif isinstance(x, int): - arg_x = np.array([float(x)]) + arg_x = xp.array([float(x)]) elif isinstance(x, list): - arg_x = np.array(x) - elif isinstance(x, np.ndarray): + arg_x = xp.array(x) + elif isinstance(x, xp.ndarray): arg_x = x.copy() else: raise ValueError(f"data type {type(x)} not supported") if isinstance(y, float): - arg_y = np.array([y]) + arg_y = xp.array([y]) elif isinstance(y, int): - arg_y = np.array([float(y)]) + arg_y = xp.array([float(y)]) elif isinstance(y, list): - arg_y = np.array(y) - elif isinstance(y, np.ndarray): + arg_y = xp.array(y) + elif isinstance(y, xp.ndarray): arg_y = y.copy() else: raise ValueError(f"data type {type(y)} not supported") if isinstance(z, float): - arg_z = np.array([z]) + arg_z = xp.array([z]) elif isinstance(z, int): - arg_z = np.array([float(z)]) + arg_z = xp.array([float(z)]) elif isinstance(z, list): - arg_z = np.array(z) - elif isinstance(z, np.ndarray): + arg_z = xp.array(z) + elif isinstance(z, xp.ndarray): arg_z = z.copy() else: raise ValueError(f"data type {type(z)} not supported") # tensor-product for given three 1D arrays if arg_x.ndim == 1 and arg_y.ndim == 1 and arg_z.ndim == 1: - E1, E2, E3 = np.meshgrid(arg_x, arg_y, arg_z, indexing="ij") + E1, E2, E3 = xp.meshgrid(arg_x, arg_y, arg_z, indexing="ij") # given xy-plane at point z: elif arg_x.ndim == 2 and arg_y.ndim == 2 and arg_z.size == 1: E1 = arg_x[:, :, None] E2 = arg_y[:, :, None] - E3 = arg_z * np.ones(E1.shape) + E3 = arg_z * xp.ones(E1.shape) # given xz-plane at point y: elif arg_x.ndim == 2 and arg_y.size == 1 and arg_z.ndim == 2: E1 = arg_x[:, None, :] - E2 = arg_y * np.ones(E1.shape) + E2 = arg_y * xp.ones(E1.shape) E3 = arg_z[:, None, :] # given yz-plane at point x: elif arg_x.size == 1 and arg_y.ndim == 2 and arg_z.ndim == 2: E2 = arg_y[None, :, :] E3 = arg_z[None, :, :] - E1 = arg_x * np.ones(E2.shape) + E1 = arg_x * xp.ones(E2.shape) # given three 3D arrays elif arg_x.ndim == 3 and arg_y.ndim == 3 and arg_z.ndim == 3: # Distinguish if input coordinates are from sparse or dense meshgrid. @@ -1224,7 +1224,7 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): # float (point-wise, scalar function) if isinstance(a_in, float): - a_out = np.array([[[[a_in]]]]) + a_out = xp.array([[[[a_in]]]]) # single callable: # scalar function -> must return a 3d array for 3d evaluation points @@ -1237,7 +1237,7 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): else: if is_sparse_meshgrid: a_out = a_in( - *np.meshgrid(Xs[0][:, 0, 0], Xs[1][0, :, 0], Xs[2][0, 0, :], indexing="ij"), + *xp.meshgrid(Xs[0][:, 0, 0], Xs[1][0, :, 0], Xs[2][0, 0, :], indexing="ij"), **a_kwargs, ) else: @@ -1245,7 +1245,7 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): # case of Field.__call__ if isinstance(a_out, list): - a_out = np.array(a_out) + a_out = xp.array(a_out) if a_out.ndim == 3: a_out = a_out[None, :, :, :] @@ -1273,7 +1273,7 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): if is_sparse_meshgrid: a_out += [ component( - *np.meshgrid( + *xp.meshgrid( Xs[0][:, 0, 0], Xs[1][0, :, 0], Xs[2][0, 0, :], @@ -1285,7 +1285,7 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): else: a_out += [component(*Xs, **a_kwargs)] - elif isinstance(component, np.ndarray): + elif isinstance(component, xp.ndarray): if flat_eval: assert component.ndim == 1, print(f"{component.ndim = }") else: @@ -1294,16 +1294,16 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): a_out += [component] elif isinstance(component, float): - a_out += [np.array([component])[:, None, None]] + a_out += [xp.array([component])[:, None, None]] - a_out = np.array(a_out, dtype=float) + a_out = xp.array(a_out, dtype=float) # numpy array: # 1d array (flat_eval=True and scalar input or flat_eval=False and length 1 (scalar) or length 3 (vector)) # 2d array (flat_eval=True and vector-valued input of shape (3,:)) # 3d array (flat_eval=False and scalar input) # 4d array (flat_eval=False and vector-valued input of shape (3,:,:,:)) - elif isinstance(a_in, np.ndarray): + elif isinstance(a_in, xp.ndarray): if flat_eval: if a_in.ndim == 1: a_out = a_in[None, :] @@ -1344,26 +1344,26 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): if flat_eval: assert a_out.ndim == 2 assert a_out.shape[0] == 1 or a_out.shape[0] == 3 - a_out = np.ascontiguousarray(np.transpose(a_out, axes=(1, 0))).copy() # Make sure we don't have stride 0 + a_out = xp.ascontiguousarray(xp.transpose(a_out, axes=(1, 0))).copy() # Make sure we don't have stride 0 # make sure that output array is 4d and of shape (:,:,:, 1) or (:,:,:, 3) for tensor-product/slice evaluation else: assert a_out.ndim == 4 assert a_out.shape[0] == 1 or a_out.shape[0] == 3 - a_out = np.ascontiguousarray( - np.transpose(a_out, axes=(1, 2, 3, 0)), + a_out = xp.ascontiguousarray( + xp.transpose(a_out, axes=(1, 2, 3, 0)), ).copy() # Make sure we don't have stride 0 return a_out # ================================ - def get_params_numpy(self) -> np.ndarray: + def get_params_numpy(self) -> xp.ndarray: """Convert parameter dict into numpy array.""" params_numpy = [] for k, v in self.params.items(): params_numpy.append(v) - return np.array(params_numpy) + return xp.array(params_numpy) def show( self, @@ -1414,12 +1414,12 @@ def show( # plot domain without MPI decomposition and high resolution if grid_info is None: - e1 = np.linspace(0.0, 1.0, 16) - e2 = np.linspace(0.0, 1.0, 65) + e1 = xp.linspace(0.0, 1.0, 16) + e2 = xp.linspace(0.0, 1.0, 65) if logical: - E1, E2 = np.meshgrid(e1, e2, indexing="ij") - X = np.stack((E1, E2), axis=0) + E1, E2 = xp.meshgrid(e1, e2, indexing="ij") + X = xp.stack((E1, E2), axis=0) else: XYZ = self(e1, e2, 0.0, squeeze_out=True) @@ -1459,11 +1459,11 @@ def show( ) # top view - e3 = np.linspace(0.0, 1.0, 65) + e3 = xp.linspace(0.0, 1.0, 65) if logical: - E1, E2 = np.meshgrid(e1, e2, indexing="ij") - X = np.stack((E1, E2), axis=0) + E1, E2 = xp.meshgrid(e1, e2, indexing="ij") + X = xp.stack((E1, E2), axis=0) else: theta_0 = self(e1, 0.0, e3, squeeze_out=True) theta_pi = self(e1, 0.5, e3, squeeze_out=True) @@ -1524,7 +1524,7 @@ def show( # coordinates # e3 = [0., .25, .5, .75] # x, y, z = self(e1, e2, e3) - # R = np.sqrt(x**2 + y**2) + # R = xp.sqrt(x**2 + y**2) # fig = plt.figure(figsize=(13, 13)) # for n in range(4): @@ -1551,14 +1551,14 @@ def show( elif isinstance(grid_info, list): assert len(grid_info) > 1 - e1 = np.linspace(0.0, 1.0, grid_info[0] + 1) - e2 = np.linspace(0.0, 1.0, grid_info[1] + 1) + e1 = xp.linspace(0.0, 1.0, grid_info[0] + 1) + e2 = xp.linspace(0.0, 1.0, grid_info[1] + 1) fig = plt.figure(figsize=figsize) ax = fig.add_subplot(1, 1, 1) if logical: - E1, E2 = np.meshgrid(e1, e2, indexing="ij") + E1, E2 = xp.meshgrid(e1, e2, indexing="ij") # eta1-isolines for i in range(e1.size): @@ -1586,7 +1586,7 @@ def show( ax.plot(X[co1, :, j], X[co2, :, j], "tab:blue", alpha=0.5) # plot domain with MPI decomposition - elif isinstance(grid_info, np.ndarray): + elif isinstance(grid_info, xp.ndarray): assert grid_info.ndim == 2 assert grid_info.shape[1] > 5 @@ -1594,7 +1594,7 @@ def show( ax = fig.add_subplot(1, 1, 1) for i in range(grid_info.shape[0]): - e1 = np.linspace( + e1 = xp.linspace( grid_info[i, 0], grid_info[i, 1], int( @@ -1602,7 +1602,7 @@ def show( ) + 1, ) - e2 = np.linspace( + e2 = xp.linspace( grid_info[i, 3], grid_info[i, 4], int( @@ -1612,7 +1612,7 @@ def show( ) if logical: - E1, E2 = np.meshgrid(e1, e2, indexing="ij") + E1, E2 = xp.meshgrid(e1, e2, indexing="ij") # eta1-isolines first_line = ax.plot( @@ -1737,7 +1737,7 @@ def show( ax.axis("equal") - if isinstance(grid_info, np.ndarray): + if isinstance(grid_info, xp.ndarray): plt.legend() if self.__class__.__name__ in torus_mappings: @@ -1772,9 +1772,9 @@ def __init__( Nel: tuple[int] = (8, 24, 6), p: tuple[int] = (2, 3, 1), spl_kind: tuple[bool] = (False, True, True), - cx: np.ndarray = None, - cy: np.ndarray = None, - cz: np.ndarray = None, + cx: xp.ndarray = None, + cy: xp.ndarray = None, + cz: xp.ndarray = None, ): self.kind_map = 0 @@ -1805,7 +1805,7 @@ def __init__( assert self.cz.shape == expected_shape # identify polar singularity at eta1=0 - if np.all(self.cx[0, :, 0] == self.cx[0, 0, 0]): + if xp.all(self.cx[0, :, 0] == self.cx[0, 0, 0]): self.pole = True else: self.pole = False @@ -1836,17 +1836,17 @@ def __init__( Nel: tuple[int] = (8, 24), p: tuple[int] = (2, 3), spl_kind: tuple[bool] = (False, True), - cx: np.ndarray = None, - cy: np.ndarray = None, + cx: xp.ndarray = None, + cy: xp.ndarray = None, ): # get default control points if cx is None or cy is None: def X(eta1, eta2): - return eta1 * np.cos(2 * np.pi * eta2) + 3.0 + return eta1 * xp.cos(2 * xp.pi * eta2) + 3.0 def Y(eta1, eta2): - return eta1 * np.sin(2 * np.pi * eta2) + return eta1 * xp.sin(2 * xp.pi * eta2) cx, cy = interp_mapping(Nel, p, spl_kind, X, Y) @@ -1869,7 +1869,7 @@ def Y(eta1, eta2): assert self.cy.shape == expected_shape # identify polar singularity at eta1=0 - if np.all(self.cx[0, :] == self.cx[0, 0]): + if xp.all(self.cx[0, :] == self.cx[0, 0]): self.pole = True else: self.pole = False @@ -1877,7 +1877,7 @@ def Y(eta1, eta2): # reshape control points to 3D self._cx = self.cx[:, :, None] self._cy = self.cy[:, :, None] - self._cz = np.zeros((1, 1, 1), dtype=float) + self._cz = xp.zeros((1, 1, 1), dtype=float) # init base class super().__init__(Nel=Nel, p=p, spl_kind=spl_kind) @@ -1902,8 +1902,8 @@ def __init__( Nel: tuple[int] = (8, 24), p: tuple[int] = (2, 3), spl_kind: tuple[bool] = (False, True), - cx: np.ndarray = None, - cy: np.ndarray = None, + cx: xp.ndarray = None, + cy: xp.ndarray = None, Lz: float = 4.0, ): self.kind_map = 1 @@ -1912,10 +1912,10 @@ def __init__( if cx is None or cy is None: def X(eta1, eta2): - return eta1 * np.cos(2 * np.pi * eta2) + return eta1 * xp.cos(2 * xp.pi * eta2) def Y(eta1, eta2): - return eta1 * np.sin(2 * np.pi * eta2) + return eta1 * xp.sin(2 * xp.pi * eta2) cx, cy = interp_mapping(Nel, p, spl_kind, X, Y) @@ -1923,7 +1923,7 @@ def Y(eta1, eta2): cx[0] = 0.0 cy[0] = 0.0 - self.params_numpy = np.array([Lz]) + self.params_numpy = xp.array([Lz]) self.periodic_eta3 = False # init base class @@ -1954,7 +1954,7 @@ class PoloidalSplineTorus(PoloidalSpline): spl_kind : tuple[bool] Kind of spline in each poloidal direction (True=periodic, False=clamped). - cx, cy : np.ndarray + cx, cy : xp.ndarray Control points (spline coefficients) of the poloidal mapping. If None, a default square-to-disc mapping of radius 1 centered around (x, y) = (3, 0) is interpolated. @@ -1967,23 +1967,23 @@ def __init__( Nel: tuple[int] = (8, 24), p: tuple[int] = (2, 3), spl_kind: tuple[bool] = (False, True), - cx: np.ndarray = None, - cy: np.ndarray = None, + cx: xp.ndarray = None, + cy: xp.ndarray = None, tor_period: int = 3, ): # use setters for mapping attributes self.kind_map = 2 - self.params_numpy = np.array([float(tor_period)]) + self.params_numpy = xp.array([float(tor_period)]) self.periodic_eta3 = True # get default control points if cx is None or cy is None: def X(eta1, eta2): - return eta1 * np.cos(2 * np.pi * eta2) + 3.0 + return eta1 * xp.cos(2 * xp.pi * eta2) + 3.0 def Y(eta1, eta2): - return eta1 * np.sin(2 * np.pi * eta2) + return eta1 * xp.sin(2 * xp.pi * eta2) cx, cy = interp_mapping(Nel, p, spl_kind, X, Y) @@ -2025,7 +2025,7 @@ def interp_mapping(Nel, p, spl_kind, X, Y, Z=None): NbaseN = [Nel + p - kind * p for Nel, p, kind in zip(Nel, p, spl_kind)] # element boundaries - el_b = [np.linspace(0.0, 1.0, Nel + 1) for Nel in Nel] + el_b = [xp.linspace(0.0, 1.0, Nel + 1) for Nel in Nel] # spline knot vectors T = [bsp.make_knots(el_b, p, kind) for el_b, p, kind in zip(el_b, p, spl_kind)] @@ -2040,7 +2040,7 @@ def interp_mapping(Nel, p, spl_kind, X, Y, Z=None): if len(Nel) == 2: I = kron(I_mat[0], I_mat[1], format="csc") - I_pts = np.meshgrid(I_pts[0], I_pts[1], indexing="ij") + I_pts = xp.meshgrid(I_pts[0], I_pts[1], indexing="ij") cx = spsolve(I, X(I_pts[0], I_pts[1]).flatten()).reshape( NbaseN[0], @@ -2073,7 +2073,7 @@ def interp_mapping(Nel, p, spl_kind, X, Y, Z=None): return 0.0 -def spline_interpolation_nd(p: list, spl_kind: list, grids_1d: list, values: np.ndarray): +def spline_interpolation_nd(p: list, spl_kind: list, grids_1d: list, values: xp.ndarray): """n-dimensional tensor-product spline interpolation with discrete input. The interpolation points are passed as a list of 1d arrays, each array with increasing entries g[0]=0 < g[1] < ... @@ -2095,7 +2095,7 @@ def spline_interpolation_nd(p: list, spl_kind: list, grids_1d: list, values: np. Returns -------- - coeffs : np.array + coeffs : xp.array spline coefficients as nd array. T : list[array] @@ -2110,11 +2110,11 @@ def spline_interpolation_nd(p: list, spl_kind: list, grids_1d: list, values: np. I_mat = [] I_LU = [] for sh, x_grid, p_i, kind_i in zip(values.shape, grids_1d, p, spl_kind): - assert isinstance(x_grid, np.ndarray) + assert isinstance(x_grid, xp.ndarray) assert sh == x_grid.size assert ( - np.all( - np.roll(x_grid, 1)[1:] < x_grid[1:], + xp.all( + xp.roll(x_grid, 1)[1:] < x_grid[1:], ) and x_grid[-1] > x_grid[-2] ) @@ -2122,17 +2122,17 @@ def spline_interpolation_nd(p: list, spl_kind: list, grids_1d: list, values: np. if kind_i: assert x_grid[-1] < 1.0, "Interpolation points must be <1 for periodic interpolation." - breaks = np.ones(x_grid.size + 1) + breaks = xp.ones(x_grid.size + 1) if p_i % 2 == 0: - breaks[1:-1] = (x_grid[1:] + np.roll(x_grid, 1)[1:]) / 2.0 + breaks[1:-1] = (x_grid[1:] + xp.roll(x_grid, 1)[1:]) / 2.0 breaks[0] = 0.0 else: breaks[:-1] = x_grid else: assert ( - np.abs( + xp.abs( x_grid[-1] - 1.0, ) < 1e-14 @@ -2149,12 +2149,12 @@ def spline_interpolation_nd(p: list, spl_kind: list, grids_1d: list, values: np. breaks[0] = 0.0 breaks[-1] = 1.0 - # breaks = np.linspace(0., 1., x_grid.size - (not kind_i)*p_i + 1) + # breaks = xp.linspace(0., 1., x_grid.size - (not kind_i)*p_i + 1) T += [bsp.make_knots(breaks, p_i, periodic=kind_i)] indN += [ - (np.indices((breaks.size - 1, p_i + 1))[1] + np.arange(breaks.size - 1)[:, None]) % x_grid.size, + (xp.indices((breaks.size - 1, p_i + 1))[1] + xp.arange(breaks.size - 1)[:, None]) % x_grid.size, ] I_mat += [bsp.collocation_matrix(T[-1], p_i, x_grid, periodic=kind_i)] diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index da8eb2fe2..5e1649dc7 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -12,7 +12,7 @@ interp_mapping, ) from struphy.geometry.utilities import field_line_tracing -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Tokamak(PoloidalSplineTorus): @@ -156,8 +156,8 @@ def __init__(self, gvec_equil=None): def XYZ(e1, e2, e3): rho = _rmin + e1 * (1.0 - _rmin) - theta = 2 * np.pi * e2 - zeta = 2 * np.pi * e3 / gvec_equil._nfp + theta = 2 * xp.pi * e2 + zeta = 2 * xp.pi * e3 / gvec_equil._nfp if gvec_equil.params["use_boozer"]: ev = gvec.EvaluationsBoozer(rho=rho, theta_B=theta, zeta_B=zeta, state=gvec_equil.state) else: @@ -286,10 +286,10 @@ def __init__( # get control points def X(eta1, eta2): - return a * eta1 * np.cos(2 * np.pi * eta2) + return a * eta1 * xp.cos(2 * xp.pi * eta2) def Y(eta1, eta2): - return a * eta1 * np.sin(2 * np.pi * eta2) + return a * eta1 * xp.sin(2 * xp.pi * eta2) spl_kind = (False, True) @@ -363,17 +363,17 @@ def __init__( if sfl: def theta(eta1, eta2): - return 2 * np.arctan(np.sqrt((1 + a * eta1 / R0) / (1 - a * eta1 / R0)) * np.tan(np.pi * eta2)) + return 2 * xp.arctan(xp.sqrt((1 + a * eta1 / R0) / (1 - a * eta1 / R0)) * xp.tan(xp.pi * eta2)) else: def theta(eta1, eta2): - return 2 * np.pi * eta2 + return 2 * xp.pi * eta2 def R(eta1, eta2): - return a * eta1 * np.cos(theta(eta1, eta2)) + R0 + return a * eta1 * xp.cos(theta(eta1, eta2)) + R0 def Z(eta1, eta2): - return a * eta1 * np.sin(theta(eta1, eta2)) + return a * eta1 * xp.sin(theta(eta1, eta2)) spl_kind = (False, True) @@ -762,24 +762,24 @@ def __init__( def inverse_map(self, x, y, z, bounded=True, change_out_order=False): """Analytical inverse map of HollowTorus""" - mr = np.sqrt(x**2 + y**2) - self.params["R0"] + mr = xp.sqrt(x**2 + y**2) - self.params["R0"] - eta3 = np.arctan2(-y, x) % (2 * np.pi / self.params["tor_period"]) / (2 * np.pi) * self.params["tor_period"] - eta2 = np.arctan2(z, mr) % (2 * np.pi / self.params["pol_period"]) / (2 * np.pi / self.params["pol_period"]) - eta1 = (z / np.sin(2 * np.pi * eta2 / self.params["pol_period"]) - self.params["a1"]) / ( + eta3 = xp.arctan2(-y, x) % (2 * xp.pi / self.params["tor_period"]) / (2 * xp.pi) * self.params["tor_period"] + eta2 = xp.arctan2(z, mr) % (2 * xp.pi / self.params["pol_period"]) / (2 * xp.pi / self.params["pol_period"]) + eta1 = (z / xp.sin(2 * xp.pi * eta2 / self.params["pol_period"]) - self.params["a1"]) / ( self.params["a2"] - self.params["a1"] ) if bounded: eta1[eta1 > 1] = 1.0 eta1[eta1 < 0] = 0.0 - assert np.all(np.logical_and(eta1 >= 0, eta1 <= 1)) + assert xp.all(xp.logical_and(eta1 >= 0, eta1 <= 1)) - assert np.all(np.logical_and(eta2 >= 0, eta2 <= 1)) - assert np.all(np.logical_and(eta3 >= 0, eta3 <= 1)) + assert xp.all(xp.logical_and(eta2 >= 0, eta2 <= 1)) + assert xp.all(xp.logical_and(eta3 >= 0, eta3 <= 1)) if change_out_order: - return np.transpose((eta1, eta2, eta3)) + return xp.transpose((eta1, eta2, eta3)) else: return eta1, eta2, eta3 diff --git a/src/struphy/geometry/evaluation_kernels.py b/src/struphy/geometry/evaluation_kernels.py index a357bbea1..4f97b9ce9 100644 --- a/src/struphy/geometry/evaluation_kernels.py +++ b/src/struphy/geometry/evaluation_kernels.py @@ -31,7 +31,7 @@ def f( args: DomainArguments Arguments for the mapping. - f_out : np.array + f_out : xp.array Output array of shape (3,). """ @@ -196,7 +196,7 @@ def df( args: DomainArguments Arguments for the mapping. - df_out : np.array + df_out : xp.array Output array of shape (3, 3). """ @@ -354,7 +354,7 @@ def det_df( args: DomainArguments Arguments for the mapping. - tmp1 : np.array + tmp1 : xp.array Temporary array of shape (3, 3). """ @@ -388,13 +388,13 @@ def df_inv( args: DomainArguments Arguments for the mapping. - tmp1: np.array + tmp1: xp.array Temporary array of shape (3, 3). avoid_round_off: bool Whether to manually set exact zeros in arrays. - dfinv_out: np.array + dfinv_out: xp.array Output array of shape (3, 3). """ @@ -484,13 +484,13 @@ def g( args: DomainArguments Arguments for the mapping. - tmp1, tmp2: np.array + tmp1, tmp2: xp.array Temporary arrays of shape (3, 3). avoid_round_off: bool Whether to manually set exact zeros in arrays. - g_out: np.array + g_out: xp.array Output array of shape (3, 3). """ df( @@ -601,13 +601,13 @@ def g_inv( args: DomainArguments Arguments for the mapping. - tmp1, tmp2, tmp3: np.array + tmp1, tmp2, tmp3: xp.array Temporary arrays of shape (3, 3). avoid_round_off: bool Whether to manually set exact zeros in arrays. - ginv_out: np.array + ginv_out: xp.array Output array of shape (3, 3). """ g( @@ -732,16 +732,16 @@ def select_metric_coeff( args: DomainArguments Arguments for the mapping. - tmp0: np.array + tmp0: xp.array Temporary array of shape (3,). - tmp1, tmp2, tmp3: np.array + tmp1, tmp2, tmp3: xp.array Temporary arrays of shape (3, 3). avoid_round_off: bool Whether to manually set exact zeros in arrays. - out: np.array + out: xp.array Output array of shape (3, 3). """ # identity map diff --git a/src/struphy/geometry/tests/test_domain.py b/src/struphy/geometry/tests/test_domain.py index 48348958b..76db82315 100644 --- a/src/struphy/geometry/tests/test_domain.py +++ b/src/struphy/geometry/tests/test_domain.py @@ -5,7 +5,7 @@ def test_prepare_arg(): """Tests prepare_arg static method in domain base class.""" from struphy.geometry.base import Domain - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp def a1(e1, e2, e3): return e1 * e2 @@ -21,12 +21,12 @@ def a_vec(e1, e2, e3): a_2 = e2 * e3 a_3 = e3 * e1 - return np.stack((a_1, a_2, a_3), axis=0) + return xp.stack((a_1, a_2, a_3), axis=0) # ========== tensor-product/slice evaluation =============== - e1 = np.random.rand(4) - e2 = np.random.rand(5) - e3 = np.random.rand(6) + e1 = xp.random.rand(4) + e2 = xp.random.rand(5) + e3 = xp.random.rand(6) E1, E2, E3, is_sparse_meshgrid = Domain.prepare_eval_pts(e1, e2, e3, flat_eval=False) @@ -84,7 +84,7 @@ def a_vec(e1, e2, e3): assert Domain.prepare_arg([A1, A2, A3], E1, E2, E3).shape == shape_vector # ============== markers evaluation ========================== - markers = np.random.rand(10, 6) + markers = xp.random.rand(10, 6) shape_scalar = (markers.shape[0], 1) shape_vector = (markers.shape[0], 3) @@ -160,13 +160,13 @@ def test_evaluation_mappings(mapping): from struphy.geometry import domains from struphy.geometry.base import Domain - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # arrays: - arr1 = np.linspace(0.0, 1.0, 4) - arr2 = np.linspace(0.0, 1.0, 5) - arr3 = np.linspace(0.0, 1.0, 6) - arrm = np.random.rand(10, 8) + arr1 = xp.linspace(0.0, 1.0, 4) + arr2 = xp.linspace(0.0, 1.0, 5) + arr3 = xp.linspace(0.0, 1.0, 6) + arrm = xp.random.rand(10, 8) print() print('Testing "evaluate"...') print("array shapes:", arr1.shape, arr2.shape, arr3.shape, arrm.shape) @@ -262,9 +262,9 @@ def test_evaluation_mappings(mapping): assert domain.metric_inv(arr1, arr2, arr3).shape == (3, 3) + arr1.shape + arr2.shape + arr3.shape # matrix evaluations at one point in third direction - mat12_x, mat12_y = np.meshgrid(arr1, arr2, indexing="ij") - mat13_x, mat13_z = np.meshgrid(arr1, arr3, indexing="ij") - mat23_y, mat23_z = np.meshgrid(arr2, arr3, indexing="ij") + mat12_x, mat12_y = xp.meshgrid(arr1, arr2, indexing="ij") + mat13_x, mat13_z = xp.meshgrid(arr1, arr3, indexing="ij") + mat23_y, mat23_z = xp.meshgrid(arr2, arr3, indexing="ij") # eta1-eta2 matrix evaluation: print("eta1-eta2 matrix evaluation, shape:", domain(mat12_x, mat12_y, 0.5, squeeze_out=True).shape) @@ -294,7 +294,7 @@ def test_evaluation_mappings(mapping): assert domain.metric_inv(0.5, mat23_y, mat23_z, squeeze_out=True).shape == (3, 3) + mat23_y.shape # matrix evaluations for sparse meshgrid - mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing="ij", sparse=True) + mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing="ij", sparse=True) print("sparse meshgrid matrix evaluation, shape:", domain(mat_x, mat_y, mat_z).shape) assert domain(mat_x, mat_y, mat_z).shape == (3,) + (mat_x.shape[0], mat_y.shape[1], mat_z.shape[2]) assert domain.jacobian(mat_x, mat_y, mat_z).shape == (3, 3) + (mat_x.shape[0], mat_y.shape[1], mat_z.shape[2]) @@ -304,7 +304,7 @@ def test_evaluation_mappings(mapping): assert domain.metric_inv(mat_x, mat_y, mat_z).shape == (3, 3) + (mat_x.shape[0], mat_y.shape[1], mat_z.shape[2]) # matrix evaluations - mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing="ij") + mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing="ij") print("matrix evaluation, shape:", domain(mat_x, mat_y, mat_z).shape) assert domain(mat_x, mat_y, mat_z).shape == (3,) + mat_x.shape assert domain.jacobian(mat_x, mat_y, mat_z).shape == (3, 3) + mat_x.shape @@ -319,21 +319,21 @@ def test_pullback(): from struphy.geometry import domains from struphy.geometry.base import Domain - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # arrays: - arr1 = np.linspace(0.0, 1.0, 4) - arr2 = np.linspace(0.0, 1.0, 5) - arr3 = np.linspace(0.0, 1.0, 6) + arr1 = xp.linspace(0.0, 1.0, 4) + arr2 = xp.linspace(0.0, 1.0, 5) + arr3 = xp.linspace(0.0, 1.0, 6) print() print('Testing "pull"...') print("array shapes:", arr1.shape, arr2.shape, arr3.shape) - markers = np.random.rand(13, 6) + markers = xp.random.rand(13, 6) # physical function to pull back (used as components of forms too): def fun(x, y, z): - return np.exp(x) * np.sin(y) * np.cos(z) + return xp.exp(x) * xp.sin(y) * xp.cos(z) domain_class = getattr(domains, "Colella") domain = domain_class() @@ -421,9 +421,9 @@ def fun(x, y, z): ) # matrix pullbacks at one point in third direction - mat12_x, mat12_y = np.meshgrid(arr1, arr2, indexing="ij") - mat13_x, mat13_z = np.meshgrid(arr1, arr3, indexing="ij") - mat23_y, mat23_z = np.meshgrid(arr2, arr3, indexing="ij") + mat12_x, mat12_y = xp.meshgrid(arr1, arr2, indexing="ij") + mat13_x, mat13_z = xp.meshgrid(arr1, arr3, indexing="ij") + mat23_y, mat23_z = xp.meshgrid(arr2, arr3, indexing="ij") # eta1-eta2 matrix pullback: if p_str == "0" or p_str == "3": @@ -450,7 +450,7 @@ def fun(x, y, z): ) # matrix pullbacks for sparse meshgrid - mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing="ij", sparse=True) + mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing="ij", sparse=True) if p_str == "0" or p_str == "3": assert domain.pull(fun_form, mat_x, mat_y, mat_z, kind=p_str).shape == ( mat_x.shape[0], @@ -466,7 +466,7 @@ def fun(x, y, z): ) # matrix pullbacks - mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing="ij") + mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing="ij") if p_str == "0" or p_str == "3": assert domain.pull(fun_form, mat_x, mat_y, mat_z, kind=p_str).shape == mat_x.shape else: @@ -478,21 +478,21 @@ def test_pushforward(): from struphy.geometry import domains from struphy.geometry.base import Domain - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # arrays: - arr1 = np.linspace(0.0, 1.0, 4) - arr2 = np.linspace(0.0, 1.0, 5) - arr3 = np.linspace(0.0, 1.0, 6) + arr1 = xp.linspace(0.0, 1.0, 4) + arr2 = xp.linspace(0.0, 1.0, 5) + arr3 = xp.linspace(0.0, 1.0, 6) print() print('Testing "push"...') print("array shapes:", arr1.shape, arr2.shape, arr3.shape) - markers = np.random.rand(13, 6) + markers = xp.random.rand(13, 6) # logical function to push (used as components of forms too): def fun(e1, e2, e3): - return np.exp(e1) * np.sin(e2) * np.cos(e3) + return xp.exp(e1) * xp.sin(e2) * xp.cos(e3) domain_class = getattr(domains, "Colella") domain = domain_class() @@ -580,9 +580,9 @@ def fun(e1, e2, e3): ) # matrix pushs at one point in third direction - mat12_x, mat12_y = np.meshgrid(arr1, arr2, indexing="ij") - mat13_x, mat13_z = np.meshgrid(arr1, arr3, indexing="ij") - mat23_y, mat23_z = np.meshgrid(arr2, arr3, indexing="ij") + mat12_x, mat12_y = xp.meshgrid(arr1, arr2, indexing="ij") + mat13_x, mat13_z = xp.meshgrid(arr1, arr3, indexing="ij") + mat23_y, mat23_z = xp.meshgrid(arr2, arr3, indexing="ij") # eta1-eta2 matrix push: if p_str == "0" or p_str == "3": @@ -609,7 +609,7 @@ def fun(e1, e2, e3): ) # matrix pushs for sparse meshgrid - mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing="ij", sparse=True) + mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing="ij", sparse=True) if p_str == "0" or p_str == "3": assert domain.push(fun_form, mat_x, mat_y, mat_z, kind=p_str).shape == ( mat_x.shape[0], @@ -625,7 +625,7 @@ def fun(e1, e2, e3): ) # matrix pushs - mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing="ij") + mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing="ij") if p_str == "0" or p_str == "3": assert domain.push(fun_form, mat_x, mat_y, mat_z, kind=p_str).shape == mat_x.shape else: @@ -637,21 +637,21 @@ def test_transform(): from struphy.geometry import domains from struphy.geometry.base import Domain - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # arrays: - arr1 = np.linspace(0.0, 1.0, 4) - arr2 = np.linspace(0.0, 1.0, 5) - arr3 = np.linspace(0.0, 1.0, 6) + arr1 = xp.linspace(0.0, 1.0, 4) + arr2 = xp.linspace(0.0, 1.0, 5) + arr3 = xp.linspace(0.0, 1.0, 6) print() print('Testing "transform"...') print("array shapes:", arr1.shape, arr2.shape, arr3.shape) - markers = np.random.rand(13, 6) + markers = xp.random.rand(13, 6) # logical function to push (used as components of forms too): def fun(e1, e2, e3): - return np.exp(e1) * np.sin(e2) * np.cos(e3) + return xp.exp(e1) * xp.sin(e2) * xp.cos(e3) domain_class = getattr(domains, "Colella") domain = domain_class() @@ -751,9 +751,9 @@ def fun(e1, e2, e3): ) # matrix transforms at one point in third direction - mat12_x, mat12_y = np.meshgrid(arr1, arr2, indexing="ij") - mat13_x, mat13_z = np.meshgrid(arr1, arr3, indexing="ij") - mat23_y, mat23_z = np.meshgrid(arr2, arr3, indexing="ij") + mat12_x, mat12_y = xp.meshgrid(arr1, arr2, indexing="ij") + mat13_x, mat13_z = xp.meshgrid(arr1, arr3, indexing="ij") + mat23_y, mat23_z = xp.meshgrid(arr2, arr3, indexing="ij") # eta1-eta2 matrix transform: if p_str == "0_to_3" or p_str == "3_to_0": @@ -789,7 +789,7 @@ def fun(e1, e2, e3): ) # matrix transforms for sparse meshgrid - mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing="ij", sparse=True) + mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing="ij", sparse=True) if p_str == "0_to_3" or p_str == "3_to_0": assert domain.transform(fun_form, mat_x, mat_y, mat_z, kind=p_str).shape == ( mat_x.shape[0], @@ -805,7 +805,7 @@ def fun(e1, e2, e3): ) # matrix transforms - mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing="ij") + mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing="ij") if p_str == "0_to_3" or p_str == "3_to_0": assert domain.transform(fun_form, mat_x, mat_y, mat_z, kind=p_str).shape == mat_x.shape else: @@ -817,18 +817,18 @@ def fun(e1, e2, e3): # """ # # from struphy.geometry import domains -# from struphy.utils.arrays import xp as np +# from struphy.utils.arrays import xp # # # arrays: -# arr1 = np.linspace(0., 1., 4) -# arr2 = np.linspace(0., 1., 5) -# arr3 = np.linspace(0., 1., 6) +# arr1 = xp.linspace(0., 1., 4) +# arr2 = xp.linspace(0., 1., 5) +# arr3 = xp.linspace(0., 1., 6) # print() # print('Testing "transform"...') # print('array shapes:', arr1.shape, arr2.shape, arr3.shape) # # # logical function to tranform (used as components of forms too): -# fun = lambda eta1, eta2, eta3: np.exp(eta1)*np.sin(eta2)*np.cos(eta3) +# fun = lambda eta1, eta2, eta3: xp.exp(eta1)*xp.sin(eta2)*xp.cos(eta3) # # domain_class = getattr(domains, 'Colella') # domain = domain_class() @@ -885,9 +885,9 @@ def fun(e1, e2, e3): # assert a.shape[0] == arr1.size and a.shape[1] == arr2.size and a.shape[2] == arr3.size # # # matrix transformation at one point in third direction -# mat12_x, mat12_y = np.meshgrid(arr1, arr2, indexing='ij') -# mat13_x, mat13_z = np.meshgrid(arr1, arr3, indexing='ij') -# mat23_y, mat23_z = np.meshgrid(arr2, arr3, indexing='ij') +# mat12_x, mat12_y = xp.meshgrid(arr1, arr2, indexing='ij') +# mat13_x, mat13_z = xp.meshgrid(arr1, arr3, indexing='ij') +# mat23_y, mat23_z = xp.meshgrid(arr2, arr3, indexing='ij') # # # eta1-eta2 matrix transformation: # a = domain.transform(fun_form, mat12_x, mat12_y, .5, p_str) @@ -903,13 +903,13 @@ def fun(e1, e2, e3): # assert a.shape == mat23_y.shape # # # matrix transformation for sparse meshgrid -# mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing='ij', sparse=True) +# mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing='ij', sparse=True) # a = domain.transform(fun_form, mat_x, mat_y, mat_z, p_str) # #print('sparse meshgrid matrix transformation, shape:', a.shape) # assert a.shape[0] == mat_x.shape[0] and a.shape[1] == mat_y.shape[1] and a.shape[2] == mat_z.shape[2] # # # matrix transformation -# mat_x, mat_y, mat_z = np.meshgrid(arr1, arr2, arr3, indexing='ij') +# mat_x, mat_y, mat_z = xp.meshgrid(arr1, arr2, arr3, indexing='ij') # a = domain.transform(fun_form, mat_x, mat_y, mat_z, p_str) # #print('matrix transformation, shape:', a.shape) # assert a.shape == mat_x.shape diff --git a/src/struphy/geometry/utilities.py b/src/struphy/geometry/utilities.py index 218a0467a..0d5c07d93 100644 --- a/src/struphy/geometry/utilities.py +++ b/src/struphy/geometry/utilities.py @@ -17,7 +17,7 @@ from struphy.geometry.utilities_kernels import weighted_arc_lengths_flux_surface from struphy.io.options import GivenInBasis from struphy.linear_algebra.linalg_kron import kron_lusolve_2d -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def field_line_tracing( @@ -129,10 +129,10 @@ def field_line_tracing( Returns ------- - cR : np.ndarray + cR : xp.ndarray Control points (2d) of flux aligned spline mapping (R-component). - cZ : np.ndarray + cZ : xp.ndarray Control points (2d) of flux aligned spline mapping (Z-component). """ @@ -145,8 +145,8 @@ def field_line_tracing( ps, px = p_pre # spline knots - Ts = bsp.make_knots(np.linspace(0.0, 1.0, ns + 1), ps, False) - Tx = bsp.make_knots(np.linspace(0.0, 1.0, nx + 1), px, True) + Ts = bsp.make_knots(xp.linspace(0.0, 1.0, ns + 1), ps, False) + Tx = bsp.make_knots(xp.linspace(0.0, 1.0, nx + 1), px, True) # interpolation (Greville) points s_gr = bsp.greville(Ts, ps, False) @@ -165,13 +165,13 @@ def field_line_tracing( ] # check if pole is included - if np.abs(psi(psi_axis_R, psi_axis_Z) - psi0) < 1e-14: + if xp.abs(psi(psi_axis_R, psi_axis_Z) - psi0) < 1e-14: pole = True else: pole = False - R = np.zeros((s_gr.size, x_gr.size), dtype=float) - Z = np.zeros((s_gr.size, x_gr.size), dtype=float) + R = xp.zeros((s_gr.size, x_gr.size), dtype=float) + Z = xp.zeros((s_gr.size, x_gr.size), dtype=float) # function whose root must be found for j, x in enumerate(x_gr): @@ -188,8 +188,8 @@ def field_line_tracing( # function whose root must be found def f(r): - _R = psi_axis_R + r * np.cos(2 * np.pi * x) - _Z = psi_axis_Z + r * np.sin(2 * np.pi * x) + _R = psi_axis_R + r * xp.cos(2 * xp.pi * x) + _Z = psi_axis_Z + r * xp.sin(2 * xp.pi * x) psi_norm = (psi(_R, _Z) - psi0) / (psi1 - psi0) @@ -200,8 +200,8 @@ def f(r): r_flux_surface = newton(f, x0=r_guess) - R[i, j] = psi_axis_R + r_flux_surface * np.cos(2 * np.pi * x) - Z[i, j] = psi_axis_Z + r_flux_surface * np.sin(2 * np.pi * x) + R[i, j] = psi_axis_R + r_flux_surface * xp.cos(2 * xp.pi * x) + Z[i, j] = psi_axis_Z + r_flux_surface * xp.sin(2 * xp.pi * x) # get control points cR_equal_angle = kron_lusolve_2d(ILUs, R) @@ -227,8 +227,8 @@ def f(r): ps, px = p # spline knots - Ts = bsp.make_knots(np.linspace(0.0, 1.0, ns + 1), ps, False) - Tx = bsp.make_knots(np.linspace(0.0, 1.0, nx + 1), px, True) + Ts = bsp.make_knots(xp.linspace(0.0, 1.0, ns + 1), ps, False) + Tx = bsp.make_knots(xp.linspace(0.0, 1.0, nx + 1), px, True) # interpolation (Greville) points s_gr = bsp.greville(Ts, ps, False) @@ -255,10 +255,10 @@ def f(r): # target function for xi parametrization def f_angles(xis, s_val): - assert np.all(np.logical_and(xis > 0.0, xis < 1.0)) + assert xp.all(xp.logical_and(xis > 0.0, xis < 1.0)) # add 0 and 1 to angles array - xis_extended = np.array([0.0] + list(xis) + [1.0]) + xis_extended = xp.array([0.0] + list(xis) + [1.0]) # compute (R, Z) coordinates for given xis on fixed flux surface corresponding to s_val _RZ = domain_eq_angle(s_val, xis_extended, 0.0) @@ -267,17 +267,17 @@ def f_angles(xis, s_val): _Z = _RZ[2] # |grad(psi)| at xis - gp = np.sqrt(psi(_R, _Z, dR=1) ** 2 + psi(_R, _Z, dZ=1) ** 2) + gp = xp.sqrt(psi(_R, _Z, dR=1) ** 2 + psi(_R, _Z, dZ=1) ** 2) # compute weighted arc_lengths between two successive points in xis_extended array - dl = np.zeros(xis_extended.size - 1, dtype=float) + dl = xp.zeros(xis_extended.size - 1, dtype=float) weighted_arc_lengths_flux_surface(_R, _Z, gp, dl, xi_param_dict[xi_param]) # total length of the flux surface - l = np.sum(dl) + l = xp.sum(dl) # cumulative sum of arc lengths, start with 0! - l_cum = np.cumsum(dl) + l_cum = xp.cumsum(dl) # odd spline degree if px % 2 == 1: @@ -289,8 +289,8 @@ def f_angles(xis, s_val): return xi_diff # loop over flux surfaces and find xi parametrization - R = np.zeros((s_gr.size, x_gr.size), dtype=float) - Z = np.zeros((s_gr.size, x_gr.size), dtype=float) + R = xp.zeros((s_gr.size, x_gr.size), dtype=float) + Z = xp.zeros((s_gr.size, x_gr.size), dtype=float) if px % 2 == 1: xis0 = x_gr[1:].copy() diff --git a/src/struphy/geometry/utilities_kernels.py b/src/struphy/geometry/utilities_kernels.py index 85ccdbcf9..1c26b25c5 100644 --- a/src/struphy/geometry/utilities_kernels.py +++ b/src/struphy/geometry/utilities_kernels.py @@ -18,16 +18,16 @@ def weighted_arc_lengths_flux_surface(r: "float[:]", z: "float[:]", grad_psi: "f Parameters ---------- - r : np.ndarray + r : xp.ndarray R coordinates of the flux surface. - z : np.ndarray + z : xp.ndarray Z coordinates of the flux surface. - grad_psi : np.ndarray + grad_psi : xp.ndarray Absolute values of the flux function gradient on the flux surface: |grad(psi)| = sqrt[ (d_R psi)**2 + (d_Z psi)**2 ]. - dwls : np.ndarray + dwls : xp.ndarray The weighted arc lengths will be written into this array. Length must be one smaller than lengths of r, z and grad_psi. kind : int diff --git a/src/struphy/initial/eigenfunctions.py b/src/struphy/initial/eigenfunctions.py index 172386bdf..093fe21ac 100644 --- a/src/struphy/initial/eigenfunctions.py +++ b/src/struphy/initial/eigenfunctions.py @@ -5,7 +5,7 @@ from sympde.topology import Derham, Line from struphy.fields_background.equils import set_defaults -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class InitialMHDAxisymHdivEigFun: @@ -54,11 +54,11 @@ def __init__(self, derham, **params): spec_path = params["spec_abs"] # load eigenvector for velocity field - omega2, U2_eig = np.split(np.load(spec_path), [1], axis=0) + omega2, U2_eig = xp.split(xp.load(spec_path), [1], axis=0) omega2 = omega2.flatten() # find eigenvector corresponding to given squared eigenfrequency range - mode = np.where((np.real(omega2) < params["eig_freq_upper"]) & (np.real(omega2) > params["eig_freq_lower"]))[0] + mode = xp.where((xp.real(omega2) < params["eig_freq_upper"]) & (xp.real(omega2) > params["eig_freq_lower"]))[0] assert mode.size == 1 mode = mode[0] @@ -89,28 +89,28 @@ def __init__(self, derham, **params): n_tor = int(os.path.split(spec_path)[-1][-6:-4]) - N_cos = p0(lambda phi: np.cos(2 * np.pi * n_tor * phi)).coeffs.toarray() - N_sin = p0(lambda phi: np.sin(2 * np.pi * n_tor * phi)).coeffs.toarray() + N_cos = p0(lambda phi: xp.cos(2 * xp.pi * n_tor * phi)).coeffs.toarray() + N_sin = p0(lambda phi: xp.sin(2 * xp.pi * n_tor * phi)).coeffs.toarray() - D_cos = p1(lambda phi: np.cos(2 * np.pi * n_tor * phi)).coeffs.toarray() - D_sin = p1(lambda phi: np.sin(2 * np.pi * n_tor * phi)).coeffs.toarray() + D_cos = p1(lambda phi: xp.cos(2 * xp.pi * n_tor * phi)).coeffs.toarray() + D_sin = p1(lambda phi: xp.sin(2 * xp.pi * n_tor * phi)).coeffs.toarray() # select real part or imaginary part assert params["kind"] == "r" or params["kind"] == "i" if params["kind"] == "r": - eig_vec_1 = (np.outer(np.real(eig_vec_1), D_cos) - np.outer(np.imag(eig_vec_1), D_sin)).flatten() - eig_vec_2 = (np.outer(np.real(eig_vec_2), D_cos) - np.outer(np.imag(eig_vec_2), D_sin)).flatten() - eig_vec_3 = (np.outer(np.real(eig_vec_3), N_cos) - np.outer(np.imag(eig_vec_3), N_sin)).flatten() + eig_vec_1 = (xp.outer(xp.real(eig_vec_1), D_cos) - xp.outer(xp.imag(eig_vec_1), D_sin)).flatten() + eig_vec_2 = (xp.outer(xp.real(eig_vec_2), D_cos) - xp.outer(xp.imag(eig_vec_2), D_sin)).flatten() + eig_vec_3 = (xp.outer(xp.real(eig_vec_3), N_cos) - xp.outer(xp.imag(eig_vec_3), N_sin)).flatten() else: - eig_vec_1 = (np.outer(np.imag(eig_vec_1), D_cos) + np.outer(np.real(eig_vec_1), D_sin)).flatten() - eig_vec_2 = (np.outer(np.imag(eig_vec_2), D_cos) + np.outer(np.real(eig_vec_2), D_sin)).flatten() - eig_vec_3 = (np.outer(np.imag(eig_vec_3), N_cos) + np.outer(np.real(eig_vec_3), N_sin)).flatten() + eig_vec_1 = (xp.outer(xp.imag(eig_vec_1), D_cos) + xp.outer(xp.real(eig_vec_1), D_sin)).flatten() + eig_vec_2 = (xp.outer(xp.imag(eig_vec_2), D_cos) + xp.outer(xp.real(eig_vec_2), D_sin)).flatten() + eig_vec_3 = (xp.outer(xp.imag(eig_vec_3), N_cos) + xp.outer(xp.real(eig_vec_3), N_sin)).flatten() # set coefficients in full space - eigvec_1_ten = np.zeros(derham.nbasis["2"][0], dtype=float) - eigvec_2_ten = np.zeros(derham.nbasis["2"][1], dtype=float) - eigvec_3_ten = np.zeros(derham.nbasis["2"][2], dtype=float) + eigvec_1_ten = xp.zeros(derham.nbasis["2"][0], dtype=float) + eigvec_2_ten = xp.zeros(derham.nbasis["2"][1], dtype=float) + eigvec_3_ten = xp.zeros(derham.nbasis["2"][2], dtype=float) bc1_1 = derham.dirichlet_bc[0][0] bc1_2 = derham.dirichlet_bc[0][1] @@ -138,19 +138,19 @@ def __init__(self, derham, **params): else: # split into polar/tensor product parts - eig_vec_1 = np.split( + eig_vec_1 = xp.split( eig_vec_1, [ derham.Vh_pol["2"].n_polar[0] * nnz_tor[0], ], ) - eig_vec_2 = np.split( + eig_vec_2 = xp.split( eig_vec_2, [ derham.Vh_pol["2"].n_polar[1] * nnz_tor[1], ], ) - eig_vec_3 = np.split( + eig_vec_3 = xp.split( eig_vec_3, [ derham.Vh_pol["2"].n_polar[2] * nnz_tor[2], diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 1be8d4191..a062c3f61 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -8,7 +8,7 @@ from struphy.initial.base import Perturbation from struphy.io.options import GivenInBasis, NoiseDirections, check_option -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @dataclass @@ -168,7 +168,7 @@ def __init__( self._pfuns += [lambda eta3: 1.0] elif pfun == "localize": self._pfuns += [ - lambda eta3: np.tanh((eta3 - 0.5) / params) / np.cosh((eta3 - 0.5) / params), + lambda eta3: xp.tanh((eta3 - 0.5) / params) / xp.cosh((eta3 - 0.5) / params), ] else: raise ValueError(f"Profile function {pfun} is not defined..") @@ -193,10 +193,10 @@ def __call__(self, x, y, z): val += ( amp * pfun(z) - * np.sin( - l * 2.0 * np.pi / self._Lx * x - + m * 2.0 * np.pi / self._Ly * y - + n * 2.0 * np.pi / self._Lz * z + * xp.sin( + l * 2.0 * xp.pi / self._Lx * x + + m * 2.0 * xp.pi / self._Ly * y + + n * 2.0 * xp.pi / self._Lz * z + t, ) ) @@ -296,8 +296,8 @@ def __call__(self, x, y, z): val = 0.0 for amp, l, m, n in zip(self._amps, self._ls, self._ms, self._ns): - val += amp * np.cos( - l * 2.0 * np.pi / self._Lx * x + m * 2.0 * np.pi / self._Ly * y + n * 2.0 * np.pi / self._Lz * z, + val += amp * xp.cos( + l * 2.0 * xp.pi / self._Lx * x + m * 2.0 * xp.pi / self._Ly * y + n * 2.0 * xp.pi / self._Lz * z, ) # print( "Cos max value", val.max()) return val @@ -333,12 +333,12 @@ def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28): def __call__(self, eta1, eta2, eta3): val = 0.0 r = eta1 * (self._r2 - self._r1) + self._r1 - theta = eta2 * 2.0 * np.pi + theta = eta2 * 2.0 * xp.pi val += ( -self._m / r - * np.cos(self._m * theta) + * xp.cos(self._m * theta) * (self._a * scipy.special.jv(self._m, r) + self._b * scipy.special.yn(self._m, r)) ) return val @@ -375,12 +375,12 @@ def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28): def __call__(self, eta1, eta2, eta3): val = 0.0 r = eta1 * (self._r2 - self._r1) + self._r1 - theta = eta2 * 2.0 * np.pi + theta = eta2 * 2.0 * xp.pi val += ( self._a * ((self._m / r) * scipy.special.jv(self._m, r) - scipy.special.jv(self._m + 1, r)) + (self._b * ((self._m / r) * scipy.special.yn(self._m, r) - scipy.special.yn(self._m + 1, r))) - ) * np.sin(self._m * theta) + ) * xp.sin(self._m * theta) return val @@ -414,10 +414,10 @@ def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28): def __call__(self, eta1, eta2, eta3): val = 0.0 r = eta1 * (self._r2 - self._r1) + self._r1 - theta = eta2 * 2.0 * np.pi + theta = eta2 * 2.0 * xp.pi z = eta3 - val += (self._a * scipy.special.jv(self._m, r) + self._b * scipy.special.yn(self._m, r)) * np.cos( + val += (self._a * scipy.special.jv(self._m, r) + self._b * scipy.special.yn(self._m, r)) * xp.cos( self._m * theta ) return val @@ -509,7 +509,7 @@ def __init__( if pfun == "Id": self._pfuns += [lambda z: 1.0] elif pfun == "localize": - self._pfuns += [lambda z, p=params: np.tanh((z - 0.5) / p) / np.cosh((z - 0.5) / p)] + self._pfuns += [lambda z, p=params: xp.tanh((z - 0.5) / p) / xp.cosh((z - 0.5) / p)] else: raise ValueError(f"Profile function {pfun} is not defined..") @@ -523,8 +523,8 @@ def __call__(self, x, y, z): val += ( amp * pfun(z) - * np.cos(l * 2.0 * np.pi / self._Lx * x + thx) - * np.cos(m * 2.0 * np.pi / self._Ly * y + thy) + * xp.cos(l * 2.0 * xp.pi / self._Lx * x + thx) + * xp.cos(m * 2.0 * xp.pi / self._Ly * y + thy) ) return val @@ -614,7 +614,7 @@ def __init__( if pfun == "Id": self._pfuns += [lambda z: 1.0] elif pfun == "localize": - self._pfuns += [lambda z, p=params: np.tanh((z - 0.5) / p) / np.cosh((z - 0.5) / p)] + self._pfuns += [lambda z, p=params: xp.tanh((z - 0.5) / p) / xp.cosh((z - 0.5) / p)] else: raise ValueError(f"Profile function {pfun} is not defined..") @@ -628,8 +628,8 @@ def __call__(self, x, y, z): val += ( amp * pfun(z) - * np.sin(l * 2.0 * np.pi / self._Lx * x + thx) - * np.sin(m * 2.0 * np.pi / self._Ly * y + thy) + * xp.sin(l * 2.0 * xp.pi / self._Lx * x + thx) + * xp.sin(m * 2.0 * xp.pi / self._Ly * y + thy) ) return val @@ -721,7 +721,7 @@ def __init__( if pfun == "Id": self._pfuns += [lambda z: 1.0] elif pfun == "localize": - self._pfuns += [lambda z, p=params: np.tanh((z - 0.5) / p) / np.cosh((z - 0.5) / p)] + self._pfuns += [lambda z, p=params: xp.tanh((z - 0.5) / p) / xp.cosh((z - 0.5) / p)] else: raise ValueError(f"Profile function {pfun} is not defined..") @@ -735,8 +735,8 @@ def __call__(self, x, y, z): val += ( amp * pfun(z) - * np.sin(l * 2.0 * np.pi / self._Lx * x + thx) - * np.cos(m * 2.0 * np.pi / self._Ly * y + thy) + * xp.sin(l * 2.0 * xp.pi / self._Lx * x + thx) + * xp.cos(m * 2.0 * xp.pi / self._Ly * y + thy) ) return val @@ -828,7 +828,7 @@ def __init__( if pfun == "Id": self._pfuns += [lambda z: 1.0] elif pfun == "localize": - self._pfuns += [lambda z, p=params: np.tanh((z - 0.5) / p) / np.cosh((z - 0.5) / p)] + self._pfuns += [lambda z, p=params: xp.tanh((z - 0.5) / p) / xp.cosh((z - 0.5) / p)] else: raise ValueError(f"Profile function {pfun} is not defined..") @@ -842,8 +842,8 @@ def __call__(self, x, y, z): val += ( amp * pfun(z) - * np.cos(l * 2.0 * np.pi / self._Lx * x + thx) - * np.sin(m * 2.0 * np.pi / self._Ly * y + thy) + * xp.cos(l * 2.0 * xp.pi / self._Lx * x + thx) + * xp.sin(m * 2.0 * xp.pi / self._Ly * y + thy) ) return val @@ -946,18 +946,18 @@ def __init__( ls = 1 else: ls = params - self._pfuns += [lambda eta1: np.sin(ls * np.pi * eta1)] + self._pfuns += [lambda eta1: xp.sin(ls * xp.pi * eta1)] elif pfun == "exp": self._pfuns += [ - lambda eta1: np.exp(-((eta1 - params[0]) ** 2) / (2 * params[1] ** 2)) - / np.sqrt(2 * np.pi * params[1] ** 2), + lambda eta1: xp.exp(-((eta1 - params[0]) ** 2) / (2 * params[1] ** 2)) + / xp.sqrt(2 * xp.pi * params[1] ** 2), ] elif pfun == "d_exp": self._pfuns += [ lambda eta1: -(eta1 - params[0]) / params[1] ** 2 - * np.exp(-((eta1 - params[0]) ** 2) / (2 * params[1] ** 2)) - / np.sqrt(2 * np.pi * params[1] ** 2), + * xp.exp(-((eta1 - params[0]) ** 2) / (2 * params[1] ** 2)) + / xp.sqrt(2 * xp.pi * params[1] ** 2), ] else: raise ValueError(f"Profile function {pfun} is not defined..") @@ -972,8 +972,8 @@ def __call__(self, eta1, eta2, eta3): val += ( amp * pfun(eta1) - * np.sin( - mi * 2.0 * np.pi * eta2 + ni * 2.0 * np.pi * eta3, + * xp.sin( + mi * 2.0 * xp.pi * eta2 + ni * 2.0 * xp.pi * eta3, ) ) @@ -1078,20 +1078,20 @@ def __init__( ls = 1 else: ls = params - self._pfuns += [lambda eta1: np.sin(ls * np.pi * eta1)] + self._pfuns += [lambda eta1: xp.sin(ls * xp.pi * eta1)] elif pfun == "cos": - self._pfuns += [lambda eta1: np.cos(np.pi * eta1)] + self._pfuns += [lambda eta1: xp.cos(xp.pi * eta1)] elif pfun == "exp": self._pfuns += [ - lambda eta1: np.exp(-((eta1 - params[0]) ** 2) / (2 * params[1] ** 2)) - / np.sqrt(2 * np.pi * params[1] ** 2), + lambda eta1: xp.exp(-((eta1 - params[0]) ** 2) / (2 * params[1] ** 2)) + / xp.sqrt(2 * xp.pi * params[1] ** 2), ] elif pfun == "d_exp": self._pfuns += [ lambda eta1: -(eta1 - params[0]) / params[1] ** 2 - * np.exp(-((eta1 - params[0]) ** 2) / (2 * params[1] ** 2)) - / np.sqrt(2 * np.pi * params[1] ** 2), + * xp.exp(-((eta1 - params[0]) ** 2) / (2 * params[1] ** 2)) + / xp.sqrt(2 * xp.pi * params[1] ** 2), ] else: raise ValueError( @@ -1108,8 +1108,8 @@ def __call__(self, eta1, eta2, eta3): val += ( amp * pfun(eta1) - * np.cos( - mi * 2.0 * np.pi * eta2 + ni * 2.0 * np.pi * eta3, + * xp.cos( + mi * 2.0 * xp.pi * eta2 + ni * 2.0 * xp.pi * eta3, ) ) @@ -1157,7 +1157,7 @@ def __init__( self.comp = comp def __call__(self, e1, e2, e3): - val = self._amp * (-np.tanh((e1 - 0.75) / self._delta) + np.tanh((e1 - 0.25) / self._delta) - 1) + val = self._amp * (-xp.tanh((e1 - 0.75) / self._delta) + xp.tanh((e1 - 0.25) / self._delta) - 1) return val @@ -1203,7 +1203,7 @@ def __init__( self.comp = comp def __call__(self, e1, e2, e3): - val = self._amp * (-np.tanh((e2 - 0.75) / self._delta) + np.tanh((e2 - 0.25) / self._delta) - 1) + val = self._amp * (-xp.tanh((e2 - 0.75) / self._delta) + xp.tanh((e2 - 0.25) / self._delta) - 1) return val @@ -1249,7 +1249,7 @@ def __init__( self.comp = comp def __call__(self, e1, e2, e3): - val = self._amp * (-np.tanh((e3 - 0.75) / self._delta) + np.tanh((e3 - 0.25) / self._delta) - 1) + val = self._amp * (-xp.tanh((e3 - 0.75) / self._delta) + xp.tanh((e3 - 0.25) / self._delta) - 1) return val @@ -1376,9 +1376,9 @@ def __init__( # equilibrium ion velocity def __call__(self, x, y, z): """Velocity of ions and electrons.""" - R = np.sqrt(x**2 + y**2) - R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(-y, x) + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) ustarR = ( self._alpha * R / (self._a * self._R0) * (-z) + self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * z @@ -1396,10 +1396,10 @@ def __call__(self, x, y, z): # from cylindrical to cartesian: if self.comp == 0: - ux = np.cos(phi) * uR - R * np.sin(phi) * uphi + ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi return ux elif self.comp == 1: - uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi + uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi return uy elif self.comp == 2: uz = uZ @@ -1482,9 +1482,9 @@ def __init__( # equilibrium ion velocity def __call__(self, x, y, z): """Velocity of ions and electrons.""" - R = np.sqrt(x**2 + y**2) - R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(-y, x) + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) ustarR = ( self._alpha * R / (self._a * self._R0) * (-z) + self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * z @@ -1502,10 +1502,10 @@ def __call__(self, x, y, z): # from cylindrical to cartesian: if self.comp == 0: - ux = np.cos(phi) * uR - R * np.sin(phi) * uphi + ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi return ux elif self.comp == 1: - uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi + uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi return uy elif self.comp == 2: uz = uZ @@ -1588,9 +1588,9 @@ def __init__( # equilibrium ion velocity def __call__(self, x, y, z): """Velocity of ions and electrons.""" - R = np.sqrt(x**2 + y**2) - R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(-y, x) + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) ustarR = ( self._alpha * R / (self._a * self._R0) * (-z) + self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * z @@ -1608,10 +1608,10 @@ def __call__(self, x, y, z): # from cylindrical to cartesian: if self.comp == 0: - ux = np.cos(phi) * uR - R * np.sin(phi) * uphi + ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi return ux elif self.comp == 1: - uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi + uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi return uy elif self.comp == 2: uz = uZ @@ -1682,7 +1682,7 @@ def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0): # equilibrium potential def __call__(self, x, y, z): """Equilibrium potential.""" - R = np.sqrt(x**2 + y**2) + R = xp.sqrt(x**2 + y**2) pp = 0.5 * self._a * self._B0 * self._alpha * (((R - self._R0) ** 2 + z**2) / self._a**2 - 2.0 / 3.0) return pp @@ -1745,15 +1745,15 @@ def __call__(self, x, y, z): """Velocity of ions.""" """x component""" if self._dimension == "2D": - ux = -np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y) + ux = -xp.sin(2 * xp.pi * x) * xp.sin(2 * xp.pi * y) elif self._dimension == "1D": - ux = np.sin(2 * np.pi * x) + 1.0 + ux = xp.sin(2 * xp.pi * x) + 1.0 """y component""" if self._dimension == "2D": - uy = -np.cos(2 * np.pi * x) * np.cos(2 * np.pi * y) + uy = -xp.cos(2 * xp.pi * x) * xp.cos(2 * xp.pi * y) elif self._dimension == "1D": - uy = np.cos(2 * np.pi * x) + uy = xp.cos(2 * xp.pi * x) """z component""" uz = 0.0 * x @@ -1771,15 +1771,15 @@ def __call__(self, x, y, z): """Velocity of electrons.""" """x component""" if self._dimension == "2D": - ux = -np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y) + ux = -xp.sin(4 * xp.pi * x) * xp.sin(4 * xp.pi * y) elif self._dimension == "1D": - ux = np.sin(2.0 * np.pi * x) + ux = xp.sin(2.0 * xp.pi * x) """y component""" if self._dimension == "2D": - uy = -np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y) + uy = -xp.cos(4 * xp.pi * x) * xp.cos(4 * xp.pi * y) elif self._dimension == "1D": - uy = np.cos(2 * np.pi * x) + uy = xp.cos(2 * xp.pi * x) """z component""" uz = 0.0 * x @@ -1844,9 +1844,9 @@ def __init__(self, dimension="1D", b0=1.0): def __call__(self, x, y, z): """Potential.""" if self._dimension == "2D": - phi = np.cos(2 * np.pi * x) + np.sin(2 * np.pi * y) + phi = xp.cos(2 * xp.pi * x) + xp.sin(2 * xp.pi * y) elif self._dimension == "1D": - phi = np.sin(2.0 * np.pi * x) + phi = xp.sin(2.0 * xp.pi * x) return phi @@ -1906,15 +1906,15 @@ def __call__(self, x, y, z): """Velocity of ions.""" """x component""" if self._dimension == "2D": - ux = -np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y) + ux = -xp.sin(2 * xp.pi * x) * xp.sin(2 * xp.pi * y) elif self._dimension == "1D": - ux = np.sin(2 * np.pi * x) + 1.0 + ux = xp.sin(2 * xp.pi * x) + 1.0 """y component""" if self._dimension == "2D": - uy = -np.cos(2 * np.pi * x) * np.cos(2 * np.pi * y) + uy = -xp.cos(2 * xp.pi * x) * xp.cos(2 * xp.pi * y) elif self._dimension == "1D": - uy = np.cos(2 * np.pi * x) + uy = xp.cos(2 * xp.pi * x) """z component""" uz = 0.0 * x @@ -1932,15 +1932,15 @@ def __call__(self, x, y, z): """Velocity of electrons.""" """x component""" if self._dimension == "2D": - ux = -np.sin(4 * np.pi * x) * np.sin(4 * np.pi * y) + ux = -xp.sin(4 * xp.pi * x) * xp.sin(4 * xp.pi * y) elif self._dimension == "1D": - ux = np.sin(2.0 * np.pi * x) + ux = xp.sin(2.0 * xp.pi * x) """y component""" if self._dimension == "2D": - uy = -np.cos(4 * np.pi * x) * np.cos(4 * np.pi * y) + uy = -xp.cos(4 * xp.pi * x) * xp.cos(4 * xp.pi * y) elif self._dimension == "1D": - uy = np.cos(2 * np.pi * x) + uy = xp.cos(2 * xp.pi * x) """z component""" uz = 0.0 * x @@ -2001,8 +2001,8 @@ def __call__(self, eta1, eta2=None, eta3=None): val = ( self._n0 * self._c[3] - * np.exp( - -self._c[2] / self._c[1] * np.tanh((eta1 - self._c[0]) / self._c[2]), + * xp.exp( + -self._c[2] / self._c[1] * xp.tanh((eta1 - self._c[0]) / self._c[2]), ) ) @@ -2088,9 +2088,9 @@ def __init__( # equilibrium ion velocity def __call__(self, x, y, z): """Velocity of ions and electrons.""" - R = np.sqrt(x**2 + y**2) - R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(-y, x) + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) A = self._alpha / (self._a * self._R0) C = self._beta * self._Bp * self._R0 / (self._B0 * self._a) @@ -2101,10 +2101,10 @@ def __call__(self, x, y, z): # from cylindrical to cartesian: if self.comp == 0: - ux = np.cos(phi) * uR - R * np.sin(phi) * uphi + ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi return ux elif self.comp == 1: - uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi + uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi return uy elif self.comp == 2: uz = uZ @@ -2192,9 +2192,9 @@ def __init__( # equilibrium ion velocity def __call__(self, x, y, z): """Velocity of ions and electrons.""" - R = np.sqrt(x**2 + y**2) - R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(-y, x) + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) A = self._alpha / (self._a * self._R0) C = self._beta * self._Bp * self._R0 / (self._B0 * self._a) @@ -2205,10 +2205,10 @@ def __call__(self, x, y, z): # from cylindrical to cartesian: if self.comp == 0: - ux = np.cos(phi) * uR - R * np.sin(phi) * uphi + ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi return ux elif self.comp == 1: - uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi + uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi return uy elif self.comp == 2: uz = uZ @@ -2296,9 +2296,9 @@ def __init__( # equilibrium ion velocity def __call__(self, x, y, z): """Velocity of ions and electrons.""" - R = np.sqrt(x**2 + y**2) - R = np.where(R == 0.0, 1e-9, R) - phi = np.arctan2(-y, x) + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) A = self._alpha / (self._a * self._R0) C = self._beta * self._Bp * self._R0 / (self._B0 * self._a) @@ -2309,10 +2309,10 @@ def __call__(self, x, y, z): # from cylindrical to cartesian: if self.comp == 0: - ux = np.cos(phi) * uR - R * np.sin(phi) * uphi + ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi return ux elif self.comp == 1: - uy = -np.sin(phi) * uR - R * np.cos(phi) * uphi + uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi return uy elif self.comp == 2: uz = uZ diff --git a/src/struphy/initial/tests/test_init_perturbations.py b/src/struphy/initial/tests/test_init_perturbations.py index d2a39debd..ee71c22c8 100644 --- a/src/struphy/initial/tests/test_init_perturbations.py +++ b/src/struphy/initial/tests/test_init_perturbations.py @@ -29,7 +29,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False from struphy.initial import perturbations from struphy.initial.base import Perturbation from struphy.models.variables import FEECVariable - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -50,10 +50,10 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False form_vector = ["1", "2", "v", "norm", "physical_at_eta"] # evaluation points - e1 = np.linspace(0.0, 1.0, 30) - e2 = np.linspace(0.0, 1.0, 40) - e3 = np.linspace(0.0, 1.0, 50) - eee1, eee2, eee3 = np.meshgrid(e1, e2, e3, indexing="ij") + e1 = xp.linspace(0.0, 1.0, 30) + e2 = xp.linspace(0.0, 1.0, 40) + e3 = xp.linspace(0.0, 1.0, 50) + eee1, eee2, eee3 = xp.meshgrid(e1, e2, e3, indexing="ij") # mode paramters kwargs = {} @@ -130,7 +130,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False field_vals_xyz = domain.push(field, e1, e2, e3, kind=form) x, y, z = domain(e1, e2, e3) - r = np.sqrt(x**2 + y**2) + r = xp.sqrt(x**2 + y**2) if fun_form == "physical": fun_vals_xyz = perturbation_xyz(x, y, z) @@ -139,7 +139,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False else: fun_vals_xyz = domain.push(perturbation, eee1, eee2, eee3, kind=fun_form) - error = np.max(np.abs(field_vals_xyz - fun_vals_xyz)) / np.max(np.abs(fun_vals_xyz)) + error = xp.max(xp.abs(field_vals_xyz - fun_vals_xyz)) / xp.max(xp.abs(fun_vals_xyz)) print(f"{rank=}, {key=}, {form=}, {fun_form=}, {error=}") assert error < 0.02 @@ -254,7 +254,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False f_xyz = [f1_xyz, f2_xyz, f3_xyz] x, y, z = domain(e1, e2, e3) - r = np.sqrt(x**2 + y**2) + r = xp.sqrt(x**2 + y**2) # exact values if fun_form == "physical": @@ -279,7 +279,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False error = 0.0 for fi, funi in zip(f_xyz, fun_xyz_vec): - error += np.max(np.abs(fi - funi)) / np.max(np.abs(funi)) + error += xp.max(xp.abs(fi - funi)) / xp.max(xp.abs(funi)) error /= 3.0 print(f"{rank=}, {key=}, {form=}, {fun_form=}, {error=}") assert error < 0.02 diff --git a/src/struphy/initial/utilities.py b/src/struphy/initial/utilities.py index da5edd45a..274b73ba6 100644 --- a/src/struphy/initial/utilities.py +++ b/src/struphy/initial/utilities.py @@ -4,7 +4,7 @@ from struphy.fields_background.equils import set_defaults from struphy.io.output_handling import DataContainer -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class InitFromOutput: @@ -98,6 +98,6 @@ def __init__( self._amp = amp def __call__(self, x, y, z): - val = self._amp * np.random.rand(*x.shape).squeeze() + val = self._amp * xp.random.rand(*x.shape).squeeze() return val diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 3a9cd8e61..fff19cb48 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -5,7 +5,7 @@ from psydac.ddm.mpi import mpi as MPI from struphy.physics.physics import ConstantsOfNature -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp ## Literal options @@ -189,7 +189,7 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk elif velocity_scale == "alfvén": assert A_bulk is not None, 'Need bulk species to choose velocity scale "alfvén".' - self._v = self.B / np.sqrt(self.n * A_bulk * con.mH * con.mu0) + self._v = self.B / xp.sqrt(self.n * A_bulk * con.mH * con.mu0) elif velocity_scale == "cyclotron": assert Z_bulk is not None, 'Need bulk species to choose velocity scale "cyclotron".' @@ -199,7 +199,7 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk elif velocity_scale == "thermal": assert A_bulk is not None, 'Need bulk species to choose velocity scale "thermal".' assert self.kBT is not None - self._v = np.sqrt(self.kBT * 1000 * con.e / (con.mH * A_bulk)) + self._v = xp.sqrt(self.kBT * 1000 * con.e / (con.mH * A_bulk)) # time (s) self._t = self.x / self.v diff --git a/src/struphy/io/output_handling.py b/src/struphy/io/output_handling.py index 79e23692b..ed9c563d4 100644 --- a/src/struphy/io/output_handling.py +++ b/src/struphy/io/output_handling.py @@ -3,7 +3,7 @@ import h5py -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class DataContainer: @@ -83,11 +83,11 @@ def add_data(self, data_dict): Parameters ---------- data_dict : dict - Name-object pairs to save during time stepping, e.g. {key : val}. key must be a string and val must be a np.array of fixed shape. Scalar values (floats) must therefore be passed as 1d arrays of size 1. + Name-object pairs to save during time stepping, e.g. {key : val}. key must be a string and val must be a xp.array of fixed shape. Scalar values (floats) must therefore be passed as 1d arrays of size 1. """ for key, val in data_dict.items(): - assert isinstance(val, np.ndarray) + assert isinstance(val, xp.ndarray) # if dataset already exists, check for compatibility with given array if key in self._dset_dict: diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index 375a76b4a..58bcdccf0 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -10,7 +10,7 @@ from struphy.geometry.base import Domain from struphy.io.options import DerhamOptions from struphy.topology.grids import TensorProductGrid -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def import_parameters_py(params_path: str) -> ModuleType: diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 9ffe6526e..5da0a1568 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -5,7 +5,7 @@ from struphy.fields_background.base import FluidEquilibriumWithB from struphy.initial.base import Perturbation -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class KineticBackground(metaclass=ABCMeta): @@ -103,7 +103,7 @@ def __call__(self, *args): Returns ------- - f0 : np.ndarray + f0 : xp.ndarray The evaluated background. """ pass @@ -118,12 +118,12 @@ def __rmul__(self, a): return ScalarMultiplyKineticBackground(self, a) def __div__(self, a): - assert isinstance(a, float) or isinstance(a, int) or isinstance(a, np.int64) + assert isinstance(a, float) or isinstance(a, int) or isinstance(a, xp.int64) assert a != 0, "Cannot divide by zero!" return ScalarMultiplyKineticBackground(self, 1 / a) def __rdiv__(self, a): - assert isinstance(a, float) or isinstance(a, int) or isinstance(a, np.int64) + assert isinstance(a, float) or isinstance(a, int) or isinstance(a, xp.int64) assert a != 0, "Cannot divide by zero!" return ScalarMultiplyKineticBackground(self, 1 / a) @@ -231,7 +231,7 @@ def __call__(self, *args): Returns ------- - f0 : np.ndarray + f0 : xp.ndarray The evaluated background. """ return self._f1(*args) + self._f2(*args) @@ -240,7 +240,7 @@ def __call__(self, *args): class ScalarMultiplyKineticBackground(KineticBackground): def __init__(self, f0, a): assert isinstance(f0, KineticBackground) - assert isinstance(a, float) or isinstance(a, int) or isinstance(a, np.int64) + assert isinstance(a, float) or isinstance(a, int) or isinstance(a, xp.int64) self._f = f0 self._a = a @@ -317,7 +317,7 @@ def __call__(self, *args): Returns ------- - f0 : np.ndarray + f0 : xp.ndarray The evaluated background. """ return self._a * self._f(*args) @@ -393,14 +393,14 @@ def gaussian(self, v, u=0.0, vth=1.0, polar=False, volume_form=False): An array of size(v). """ - if isinstance(v, np.ndarray) and isinstance(u, np.ndarray): + if isinstance(v, xp.ndarray) and isinstance(u, xp.ndarray): assert v.shape == u.shape, f"{v.shape = } but {u.shape = }" if not polar: - out = 1.0 / vth * 1.0 / np.sqrt(2.0 * np.pi) * np.exp(-((v - u) ** 2) / (2.0 * vth**2)) + out = 1.0 / vth * 1.0 / xp.sqrt(2.0 * xp.pi) * xp.exp(-((v - u) ** 2) / (2.0 * vth**2)) else: - assert np.all(v >= 0.0) - out = 1.0 / vth**2 * np.exp(-((v - u) ** 2) / (2.0 * vth**2)) + assert xp.all(v >= 0.0) + out = 1.0 / vth**2 * xp.exp(-((v - u) ** 2) / (2.0 * vth**2)) if volume_form: out *= v @@ -426,16 +426,16 @@ def __call__(self, *args): Returns ------- - f : np.ndarray + f : xp.ndarray The evaluated Maxwellian. """ # Check that all args have the same shape - shape0 = np.shape(args[0]) + shape0 = xp.shape(args[0]) for i, arg in enumerate(args): - assert np.shape(arg) == shape0, f"Argument {i} has {np.shape(arg) = }, but must be {shape0 = }." - assert np.ndim(arg) == 1 or np.ndim(arg) == 3 + self.vdim, ( - f"{np.ndim(arg) = } not allowed for Maxwellian evaluation." + assert xp.shape(arg) == shape0, f"Argument {i} has {xp.shape(arg) = }, but must be {shape0 = }." + assert xp.ndim(arg) == 1 or xp.ndim(arg) == 3 + self.vdim, ( + f"{xp.ndim(arg) = } not allowed for Maxwellian evaluation." ) # flat or meshgrid evaluation # Get result evaluated at eta's @@ -444,33 +444,33 @@ def __call__(self, *args): vths = self.vth(*args[: -self.vdim]) # take care of correct broadcasting, assuming args come from phase space meshgrid - if np.ndim(args[0]) > 3: + if xp.ndim(args[0]) > 3: # move eta axes to the back - arg_t = np.moveaxis(args[0], 0, -1) - arg_t = np.moveaxis(arg_t, 0, -1) - arg_t = np.moveaxis(arg_t, 0, -1) + arg_t = xp.moveaxis(args[0], 0, -1) + arg_t = xp.moveaxis(arg_t, 0, -1) + arg_t = xp.moveaxis(arg_t, 0, -1) # broadcast res_broad = res + 0.0 * arg_t # move eta axes to the front - res = np.moveaxis(res_broad, -1, 0) - res = np.moveaxis(res, -1, 0) - res = np.moveaxis(res, -1, 0) + res = xp.moveaxis(res_broad, -1, 0) + res = xp.moveaxis(res, -1, 0) + res = xp.moveaxis(res, -1, 0) # Multiply result with gaussian in v's for i, v in enumerate(args[-self.vdim :]): # correct broadcasting - if np.ndim(args[0]) > 3: + if xp.ndim(args[0]) > 3: u_broad = us[i] + 0.0 * arg_t - u = np.moveaxis(u_broad, -1, 0) - u = np.moveaxis(u, -1, 0) - u = np.moveaxis(u, -1, 0) + u = xp.moveaxis(u_broad, -1, 0) + u = xp.moveaxis(u, -1, 0) + u = xp.moveaxis(u, -1, 0) vth_broad = vths[i] + 0.0 * arg_t - vth = np.moveaxis(vth_broad, -1, 0) - vth = np.moveaxis(vth, -1, 0) - vth = np.moveaxis(vth, -1, 0) + vth = xp.moveaxis(vth_broad, -1, 0) + vth = xp.moveaxis(vth, -1, 0) + vth = xp.moveaxis(vth, -1, 0) else: u = us[i] vth = vths[i] @@ -499,9 +499,9 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name: str = "n", add_perturbatio """ # collect arguments - assert isinstance(eta1, np.ndarray) - assert isinstance(eta2, np.ndarray) - assert isinstance(eta3, np.ndarray) + assert isinstance(eta1, xp.ndarray) + assert isinstance(eta2, xp.ndarray) + assert isinstance(eta3, xp.ndarray) assert eta1.shape == eta2.shape == eta3.shape params = self.maxw_params[name] @@ -511,7 +511,7 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name: str = "n", add_perturbatio # flat evaluation for markers if eta1.ndim == 1: etas = [ - np.concatenate( + xp.concatenate( (eta1[:, None], eta2[:, None], eta3[:, None]), axis=1, ), diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index 1f6c16b22..ea541eb70 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -6,7 +6,7 @@ from struphy.fields_background.equils import set_defaults from struphy.initial.base import Perturbation from struphy.kinetic_background.base import Maxwellian -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Maxwellian3D(Maxwellian): @@ -250,7 +250,7 @@ def velocity_jacobian_det(self, eta1, eta2, eta3, *v): assert len(v) == 2 # call equilibrium - etas = (np.vstack((eta1, eta2, eta3)).T).copy() + etas = (xp.vstack((eta1, eta2, eta3)).T).copy() absB0 = self.equil.absB0(etas) # J = v_perp/B @@ -404,14 +404,14 @@ def velocity_jacobian_det(self, eta1, eta2, eta3, energy): assert eta3.ndim == 1 if self.maxw_params["type"] == "Particles6D": - return np.sqrt(2.0 * energy) * 4.0 * np.pi + return xp.sqrt(2.0 * energy) * 4.0 * xp.pi else: # call equilibrium - etas = (np.vstack((eta1, eta2, eta3)).T).copy() + etas = (xp.vstack((eta1, eta2, eta3)).T).copy() absB0 = self.equil.absB0(etas) - return np.sqrt(energy) * 2.0 * np.sqrt(2.0) / absB0 + return xp.sqrt(energy) * 2.0 * xp.sqrt(2.0) / absB0 def gaussian(self, e, vth=1.0): """3-dim. normal distribution, to which array-valued thermal velocities can be passed. @@ -429,10 +429,10 @@ def gaussian(self, e, vth=1.0): An array of size(e). """ - if isinstance(vth, np.ndarray): + if isinstance(vth, xp.ndarray): assert e.shape == vth.shape, f"{e.shape = } but {vth.shape = }" - return 2.0 * np.sqrt(e / np.pi) / vth**3 * np.exp(-e / vth**2) + return 2.0 * xp.sqrt(e / xp.pi) / vth**3 * xp.exp(-e / vth**2) def __call__(self, *args): """Evaluates the canonical Maxwellian distribution function. @@ -454,16 +454,16 @@ def __call__(self, *args): Returns ------- - f : np.ndarray + f : xp.ndarray The evaluated Maxwellian. """ # Check that all args have the same shape - shape0 = np.shape(args[0]) + shape0 = xp.shape(args[0]) for i, arg in enumerate(args): - assert np.shape(arg) == shape0, f"Argument {i} has {np.shape(arg) = }, but must be {shape0 = }." - assert np.ndim(arg) == 1 or np.ndim(arg) == 3, ( - f"{np.ndim(arg) = } not allowed for canonical Maxwellian evaluation." + assert xp.shape(arg) == shape0, f"Argument {i} has {xp.shape(arg) = }, but must be {shape0 = }." + assert xp.ndim(arg) == 1 or xp.ndim(arg) == 3, ( + f"{xp.ndim(arg) = } not allowed for canonical Maxwellian evaluation." ) # flat or meshgrid evaluation # Get result evaluated with each particles' psic @@ -471,26 +471,26 @@ def __call__(self, *args): vths = self.vth(args[2]) # take care of correct broadcasting, assuming args come from phase space meshgrid - if np.ndim(args[0]) == 3: + if xp.ndim(args[0]) == 3: # move eta axes to the back - arg_t = np.moveaxis(args[0], 0, -1) - arg_t = np.moveaxis(arg_t, 0, -1) - arg_t = np.moveaxis(arg_t, 0, -1) + arg_t = xp.moveaxis(args[0], 0, -1) + arg_t = xp.moveaxis(arg_t, 0, -1) + arg_t = xp.moveaxis(arg_t, 0, -1) # broadcast res_broad = res + 0.0 * arg_t # move eta axes to the front - res = np.moveaxis(res_broad, -1, 0) - res = np.moveaxis(res, -1, 0) - res = np.moveaxis(res, -1, 0) + res = xp.moveaxis(res_broad, -1, 0) + res = xp.moveaxis(res, -1, 0) + res = xp.moveaxis(res, -1, 0) # Multiply result with gaussian in energy - if np.ndim(args[0]) == 3: + if xp.ndim(args[0]) == 3: vth_broad = vths + 0.0 * arg_t - vth = np.moveaxis(vth_broad, -1, 0) - vth = np.moveaxis(vth, -1, 0) - vth = np.moveaxis(vth, -1, 0) + vth = xp.moveaxis(vth_broad, -1, 0) + vth = xp.moveaxis(vth, -1, 0) + vth = xp.moveaxis(vth, -1, 0) else: vth = vths @@ -543,13 +543,13 @@ def rc(self, psic): rc_squared = (psic - self.equil.psi_range[0]) / (self.equil.psi_range[1] - self.equil.psi_range[0]) # sorting out indices of negative rc² - neg_index = np.logical_not(rc_squared >= 0) + neg_index = xp.logical_not(rc_squared >= 0) # make them positive rc_squared[neg_index] *= -1 # calculate rc - rc = np.sqrt(rc_squared) + rc = xp.sqrt(rc_squared) rc[neg_index] *= -1 return rc @@ -567,7 +567,7 @@ def n(self, psic, add_perturbation: bool = None): A float (background value) or a numpy.array of the evaluated density. """ # collect arguments - assert isinstance(psic, np.ndarray) + assert isinstance(psic, xp.ndarray) # assuming that input comes from meshgrid. if psic.ndim == 3: @@ -611,7 +611,7 @@ def vth(self, psic): """ # collect arguments - assert isinstance(psic, np.ndarray) + assert isinstance(psic, xp.ndarray) # assuming that input comes from meshgrid. if psic.ndim == 3: diff --git a/src/struphy/kinetic_background/moment_functions.py b/src/struphy/kinetic_background/moment_functions.py index 8d17f670c..b573e657a 100644 --- a/src/struphy/kinetic_background/moment_functions.py +++ b/src/struphy/kinetic_background/moment_functions.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 "Analytical moment functions." -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class ITPA_density: @@ -46,8 +46,8 @@ def __call__(self, eta1, eta2=None, eta3=None): val = ( self._n0 * self._c[3] - * np.exp( - -self._c[2] / self._c[1] * np.tanh((eta1 - self._c[0]) / self._c[2]), + * xp.exp( + -self._c[2] / self._c[1] * xp.tanh((eta1 - self._c[0]) / self._c[2]), ) ) diff --git a/src/struphy/kinetic_background/tests/test_base.py b/src/struphy/kinetic_background/tests/test_base.py index a8f6a1317..5dc077c24 100644 --- a/src/struphy/kinetic_background/tests/test_base.py +++ b/src/struphy/kinetic_background/tests/test_base.py @@ -4,13 +4,13 @@ def test_kinetic_background_magics(show_plot=False): import matplotlib.pyplot as plt from struphy.kinetic_background.maxwellians import Maxwellian3D - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp Nel = [32, 1, 1] - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) - v1 = np.linspace(-7.0, 7.0, 128) + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) + v1 = xp.linspace(-7.0, 7.0, 128) m1_params = {"n": 0.5, "u1": 3.0} m2_params = {"n": 0.5, "u1": -3.0} @@ -22,11 +22,11 @@ def test_kinetic_background_magics(show_plot=False): m_rmul_int = 2 * m1 m_mul_int = m1 * 2 m_mul_float = 2.0 * m1 - m_mul_npint = np.ones(1, dtype=int)[0] * m1 + m_mul_npint = xp.ones(1, dtype=int)[0] * m1 m_sub = m1 - m2 # compare distribution function - meshgrids = np.meshgrid(e1, e2, e3, v1, [0.0], [0.0]) + meshgrids = xp.meshgrid(e1, e2, e3, v1, [0.0], [0.0]) m1_vals = m1(*meshgrids) m2_vals = m2(*meshgrids) @@ -38,15 +38,15 @@ def test_kinetic_background_magics(show_plot=False): m_mul_npint_vals = m_mul_npint(*meshgrids) m_sub_vals = m_sub(*meshgrids) - assert np.allclose(m1_vals + m2_vals, m_add_vals) - assert np.allclose(2 * m1_vals, m_rmul_int_vals) - assert np.allclose(2 * m1_vals, m_mul_int_vals) - assert np.allclose(2.0 * m1_vals, m_mul_float_vals) - assert np.allclose(np.ones(1, dtype=int)[0] * m1_vals, m_mul_npint_vals) - assert np.allclose(m1_vals - m2_vals, m_sub_vals) + assert xp.allclose(m1_vals + m2_vals, m_add_vals) + assert xp.allclose(2 * m1_vals, m_rmul_int_vals) + assert xp.allclose(2 * m1_vals, m_mul_int_vals) + assert xp.allclose(2.0 * m1_vals, m_mul_float_vals) + assert xp.allclose(xp.ones(1, dtype=int)[0] * m1_vals, m_mul_npint_vals) + assert xp.allclose(m1_vals - m2_vals, m_sub_vals) # compare first two moments - meshgrids = np.meshgrid(e1, e2, e3) + meshgrids = xp.meshgrid(e1, e2, e3) n1_vals = m1.n(*meshgrids) n2_vals = m2.n(*meshgrids) @@ -57,11 +57,11 @@ def test_kinetic_background_magics(show_plot=False): u_add1, u_add2, u_add3 = m_add.u(*meshgrids) n_sub_vals = m_sub.n(*meshgrids) - assert np.allclose(n1_vals + n2_vals, n_add_vals) - assert np.allclose(u11 + u21, u_add1) - assert np.allclose(u12 + u22, u_add2) - assert np.allclose(u13 + u23, u_add3) - assert np.allclose(n1_vals - n2_vals, n_sub_vals) + assert xp.allclose(n1_vals + n2_vals, n_add_vals) + assert xp.allclose(u11 + u21, u_add1) + assert xp.allclose(u12 + u22, u_add2) + assert xp.allclose(u13 + u23, u_add3) + assert xp.allclose(n1_vals - n2_vals, n_sub_vals) if show_plot: plt.figure(figsize=(12, 8)) diff --git a/src/struphy/kinetic_background/tests/test_maxwellians.py b/src/struphy/kinetic_background/tests/test_maxwellians.py index fb14fd7ed..44e3679d4 100644 --- a/src/struphy/kinetic_background/tests/test_maxwellians.py +++ b/src/struphy/kinetic_background/tests/test_maxwellians.py @@ -11,28 +11,28 @@ def test_maxwellian_3d_uniform(Nel, show_plot=False): import matplotlib.pyplot as plt from struphy.kinetic_background.maxwellians import Maxwellian3D - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) # ========================================================== # ==== Test uniform non-shifted, isothermal Maxwellian ===== # ========================================================== maxwellian = Maxwellian3D(n=(2.0, None)) - meshgrids = np.meshgrid(e1, e2, e3, [0.0], [0.0], [0.0]) + meshgrids = xp.meshgrid(e1, e2, e3, [0.0], [0.0], [0.0]) # Test constant value at v=0 res = maxwellian(*meshgrids).squeeze() - assert np.allclose(res, 2.0 / (2 * np.pi) ** (3 / 2) + 0 * e1, atol=10e-10), ( - f"{res=},\n {2.0 / (2 * np.pi) ** (3 / 2)}" + assert xp.allclose(res, 2.0 / (2 * xp.pi) ** (3 / 2) + 0 * e1, atol=10e-10), ( + f"{res=},\n {2.0 / (2 * xp.pi) ** (3 / 2)}" ) # test Maxwellian profile in v - v1 = np.linspace(-5, 5, 128) - meshgrids = np.meshgrid( + v1 = xp.linspace(-5, 5, 128) + meshgrids = xp.meshgrid( [0.0], [0.0], [0.0], @@ -41,8 +41,8 @@ def test_maxwellian_3d_uniform(Nel, show_plot=False): [0.0], ) res = maxwellian(*meshgrids).squeeze() - res_ana = 2.0 * np.exp(-(v1**2) / 2.0) / (2 * np.pi) ** (3 / 2) - assert np.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" + res_ana = 2.0 * xp.exp(-(v1**2) / 2.0) / (2 * xp.pi) ** (3 / 2) + assert xp.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" # ======================================================= # ===== Test non-zero shifts and thermal velocities ===== @@ -68,14 +68,14 @@ def test_maxwellian_3d_uniform(Nel, show_plot=False): # test Maxwellian profile in v for i in range(3): vs = [0, 0, 0] - vs[i] = np.linspace(-5, 5, 128) - meshgrids = np.meshgrid([0.0], [0.0], [0.0], *vs) + vs[i] = xp.linspace(-5, 5, 128) + meshgrids = xp.meshgrid([0.0], [0.0], [0.0], *vs) res = maxwellian(*meshgrids).squeeze() - res_ana = np.exp(-((vs[0] - u1) ** 2) / (2 * vth1**2)) - res_ana *= np.exp(-((vs[1] - u2) ** 2) / (2 * vth2**2)) - res_ana *= np.exp(-((vs[2] - u3) ** 2) / (2 * vth3**2)) - res_ana *= n / ((2 * np.pi) ** (3 / 2) * vth1 * vth2 * vth3) + res_ana = xp.exp(-((vs[0] - u1) ** 2) / (2 * vth1**2)) + res_ana *= xp.exp(-((vs[1] - u2) ** 2) / (2 * vth2**2)) + res_ana *= xp.exp(-((vs[2] - u3) ** 2) / (2 * vth3**2)) + res_ana *= n / ((2 * xp.pi) ** (3 / 2) * vth1 * vth2 * vth3) if show_plot: plt.plot(vs[i], res_ana, label="analytical") @@ -86,7 +86,7 @@ def test_maxwellian_3d_uniform(Nel, show_plot=False): plt.xlabel("v_" + str(i + 1)) plt.show() - assert np.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana =}" + assert xp.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana =}" @pytest.mark.parametrize("Nel", [[64, 1, 1]]) @@ -97,10 +97,10 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import Maxwellian3D - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - e1 = np.linspace(0.0, 1.0, Nel[0]) - v1 = np.linspace(-5.0, 5.0, 128) + e1 = xp.linspace(0.0, 1.0, Nel[0]) + v1 = xp.linspace(-5.0, 5.0, 128) # =============================================== # ===== Test cosine perturbation in density ===== @@ -112,10 +112,10 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): maxwellian = Maxwellian3D(n=(2.0, pert)) - meshgrids = np.meshgrid(e1, [0.0], [0.0], [0.0], [0.0], [0.0]) + meshgrids = xp.meshgrid(e1, [0.0], [0.0], [0.0], [0.0], [0.0]) res = maxwellian(*meshgrids).squeeze() - ana_res = (2.0 + amp * np.cos(2 * np.pi * mode * e1)) / (2 * np.pi) ** (3 / 2) + ana_res = (2.0 + amp * xp.cos(2 * xp.pi * mode * e1)) / (2 * xp.pi) ** (3 / 2) if show_plot: plt.plot(e1, ana_res, label="analytical") @@ -126,7 +126,7 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): plt.ylabel("f(eta_1)") plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" # ============================================= # ===== Test cosine perturbation in shift ===== @@ -140,7 +140,7 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): maxwellian = Maxwellian3D(n=(n, None), u1=(u1, pert)) - meshgrids = np.meshgrid( + meshgrids = xp.meshgrid( e1, [0.0], [0.0], @@ -150,9 +150,9 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): ) res = maxwellian(*meshgrids).squeeze() - shift = u1 + amp * np.cos(2 * np.pi * mode * e1) - ana_res = np.exp(-((v1 - shift[:, None]) ** 2) / 2) - ana_res *= n / (2 * np.pi) ** (3 / 2) + shift = u1 + amp * xp.cos(2 * xp.pi * mode * e1) + ana_res = xp.exp(-((v1 - shift[:, None]) ** 2) / 2) + ana_res *= n / (2 * xp.pi) ** (3 / 2) if show_plot: plt.figure(1) @@ -173,7 +173,7 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" # =========================================== # ===== Test cosine perturbation in vth ===== @@ -187,7 +187,7 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): maxwellian = Maxwellian3D(n=(n, None), vth1=(vth1, pert)) - meshgrids = np.meshgrid( + meshgrids = xp.meshgrid( e1, [0.0], [0.0], @@ -197,9 +197,9 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): ) res = maxwellian(*meshgrids).squeeze() - thermal = vth1 + amp * np.cos(2 * np.pi * mode * e1) - ana_res = np.exp(-(v1**2) / (2.0 * thermal[:, None] ** 2)) - ana_res *= n / ((2 * np.pi) ** (3 / 2) * thermal[:, None]) + thermal = vth1 + amp * xp.cos(2 * xp.pi * mode * e1) + ana_res = xp.exp(-(v1**2) / (2.0 * thermal[:, None] ** 2)) + ana_res *= n / ((2 * xp.pi) ** (3 / 2) * thermal[:, None]) if show_plot: plt.figure(1) @@ -220,7 +220,7 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" # ============================================= # ===== Test ITPA perturbation in density ===== @@ -232,10 +232,10 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): maxwellian = Maxwellian3D(n=(0.0, pert)) - meshgrids = np.meshgrid(e1, [0.0], [0.0], [0.0], [0.0], [0.0]) + meshgrids = xp.meshgrid(e1, [0.0], [0.0], [0.0], [0.0], [0.0]) res = maxwellian(*meshgrids).squeeze() - ana_res = n0 * c[3] * np.exp(-c[2] / c[1] * np.tanh((e1 - c[0]) / c[2])) / (2 * np.pi) ** (3 / 2) + ana_res = n0 * c[3] * xp.exp(-c[2] / c[1] * xp.tanh((e1 - c[0]) / c[2])) / (2 * xp.pi) ** (3 / 2) if show_plot: plt.plot(e1, ana_res, label="analytical") @@ -246,7 +246,7 @@ def test_maxwellian_3d_perturbed(Nel, show_plot=False): plt.ylabel("f(eta_1)") plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" @pytest.mark.parametrize("Nel", [[8, 11, 12]]) @@ -263,27 +263,27 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): from struphy.initial import perturbations from struphy.initial.base import Perturbation from struphy.kinetic_background.maxwellians import Maxwellian3D - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) v1 = [0.0] v2 = [0.0, -1.0] v3 = [0.0, -1.0, -1.3] - meshgrids = np.meshgrid(e1, e2, e3, v1, v2, v3, indexing="ij") - e_meshgrids = np.meshgrid(e1, e2, e3, indexing="ij") + meshgrids = xp.meshgrid(e1, e2, e3, v1, v2, v3, indexing="ij") + e_meshgrids = xp.meshgrid(e1, e2, e3, indexing="ij") n_mks = 17 - e1_fl = np.random.rand(n_mks) - e2_fl = np.random.rand(n_mks) - e3_fl = np.random.rand(n_mks) - v1_fl = np.random.randn(n_mks) - v2_fl = np.random.randn(n_mks) - v3_fl = np.random.randn(n_mks) + e1_fl = xp.random.rand(n_mks) + e2_fl = xp.random.rand(n_mks) + e3_fl = xp.random.rand(n_mks) + v1_fl = xp.random.randn(n_mks) + v2_fl = xp.random.randn(n_mks) + v3_fl = xp.random.randn(n_mks) args_fl = [e1_fl, e2_fl, e3_fl, v1_fl, v2_fl, v3_fl] - e_args_fl = np.concatenate((e1_fl[:, None], e2_fl[:, None], e3_fl[:, None]), axis=1) + e_args_fl = xp.concatenate((e1_fl[:, None], e2_fl[:, None], e3_fl[:, None]), axis=1) for key, val in inspect.getmembers(equils): if inspect.isclass(val) and val.__module__ == equils.__name__: @@ -314,8 +314,8 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): elif "ShearedSlab" in key: mhd_equil.domain = domains.Cuboid( r1=mhd_equil.params["a"], - r2=mhd_equil.params["a"] * 2 * np.pi, - r3=mhd_equil.params["R0"] * 2 * np.pi, + r2=mhd_equil.params["a"] * 2 * xp.pi, + r3=mhd_equil.params["R0"] * 2 * xp.pi, ) elif "ShearFluid" in key: mhd_equil.domain = domains.Cuboid( @@ -323,7 +323,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): ) elif "ScrewPinch" in key: mhd_equil.domain = domains.HollowCylinder( - a1=1e-3, a2=mhd_equil.params["a"], Lz=mhd_equil.params["R0"] * 2 * np.pi + a1=1e-3, a2=mhd_equil.params["a"], Lz=mhd_equil.params["R0"] * 2 * xp.pi ) else: try: @@ -353,11 +353,11 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): # test meshgrid evaluation n0 = mhd_equil.n0(*e_meshgrids) - assert np.allclose( + assert xp.allclose( maxwellian(*meshgrids)[:, :, :, 0, 0, 0], n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 0, 0] ) - assert np.allclose( + assert xp.allclose( maxwellian(*meshgrids)[:, :, :, 0, 1, 2], n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 1, 2] ) @@ -365,16 +365,16 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): if "GVECequilibrium" in key: pass else: - assert np.allclose(maxwellian(*args_fl), mhd_equil.n0(e_args_fl) * maxwellian_1(*args_fl)) - assert np.allclose(maxwellian.n(e1_fl, e2_fl, e3_fl), mhd_equil.n0(e_args_fl)) + assert xp.allclose(maxwellian(*args_fl), mhd_equil.n0(e_args_fl) * maxwellian_1(*args_fl)) + assert xp.allclose(maxwellian.n(e1_fl, e2_fl, e3_fl), mhd_equil.n0(e_args_fl)) u_maxw = maxwellian.u(e1_fl, e2_fl, e3_fl) u_eq = mhd_equil.u_cart(e_args_fl)[0] - assert all([np.allclose(m, e) for m, e in zip(u_maxw, u_eq)]) + assert all([xp.allclose(m, e) for m, e in zip(u_maxw, u_eq)]) vth_maxw = maxwellian.vth(e1_fl, e2_fl, e3_fl) - vth_eq = np.sqrt(mhd_equil.p0(e_args_fl) / mhd_equil.n0(e_args_fl)) - assert all([np.allclose(v, vth_eq) for v in vth_maxw]) + vth_eq = xp.sqrt(mhd_equil.p0(e_args_fl) / mhd_equil.n0(e_args_fl)) + assert all([xp.allclose(v, vth_eq) for v in vth_maxw]) # plotting moments if show_plot: @@ -384,7 +384,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): # density plots n_cart = mhd_equil.domain.push(maxwellian.n, *e_meshgrids) - levels = np.linspace(np.min(n_cart) - 1e-10, np.max(n_cart), 20) + levels = xp.linspace(xp.min(n_cart) - 1e-10, xp.max(n_cart), 20) plt.subplot(2, 5, 1) if "Slab" in key or "Pinch" in key: @@ -420,7 +420,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): # velocity plots us = maxwellian.u(*e_meshgrids) for i, u in enumerate(us): - levels = np.linspace(np.min(u) - 1e-10, np.max(u), 20) + levels = xp.linspace(xp.min(u) - 1e-10, xp.max(u), 20) plt.subplot(2, 5, 2 + i) if "Slab" in key or "Pinch" in key: @@ -453,7 +453,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): vth = maxwellian.vth(*e_meshgrids)[0] vth_cart = mhd_equil.domain.push(vth, *e_meshgrids) - levels = np.linspace(np.min(vth_cart) - 1e-10, np.max(vth_cart), 20) + levels = xp.linspace(xp.min(vth_cart) - 1e-10, xp.max(vth_cart), 20) plt.subplot(2, 5, 5) if "Slab" in key or "Pinch" in key: @@ -528,13 +528,13 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): vth3=(0.0, pert), ) - assert np.allclose(maxwellian_zero_bckgr.n(*e_meshgrids), pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[0], pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[1], pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[2], pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[0], pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[1], pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[2], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.n(*e_meshgrids), pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[0], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[1], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[2], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[0], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[1], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[2], pert(*e_meshgrids)) # plotting perturbations if show_plot: # and 'Torus' in key_2: @@ -544,7 +544,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): # density plots n_cart = mhd_equil.domain.push(maxwellian_zero_bckgr.n, *e_meshgrids) - levels = np.linspace(np.min(n_cart) - 1e-10, np.max(n_cart), 20) + levels = xp.linspace(xp.min(n_cart) - 1e-10, xp.max(n_cart), 20) plt.subplot(2, 5, 1) if "Slab" in key or "Pinch" in key: @@ -580,7 +580,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): # velocity plots us = maxwellian_zero_bckgr.u(*e_meshgrids) for i, u in enumerate(us): - levels = np.linspace(np.min(u) - 1e-10, np.max(u), 20) + levels = xp.linspace(xp.min(u) - 1e-10, xp.max(u), 20) plt.subplot(2, 5, 2 + i) if "Slab" in key or "Pinch" in key: @@ -617,7 +617,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): vth = maxwellian_zero_bckgr.vth(*e_meshgrids)[0] vth_cart = mhd_equil.domain.push(vth, *e_meshgrids) - levels = np.linspace(np.min(vth_cart) - 1e-10, np.max(vth_cart), 20) + levels = xp.linspace(xp.min(vth_cart) - 1e-10, xp.max(vth_cart), 20) plt.subplot(2, 5, 5) if "Slab" in key or "Pinch" in key: @@ -669,31 +669,31 @@ def test_maxwellian_2d_uniform(Nel, show_plot=False): import matplotlib.pyplot as plt from struphy.kinetic_background.maxwellians import GyroMaxwellian2D - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) # =========================================================== # ===== Test uniform non-shifted, isothermal Maxwellian ===== # =========================================================== maxwellian = GyroMaxwellian2D(n=(2.0, None), volume_form=False) - meshgrids = np.meshgrid(e1, e2, e3, [0.01], [0.01]) + meshgrids = xp.meshgrid(e1, e2, e3, [0.01], [0.01]) # Test constant value at v_para = v_perp = 0.01 res = maxwellian(*meshgrids).squeeze() - assert np.allclose(res, 2.0 / (2 * np.pi) ** (1 / 2) * np.exp(-(0.01**2)) + 0 * e1, atol=10e-10), ( - f"{res=},\n {2.0 / (2 * np.pi) ** (3 / 2)}" + assert xp.allclose(res, 2.0 / (2 * xp.pi) ** (1 / 2) * xp.exp(-(0.01**2)) + 0 * e1, atol=10e-10), ( + f"{res=},\n {2.0 / (2 * xp.pi) ** (3 / 2)}" ) # test Maxwellian profile in v - v_para = np.linspace(-5, 5, 64) - v_perp = np.linspace(0, 2.5, 64) - vpara, vperp = np.meshgrid(v_para, v_perp) + v_para = xp.linspace(-5, 5, 64) + v_perp = xp.linspace(0, 2.5, 64) + vpara, vperp = xp.meshgrid(v_para, v_perp) - meshgrids = np.meshgrid( + meshgrids = xp.meshgrid( [0.0], [0.0], [0.0], @@ -702,8 +702,8 @@ def test_maxwellian_2d_uniform(Nel, show_plot=False): ) res = maxwellian(*meshgrids).squeeze() - res_ana = 2.0 / (2 * np.pi) ** (1 / 2) * np.exp(-(vpara.T**2) / 2.0 - vperp.T**2 / 2.0) - assert np.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" + res_ana = 2.0 / (2 * xp.pi) ** (1 / 2) * xp.exp(-(vpara.T**2) / 2.0 - vperp.T**2 / 2.0) + assert xp.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" # ======================================================= # ===== Test non-zero shifts and thermal velocities ===== @@ -724,16 +724,16 @@ def test_maxwellian_2d_uniform(Nel, show_plot=False): ) # test Maxwellian profile in v - v_para = np.linspace(-5, 5, 64) - v_perp = np.linspace(0, 2.5, 64) - vpara, vperp = np.meshgrid(v_para, v_perp) + v_para = xp.linspace(-5, 5, 64) + v_perp = xp.linspace(0, 2.5, 64) + vpara, vperp = xp.meshgrid(v_para, v_perp) - meshgrids = np.meshgrid([0.0], [0.0], [0.0], v_para, v_perp) + meshgrids = xp.meshgrid([0.0], [0.0], [0.0], v_para, v_perp) res = maxwellian(*meshgrids).squeeze() - res_ana = np.exp(-((vpara.T - u_para) ** 2) / (2 * vth_para**2)) - res_ana *= np.exp(-((vperp.T - u_perp) ** 2) / (2 * vth_perp**2)) - res_ana *= n / ((2 * np.pi) ** (1 / 2) * vth_para * vth_perp**2) + res_ana = xp.exp(-((vpara.T - u_para) ** 2) / (2 * vth_para**2)) + res_ana *= xp.exp(-((vperp.T - u_perp) ** 2) / (2 * vth_perp**2)) + res_ana *= n / ((2 * xp.pi) ** (1 / 2) * vth_para * vth_perp**2) if show_plot: plt.plot(v_para, res_ana[:, 32], label="analytical") @@ -752,7 +752,7 @@ def test_maxwellian_2d_uniform(Nel, show_plot=False): plt.xlabel("v_" + "perp") plt.show() - assert np.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana =}" + assert xp.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana =}" @pytest.mark.parametrize("Nel", [[6, 1, 1]]) @@ -763,11 +763,11 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import GyroMaxwellian2D - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - e1 = np.linspace(0.0, 1.0, Nel[0]) - v1 = np.linspace(-5.0, 5.0, 128) - v2 = np.linspace(0, 2.5, 128) + e1 = xp.linspace(0.0, 1.0, Nel[0]) + v1 = xp.linspace(-5.0, 5.0, 128) + v2 = xp.linspace(0, 2.5, 128) # =============================================== # ===== Test cosine perturbation in density ===== @@ -779,11 +779,11 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): maxwellian = GyroMaxwellian2D(n=(2.0, pert), volume_form=False) v_perp = 0.1 - meshgrids = np.meshgrid(e1, [0.0], [0.0], [0.0], v_perp) + meshgrids = xp.meshgrid(e1, [0.0], [0.0], [0.0], v_perp) res = maxwellian(*meshgrids).squeeze() - ana_res = (2.0 + amp * np.cos(2 * np.pi * mode * e1)) / (2 * np.pi) ** (1 / 2) - ana_res *= np.exp(-(v_perp**2) / 2) + ana_res = (2.0 + amp * xp.cos(2 * xp.pi * mode * e1)) / (2 * xp.pi) ** (1 / 2) + ana_res *= xp.exp(-(v_perp**2) / 2) if show_plot: plt.plot(e1, ana_res, label="analytical") @@ -794,7 +794,7 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): plt.ylabel("f(eta_1)") plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" # ==================================================== # ===== Test cosine perturbation in shift (para) ===== @@ -812,12 +812,12 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): ) v_perp = 0.1 - meshgrids = np.meshgrid(e1, [0.0], [0.0], v1, v_perp) + meshgrids = xp.meshgrid(e1, [0.0], [0.0], v1, v_perp) res = maxwellian(*meshgrids).squeeze() - shift = u_para + amp * np.cos(2 * np.pi * mode * e1) - ana_res = np.exp(-((v1 - shift[:, None]) ** 2) / 2.0) - ana_res *= n / (2 * np.pi) ** (1 / 2) * np.exp(-(v_perp**2) / 2.0) + shift = u_para + amp * xp.cos(2 * xp.pi * mode * e1) + ana_res = xp.exp(-((v1 - shift[:, None]) ** 2) / 2.0) + ana_res *= n / (2 * xp.pi) ** (1 / 2) * xp.exp(-(v_perp**2) / 2.0) if show_plot: plt.figure(1) @@ -838,7 +838,7 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" # ==================================================== # ===== Test cosine perturbation in shift (perp) ===== @@ -855,12 +855,12 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): volume_form=False, ) - meshgrids = np.meshgrid(e1, [0.0], [0.0], 0.0, v2) + meshgrids = xp.meshgrid(e1, [0.0], [0.0], 0.0, v2) res = maxwellian(*meshgrids).squeeze() - shift = u_perp + amp * np.cos(2 * np.pi * mode * e1) - ana_res = np.exp(-((v2 - shift[:, None]) ** 2) / 2.0) - ana_res *= n / (2 * np.pi) ** (1 / 2) + shift = u_perp + amp * xp.cos(2 * xp.pi * mode * e1) + ana_res = xp.exp(-((v2 - shift[:, None]) ** 2) / 2.0) + ana_res *= n / (2 * xp.pi) ** (1 / 2) if show_plot: plt.figure(1) @@ -881,7 +881,7 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" # ================================================== # ===== Test cosine perturbation in vth (para) ===== @@ -899,7 +899,7 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): ) v_perp = 0.1 - meshgrids = np.meshgrid( + meshgrids = xp.meshgrid( e1, [0.0], [0.0], @@ -908,10 +908,10 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): ) res = maxwellian(*meshgrids).squeeze() - thermal = vth_para + amp * np.cos(2 * np.pi * mode * e1) - ana_res = np.exp(-(v1**2) / (2.0 * thermal[:, None] ** 2)) - ana_res *= n / ((2 * np.pi) ** (1 / 2) * thermal[:, None]) - ana_res *= np.exp(-(v_perp**2) / 2.0) + thermal = vth_para + amp * xp.cos(2 * xp.pi * mode * e1) + ana_res = xp.exp(-(v1**2) / (2.0 * thermal[:, None] ** 2)) + ana_res *= n / ((2 * xp.pi) ** (1 / 2) * thermal[:, None]) + ana_res *= xp.exp(-(v_perp**2) / 2.0) if show_plot: plt.figure(1) @@ -932,7 +932,7 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" # ================================================== # ===== Test cosine perturbation in vth (perp) ===== @@ -949,7 +949,7 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): volume_form=False, ) - meshgrids = np.meshgrid( + meshgrids = xp.meshgrid( e1, [0.0], [0.0], @@ -958,9 +958,9 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): ) res = maxwellian(*meshgrids).squeeze() - thermal = vth_perp + amp * np.cos(2 * np.pi * mode * e1) - ana_res = np.exp(-(v2**2) / (2.0 * thermal[:, None] ** 2)) - ana_res *= n / ((2 * np.pi) ** (1 / 2) * thermal[:, None] ** 2) + thermal = vth_perp + amp * xp.cos(2 * xp.pi * mode * e1) + ana_res = xp.exp(-(v2**2) / (2.0 * thermal[:, None] ** 2)) + ana_res *= n / ((2 * xp.pi) ** (1 / 2) * thermal[:, None] ** 2) if show_plot: plt.figure(1) @@ -981,7 +981,7 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" # ============================================= # ===== Test ITPA perturbation in density ===== @@ -993,11 +993,11 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): maxwellian = GyroMaxwellian2D(n=(0.0, pert), volume_form=False) v_perp = 0.1 - meshgrids = np.meshgrid(e1, [0.0], [0.0], [0.0], v_perp) + meshgrids = xp.meshgrid(e1, [0.0], [0.0], [0.0], v_perp) res = maxwellian(*meshgrids).squeeze() - ana_res = n0 * c[3] * np.exp(-c[2] / c[1] * np.tanh((e1 - c[0]) / c[2])) / (2 * np.pi) ** (1 / 2) - ana_res *= np.exp(-(v_perp**2) / 2.0) + ana_res = n0 * c[3] * xp.exp(-c[2] / c[1] * xp.tanh((e1 - c[0]) / c[2])) / (2 * xp.pi) ** (1 / 2) + ana_res *= xp.exp(-(v_perp**2) / 2.0) if show_plot: plt.plot(e1, ana_res, label="analytical") @@ -1008,7 +1008,7 @@ def test_maxwellian_2d_perturbed(Nel, show_plot=False): plt.ylabel("f(eta_1)") plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" @pytest.mark.parametrize("Nel", [[8, 12, 12]]) @@ -1025,25 +1025,25 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): from struphy.initial import perturbations from struphy.initial.base import Perturbation from struphy.kinetic_background.maxwellians import GyroMaxwellian2D - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) v1 = [0.0] v2 = [0.0, 2.0] - meshgrids = np.meshgrid(e1, e2, e3, v1, v2, indexing="ij") - e_meshgrids = np.meshgrid(e1, e2, e3, indexing="ij") + meshgrids = xp.meshgrid(e1, e2, e3, v1, v2, indexing="ij") + e_meshgrids = xp.meshgrid(e1, e2, e3, indexing="ij") n_mks = 17 - e1_fl = np.random.rand(n_mks) - e2_fl = np.random.rand(n_mks) - e3_fl = np.random.rand(n_mks) - v1_fl = np.random.randn(n_mks) - v2_fl = np.random.rand(n_mks) + e1_fl = xp.random.rand(n_mks) + e2_fl = xp.random.rand(n_mks) + e3_fl = xp.random.rand(n_mks) + v1_fl = xp.random.randn(n_mks) + v2_fl = xp.random.rand(n_mks) args_fl = [e1_fl, e2_fl, e3_fl, v1_fl, v2_fl] - e_args_fl = np.concatenate((e1_fl[:, None], e2_fl[:, None], e3_fl[:, None]), axis=1) + e_args_fl = xp.concatenate((e1_fl[:, None], e2_fl[:, None], e3_fl[:, None]), axis=1) for key, val in inspect.getmembers(equils): if inspect.isclass(val) and val.__module__ == equils.__name__: @@ -1076,8 +1076,8 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): elif "ShearedSlab" in key: mhd_equil.domain = domains.Cuboid( r1=mhd_equil.params["a"], - r2=mhd_equil.params["a"] * 2 * np.pi, - r3=mhd_equil.params["R0"] * 2 * np.pi, + r2=mhd_equil.params["a"] * 2 * xp.pi, + r3=mhd_equil.params["R0"] * 2 * xp.pi, ) elif "ShearFluid" in key: mhd_equil.domain = domains.Cuboid( @@ -1085,7 +1085,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): ) elif "ScrewPinch" in key: mhd_equil.domain = domains.HollowCylinder( - a1=1e-3, a2=mhd_equil.params["a"], Lz=mhd_equil.params["R0"] * 2 * np.pi + a1=1e-3, a2=mhd_equil.params["a"], Lz=mhd_equil.params["R0"] * 2 * xp.pi ) else: try: @@ -1111,27 +1111,27 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): # test meshgrid evaluation n0 = mhd_equil.n0(*e_meshgrids) - assert np.allclose(maxwellian(*meshgrids)[:, :, :, 0, 0], n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 0]) + assert xp.allclose(maxwellian(*meshgrids)[:, :, :, 0, 0], n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 0]) - assert np.allclose(maxwellian(*meshgrids)[:, :, :, 0, 1], n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 1]) + assert xp.allclose(maxwellian(*meshgrids)[:, :, :, 0, 1], n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 1]) # test flat evaluation if "GVECequilibrium" in key: pass else: - assert np.allclose(maxwellian(*args_fl), mhd_equil.n0(e_args_fl) * maxwellian_1(*args_fl)) - assert np.allclose(maxwellian.n(e1_fl, e2_fl, e3_fl), mhd_equil.n0(e_args_fl)) + assert xp.allclose(maxwellian(*args_fl), mhd_equil.n0(e_args_fl) * maxwellian_1(*args_fl)) + assert xp.allclose(maxwellian.n(e1_fl, e2_fl, e3_fl), mhd_equil.n0(e_args_fl)) u_maxw = maxwellian.u(e1_fl, e2_fl, e3_fl) tmp_jv = mhd_equil.jv(e_args_fl) / mhd_equil.n0(e_args_fl) tmp_unit_b1 = mhd_equil.unit_b1(e_args_fl) # j_parallel = jv.b1 j_para = sum([ji * bi for ji, bi in zip(tmp_jv, tmp_unit_b1)]) - assert np.allclose(u_maxw[0], j_para) + assert xp.allclose(u_maxw[0], j_para) vth_maxw = maxwellian.vth(e1_fl, e2_fl, e3_fl) - vth_eq = np.sqrt(mhd_equil.p0(e_args_fl) / mhd_equil.n0(e_args_fl)) - assert all([np.allclose(v, vth_eq) for v in vth_maxw]) + vth_eq = xp.sqrt(mhd_equil.p0(e_args_fl) / mhd_equil.n0(e_args_fl)) + assert all([xp.allclose(v, vth_eq) for v in vth_maxw]) # plotting moments if show_plot: @@ -1141,7 +1141,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): # density plots n_cart = mhd_equil.domain.push(maxwellian.n, *e_meshgrids) - levels = np.linspace(np.min(n_cart) - 1e-10, np.max(n_cart), 20) + levels = xp.linspace(xp.min(n_cart) - 1e-10, xp.max(n_cart), 20) plt.subplot(2, 4, 1) if "Slab" in key or "Pinch" in key: @@ -1177,7 +1177,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): # velocity plots us = maxwellian.u(*e_meshgrids) for i, u in enumerate(us[:1]): - levels = np.linspace(np.min(u) - 1e-10, np.max(u), 20) + levels = xp.linspace(xp.min(u) - 1e-10, xp.max(u), 20) plt.subplot(2, 4, 2 + i) if "Slab" in key or "Pinch" in key: @@ -1210,7 +1210,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): vth = maxwellian.vth(*e_meshgrids)[0] vth_cart = mhd_equil.domain.push(vth, *e_meshgrids) - levels = np.linspace(np.min(vth_cart) - 1e-10, np.max(vth_cart), 20) + levels = xp.linspace(xp.min(vth_cart) - 1e-10, xp.max(vth_cart), 20) plt.subplot(2, 4, 4) if "Slab" in key or "Pinch" in key: @@ -1281,11 +1281,11 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): volume_form=False, ) - assert np.allclose(maxwellian_zero_bckgr.n(*e_meshgrids), pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[0], pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[1], pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[0], pert(*e_meshgrids)) - assert np.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[1], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.n(*e_meshgrids), pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[0], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.u(*e_meshgrids)[1], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[0], pert(*e_meshgrids)) + assert xp.allclose(maxwellian_zero_bckgr.vth(*e_meshgrids)[1], pert(*e_meshgrids)) # plotting perturbations if show_plot and "EQDSKequilibrium" in key: # and 'Torus' in key_2: @@ -1295,7 +1295,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): # density plots n_cart = mhd_equil.domain.push(maxwellian_zero_bckgr.n, *e_meshgrids) - levels = np.linspace(np.min(n_cart) - 1e-10, np.max(n_cart), 20) + levels = xp.linspace(xp.min(n_cart) - 1e-10, xp.max(n_cart), 20) plt.subplot(2, 4, 1) if "Slab" in key or "Pinch" in key: @@ -1331,7 +1331,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): # velocity plots us = maxwellian_zero_bckgr.u(*e_meshgrids) for i, u in enumerate(us): - levels = np.linspace(np.min(u) - 1e-10, np.max(u), 20) + levels = xp.linspace(xp.min(u) - 1e-10, xp.max(u), 20) plt.subplot(2, 4, 2 + i) if "Slab" in key or "Pinch" in key: @@ -1368,7 +1368,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): vth = maxwellian_zero_bckgr.vth(*e_meshgrids)[0] vth_cart = mhd_equil.domain.push(vth, *e_meshgrids) - levels = np.linspace(np.min(vth_cart) - 1e-10, np.max(vth_cart), 20) + levels = xp.linspace(xp.min(vth_cart) - 1e-10, xp.max(vth_cart), 20) plt.subplot(2, 4, 4) if "Slab" in key or "Pinch" in key: @@ -1423,13 +1423,13 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): from struphy.geometry import domains from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import CanonicalMaxwellian - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) - eta_meshgrid = np.meshgrid(e1, e2, e3) + eta_meshgrid = xp.meshgrid(e1, e2, e3) v_para = 0.01 v_perp = 0.01 @@ -1476,7 +1476,7 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): psi = mhd_equil.psi_r(r) psic = psi - epsilon * B0 * R0 / absB * v_para - psic += epsilon * np.sign(v_para) * np.sqrt(2 * (energy - mu * B0)) * R0 * np.heaviside(energy - mu * B0, 0) + psic += epsilon * xp.sign(v_para) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) # =========================================================== # ===== Test uniform, isothermal canonical Maxwellian ===== @@ -1490,14 +1490,14 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): res_ana = ( maxw_params["n"] * 2 - * np.sqrt(energy / np.pi) + * xp.sqrt(energy / xp.pi) / maxw_params["vth"] ** 3 - * np.exp(-energy / maxw_params["vth"] ** 2) + * xp.exp(-energy / maxw_params["vth"] ** 2) ) - assert np.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" + assert xp.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" # test canonical Maxwellian profile in v_para - v_para = np.linspace(-5, 5, 64) + v_para = xp.linspace(-5, 5, 64) v_perp = 0.1 absB = mhd_equil.absB0(0.0, 0.0, 0.0)[0, 0, 0] @@ -1514,18 +1514,18 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): psi = mhd_equil.psi_r(r) psic = psi - epsilon * B0 * R0 / absB * v_para - psic += epsilon * np.sign(v_para) * np.sqrt(2 * (energy - mu * B0)) * R0 * np.heaviside(energy - mu * B0, 0) + psic += epsilon * xp.sign(v_para) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) - com_meshgrids = np.meshgrid(energy, mu, psic) + com_meshgrids = xp.meshgrid(energy, mu, psic) res = maxwellian(*com_meshgrids).squeeze() res_ana = ( maxw_params["n"] * 2 - * np.sqrt(com_meshgrids[0] / np.pi) + * xp.sqrt(com_meshgrids[0] / xp.pi) / maxw_params["vth"] ** 3 - * np.exp(-com_meshgrids[0] / maxw_params["vth"] ** 2) + * xp.exp(-com_meshgrids[0] / maxw_params["vth"] ** 2) ) if show_plot: @@ -1537,11 +1537,11 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): plt.xlabel("v_para") plt.show() - assert np.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" + assert xp.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" # test canonical Maxwellian profile in v_perp v_para = 0.1 - v_perp = np.linspace(0, 2.5, 64) + v_perp = xp.linspace(0, 2.5, 64) absB = mhd_equil.absB0(0.5, 0.5, 0.5)[0, 0, 0] @@ -1557,18 +1557,18 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): psi = mhd_equil.psi_r(r) psic = psi - epsilon * B0 * R0 / absB * v_para - psic += epsilon * np.sign(v_para) * np.sqrt(2 * (energy - mu * B0)) * R0 * np.heaviside(energy - mu * B0, 0) + psic += epsilon * xp.sign(v_para) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) - com_meshgrids = np.meshgrid(energy, mu, psic) + com_meshgrids = xp.meshgrid(energy, mu, psic) res = maxwellian(*com_meshgrids).squeeze() res_ana = ( maxw_params["n"] * 2 - * np.sqrt(com_meshgrids[0] / np.pi) + * xp.sqrt(com_meshgrids[0] / xp.pi) / maxw_params["vth"] ** 3 - * np.exp(-com_meshgrids[0] / maxw_params["vth"] ** 2) + * xp.exp(-com_meshgrids[0] / maxw_params["vth"] ** 2) ) if show_plot: @@ -1580,7 +1580,7 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): plt.xlabel("v_perp") plt.show() - assert np.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" + assert xp.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana}" # ============================================= # ===== Test ITPA perturbation in density ===== @@ -1595,11 +1595,11 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): maxwellian = CanonicalMaxwellian(n=(0.0, pert), equil=mhd_equil, volume_form=False) - e1 = np.linspace(0.0, 1.0, Nel[0]) - e2 = np.linspace(0.0, 1.0, Nel[1]) - e3 = np.linspace(0.0, 1.0, Nel[2]) + e1 = xp.linspace(0.0, 1.0, Nel[0]) + e2 = xp.linspace(0.0, 1.0, Nel[1]) + e3 = xp.linspace(0.0, 1.0, Nel[2]) - eta_meshgrid = np.meshgrid(e1, e2, e3) + eta_meshgrid = xp.meshgrid(e1, e2, e3) v_para = 0.01 v_perp = 0.01 @@ -1618,16 +1618,16 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): psi = mhd_equil.psi_r(r[0, :, 0]) psic = psi - epsilon * B0 * R0 / absB * v_para - psic += epsilon * np.sign(v_para) * np.sqrt(2 * (energy - mu * B0)) * R0 * np.heaviside(energy - mu * B0, 0) + psic += epsilon * xp.sign(v_para) * xp.sqrt(2 * (energy - mu * B0)) * R0 * xp.heaviside(energy - mu * B0, 0) - com_meshgrids = np.meshgrid(energy, mu, psic) + com_meshgrids = xp.meshgrid(energy, mu, psic) res = maxwellian(energy, mu, psic).squeeze() # calculate rc rc = maxwellian.rc(psic) - ana_res = n0 * c[3] * np.exp(-c[2] / c[1] * np.tanh((rc - c[0]) / c[2])) - ana_res *= 2 * np.sqrt(energy / np.pi) / maxw_params["vth"] ** 3 * np.exp(-energy / maxw_params["vth"] ** 2) + ana_res = n0 * c[3] * xp.exp(-c[2] / c[1] * xp.tanh((rc - c[0]) / c[2])) + ana_res *= 2 * xp.sqrt(energy / xp.pi) / maxw_params["vth"] ** 3 * xp.exp(-energy / maxw_params["vth"] ** 2) if show_plot: plt.plot(e1, ana_res, label="analytical") @@ -1638,7 +1638,7 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): plt.ylabel("f(eta_1)") plt.show() - assert np.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" + assert xp.allclose(res, ana_res, atol=10e-10), f"{res=},\n {ana_res}" if __name__ == "__main__": diff --git a/src/struphy/linear_algebra/linalg_kron.py b/src/struphy/linear_algebra/linalg_kron.py index 68e2513f7..29be15d60 100644 --- a/src/struphy/linear_algebra/linalg_kron.py +++ b/src/struphy/linear_algebra/linalg_kron.py @@ -16,7 +16,7 @@ from scipy.linalg import solve_circulant from scipy.sparse.linalg import splu -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def kron_matvec_2d(kmat, vec2d): @@ -197,9 +197,9 @@ def kron_matmat_fft_3d(a_vec, b_vec): c_vec = [0, 0, 0] - c_vec[0] = np.fft.ifft(np.fft.fft(a_vec[0]) * np.fft.fft(b_vec[0])) - c_vec[1] = np.fft.ifft(np.fft.fft(a_vec[1]) * np.fft.fft(b_vec[1])) - c_vec[2] = np.fft.ifft(np.fft.fft(a_vec[2]) * np.fft.fft(b_vec[2])) + c_vec[0] = xp.fft.ifft(xp.fft.fft(a_vec[0]) * xp.fft.fft(b_vec[0])) + c_vec[1] = xp.fft.ifft(xp.fft.fft(a_vec[1]) * xp.fft.fft(b_vec[1])) + c_vec[2] = xp.fft.ifft(xp.fft.fft(a_vec[2]) * xp.fft.fft(b_vec[2])) return c_vec diff --git a/src/struphy/linear_algebra/saddle_point.py b/src/struphy/linear_algebra/saddle_point.py index 5da783b3a..efa4dbbbe 100644 --- a/src/struphy/linear_algebra/saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point.py @@ -7,7 +7,7 @@ from psydac.linalg.solvers import inverse from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class SaddlePointSolver: @@ -28,7 +28,7 @@ class SaddlePointSolver: } \right) using either the Uzawa iteration :math:`BA^{-1}B^{\top} y = BA^{-1} f` or using on of the solvers given in :mod:`psydac.linalg.solvers`. The prefered solver is GMRES. - The decission which variant to use is given by the type of A. If A is of type list of np.ndarrays or sc.sparse.csr_matrices, then this class uses the Uzawa algorithm. + The decission which variant to use is given by the type of A. If A is of type list of xp.ndarrays or sc.sparse.csr_matrices, then this class uses the Uzawa algorithm. If A is of type LinearOperator or BlockLinearOperator, a solver is used for the inverse. Using the Uzawa algorithm, solution is given by: @@ -41,7 +41,7 @@ class SaddlePointSolver: ---------- A : list, LinearOperator or BlockLinearOperator Upper left block. - Either the entries on the diagonals of block A are given as list of np.ndarray or sc.sparse.csr_matrix. + Either the entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. Alternative: Give whole matrice A as LinearOperator or BlockLinearOperator. list: Uzawa algorithm is used. LinearOperator: A solver given in :mod:`psydac.linalg.solvers` is used. Specified by solver_name. @@ -49,16 +49,16 @@ class SaddlePointSolver: B : list, LinearOperator or BlockLinearOperator Lower left block. - Uzwaw Algorithm: All entries of block B are given either as list of np.ndarray or sc.sparse.csr_matrix. + Uzwaw Algorithm: All entries of block B are given either as list of xp.ndarray or sc.sparse.csr_matrix. Solver: Give whole B as LinearOperator or BlocklinearOperator F : list Right hand side of the upper block. - Uzawa: Given as list of np.ndarray or sc.sparse.csr_matrix. + Uzawa: Given as list of xp.ndarray or sc.sparse.csr_matrix. Solver: Given as LinearOperator or BlockLinearOperator Apre : list - The non-inverted preconditioner for entries on the diagonals of block A are given as list of np.ndarray or sc.sparse.csr_matrix. Only required for the Uzawa algorithm. + The non-inverted preconditioner for entries on the diagonals of block A are given as list of xp.ndarray or sc.sparse.csr_matrix. Only required for the Uzawa algorithm. method_to_solve : str Method for the inverses. Choose from 'DirectNPInverse', 'ScipySparse', 'InexactNPInverse' ,'SparseSolver'. Only required for the Uzawa algorithm. @@ -98,14 +98,14 @@ def __init__( if isinstance(A, list): self._variant = "Uzawa" for i in A: - assert isinstance(i, np.ndarray) or isinstance(i, sc.sparse.csr_matrix) + assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) for i in B: - assert isinstance(i, np.ndarray) or isinstance(i, sc.sparse.csr_matrix) + assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) for i in F: - assert isinstance(i, np.ndarray) or isinstance(i, sc.sparse.csr_matrix) + assert isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) for i in Apre: assert ( - isinstance(i, np.ndarray) + isinstance(i, xp.ndarray) or isinstance(i, sc.sparse.csr_matrix) or isinstance(i, sc.sparse.csr_array) ) @@ -169,9 +169,9 @@ def __init__( self._setup_inverses() # Solution vectors numpy - self._Pnp = np.zeros(self._B1np.shape[0]) - self._Unp = np.zeros(self._A[0].shape[1]) - self._Uenp = np.zeros(self._A[1].shape[1]) + self._Pnp = xp.zeros(self._B1np.shape[0]) + self._Unp = xp.zeros(self._A[0].shape[1]) + self._Uenp = xp.zeros(self._A[1].shape[1]) # Allocate memory for matrices used in solving the system self._rhs0np = self._F[0].copy() self._rhs1np = self._F[1].copy() @@ -195,8 +195,8 @@ def A(self, a): same_A0 = (A0_old != A0_new).nnz == 0 same_A1 = (A1_old != A1_new).nnz == 0 else: - same_A0 = np.allclose(A0_old, A0_new, atol=1e-10) - same_A1 = np.allclose(A1_old, A1_new, atol=1e-10) + same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) + same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) if same_A0 and same_A1: need_update = False self._A = a @@ -240,8 +240,8 @@ def Apre(self, a): same_A0 = (A0_old != A0_new).nnz == 0 same_A1 = (A1_old != A1_new).nnz == 0 else: - same_A0 = np.allclose(A0_old, A0_new, atol=1e-10) - same_A1 = np.allclose(A1_old, A1_new, atol=1e-10) + same_A0 = xp.allclose(A0_old, A0_new, atol=1e-10) + same_A1 = xp.allclose(A1_old, A1_new, atol=1e-10) if same_A0 and same_A1: need_update = False self._Apre = a @@ -256,11 +256,11 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): Parameters ---------- - U_init : Vector, np.ndarray or sc.sparse.csr.csr_matrix, optional - Initial guess for the velocity of the ions. If None, initializes to zero. Types np.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. + U_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional + Initial guess for the velocity of the ions. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. - Ue_init : Vector, np.ndarray or sc.sparse.csr.csr_matrix, optional - Initial guess for the velocity of the electrons. If None, initializes to zero. Types np.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. + Ue_init : Vector, xp.ndarray or sc.sparse.csr.csr_matrix, optional + Initial guess for the velocity of the electrons. If None, initializes to zero. Types xp.ndarray and sc.sparse.csr.csr_matrix can only be given if system should be solved with Uzawa algorithm. P_init : Vector, optional Initial guess for the potential. If None, initializes to zero. @@ -310,7 +310,7 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): self._spectralresult = [] # Initialize P to zero or given initial guess - if isinstance(U_init, np.ndarray) or isinstance(U_init, sc.sparse.csr.csr_matrix): + if isinstance(U_init, xp.ndarray) or isinstance(U_init, sc.sparse.csr.csr_matrix): self._Pnp = P_init if P_init is not None else self._P self._Unp = U_init if U_init is not None else self._U self._Uenp = Ue_init if U_init is not None else self._Ue @@ -353,8 +353,8 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): # Step 2: Compute residual R = BU (divergence of U) R = R1 + R2 # self._B1np.dot(self._Unp) + self._B2np.dot(self._Uenp) - residual_norm = np.linalg.norm(R) - residual_normR1 = np.linalg.norm(R) + residual_norm = xp.linalg.norm(R) + residual_normR1 = xp.linalg.norm(R) self._residual_norms.append(residual_normR1) # Store residual norm # Check for convergence based on residual norm if residual_norm < self._tol: @@ -444,10 +444,10 @@ def _is_inverse_still_valid(self, inv, mat, name="", pre=None): I_approx = inv @ test_mat if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - I_exact = np.eye(test_mat.shape[0]) - if not np.allclose(I_approx, I_exact, atol=1e-6): + I_exact = xp.eye(test_mat.shape[0]) + if not xp.allclose(I_approx, I_exact, atol=1e-6): diff = I_approx - I_exact - max_abs = np.abs(diff).max() + max_abs = xp.abs(diff).max() print(f"{name} inverse is NOT valid anymore. Max diff: {max_abs:.2e}") return False print(f"{name} inverse is still valid.") @@ -455,7 +455,7 @@ def _is_inverse_still_valid(self, inv, mat, name="", pre=None): elif self._method_to_solve == "ScipySparse": I_exact = sc.sparse.identity(I_approx.shape[0], format=I_approx.format) diff = (I_approx - I_exact).tocoo() - max_abs = np.abs(diff.data).max() if diff.nnz > 0 else 0.0 + max_abs = xp.abs(diff.data).max() if diff.nnz > 0 else 0.0 if max_abs > 1e-6: print(f"{name} inverse is NOT valid anymore.") @@ -468,12 +468,12 @@ def _is_inverse_still_valid(self, inv, mat, name="", pre=None): def _compute_inverse(self, mat, which="matrix"): print(f"Computing inverse for {which} using method {self._method_to_solve}") if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - return np.linalg.inv(mat) + return xp.linalg.inv(mat) elif self._method_to_solve == "ScipySparse": return sc.sparse.linalg.inv(mat) elif self._method_to_solve == "SparseSolver": solver = SparseSolver(mat) - return solver.solve(np.eye(mat.shape[0])) + return solver.solve(xp.eye(mat.shape[0])) else: raise ValueError(f"Unknown solver method {self._method_to_solve}") @@ -481,14 +481,14 @@ def _spectral_analysis(self): # Spectral analysis # A11 before if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA11_before, eigvecs_before = np.linalg.eig(self._A[0]) - condA11_before = np.linalg.cond(self._A[0]) + eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0]) + condA11_before = xp.linalg.cond(self._A[0]) elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA11_before, eigvecs_before = np.linalg.eig(self._A[0].toarray()) - condA11_before = np.linalg.cond(self._A[0].toarray()) + eigvalsA11_before, eigvecs_before = xp.linalg.eig(self._A[0].toarray()) + condA11_before = xp.linalg.cond(self._A[0].toarray()) maxbeforeA11 = max(eigvalsA11_before) - maxbeforeA11_abs = np.max(np.abs(eigvalsA11_before)) - minbeforeA11_abs = np.min(np.abs(eigvalsA11_before)) + maxbeforeA11_abs = xp.max(xp.abs(eigvalsA11_before)) + minbeforeA11_abs = xp.min(xp.abs(eigvalsA11_before)) minbeforeA11 = min(eigvalsA11_before) specA11_bef = maxbeforeA11 / minbeforeA11 specA11_bef_abs = maxbeforeA11_abs / minbeforeA11_abs @@ -501,14 +501,14 @@ def _spectral_analysis(self): # A22 before if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA22_before, eigvecs_before = np.linalg.eig(self._A[1]) - condA22_before = np.linalg.cond(self._A[1]) + eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1]) + condA22_before = xp.linalg.cond(self._A[1]) elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA22_before, eigvecs_before = np.linalg.eig(self._A[1].toarray()) - condA22_before = np.linalg.cond(self._A[1].toarray()) + eigvalsA22_before, eigvecs_before = xp.linalg.eig(self._A[1].toarray()) + condA22_before = xp.linalg.cond(self._A[1].toarray()) maxbeforeA22 = max(eigvalsA22_before) - maxbeforeA22_abs = np.max(np.abs(eigvalsA22_before)) - minbeforeA22_abs = np.min(np.abs(eigvalsA22_before)) + maxbeforeA22_abs = xp.max(xp.abs(eigvalsA22_before)) + minbeforeA22_abs = xp.min(xp.abs(eigvalsA22_before)) minbeforeA22 = min(eigvalsA22_before) specA22_bef = maxbeforeA22 / minbeforeA22 specA22_bef_abs = maxbeforeA22_abs / minbeforeA22_abs @@ -523,13 +523,13 @@ def _spectral_analysis(self): if self._preconditioner == True: # A11 after preconditioning with its inverse if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA11_after_prec, eigvecs_after = np.linalg.eig(self._A11npinv @ self._A[0]) # Implement this + eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig(self._A11npinv @ self._A[0]) # Implement this elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA11_after_prec, eigvecs_after = np.linalg.eig((self._A11npinv @ self._A[0]).toarray()) + eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig((self._A11npinv @ self._A[0]).toarray()) maxafterA11_prec = max(eigvalsA11_after_prec) minafterA11_prec = min(eigvalsA11_after_prec) - maxafterA11_abs_prec = np.max(np.abs(eigvalsA11_after_prec)) - minafterA11_abs_prec = np.min(np.abs(eigvalsA11_after_prec)) + maxafterA11_abs_prec = xp.max(xp.abs(eigvalsA11_after_prec)) + minafterA11_abs_prec = xp.min(xp.abs(eigvalsA11_after_prec)) specA11_aft_prec = maxafterA11_prec / minafterA11_prec specA11_aft_abs_prec = maxafterA11_abs_prec / minafterA11_abs_prec # print(f'{maxafterA11_prec = }') @@ -541,15 +541,15 @@ def _spectral_analysis(self): # A22 after preconditioning with its inverse if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - eigvalsA22_after_prec, eigvecs_after = np.linalg.eig(self._A22npinv @ self._A[1]) # Implement this - condA22_after = np.linalg.cond(self._A22npinv @ self._A[1]) + eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig(self._A22npinv @ self._A[1]) # Implement this + condA22_after = xp.linalg.cond(self._A22npinv @ self._A[1]) elif self._method_to_solve in ("SparseSolver", "ScipySparse"): - eigvalsA22_after_prec, eigvecs_after = np.linalg.eig((self._A22npinv @ self._A[1]).toarray()) - condA22_after = np.linalg.cond((self._A22npinv @ self._A[1]).toarray()) + eigvalsA22_after_prec, eigvecs_after = xp.linalg.eig((self._A22npinv @ self._A[1]).toarray()) + condA22_after = xp.linalg.cond((self._A22npinv @ self._A[1]).toarray()) maxafterA22_prec = max(eigvalsA22_after_prec) minafterA22_prec = min(eigvalsA22_after_prec) - maxafterA22_abs_prec = np.max(np.abs(eigvalsA22_after_prec)) - minafterA22_abs_prec = np.min(np.abs(eigvalsA22_after_prec)) + maxafterA22_abs_prec = xp.max(xp.abs(eigvalsA22_after_prec)) + minafterA22_abs_prec = xp.min(xp.abs(eigvalsA22_after_prec)) specA22_aft_prec = maxafterA22_prec / minafterA22_prec specA22_aft_abs_prec = maxafterA22_abs_prec / minafterA22_abs_prec # print(f'{maxafterA22_prec = }') diff --git a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py index 8ab0daa03..9f196411e 100644 --- a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py +++ b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py @@ -29,7 +29,7 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m from struphy.geometry import domains from struphy.initial import perturbations from struphy.linear_algebra.saddle_point import SaddlePointSolver - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -132,12 +132,12 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m A11np = M2np / dt + nu * (Dnp.T @ M3np @ Dnp + S21np.T @ Cnp.T @ M2np @ Cnp @ S21np) - M2Bnp if method_to_solve in ("DirectNPInverse", "InexactNPInverse"): A22np = ( - stab_sigma * np.identity(A11np.shape[0]) + stab_sigma * xp.identity(A11np.shape[0]) + nue * (Dnp.T @ M3np @ Dnp + S21np.T @ Cnp.T @ M2np @ Cnp @ S21np) + M2Bnp ) # Preconditioner - _A22np_pre = stab_sigma * np.identity(A22np.shape[0]) # + nue*(Dnp.T @ M3np @ Dnp) + _A22np_pre = stab_sigma * xp.identity(A22np.shape[0]) # + nue*(Dnp.T @ M3np @ Dnp) _A11np_pre = M2np / dt # + nu * (Dnp.T @ M3np @ Dnp) elif method_to_solve in ("SparseSolver", "ScipySparse"): A22np = ( @@ -201,9 +201,9 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m - (B[0, 1].T).dot(y1_rdm) ) TestDiv = -B1.dot(x1) + B2.dot(x2) - RestDiv = np.linalg.norm(TestDiv.toarray()) - RestA = np.linalg.norm(TestA.toarray()) - RestAe = np.linalg.norm(TestAe.toarray()) + RestDiv = xp.linalg.norm(TestDiv.toarray()) + RestA = xp.linalg.norm(TestA.toarray()) + RestAe = xp.linalg.norm(TestAe.toarray()) print(f"{RestA =}") print(f"{RestAe =}") print(f"{RestDiv =}") @@ -218,10 +218,10 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m - (nue * (Dnp.T @ M3np @ Dnp + S21np.T @ Cnp.T @ M2np @ Cnp @ S21np) + M2Bnp).dot(x2np) - B2np.T.dot(ynp) ) - RestAnp = np.linalg.norm(TestAnp) - RestAenp = np.linalg.norm(TestAenp) + RestAnp = xp.linalg.norm(TestAnp) + RestAenp = xp.linalg.norm(TestAenp) TestDivnp = -B1np.dot(x1np) + B2np.dot(x2np) - RestDivnp = np.linalg.norm(TestDivnp) + RestDivnp = xp.linalg.norm(TestDivnp) print(f"{RestAnp =}") print(f"{RestAenp =}") print(f"{RestDivnp =}") @@ -296,22 +296,22 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m elapsed_time = end_time - start_time print(f"Method execution time: {elapsed_time:.6f} seconds") - if isinstance(x_uzawa[0], np.ndarray): - # Output as np.ndarray + if isinstance(x_uzawa[0], xp.ndarray): + # Output as xp.ndarray Rx1 = x1np - x_uzawa[0] Rx2 = x2np - x_uzawa[1] Ry = ynp - y_uzawa - residualx_normx1 = np.linalg.norm(Rx1) - residualx_normx2 = np.linalg.norm(Rx2) - residualy_norm = np.linalg.norm(Ry) + residualx_normx1 = xp.linalg.norm(Rx1) + residualx_normx2 = xp.linalg.norm(Rx2) + residualy_norm = xp.linalg.norm(Ry) TestRest1 = F1np - A11np.dot(x_uzawa[0]) - B1np.T.dot(y_uzawa) - TestRest1val = np.max(abs(TestRest1)) + TestRest1val = xp.max(abs(TestRest1)) Testoldy1 = F1np - A11np.dot(x_uzawa[0]) - B1np.T.dot(ynp) - Testoldy1val = np.max(abs(Testoldy1)) + Testoldy1val = xp.max(abs(Testoldy1)) TestRest2 = F2np - A22np.dot(x_uzawa[1]) - B2np.T.dot(y_uzawa) - TestRest2val = np.max(abs(TestRest2)) + TestRest2val = xp.max(abs(TestRest2)) Testoldy2 = F2np - A22np.dot(x_uzawa[1]) - B2np.T.dot(ynp) - Testoldy2val = np.max(abs(Testoldy2)) + Testoldy2val = xp.max(abs(Testoldy2)) print(f"{TestRest1val =}") print(f"{TestRest2val =}") print(f"{Testoldy1val =}") @@ -329,18 +329,18 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m Rx1 = x1 - x_uzawa[0] Rx2 = x2 - x_uzawa[1] Ry = y1_rdm - y_uzawa - residualx_normx1 = np.linalg.norm(Rx1.toarray()) - residualx_normx2 = np.linalg.norm(Rx2.toarray()) - residualy_norm = np.linalg.norm(Ry.toarray()) + residualx_normx1 = xp.linalg.norm(Rx1.toarray()) + residualx_normx2 = xp.linalg.norm(Rx2.toarray()) + residualy_norm = xp.linalg.norm(Ry.toarray()) TestRest1 = F1 - A11.dot(x_uzawa[0]) - B1T.dot(y_uzawa) - TestRest1val = np.max(abs(TestRest1.toarray())) + TestRest1val = xp.max(abs(TestRest1.toarray())) Testoldy1 = F1 - A11.dot(x_uzawa[0]) - B1T.dot(y1_rdm) - Testoldy1val = np.max(abs(Testoldy1.toarray())) + Testoldy1val = xp.max(abs(Testoldy1.toarray())) TestRest2 = F2 - A22.dot(x_uzawa[1]) - B2T.dot(y_uzawa) - TestRest2val = np.max(abs(TestRest2.toarray())) + TestRest2val = xp.max(abs(TestRest2.toarray())) Testoldy2 = F2 - A22.dot(x_uzawa[1]) - B2T.dot(y1_rdm) - Testoldy2val = np.max(abs(Testoldy2.toarray())) + Testoldy2val = xp.max(abs(Testoldy2.toarray())) # print(f"{TestRest1val =}") # print(f"{TestRest2val =}") # print(f"{Testoldy1val =}") @@ -375,13 +375,13 @@ def _plot_velocity(data_reshaped): import matplotlib import matplotlib.pyplot as plt - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp matplotlib.use("Agg") - x = np.linspace(0, 1, 30) - y = np.linspace(0, 1, 30) - X, Y = np.meshgrid(x, y) + x = xp.linspace(0, 1, 30) + y = xp.linspace(0, 1, 30) + X, Y = xp.meshgrid(x, y) plt.figure(figsize=(6, 5)) plt.imshow(data_reshaped.T, cmap="viridis", origin="lower", extent=[0, 1, 0, 1]) diff --git a/src/struphy/linear_algebra/tests/test_stencil_dot_kernels.py b/src/struphy/linear_algebra/tests/test_stencil_dot_kernels.py index f91d6872d..0475b418a 100644 --- a/src/struphy/linear_algebra/tests/test_stencil_dot_kernels.py +++ b/src/struphy/linear_algebra/tests/test_stencil_dot_kernels.py @@ -19,7 +19,7 @@ def test_1d(Nel, p, spl_kind, domain_ind, codomain_ind): from struphy.feec.psydac_derham import Derham from struphy.linear_algebra.stencil_dot_kernels import matvec_1d_kernel - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" @@ -78,8 +78,8 @@ def test_1d(Nel, p, spl_kind, domain_ind, codomain_ind): mat_pre._data[p_out + i_loc, d1] = m - i # random vector - # np.random.seed(123) - x[s_in : e_in + 1] = np.random.rand(domain.coeff_space.npts[0]) + # xp.random.seed(123) + x[s_in : e_in + 1] = xp.random.rand(domain.coeff_space.npts[0]) if rank == 0: print(f"spl_kind={spl_kind}") @@ -118,8 +118,8 @@ def test_1d(Nel, p, spl_kind, domain_ind, codomain_ind): print("\nout_ker=", out_ker._data) print("\nout_pre=", out_pre._data) - assert np.allclose(out_ker._data, out._data) - assert np.allclose(out_pre._data, out._data) + assert xp.allclose(out_ker._data, out._data) + assert xp.allclose(out_pre._data, out._data) @pytest.mark.parametrize("Nel", [[12, 16, 20]]) @@ -140,7 +140,7 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): from struphy.feec.psydac_derham import Derham from struphy.linear_algebra.stencil_dot_kernels import matvec_3d_kernel - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" @@ -177,16 +177,16 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): x = StencilVector(domain.coeff_space) out_ker = StencilVector(codomain.coeff_space) - s_out = np.array(mat.codomain.starts) - e_out = np.array(mat.codomain.ends) - p_out = np.array(mat.codomain.pads) - s_in = np.array(mat.domain.starts) - e_in = np.array(mat.domain.ends) - p_in = np.array(mat.domain.pads) + s_out = xp.array(mat.codomain.starts) + e_out = xp.array(mat.codomain.ends) + p_out = xp.array(mat.codomain.pads) + s_in = xp.array(mat.domain.starts) + e_in = xp.array(mat.domain.ends) + p_in = xp.array(mat.domain.pads) # random matrix - np.random.seed(123) - tmp1 = np.random.rand(*codomain.coeff_space.npts, *[2 * q + 1 for q in p]) + xp.random.seed(123) + tmp1 = xp.random.rand(*codomain.coeff_space.npts, *[2 * q + 1 for q in p]) mat[ s_out[0] : e_out[0] + 1, s_out[1] : e_out[1] + 1, @@ -207,7 +207,7 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): ] # random vector - tmp2 = np.random.rand(*domain.coeff_space.npts) + tmp2 = xp.random.rand(*domain.coeff_space.npts) x[ s_in[0] : e_in[0] + 1, s_in[1] : e_in[1] + 1, @@ -226,7 +226,7 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): # kernel matvec add = [int(end_in >= end_out) for end_in, end_out in zip(mat.domain.ends, mat.codomain.ends)] - add = np.array(add) + add = xp.array(add) matvec_3d_kernel(mat._data, x._data, out_ker._data, s_in, p_in, add, s_out, e_out, p_out) # precompiled .dot @@ -253,12 +253,12 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): print("\nout_ker[2]=", out_ker._data[p_out[0], p_out[1], :]) print("\nout_pre[2]=", out_pre._data[p_out[0], p_out[1], :]) - assert np.allclose( + assert xp.allclose( out_ker[s_out[0] : e_out[0] + 1, s_out[1] : e_out[1] + 1, s_out[2] : e_out[2] + 1], out[s_out[0] : e_out[0] + 1, s_out[1] : e_out[1] + 1, s_out[2] : e_out[2] + 1], ) - assert np.allclose( + assert xp.allclose( out_pre[s_out[0] : e_out[0] + 1, s_out[1] : e_out[1] + 1, s_out[2] : e_out[2] + 1], out[s_out[0] : e_out[0] + 1, s_out[1] : e_out[1] + 1, s_out[2] : e_out[2] + 1], ) diff --git a/src/struphy/linear_algebra/tests/test_stencil_transpose_kernels.py b/src/struphy/linear_algebra/tests/test_stencil_transpose_kernels.py index 0265ba741..486e302f5 100644 --- a/src/struphy/linear_algebra/tests/test_stencil_transpose_kernels.py +++ b/src/struphy/linear_algebra/tests/test_stencil_transpose_kernels.py @@ -19,7 +19,7 @@ def test_1d(Nel, p, spl_kind, domain_ind, codomain_ind): from struphy.feec.psydac_derham import Derham from struphy.linear_algebra.stencil_transpose_kernels import transpose_1d_kernel - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" @@ -112,8 +112,8 @@ def test_1d(Nel, p, spl_kind, domain_ind, codomain_ind): print("\nmatT_pre=", matT_pre._data) print("\nmatT_pre.toarray=\n", matT_pre.toarray()) - assert np.allclose(matT_ker[s_in : e_in + 1, :], matT[s_in : e_in + 1, :]) - assert np.allclose(matT_pre[s_in : e_in + 1, :], matT[s_in : e_in + 1, :]) + assert xp.allclose(matT_ker[s_in : e_in + 1, :], matT[s_in : e_in + 1, :]) + assert xp.allclose(matT_pre[s_in : e_in + 1, :], matT[s_in : e_in + 1, :]) @pytest.mark.parametrize("Nel", [[12, 16, 20]]) @@ -134,7 +134,7 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): from struphy.feec.psydac_derham import Derham from struphy.linear_algebra.stencil_transpose_kernels import transpose_3d_kernel - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" @@ -170,16 +170,16 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): mat_pre = StencilMatrix(domain.coeff_space, codomain.coeff_space, backend=PSYDAC_BACKEND_GPYCCEL, precompiled=True) matT_ker = StencilMatrix(codomain.coeff_space, domain.coeff_space) - s_out = np.array(mat.codomain.starts) - e_out = np.array(mat.codomain.ends) - p_out = np.array(mat.codomain.pads) - s_in = np.array(mat.domain.starts) - e_in = np.array(mat.domain.ends) - p_in = np.array(mat.domain.pads) + s_out = xp.array(mat.codomain.starts) + e_out = xp.array(mat.codomain.ends) + p_out = xp.array(mat.codomain.pads) + s_in = xp.array(mat.domain.starts) + e_in = xp.array(mat.domain.ends) + p_in = xp.array(mat.domain.pads) # random matrix - np.random.seed(123) - tmp1 = np.random.rand(*codomain.coeff_space.npts, *[2 * q + 1 for q in p]) + xp.random.seed(123) + tmp1 = xp.random.rand(*codomain.coeff_space.npts, *[2 * q + 1 for q in p]) mat[ s_out[0] : e_out[0] + 1, s_out[1] : e_out[1] + 1, @@ -208,7 +208,7 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): # kernel transpose add = [int(end_out >= end_in) for end_in, end_out in zip(mat.domain.ends, mat.codomain.ends)] - add = np.array(add) + add = xp.array(add) transpose_3d_kernel(mat._data, matT_ker._data, s_out, p_out, add, s_in, e_in, p_in) # precompiled transpose @@ -237,12 +237,12 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): print("\nmatT_ker[2]=", matT_ker._data[p_in[0], p_in[1], :, 1, 1, :]) print("\nmatT_pre[2]=", matT_pre._data[p_in[0], p_in[1], :, 1, 1, :]) - assert np.allclose( + assert xp.allclose( matT_ker[s_in[0] : e_in[0] + 1, s_in[1] : e_in[1] + 1, s_in[2] : e_in[2] + 1], matT[s_in[0] : e_in[0] + 1, s_in[1] : e_in[1] + 1, s_in[2] : e_in[2] + 1], ) - assert np.allclose( + assert xp.allclose( matT_pre[s_in[0] : e_in[0] + 1, s_in[1] : e_in[1] + 1, s_in[2] : e_in[2] + 1], matT[s_in[0] : e_in[0] + 1, s_in[1] : e_in[1] + 1, s_in[2] : e_in[2] + 1], ) diff --git a/src/struphy/main.py b/src/struphy/main.py index e40059537..3b06b671e 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -38,7 +38,7 @@ from struphy.profiling.profiling import ProfileManager from struphy.topology import grids from struphy.topology.grids import TensorProductGrid -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml @@ -235,9 +235,9 @@ def run( # store geometry vtk if rank == 0: grids_log = [ - np.linspace(1e-6, 1.0, 32), - np.linspace(0.0, 1.0, 32), - np.linspace(0.0, 1.0, 32), + xp.linspace(1e-6, 1.0, 32), + xp.linspace(0.0, 1.0, 32), + xp.linspace(0.0, 1.0, 32), ] tmp = model.domain(*grids_log) @@ -262,9 +262,9 @@ def run( # time quantities (current time value, value in seconds and index) time_state = {} - time_state["value"] = np.zeros(1, dtype=float) - time_state["value_sec"] = np.zeros(1, dtype=float) - time_state["index"] = np.zeros(1, dtype=int) + time_state["value"] = xp.zeros(1, dtype=float) + time_state["value_sec"] = xp.zeros(1, dtype=float) + time_state["index"] = xp.zeros(1, dtype=int) # add time quantities to data object for saving for key, val in time_state.items(): @@ -476,7 +476,7 @@ def pproc( file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") # save time grid at which post-processing data is created - np.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) + xp.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) if "feec" in file.keys(): exist_fields = True @@ -635,28 +635,28 @@ def __init__(self, path: str): self._f = {} self._spline_values = {} self._n_sph = {} - self.grids_log: list[np.ndarray] = None - self.grids_phy: list[np.ndarray] = None - self.t_grid: np.ndarray = None + self.grids_log: list[xp.ndarray] = None + self.grids_phy: list[xp.ndarray] = None + self.t_grid: xp.ndarray = None @property - def orbits(self) -> dict[str, np.ndarray]: + def orbits(self) -> dict[str, xp.ndarray]: """Keys: species name. Values: 3d arrays indexed by (n, p, a), where 'n' is the time index, 'p' the particle index and 'a' the attribute index.""" return self._orbits @property - def f(self) -> dict[str, dict[str, dict[str, np.ndarray]]]: - """Keys: species name. Values: dicts of slice names ('e1_v1' etc.) holding dicts of corresponding np.arrays for plotting.""" + def f(self) -> dict[str, dict[str, dict[str, xp.ndarray]]]: + """Keys: species name. Values: dicts of slice names ('e1_v1' etc.) holding dicts of corresponding xp.arrays for plotting.""" return self._f @property - def spline_values(self) -> dict[str, dict[str, np.ndarray]]: + def spline_values(self) -> dict[str, dict[str, xp.ndarray]]: """Keys: species name. Values: dicts of variable names with values being 3d arrays on the grid.""" return self._spline_values @property - def n_sph(self) -> dict[str, dict[str, dict[str, np.ndarray]]]: - """Keys: species name. Values: dicts of view names ('view_0' etc.) holding dicts of corresponding np.arrays for plotting.""" + def n_sph(self) -> dict[str, dict[str, dict[str, xp.ndarray]]]: + """Keys: species name. Values: dicts of view names ('view_0' etc.) holding dicts of corresponding xp.arrays for plotting.""" return self._n_sph @property @@ -704,7 +704,7 @@ def load_data(path: str) -> SimData: simdata = SimData(path) # load time grid - simdata.t_grid = np.load(os.path.join(path_pproc, "t_grid.npy")) + simdata.t_grid = xp.load(os.path.join(path_pproc, "t_grid.npy")) # data paths path_fields = os.path.join(path_pproc, "fields_data") @@ -755,9 +755,9 @@ def load_data(path: str) -> SimData: # print(f"{file = }") if ".npy" in file: step = int(file.split(".")[0].split("_")[-1]) - tmp = np.load(os.path.join(path_dat, file)) + tmp = xp.load(os.path.join(path_dat, file)) if n == 0: - simdata._orbits[spec] = np.zeros((Nt, *tmp.shape), dtype=float) + simdata._orbits[spec] = xp.zeros((Nt, *tmp.shape), dtype=float) simdata._orbits[spec][step] = tmp n += 1 @@ -772,7 +772,7 @@ def load_data(path: str) -> SimData: # print(f"{files = }") for file in files: name = file.split(".")[0] - tmp = np.load(os.path.join(path_dat, sli, file)) + tmp = xp.load(os.path.join(path_dat, sli, file)) # print(f"{name = }") simdata._f[spec][sli][name] = tmp @@ -787,7 +787,7 @@ def load_data(path: str) -> SimData: # print(f"{files = }") for file in files: name = file.split(".")[0] - tmp = np.load(os.path.join(path_dat, sli, file)) + tmp = xp.load(os.path.join(path_dat, sli, file)) # print(f"{name = }") simdata._n_sph[spec][sli][name] = tmp diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 1e5ebff3d..d910cc41b 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -35,7 +35,7 @@ from struphy.profiling.profiling import ProfileManager from struphy.propagators.base import Propagator from struphy.topology.grids import TensorProductGrid -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml, read_state @@ -478,7 +478,7 @@ def add_scalar(self, name: str, variable: PICVariable | SPHVariable = None, comp self._scalar_quantities = {} self._scalar_quantities[name] = { - "value": np.empty(1, dtype=float), + "value": xp.empty(1, dtype=float), "variable": variable, "compute": compute, "summands": summands, @@ -524,7 +524,7 @@ def update_scalar(self, name, value=None): assert isinstance(value, float) # Create a numpy array to hold the scalar value - value_array = np.array([value], dtype=np.float64) + value_array = xp.array([value], dtype=xp.float64) # Perform MPI operations based on the compute flags if "sum_world" in compute_operations and not isinstance(MPI, MockMPI): @@ -562,7 +562,7 @@ def update_scalar(self, name, value=None): if "divide_n_mks" in compute_operations: # Initialize the total number of markers - n_mks_tot = np.array([variable.particles.Np]) + n_mks_tot = xp.array([variable.particles.Np]) value_array /= n_mks_tot # Update the scalar value @@ -728,11 +728,11 @@ def update_markers_to_be_saved(self): assert isinstance(obj, Particles) if var.n_to_save > 0: - markers_on_proc = np.logical_and( + markers_on_proc = xp.logical_and( obj.markers[:, -1] >= 0.0, obj.markers[:, -1] < var.n_to_save, ) - n_markers_on_proc = np.count_nonzero(markers_on_proc) + n_markers_on_proc = xp.count_nonzero(markers_on_proc) var.saved_markers[:] = -1.0 var.saved_markers[:n_markers_on_proc] = obj.markers[markers_on_proc] @@ -776,7 +776,7 @@ def update_distr_functions(self): h2 = 1 / obj.boxes_per_dim[1] h3 = 1 / obj.boxes_per_dim[2] - ndim = np.count_nonzero([d > 1 for d in obj.boxes_per_dim]) + ndim = xp.count_nonzero([d > 1 for d in obj.boxes_per_dim]) if ndim == 0: kernel_type = "gaussian_3d" else: @@ -800,7 +800,7 @@ def print_scalar_quantities(self): sq_str = "" for key, scalar_dict in self._scalar_quantities.items(): val = scalar_dict["value"] - assert not np.isnan(val[0]), f"Scalar {key} is {val[0]}." + assert not xp.isnan(val[0]), f"Scalar {key} is {val[0]}." sq_str += key + ": {:14.11f}".format(val[0]) + " " print(sq_str) @@ -1554,15 +1554,15 @@ def compute_plasma_params(self, verbose=True): units_affix["epsilon"] = "" h = 1 / 20 - eta1 = np.linspace(h / 2.0, 1.0 - h / 2.0, 20) - eta2 = np.linspace(h / 2.0, 1.0 - h / 2.0, 20) - eta3 = np.linspace(h / 2.0, 1.0 - h / 2.0, 20) + eta1 = xp.linspace(h / 2.0, 1.0 - h / 2.0, 20) + eta2 = xp.linspace(h / 2.0, 1.0 - h / 2.0, 20) + eta3 = xp.linspace(h / 2.0, 1.0 - h / 2.0, 20) ## global parameters # plasma volume (hat x^3) det_tmp = self.domain.jacobian_det(eta1, eta2, eta3) - vol1 = np.mean(np.abs(det_tmp)) + vol1 = xp.mean(xp.abs(det_tmp)) # plasma volume (m⁻³) plasma_volume = vol1 * self.units.x**3 # transit length (m) @@ -1571,13 +1571,13 @@ def compute_plasma_params(self, verbose=True): if isinstance(self.equil, FluidEquilibriumWithB): B_tmp = self.equil.absB0(eta1, eta2, eta3) else: - B_tmp = np.zeros((eta1.size, eta2.size, eta3.size)) - magnetic_field = np.mean(B_tmp * np.abs(det_tmp)) / vol1 * self.units.B - B_max = np.max(B_tmp) * self.units.B - B_min = np.min(B_tmp) * self.units.B + B_tmp = xp.zeros((eta1.size, eta2.size, eta3.size)) + magnetic_field = xp.mean(B_tmp * xp.abs(det_tmp)) / vol1 * self.units.B + B_max = xp.max(B_tmp) * self.units.B + B_min = xp.min(B_tmp) * self.units.B if magnetic_field < 1e-14: - magnetic_field = np.nan + magnetic_field = xp.nan # print("\n+++++++ WARNING +++++++ magnetic field is zero - set to nan !!") if verbose and MPI.COMM_WORLD.Get_rank() == 0: @@ -1617,13 +1617,13 @@ def compute_plasma_params(self, verbose=True): # self._pparams[species]["charge"] = val["params"]["phys_params"]["Z"] * e # # density (m⁻³) # self._pparams[species]["density"] = ( - # np.mean( + # xp.mean( # self.equil.n0( # eta1, # eta2, # eta3, # ) - # * np.abs(det_tmp), + # * xp.abs(det_tmp), # ) # * self.units.x ** 3 # / plasma_volume @@ -1631,13 +1631,13 @@ def compute_plasma_params(self, verbose=True): # ) # # pressure (bar) # self._pparams[species]["pressure"] = ( - # np.mean( + # xp.mean( # self.equil.p0( # eta1, # eta2, # eta3, # ) - # * np.abs(det_tmp), + # * xp.abs(det_tmp), # ) # * self.units.x ** 3 # / plasma_volume @@ -1648,7 +1648,7 @@ def compute_plasma_params(self, verbose=True): # self._pparams[species]["kBT"] = self._pparams[species]["pressure"] * 1e5 / self._pparams[species]["density"] / e * 1e-3 # if len(self.kinetic) > 0: - # eta1mg, eta2mg, eta3mg = np.meshgrid( + # eta1mg, eta2mg, eta3mg = xp.meshgrid( # eta1, # eta2, # eta3, @@ -1694,11 +1694,11 @@ def compute_plasma_params(self, verbose=True): # # density (m⁻³) # self._pparams[species]["density"] = ( - # np.mean(tmp.n(psi) * np.abs(det_tmp)) * self.units.x ** 3 / plasma_volume * self.units.n + # xp.mean(tmp.n(psi) * xp.abs(det_tmp)) * self.units.x ** 3 / plasma_volume * self.units.n # ) # # thermal speed (m/s) # self._pparams[species]["v_th"] = ( - # np.mean(tmp.vth(psi) * np.abs(det_tmp)) * self.units.x ** 3 / plasma_volume * self.units.v + # xp.mean(tmp.vth(psi) * xp.abs(det_tmp)) * self.units.x ** 3 / plasma_volume * self.units.v # ) # # thermal energy (keV) # self._pparams[species]["kBT"] = self._pparams[species]["mass"] * self._pparams[species]["v_th"] ** 2 / e * 1e-3 @@ -1709,8 +1709,8 @@ def compute_plasma_params(self, verbose=True): # else: # # density (m⁻³) - # # self._pparams[species]['density'] = np.mean(tmp.n( - # # eta1mg, eta2mg, eta3mg) * np.abs(det_tmp)) * units['x']**3 / plasma_volume * units['n'] + # # self._pparams[species]['density'] = xp.mean(tmp.n( + # # eta1mg, eta2mg, eta3mg) * xp.abs(det_tmp)) * units['x']**3 / plasma_volume * units['n'] # self._pparams[species]["density"] = 99.0 # # thermal speeds (m/s) # vth = [] @@ -1718,11 +1718,11 @@ def compute_plasma_params(self, verbose=True): # vths = [99.0] # for k in range(len(vths)): # vth += [ - # vths[k] * np.abs(det_tmp) * self.units.x ** 3 / plasma_volume * self.units.v, + # vths[k] * xp.abs(det_tmp) * self.units.x ** 3 / plasma_volume * self.units.v, # ] # thermal_speed = 0.0 # for dir in range(val["obj"].vdim): - # # self._pparams[species]['vth' + str(dir + 1)] = np.mean(vth[dir]) + # # self._pparams[species]['vth' + str(dir + 1)] = xp.mean(vth[dir]) # self._pparams[species]["vth" + str(dir + 1)] = 99.0 # thermal_speed += self._pparams[species]["vth" + str(dir + 1)] # # TODO: here it is assumed that background density parameter is called "n", @@ -1741,11 +1741,11 @@ def compute_plasma_params(self, verbose=True): # for species in self._pparams: # # alfvén speed (m/s) - # self._pparams[species]["v_A"] = magnetic_field / np.sqrt( + # self._pparams[species]["v_A"] = magnetic_field / xp.sqrt( # mu0 * self._pparams[species]["mass"] * self._pparams[species]["density"], # ) # # thermal speed (m/s) - # self._pparams[species]["v_th"] = np.sqrt( + # self._pparams[species]["v_th"] = xp.sqrt( # self._pparams[species]["kBT"] * 1e3 * e / self._pparams[species]["mass"], # ) # # thermal frequency (Mrad/s) @@ -1754,7 +1754,7 @@ def compute_plasma_params(self, verbose=True): # self._pparams[species]["Omega_c"] = self._pparams[species]["charge"] * magnetic_field / self._pparams[species]["mass"] * 1e-6 # # plasma frequency (Mrad/s) # self._pparams[species]["Omega_p"] = ( - # np.sqrt( + # xp.sqrt( # self._pparams[species]["density"] * (self._pparams[species]["charge"]) ** 2 / eps0 / self._pparams[species]["mass"], # ) # * 1e-6 @@ -1764,7 +1764,7 @@ def compute_plasma_params(self, verbose=True): # # Larmor radius (m) # self._pparams[species]["rho_th"] = self._pparams[species]["v_th"] / (self._pparams[species]["Omega_c"] * 1e6) # # MHD length scale (m) - # self._pparams[species]["v_A/Omega_c"] = self._pparams[species]["v_A"] / (np.abs(self._pparams[species]["Omega_c"]) * 1e6) + # self._pparams[species]["v_A/Omega_c"] = self._pparams[species]["v_A"] / (xp.abs(self._pparams[species]["Omega_c"]) * 1e6) # # dim-less ratios # self._pparams[species]["rho_th/L"] = self._pparams[species]["rho_th"] / transit_length diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 2ddb72972..98983da4b 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -9,7 +9,7 @@ from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.polar.basic import PolarVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp rank = MPI.COMM_WORLD.Get_rank() @@ -577,7 +577,7 @@ def allocate_helpers(self): def f(e1, e2, e3): return 1 - f = np.vectorize(f) + f = xp.vectorize(f) self._integrator = projV3(f) self._energy_evaluator = InternalEnergyEvaluator(self.derham, self.propagators.variat_ent.options.gamma) @@ -777,7 +777,7 @@ def allocate_helpers(self): def f(e1, e2, e3): return 1 - f = np.vectorize(f) + f = xp.vectorize(f) self._integrator = projV3(f) self._energy_evaluator = InternalEnergyEvaluator(self.derham, self.propagators.variat_ent.options.gamma) @@ -996,7 +996,7 @@ def allocate_helpers(self): def f(e1, e2, e3): return 1 - f = np.vectorize(f) + f = xp.vectorize(f) self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() @@ -1195,7 +1195,7 @@ def allocate_helpers(self): def f(e1, e2, e3): return 1 - f = np.vectorize(f) + f = xp.vectorize(f) self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() @@ -1433,7 +1433,7 @@ def allocate_helpers(self): def f(e1, e2, e3): return 1 - f = np.vectorize(f) + f = xp.vectorize(f) self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() @@ -1659,7 +1659,7 @@ def allocate_helpers(self): def f(e1, e2, e3): return 1 - f = np.vectorize(f) + f = xp.vectorize(f) self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() @@ -1862,7 +1862,7 @@ def allocate_helpers(self): def f(e1, e2, e3): return 1 - f = np.vectorize(f) + f = xp.vectorize(f) self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() @@ -2087,7 +2087,7 @@ def allocate_helpers(self): def f(e1, e2, e3): return 1 - f = np.vectorize(f) + f = xp.vectorize(f) self._integrator = projV3(f) self._ones = self.derham.Vh_pol["3"].zeros() diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index 70968d643..0d765981d 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -6,7 +6,7 @@ from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.polar.basic import PolarVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel rank = MPI.COMM_WORLD.Get_rank() @@ -236,8 +236,8 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("n_lost_particles", compute="from_particles", species="energetic_ions") # temporary vectors for scalar quantities: - self._tmp = np.empty(1, dtype=float) - self._n_lost_particles = np.empty(1, dtype=float) + self._tmp = xp.empty(1, dtype=float) + self._n_lost_particles = xp.empty(1, dtype=float) def update_scalar_quantities(self): # perturbed fields @@ -500,8 +500,8 @@ def __init__(self, params, comm, clone_config=None): # temporary vectors for scalar quantities self._tmp_u = self.derham.Vh["2"].zeros() self._tmp_b1 = self.derham.Vh["2"].zeros() - self._tmp = np.empty(1, dtype=float) - self._n_lost_particles = np.empty(1, dtype=float) + self._tmp = xp.empty(1, dtype=float) + self._n_lost_particles = xp.empty(1, dtype=float) def update_scalar_quantities(self): # perturbed fields @@ -715,10 +715,10 @@ def allocate_helpers(self): else: self._ones[:] = 1.0 - self._en_fv = np.empty(1, dtype=float) - self._en_fB = np.empty(1, dtype=float) - self._en_tot = np.empty(1, dtype=float) - self._n_lost_particles = np.empty(1, dtype=float) + self._en_fv = xp.empty(1, dtype=float) + self._en_fB = xp.empty(1, dtype=float) + self._en_tot = xp.empty(1, dtype=float) + self._n_lost_particles = xp.empty(1, dtype=float) self._PB = getattr(self.basis_ops, "PB") self._PBb = self._PB.codomain.zeros() @@ -771,7 +771,7 @@ def update_scalar_quantities(self): self.update_scalar("en_tot") # print number of lost particles - n_lost_markers = np.array(particles.n_lost_markers) + n_lost_markers = xp.array(particles.n_lost_markers) if self.derham.comm is not None: self.derham.comm.Allreduce( @@ -956,7 +956,7 @@ def __init__(self, params, comm, clone_config=None): hot_params = params["kinetic"]["hot_electrons"] # model parameters - self._alpha = np.abs( + self._alpha = xp.abs( self.equation_params["cold_electrons"]["alpha"], ) self._epsilon_cold = self.equation_params["cold_electrons"]["epsilon"] @@ -1022,7 +1022,7 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("en_tot") # temporaries - self._tmp = np.empty(1, dtype=float) + self._tmp = xp.empty(1, dtype=float) def initialize_from_params(self): """:meta private:""" @@ -1044,8 +1044,8 @@ def initialize_from_params(self): charge_accum() # Locally subtract mean charge for solvability with periodic bc - if np.all(charge_accum.vectors[0].space.periods): - charge_accum._vectors[0][:] -= np.mean( + if xp.all(charge_accum.vectors[0].space.periods): + charge_accum._vectors[0][:] -= xp.mean( charge_accum.vectors[0].toarray()[charge_accum.vectors[0].toarray() != 0], ) @@ -1082,7 +1082,7 @@ def update_scalar_quantities(self): * self._epsilon_hot / self._epsilon_cold / (2 * self.pointer["hot_electrons"].Np) - * np.dot( + * xp.dot( self.pointer["hot_electrons"].markers_wo_holes[:, 3] ** 2 + self.pointer["hot_electrons"].markers_wo_holes[:, 4] ** 2 + self.pointer["hot_electrons"].markers_wo_holes[:, 5] ** 2, diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 792bc14cb..791477419 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -7,7 +7,7 @@ from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.pic.accumulation.particles_to_grid import AccumulatorVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel rank = MPI.COMM_WORLD.Get_rank() @@ -144,7 +144,7 @@ def velocity_scale(self): return "light" def allocate_helpers(self): - self._tmp = np.empty(1, dtype=float) + self._tmp = xp.empty(1, dtype=float) def update_scalar_quantities(self): # e*M1*e/2 @@ -158,7 +158,7 @@ def update_scalar_quantities(self): self._tmp[0] = ( alpha**2 / (2 * particles.Np) - * np.dot( + * xp.dot( particles.markers_wo_holes[:, 3] ** 2 + particles.markers_wo_holes[:, 4] ** 2 + particles.markers_wo_holes[:, 5] ** 2, @@ -188,7 +188,7 @@ def allocate_propagators(self): # sanity check # self.pointer['species1'].show_distribution_function( - # [True] + [False]*5, [np.linspace(0, 1, 32)]) + # [True] + [False]*5, [xp.linspace(0, 1, 32)]) # accumulate charge density charge_accum = AccumulatorVector( @@ -435,7 +435,7 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("en_tot") # temporaries - self._tmp = np.empty(1, dtype=float) + self._tmp = xp.empty(1, dtype=float) def initialize_from_params(self): """:meta private:""" @@ -453,7 +453,7 @@ def initialize_from_params(self): # sanity check # self.pointer['species1'].show_distribution_function( - # [True] + [False]*5, [np.linspace(0, 1, 32)]) + # [True] + [False]*5, [xp.linspace(0, 1, 32)]) # accumulate charge density charge_accum = AccumulatorVector( @@ -500,7 +500,7 @@ def update_scalar_quantities(self): self._tmp[0] = ( self._alpha**2 / (2 * self.pointer["species1"].Np) - * np.dot( + * xp.dot( self.pointer["species1"].markers_wo_holes[:, 3] ** 2 + self.pointer["species1"].markers_wo_holes[:, 4] ** 2 + self.pointer["species1"].markers_wo_holes[:, 5] ** 2, @@ -685,7 +685,7 @@ def __init__(self, params, comm, clone_config=None, baseclass=False): self.alpha = self.equation_params["species1"]["alpha"] # allocate memory for evaluating f0 in energy computation - self._f0_values = np.zeros( + self._f0_values = xp.zeros( self.pointer["species1"].markers.shape[0], dtype=float, ) @@ -752,7 +752,7 @@ def __init__(self, params, comm, clone_config=None, baseclass=False): self.add_scalar("en_tot") # temporaries - self._tmp = np.empty(1, dtype=float) + self._tmp = xp.empty(1, dtype=float) self.en_E = 0.0 def initialize_from_params(self): @@ -808,7 +808,7 @@ def update_scalar_quantities(self): self.alpha**2 * self.vth**2 / (2 * self.pointer["species1"].Np) - * np.dot( + * xp.dot( self.pointer["species1"].weights ** 2, # w_p^2 self.pointer["species1"].sampling_density / self._f0_values[self.pointer["species1"].valid_mks], # s_{0,p} / f_{0,p} @@ -1137,7 +1137,7 @@ def __init__(self, params, comm, clone_config=None): self.add_scalar("en_tot") # MPI operations needed for scalar variables - self._tmp3 = np.empty(1, dtype=float) + self._tmp3 = xp.empty(1, dtype=float) self._e_field = self.derham.Vh["1"].zeros() def update_scalar_quantities(self): @@ -1158,7 +1158,7 @@ def update_scalar_quantities(self): self._tmp3[0] = ( 1 / self.pointer["ions"].Np - * np.sum( + * xp.sum( self.pointer["ions"].weights * self.pointer["ions"].velocities[:, 0] ** 2 / 2.0 + self.pointer["ions"].markers_wo_holes_and_ghost[:, 8], ) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 3c33d1b71..4ad4c14d6 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -13,7 +13,7 @@ LoadingParameters, WeightsParameters, ) -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Species(metaclass=ABCMeta): @@ -82,7 +82,7 @@ def __init__( con = ConstantsOfNature() # relevant frequencies - om_p = np.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_p = xp.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters diff --git a/src/struphy/models/tests/test_verif_EulerSPH.py b/src/struphy/models/tests/test_verif_EulerSPH.py index d16e76569..fa9b2bcfa 100644 --- a/src/struphy/models/tests/test_verif_EulerSPH.py +++ b/src/struphy/models/tests/test_verif_EulerSPH.py @@ -19,7 +19,7 @@ WeightsParameters, ) from struphy.topology import grids -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") @@ -144,7 +144,7 @@ def test_soundwave_1d(nx: int, plot_pts: int, do_plot: bool = False): plt.plot(x.squeeze(), n_sph[i, :, 0, 0], style, label=f"time={i * dt:4.2f}") plt.xlim(0, 2.5) plt.legend() - ax.set_xticks(np.linspace(0, 2.5, nx + 1)) + ax.set_xticks(xp.linspace(0, 2.5, nx + 1)) ax.xaxis.set_major_formatter(FormatStrFormatter("%.2f")) plt.grid(c="k") plt.xlabel("x") @@ -156,7 +156,7 @@ def test_soundwave_1d(nx: int, plot_pts: int, do_plot: bool = False): plt.show() - error = np.max(np.abs(n_sph[0] - n_sph[-1])) + error = xp.max(xp.abs(n_sph[0] - n_sph[-1])) print(f"SPH sound wave {error = }.") assert error < 6e-4 print("Assertion passed.") diff --git a/src/struphy/models/tests/test_verif_LinearMHD.py b/src/struphy/models/tests/test_verif_LinearMHD.py index faa59fc1a..3ed92aca0 100644 --- a/src/struphy/models/tests/test_verif_LinearMHD.py +++ b/src/struphy/models/tests/test_verif_LinearMHD.py @@ -11,7 +11,7 @@ from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time from struphy.kinetic_background import maxwellians from struphy.topology import grids -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "verification_tests") @@ -113,10 +113,10 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): ) # assert - vA = np.sqrt(Bsquare / n0) - v_alfven = vA * B0z / np.sqrt(Bsquare) + vA = xp.sqrt(Bsquare / n0) + v_alfven = vA * B0z / xp.sqrt(Bsquare) print(f"{v_alfven = }") - assert np.abs(coeffs[0][0] - v_alfven) < 0.07 + assert xp.abs(coeffs[0][0] - v_alfven) < 0.07 # second fft p_of_t = simdata.spline_values["mhd"]["pressure_log"] @@ -139,15 +139,15 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): # assert gamma = 5 / 3 - cS = np.sqrt(gamma * p0 / n0) + cS = xp.sqrt(gamma * p0 / n0) delta = (4 * B0z**2 * cS**2 * vA**2) / ((cS**2 + vA**2) ** 2 * Bsquare) - v_slow = np.sqrt(1 / 2 * (cS**2 + vA**2) * (1 - np.sqrt(1 - delta))) - v_fast = np.sqrt(1 / 2 * (cS**2 + vA**2) * (1 + np.sqrt(1 - delta))) + v_slow = xp.sqrt(1 / 2 * (cS**2 + vA**2) * (1 - xp.sqrt(1 - delta))) + v_fast = xp.sqrt(1 / 2 * (cS**2 + vA**2) * (1 + xp.sqrt(1 - delta))) print(f"{v_slow = }") print(f"{v_fast = }") - assert np.abs(coeffs[0][0] - v_slow) < 0.05 - assert np.abs(coeffs[1][0] - v_fast) < 0.19 + assert xp.abs(coeffs[0][0] - v_slow) < 0.05 + assert xp.abs(coeffs[1][0] - v_fast) < 0.19 if __name__ == "__main__": diff --git a/src/struphy/models/tests/test_verif_Maxwell.py b/src/struphy/models/tests/test_verif_Maxwell.py index 0cb4e778e..dc1a4794c 100644 --- a/src/struphy/models/tests/test_verif_Maxwell.py +++ b/src/struphy/models/tests/test_verif_Maxwell.py @@ -14,7 +14,7 @@ from struphy.kinetic_background import maxwellians from struphy.models.toy import Maxwell from struphy.topology import grids -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") @@ -97,7 +97,7 @@ def test_light_wave_1d(algo: str, do_plot: bool = False): # assert c_light_speed = 1.0 - assert np.abs(coeffs[0][0] - c_light_speed) < 0.02 + assert xp.abs(coeffs[0][0] - c_light_speed) < 0.02 @pytest.mark.mpi(min_size=4) @@ -189,32 +189,32 @@ def test_coaxial(do_plot: bool = False): def B_z(X, Y, Z, m, t): """Magnetic field in z direction of coaxial cabel""" r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return (jv(m, r) - 0.28 * yn(m, r)) * np.cos(m * theta - t) + theta = xp.arctan2(Y, X) + return (jv(m, r) - 0.28 * yn(m, r)) * xp.cos(m * theta - t) def E_r(X, Y, Z, m, t): """Electrical field in radial direction of coaxial cabel""" r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return -m / r * (jv(m, r) - 0.28 * yn(m, r)) * np.cos(m * theta - t) + theta = xp.arctan2(Y, X) + return -m / r * (jv(m, r) - 0.28 * yn(m, r)) * xp.cos(m * theta - t) def E_theta(X, Y, Z, m, t): """Electrical field in azimuthal direction of coaxial cabel""" r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return ((m / r * jv(m, r) - jv(m + 1, r)) - 0.28 * (m / r * yn(m, r) - yn(m + 1, r))) * np.sin( + theta = xp.arctan2(Y, X) + return ((m / r * jv(m, r) - jv(m + 1, r)) - 0.28 * (m / r * yn(m, r) - yn(m + 1, r))) * xp.sin( m * theta - t ) def to_E_r(X, Y, E_x, E_y): r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return np.cos(theta) * E_x + np.sin(theta) * E_y + theta = xp.arctan2(Y, X) + return xp.cos(theta) * E_x + xp.sin(theta) * E_y def to_E_theta(X, Y, E_x, E_y): r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return -np.sin(theta) * E_x + np.cos(theta) * E_y + theta = xp.arctan2(Y, X) + return -xp.sin(theta) * E_x + xp.cos(theta) * E_y # plot if do_plot: @@ -247,13 +247,13 @@ def to_E_theta(X, Y, E_x, E_y): Bz_tend = b_field_phy[t_grid[-1]][2][:, :, 0] Bz_exact = B_z(X, Y, grids_phy[0], modes, t_grid[-1]) - error_Er = np.max(np.abs((to_E_r(X, Y, Ex_tend, Ey_tend) - Er_exact))) - error_Etheta = np.max(np.abs((to_E_theta(X, Y, Ex_tend, Ey_tend) - Etheta_exact))) - error_Bz = np.max(np.abs((Bz_tend - Bz_exact))) + error_Er = xp.max(xp.abs((to_E_r(X, Y, Ex_tend, Ey_tend) - Er_exact))) + error_Etheta = xp.max(xp.abs((to_E_theta(X, Y, Ex_tend, Ey_tend) - Etheta_exact))) + error_Bz = xp.max(xp.abs((Bz_tend - Bz_exact))) - rel_err_Er = error_Er / np.max(np.abs(Er_exact)) - rel_err_Etheta = error_Etheta / np.max(np.abs(Etheta_exact)) - rel_err_Bz = error_Bz / np.max(np.abs(Bz_exact)) + rel_err_Er = error_Er / xp.max(xp.abs(Er_exact)) + rel_err_Etheta = error_Etheta / xp.max(xp.abs(Etheta_exact)) + rel_err_Bz = error_Bz / xp.max(xp.abs(Bz_exact)) print("") assert rel_err_Bz < 0.0021, f"Assertion for magnetic field Maxwell failed: {rel_err_Bz = }" diff --git a/src/struphy/models/tests/test_verif_Poisson.py b/src/struphy/models/tests/test_verif_Poisson.py index 0264d4b6d..289886cb3 100644 --- a/src/struphy/models/tests/test_verif_Poisson.py +++ b/src/struphy/models/tests/test_verif_Poisson.py @@ -18,7 +18,7 @@ WeightsParameters, ) from struphy.topology import grids -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") @@ -59,7 +59,7 @@ def test_poisson_1d(do_plot=False): model = Poisson() # propagator options - omega = 2 * np.pi + omega = 2 * xp.pi model.propagators.source.options = model.propagators.source.Options(omega=omega) model.propagators.poisson.options = model.propagators.poisson.Options(rho=model.em_fields.source) @@ -71,9 +71,9 @@ def test_poisson_1d(do_plot=False): # analytical solution Lx = r1 - l1 - rhs_exact = lambda e1, e2, e3, t: amp * np.cos(l * 2 * np.pi / Lx * e1) * np.cos(omega * t) + rhs_exact = lambda e1, e2, e3, t: amp * xp.cos(l * 2 * xp.pi / Lx * e1) * xp.cos(omega * t) phi_exact = ( - lambda e1, e2, e3, t: amp / (l * 2 * np.pi / Lx) ** 2 * np.cos(l * 2 * np.pi / Lx * e1) * np.cos(omega * t) + lambda e1, e2, e3, t: amp / (l * 2 * xp.pi / Lx) ** 2 * xp.cos(l * 2 * xp.pi / Lx * e1) * xp.cos(omega * t) ) # start run @@ -116,7 +116,7 @@ def test_poisson_1d(do_plot=False): for i, t in enumerate(phi): phi_h = phi[t][0][:, 0, 0] phi_e = phi_exact(x, 0, 0, t) - new_err = np.abs(np.max(phi_h - phi_e)) / (amp / (l * 2 * np.pi / Lx) ** 2) + new_err = xp.abs(xp.max(phi_h - phi_e)) / (amp / (l * 2 * xp.pi / Lx) ** 2) if new_err > err: err = new_err @@ -125,7 +125,7 @@ def test_poisson_1d(do_plot=False): plt.plot(x, phi_h, label="phi") plt.plot(x, phi_e, "r--", label="exact") plt.title(f"phi at {t = }") - plt.ylim(-amp / (l * 2 * np.pi / Lx) ** 2, amp / (l * 2 * np.pi / Lx) ** 2) + plt.ylim(-amp / (l * 2 * xp.pi / Lx) ** 2, amp / (l * 2 * xp.pi / Lx) ** 2) plt.legend() plt.subplot(5, 2, 2 * c + 2) diff --git a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py index 75a6f1ddc..980352a04 100644 --- a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py @@ -19,7 +19,7 @@ WeightsParameters, ) from struphy.topology import grids -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") @@ -116,7 +116,7 @@ def E_exact(t): r = 0.3677 omega = 1.4156 phi = 0.5362 - return 2 * eps**2 * np.pi / k**2 * r**2 * np.exp(2 * gamma * t) * np.cos(omega * t - phi) ** 2 + return 2 * eps**2 * xp.pi / k**2 * r**2 * xp.exp(2 * gamma * t) * xp.cos(omega * t - phi) ** 2 # get parameters dt = time_opts.dt @@ -130,12 +130,12 @@ def E_exact(t): with h5py.File(os.path.join(pa_data, "data_proc0.hdf5"), "r") as f: time = f["time"]["value"][()] E = f["scalar"]["en_E"][()] - logE = np.log10(E) + logE = xp.log10(E) # find where time derivative of E is zero - dEdt = (np.roll(logE, -1) - np.roll(logE, 1))[1:-1] / (2.0 * dt) - zeros = dEdt * np.roll(dEdt, -1) < 0.0 - maxima_inds = np.logical_and(zeros, dEdt > 0.0) + dEdt = (xp.roll(logE, -1) - xp.roll(logE, 1))[1:-1] / (2.0 * dt) + zeros = dEdt * xp.roll(dEdt, -1) < 0.0 + maxima_inds = xp.logical_and(zeros, dEdt > 0.0) maxima = logE[1:-1][maxima_inds] t_maxima = time[1:-1][maxima_inds] @@ -143,7 +143,7 @@ def E_exact(t): if do_plot: plt.figure(figsize=(18, 12)) plt.plot(time, logE, label="numerical") - plt.plot(time, np.log10(E_exact(time)), label="exact") + plt.plot(time, xp.log10(E_exact(time)), label="exact") plt.legend() plt.title(f"{dt=}, {algo=}, {Nel=}, {p=}, {ppc=}") plt.xlabel("time [m/c]") @@ -154,11 +154,11 @@ def E_exact(t): plt.show() # linear fit - linfit = np.polyfit(t_maxima[:5], maxima[:5], 1) + linfit = xp.polyfit(t_maxima[:5], maxima[:5], 1) gamma_num = linfit[0] # assert - rel_error = np.abs(gamma_num - gamma) / np.abs(gamma) + rel_error = xp.abs(gamma_num - gamma) / xp.abs(gamma) assert rel_error < 0.22, f"Assertion for weak Landau damping failed: {gamma_num = } vs. {gamma = }." print(f"Assertion for weak Landau damping passed ({rel_error = }).") diff --git a/src/struphy/models/tests/verification.py b/src/struphy/models/tests/verification.py index 0bfd53319..d54263ee5 100644 --- a/src/struphy/models/tests/verification.py +++ b/src/struphy/models/tests/verification.py @@ -11,7 +11,7 @@ import struphy from struphy.post_processing import pproc_struphy -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def VlasovAmpereOneSpecies_weakLandau( @@ -40,7 +40,7 @@ def E_exact(t): r = 0.3677 omega = 1.4156 phi = 0.5362 - return 2 * eps**2 * np.pi / k**2 * r**2 * np.exp(2 * gamma * t) * np.cos(omega * t - phi) ** 2 + return 2 * eps**2 * xp.pi / k**2 * r**2 * xp.exp(2 * gamma * t) * xp.cos(omega * t - phi) ** 2 # get parameters with open(os.path.join(path_out, "parameters.yml")) as f: @@ -56,24 +56,24 @@ def E_exact(t): with h5py.File(os.path.join(pa_data, "data_proc0.hdf5"), "r") as f: time = f["time"]["value"][()] E = f["scalar"]["en_E"][()] - logE = np.log10(E) + logE = xp.log10(E) # find where time derivative of E is zero - dEdt = (np.roll(logE, -1) - np.roll(logE, 1))[1:-1] / (2.0 * dt) - zeros = dEdt * np.roll(dEdt, -1) < 0.0 - maxima_inds = np.logical_and(zeros, dEdt > 0.0) + dEdt = (xp.roll(logE, -1) - xp.roll(logE, 1))[1:-1] / (2.0 * dt) + zeros = dEdt * xp.roll(dEdt, -1) < 0.0 + maxima_inds = xp.logical_and(zeros, dEdt > 0.0) maxima = logE[1:-1][maxima_inds] t_maxima = time[1:-1][maxima_inds] # linear fit - linfit = np.polyfit(t_maxima[:5], maxima[:5], 1) + linfit = xp.polyfit(t_maxima[:5], maxima[:5], 1) gamma_num = linfit[0] # plot if show_plots and rank == 0: plt.figure(figsize=(18, 12)) plt.plot(time, logE, label="numerical") - plt.plot(time, np.log10(E_exact(time)), label="exact") + plt.plot(time, xp.log10(E_exact(time)), label="exact") plt.legend() plt.title(f"{dt=}, {algo=}, {Nel=}, {p=}, {ppc=}") plt.xlabel("time [m/c]") @@ -84,7 +84,7 @@ def E_exact(t): plt.show() # assert - rel_error = np.abs(gamma_num - gamma) / np.abs(gamma) + rel_error = xp.abs(gamma_num - gamma) / xp.abs(gamma) assert rel_error < 0.25, f"{rank = }: Assertion for weak Landau damping failed: {gamma_num = } vs. {gamma = }." print(f"{rank = }: Assertion for weak Landau damping passed ({rel_error = }).") @@ -115,7 +115,7 @@ def E_exact(t): r = 0.3677 omega = 1.4156 phi = 0.5362 - return 2 * eps**2 * np.pi / k**2 * r**2 * np.exp(2 * gamma * t) * np.cos(omega * t - phi) ** 2 + return 2 * eps**2 * xp.pi / k**2 * r**2 * xp.exp(2 * gamma * t) * xp.cos(omega * t - phi) ** 2 # get parameters with open(os.path.join(path_out, "parameters.yml")) as f: @@ -131,24 +131,24 @@ def E_exact(t): with h5py.File(os.path.join(pa_data, "data_proc0.hdf5"), "r") as f: time = f["time"]["value"][()] E = f["scalar"]["en_E"][()] - logE = np.log10(E) + logE = xp.log10(E) # find where time derivative of E is zero - dEdt = (np.roll(logE, -1) - np.roll(logE, 1))[1:-1] / (2.0 * dt) - zeros = dEdt * np.roll(dEdt, -1) < 0.0 - maxima_inds = np.logical_and(zeros, dEdt > 0.0) + dEdt = (xp.roll(logE, -1) - xp.roll(logE, 1))[1:-1] / (2.0 * dt) + zeros = dEdt * xp.roll(dEdt, -1) < 0.0 + maxima_inds = xp.logical_and(zeros, dEdt > 0.0) maxima = logE[1:-1][maxima_inds] t_maxima = time[1:-1][maxima_inds] # linear fit - linfit = np.polyfit(t_maxima[:5], maxima[:5], 1) + linfit = xp.polyfit(t_maxima[:5], maxima[:5], 1) gamma_num = linfit[0] # plot if show_plots and rank == 0: plt.figure(figsize=(18, 12)) plt.plot(time, logE, label="numerical") - plt.plot(time, np.log10(E_exact(time)), label="exact") + plt.plot(time, xp.log10(E_exact(time)), label="exact") plt.legend() plt.title(f"{dt=}, {algo=}, {Nel=}, {p=}, {ppc=}") plt.xlabel("time [m/c]") @@ -160,7 +160,7 @@ def E_exact(t): # plt.show() # assert - rel_error = np.abs(gamma_num - gamma) / np.abs(gamma) + rel_error = xp.abs(gamma_num - gamma) / xp.abs(gamma) assert rel_error < 0.25, f"{rank = }: Assertion for weak Landau damping failed: {gamma_num = } vs. {gamma = }." print(f"{rank = }: Assertion for weak Landau damping passed ({rel_error = }).") @@ -190,8 +190,8 @@ def IsothermalEulerSPH_soundwave( MPI.COMM_WORLD.Barrier() path_n_sph = os.path.join(path_pp, "kinetic_data/euler_fluid/n_sph/view_0/") - ee1, ee2, ee3 = np.load(os.path.join(path_n_sph, "grid_n_sph.npy")) - n_sph = np.load(os.path.join(path_n_sph, "n_sph.npy")) + ee1, ee2, ee3 = xp.load(os.path.join(path_n_sph, "grid_n_sph.npy")) + n_sph = xp.load(os.path.join(path_n_sph, "n_sph.npy")) # print(f'{ee1.shape = }, {n_sph.shape = }') if show_plots and rank == 0: @@ -218,7 +218,7 @@ def IsothermalEulerSPH_soundwave( plt.plot(x.squeeze(), n_sph[i, :, 0, 0], style, label=f"time={i * dt:4.2f}") plt.xlim(0, 2.5) plt.legend() - ax.set_xticks(np.linspace(0, 2.5, nx + 1)) + ax.set_xticks(xp.linspace(0, 2.5, nx + 1)) ax.xaxis.set_major_formatter(FormatStrFormatter("%.2f")) plt.grid(c="k") plt.xlabel("x") @@ -231,7 +231,7 @@ def IsothermalEulerSPH_soundwave( plt.show() # assert - error = np.max(np.abs(n_sph[0] - n_sph[-1])) + error = xp.max(xp.abs(n_sph[0] - n_sph[-1])) print(f"{rank = }: Assertion for SPH sound wave passed ({error = }).") assert error < 1.3e-3 @@ -262,30 +262,30 @@ def Maxwell_coaxial( def B_z(X, Y, Z, m, t): """Magnetic field in z direction of coaxial cabel""" r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return (jv(m, r) - 0.28 * yn(m, r)) * np.cos(m * theta - t) + theta = xp.arctan2(Y, X) + return (jv(m, r) - 0.28 * yn(m, r)) * xp.cos(m * theta - t) def E_r(X, Y, Z, m, t): """Electrical field in radial direction of coaxial cabel""" r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return -m / r * (jv(m, r) - 0.28 * yn(m, r)) * np.cos(m * theta - t) + theta = xp.arctan2(Y, X) + return -m / r * (jv(m, r) - 0.28 * yn(m, r)) * xp.cos(m * theta - t) def E_theta(X, Y, Z, m, t): """Electrical field in azimuthal direction of coaxial cabel""" r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return ((m / r * jv(m, r) - jv(m + 1, r)) - 0.28 * (m / r * yn(m, r) - yn(m + 1, r))) * np.sin(m * theta - t) + theta = xp.arctan2(Y, X) + return ((m / r * jv(m, r) - jv(m + 1, r)) - 0.28 * (m / r * yn(m, r) - yn(m + 1, r))) * xp.sin(m * theta - t) def to_E_r(X, Y, E_x, E_y): r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return np.cos(theta) * E_x + np.sin(theta) * E_y + theta = xp.arctan2(Y, X) + return xp.cos(theta) * E_x + xp.sin(theta) * E_y def to_E_theta(X, Y, E_x, E_y): r = (X**2 + Y**2) ** 0.5 - theta = np.arctan2(Y, X) - return -np.sin(theta) * E_x + np.cos(theta) * E_y + theta = xp.arctan2(Y, X) + return -xp.sin(theta) * E_x + xp.cos(theta) * E_y # get parameters with open(os.path.join(path_out, "parameters.yml")) as f: @@ -297,7 +297,7 @@ def to_E_theta(X, Y, E_x, E_y): pproc_path = os.path.join(path_out, "post_processing/") em_fields_path = os.path.join(pproc_path, "fields_data/em_fields/") - t_grid = np.load(os.path.join(pproc_path, "t_grid.npy")) + t_grid = xp.load(os.path.join(pproc_path, "t_grid.npy")) grids_phy = pickle.loads(Path(os.path.join(pproc_path, "fields_data/grids_phy.bin")).read_bytes()) b_field_phy = pickle.loads(Path(os.path.join(em_fields_path, "b_field_phy.bin")).read_bytes()) e_field_phy = pickle.loads(Path(os.path.join(em_fields_path, "e_field_phy.bin")).read_bytes()) @@ -336,13 +336,13 @@ def to_E_theta(X, Y, E_x, E_y): Bz_tend = b_field_phy[t_grid[-1]][2][:, :, 0] Bz_exact = B_z(X, Y, grids_phy[0], modes, t_grid[-1]) - error_Er = np.max(np.abs((to_E_r(X, Y, Ex_tend, Ey_tend) - Er_exact))) - error_Etheta = np.max(np.abs((to_E_theta(X, Y, Ex_tend, Ey_tend) - Etheta_exact))) - error_Bz = np.max(np.abs((Bz_tend - Bz_exact))) + error_Er = xp.max(xp.abs((to_E_r(X, Y, Ex_tend, Ey_tend) - Er_exact))) + error_Etheta = xp.max(xp.abs((to_E_theta(X, Y, Ex_tend, Ey_tend) - Etheta_exact))) + error_Bz = xp.max(xp.abs((Bz_tend - Bz_exact))) - rel_err_Er = error_Er / np.max(np.abs(Er_exact)) - rel_err_Etheta = error_Etheta / np.max(np.abs(Etheta_exact)) - rel_err_Bz = error_Bz / np.max(np.abs(Bz_exact)) + rel_err_Er = error_Er / xp.max(xp.abs(Er_exact)) + rel_err_Etheta = error_Etheta / xp.max(xp.abs(Etheta_exact)) + rel_err_Bz = error_Bz / xp.max(xp.abs(Bz_exact)) print(f"{rel_err_Er = }") print(f"{rel_err_Etheta = }") diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index fdeac6df8..666e87649 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -6,7 +6,7 @@ from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp rank = MPI.COMM_WORLD.Get_rank() @@ -155,7 +155,7 @@ def velocity_scale(self): return "cyclotron" def allocate_helpers(self): - self._tmp = np.empty(1, dtype=float) + self._tmp = xp.empty(1, dtype=float) def update_scalar_quantities(self): particles = self.kinetic_ions.var.particles @@ -245,10 +245,10 @@ def velocity_scale(self): return "alfvén" def allocate_helpers(self): - self._en_fv = np.empty(1, dtype=float) - self._en_fB = np.empty(1, dtype=float) - self._en_tot = np.empty(1, dtype=float) - self._n_lost_particles = np.empty(1, dtype=float) + self._en_fv = xp.empty(1, dtype=float) + self._en_fB = xp.empty(1, dtype=float) + self._en_tot = xp.empty(1, dtype=float) + self._n_lost_particles = xp.empty(1, dtype=float) def update_scalar_quantities(self): particles = self.kinetic_ions.var.particles @@ -678,7 +678,7 @@ def allocate_helpers(self): def f(e1, e2, e3): return 1 - f = np.vectorize(f) + f = xp.vectorize(f) self._integrator = projV3(f) self._energy_evaluator = InternalEnergyEvaluator(self.derham, self.propagators.variat_ent.options.gamma) @@ -754,7 +754,7 @@ def update_thermo_energy(self): def __ener(self, rho, s): """Themodynamical energy as a function of rho and s, usign the perfect gaz hypothesis E(rho, s) = rho^gamma*exp(s/rho)""" - return np.power(rho, self.propagators.variat_ent.options.gamma) * np.exp(s / rho) + return xp.power(rho, self.propagators.variat_ent.options.gamma) * xp.exp(s / rho) class Poisson(StruphyModel): diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index b39fcc38b..3985f4b35 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -21,7 +21,7 @@ from struphy.pic import particles from struphy.pic.base import Particles from struphy.pic.particles import ParticlesSPH -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig if TYPE_CHECKING: @@ -257,7 +257,7 @@ def allocate( f"The number of markers for which data should be stored (={self._n_to_save}) murst be <= than the total number of markers (={obj.Np})" ) if self._n_to_save > 0: - self._saved_markers = np.zeros( + self._saved_markers = xp.zeros( (self._n_to_save, self.particles.markers.shape[1]), dtype=float, ) @@ -270,7 +270,7 @@ def n_to_save(self) -> int: return self._n_to_save @property - def saved_markers(self) -> np.ndarray: + def saved_markers(self) -> xp.ndarray: return self._saved_markers @@ -398,7 +398,7 @@ def allocate( f"The number of markers for which data should be stored (={self._n_to_save}) murst be <= than the total number of markers (={obj.Np})" ) if self._n_to_save > 0: - self._saved_markers = np.zeros( + self._saved_markers = xp.zeros( (self._n_to_save, self.particles.markers.shape[1]), dtype=float, ) @@ -411,5 +411,5 @@ def n_to_save(self) -> int: return self._n_to_save @property - def saved_markers(self) -> np.ndarray: + def saved_markers(self) -> xp.ndarray: return self._saved_markers diff --git a/src/struphy/ode/solvers.py b/src/struphy/ode/solvers.py index 96da5b790..4216198ff 100644 --- a/src/struphy/ode/solvers.py +++ b/src/struphy/ode/solvers.py @@ -4,7 +4,7 @@ from psydac.linalg.stencil import StencilVector from struphy.ode.utils import ButcherTableau -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class ODEsolverFEEC: diff --git a/src/struphy/ode/tests/test_ode_feec.py b/src/struphy/ode/tests/test_ode_feec.py index afd76fc3d..9d1a49be5 100644 --- a/src/struphy/ode/tests/test_ode_feec.py +++ b/src/struphy/ode/tests/test_ode_feec.py @@ -28,7 +28,7 @@ def test_exp_growth(spaces, algo, show_plots=False): from struphy.feec.psydac_derham import Derham from struphy.ode.solvers import ODEsolverFEEC from struphy.ode.utils import ButcherTableau - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -40,7 +40,7 @@ def test_exp_growth(spaces, algo, show_plots=False): c0 = 1.2 omega = 2.3 - y_exact = lambda t: c0 * np.exp(omega * t) + y_exact = lambda t: c0 * xp.exp(omega * t) vector_field = {} for i, space in enumerate(spaces): @@ -117,7 +117,7 @@ def f(t, y1, y2, y3, out=out): errors = {} for i, h in enumerate(hs): errors[h] = {} - time = np.linspace(0, Tend, int(Tend / h) + 1) + time = xp.linspace(0, Tend, int(Tend / h) + 1) print(f"{h = }, {time.size = }") yvec = y_exact(time) ymax = {} @@ -129,16 +129,16 @@ def f(t, y1, y2, y3, out=out): for b in var.blocks: b[:] = c0 var.update_ghost_regions() - ymax[var] = c0 * np.ones_like(time) + ymax[var] = c0 * xp.ones_like(time) for n in range(time.size - 1): tn = h * n solver(tn, h) for var in vector_field: - ymax[var][n + 1] = np.max(var.toarray()) + ymax[var][n + 1] = xp.max(var.toarray()) # checks for var in vector_field: - errors[h][var] = h * np.sum(np.abs(yvec - ymax[var])) / (h * np.sum(np.abs(yvec))) + errors[h][var] = h * xp.sum(xp.abs(yvec - ymax[var])) / (h * xp.sum(xp.abs(yvec))) print(f"{errors[h][var] = }") assert errors[h][var] < 0.31 @@ -161,9 +161,9 @@ def f(t, y1, y2, y3, out=out): h_vec += [h] err_vec += [dct[var]] - m, _ = np.polyfit(np.log(h_vec), np.log(err_vec), deg=1) + m, _ = xp.polyfit(xp.log(h_vec), xp.log(err_vec), deg=1) print(f"{spaces[j]}-space, fitted convergence rate = {m} for {algo = } with {solver.butcher.conv_rate = }") - assert np.abs(m - solver.butcher.conv_rate) < 0.1 + assert xp.abs(m - solver.butcher.conv_rate) < 0.1 print(f"Convergence check passed on {rank = }.") if rank == 0: diff --git a/src/struphy/ode/utils.py b/src/struphy/ode/utils.py index 1dfa881b3..997225d42 100644 --- a/src/struphy/ode/utils.py +++ b/src/struphy/ode/utils.py @@ -1,7 +1,7 @@ from dataclasses import dataclass from typing import Literal, get_args -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp OptsButcher = Literal[ "rk4", @@ -67,14 +67,14 @@ def __post_init__(self): else: raise NotImplementedError(f"Chosen algorithm {self.algo} is not implemented.") - self._b = np.array(b) - self._c = np.array(c) + self._b = xp.array(b) + self._c = xp.array(c) assert self._b.size == self._c.size self._n_stages = self._b.size assert len(a) == self.n_stages - 1 - self._a = np.tri(self.n_stages, k=-1) + self._a = xp.tri(self.n_stages, k=-1) for l, st in enumerate(a): assert len(st) == l + 1 self._a[l + 1, : l + 1] = st diff --git a/src/struphy/pic/accumulation/filter.py b/src/struphy/pic/accumulation/filter.py index b80da8c7a..2c73e3a06 100644 --- a/src/struphy/pic/accumulation/filter.py +++ b/src/struphy/pic/accumulation/filter.py @@ -127,11 +127,11 @@ def _apply_three_point(self, vec, repeat: int, alpha: float): comp._data, axis, self.form_int, - np.array(self.derham.Nel), - np.array(self.derham.spl_kind), - np.array(self.derham.p), - np.array(starts), - np.array(ends), + xp.array(self.derham.Nel), + xp.array(self.derham.spl_kind), + xp.array(self.derham.p), + xp.array(starts), + xp.array(ends), alpha=alpha, ) @@ -150,19 +150,19 @@ def _apply_toroidal_fourier_filter(self, vec, modes: tuple[int, ...]): """ tor_Nel = self.derham.Nel[2] - modes = np.asarray(modes, dtype=int) + modes = xp.asarray(modes, dtype=int) - assert tor_Nel >= 2 * int(np.max(modes)), "Nel[2] must be at least 2*max(modes)" + assert tor_Nel >= 2 * int(xp.max(modes)), "Nel[2] must be at least 2*max(modes)" assert self.derham.domain_decomposition.nprocs[2] == 1, "No domain decomposition along toroidal direction" - pn = np.asarray(self.derham.p, dtype=int) - ir = np.empty(3, dtype=int) + pn = xp.asarray(self.derham.p, dtype=int) + ir = xp.empty(3, dtype=int) # rfft output length if (tor_Nel % 2) == 0: - vec_temp = np.zeros(int(tor_Nel / 2) + 1, dtype=complex) + vec_temp = xp.zeros(int(tor_Nel / 2) + 1, dtype=complex) else: - vec_temp = np.zeros(int((tor_Nel - 1) / 2) + 1, dtype=complex) + vec_temp = xp.zeros(int((tor_Nel - 1) / 2) + 1, dtype=complex) for axis, comp, starts, ends in self._yield_dir_components(vec): for i in range(3): diff --git a/src/struphy/pic/accumulation/particles_to_grid.py b/src/struphy/pic/accumulation/particles_to_grid.py index 17de1a29a..d694c017b 100644 --- a/src/struphy/pic/accumulation/particles_to_grid.py +++ b/src/struphy/pic/accumulation/particles_to_grid.py @@ -12,7 +12,7 @@ from struphy.pic.accumulation.filter import AccumFilter, FilterParameters from struphy.pic.base import Particles from struphy.profiling.profiling import ProfileManager -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel @@ -191,7 +191,7 @@ def __call__(self, *optional_args, **args_control): Entries must be pyccel-conform types. args_control : any - Keyword arguments for an analytical control variate correction in the accumulation step. Possible keywords are 'control_vec' for a vector correction or 'control_mat' for a matrix correction. Values are a 1d (vector) or 2d (matrix) list with callables or np.ndarrays used for the correction. + Keyword arguments for an analytical control variate correction in the accumulation step. Possible keywords are 'control_vec' for a vector correction or 'control_mat' for a matrix correction. Values are a 1d (vector) or 2d (matrix) list with callables or xp.ndarrays used for the correction. """ # flags for break @@ -388,7 +388,7 @@ def show_accumulated_spline_field(self, mass_ops: WeightedMassOperators, eta_dir field.vector = a # plot field - eta = np.linspace(0, 1, 100) + eta = xp.linspace(0, 1, 100) if eta_direction == 0: args = (eta, 0.5, 0.5) elif eta_direction == 1: @@ -510,7 +510,7 @@ def __call__(self, *optional_args, **args_control): args_control : any Keyword arguments for an analytical control variate correction in the accumulation step. Possible keywords are 'control_vec' for a vector correction or 'control_mat' for a matrix correction. - Values are a 1d (vector) or 2d (matrix) list with callables or np.ndarrays used for the correction. + Values are a 1d (vector) or 2d (matrix) list with callables or xp.ndarrays used for the correction. """ # flags for break @@ -644,7 +644,7 @@ def show_accumulated_spline_field(self, mass_ops, eta_direction=0): field.vector = a # plot field - eta = np.linspace(0, 1, 100) + eta = xp.linspace(0, 1, 100) if eta_direction == 0: args = (eta, 0.5, 0.5) elif eta_direction == 1: diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index e9ab1ba7b..824c51c0a 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -53,7 +53,7 @@ class Intracomm: WeightsParameters, ) from struphy.utils import utils -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig from struphy.utils.pyccel import Pyccelkernel @@ -221,7 +221,7 @@ def __init__( self._nprocs = domain_decomp[1] # total number of cells (equal to mpi_size if no grid) - n_cells = np.sum(np.prod(self.domain_array[:, 2::3], axis=1, dtype=int)) * self.num_clones + n_cells = xp.sum(xp.prod(self.domain_array[:, 2::3], axis=1, dtype=int)) * self.num_clones # if verbose: # print(f"\n{self.mpi_rank = }, {n_cells = }") @@ -235,7 +235,7 @@ def __init__( assert all([nboxes % nproc == 0 for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)]), ( f"Number of boxes {self.boxes_per_dim = } must be divisible by number of processes {self.nprocs = } in each direction." ) - n_boxes = np.prod(self.boxes_per_dim, dtype=int) * self.num_clones + n_boxes = xp.prod(self.boxes_per_dim, dtype=int) * self.num_clones # if verbose: # print(f"\n{self.mpi_rank = }, {n_boxes = }") @@ -342,9 +342,9 @@ def __init__( self._generate_sampling_moments() # create buffers for mpi_sort_markers - self._sorting_etas = np.zeros(self.markers.shape, dtype=float) - self._is_on_proc_domain = np.zeros((self.markers.shape[0], 3), dtype=bool) - self._can_stay = np.zeros(self.markers.shape[0], dtype=bool) + self._sorting_etas = xp.zeros(self.markers.shape, dtype=float) + self._is_on_proc_domain = xp.zeros((self.markers.shape[0], 3), dtype=bool) + self._can_stay = xp.zeros(self.markers.shape[0], dtype=bool) self._reqs = [None] * self.mpi_size self._recvbufs = [None] * self.mpi_size self._send_to_i = [None] * self.mpi_size @@ -726,16 +726,16 @@ def index(self): def valid_mks(self): """Array of booleans stating if an entry in the markers array is a true local particle (not a hole or ghost).""" if not hasattr(self, "_valid_mks"): - self._valid_mks = ~np.logical_or(self.holes, self.ghost_particles) + self._valid_mks = ~xp.logical_or(self.holes, self.ghost_particles) return self._valid_mks def update_valid_mks(self): - self._valid_mks[:] = ~np.logical_or(self.holes, self.ghost_particles) + self._valid_mks[:] = ~xp.logical_or(self.holes, self.ghost_particles) @property def n_mks_loc(self): """Number of valid markers on process (without holes and ghosts).""" - return np.count_nonzero(self.valid_mks) + return xp.count_nonzero(self.valid_mks) @property def n_mks_on_each_proc(self): @@ -745,7 +745,7 @@ def n_mks_on_each_proc(self): @property def n_mks_on_clone(self): """Number of valid markers on current clone (without holes and ghosts).""" - return np.sum(self.n_mks_on_each_proc) + return xp.sum(self.n_mks_on_each_proc) @property def n_mks_on_each_clone(self): @@ -755,7 +755,7 @@ def n_mks_on_each_clone(self): @property def n_mks_global(self): """Number of valid markers on current clone (without holes and ghosts).""" - return np.sum(self.n_mks_on_each_clone) + return xp.sum(self.n_mks_on_each_clone) @property def positions(self): @@ -764,7 +764,7 @@ def positions(self): @positions.setter def positions(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc, 3) self._markers[self.valid_mks, self.index["pos"]] = new @@ -775,7 +775,7 @@ def velocities(self): @velocities.setter def velocities(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc, self.vdim), f"{self.n_mks_loc = } and {self.vdim = } but {new.shape = }" self._markers[self.valid_mks, self.index["vel"]] = new @@ -786,7 +786,7 @@ def phasespace_coords(self): @phasespace_coords.setter def phasespace_coords(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc, 3 + self.vdim) self._markers[self.valid_mks, self.index["coords"]] = new @@ -797,7 +797,7 @@ def weights(self): @weights.setter def weights(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc,) self._markers[self.valid_mks, self.index["weights"]] = new @@ -808,7 +808,7 @@ def sampling_density(self): @sampling_density.setter def sampling_density(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc,) self._markers[self.valid_mks, self.index["s0"]] = new @@ -819,7 +819,7 @@ def weights0(self): @weights0.setter def weights0(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc,) self._markers[self.valid_mks, self.index["w0"]] = new @@ -830,7 +830,7 @@ def marker_ids(self): @marker_ids.setter def marker_ids(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc,) self._markers[self.valid_mks, self.index["ids"]] = new @@ -861,7 +861,7 @@ def f_coords(self): @f_coords.setter def f_coords(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) self.markers[self.valid_mks, self.f_coords_index] = new @property @@ -873,16 +873,16 @@ def args_markers(self): def f_jacobian_coords(self): """Coordinates of the velocity jacobian determinant of the distribution fuction.""" if isinstance(self.f_jacobian_coords_index, list): - return self.markers[np.ix_(~self.holes, self.f_jacobian_coords_index)] + return self.markers[xp.ix_(~self.holes, self.f_jacobian_coords_index)] else: return self.markers[~self.holes, self.f_jacobian_coords_index] @f_jacobian_coords.setter def f_jacobian_coords(self, new): - assert isinstance(new, np.ndarray) + assert isinstance(new, xp.ndarray) if isinstance(self.f_jacobian_coords_index, list): self.markers[ - np.ix_( + xp.ix_( ~self.holes, self.f_jacobian_coords_index, ) @@ -929,7 +929,7 @@ def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): Returns ------- - dom_arr : np.ndarray + dom_arr : xp.ndarray A 2d array of shape (#MPI processes, 9). The row index denotes the process rank. The columns are for n=0,1,2: - arr[i, 3*n + 0] holds the LEFT domain boundary of process i in direction eta_(n+1). - arr[i, 3*n + 1] holds the RIGHT domain boundary of process i in direction eta_(n+1). @@ -941,7 +941,7 @@ def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): if mpi_dims_mask is None: mpi_dims_mask = [True, True, True] - dom_arr = np.zeros((self.mpi_size, 9), dtype=float) + dom_arr = xp.zeros((self.mpi_size, 9), dtype=float) # factorize mpi size factors = factorint(self.mpi_size) @@ -965,10 +965,10 @@ def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): mm = (mm + 1) % 3 nprocs[mm] *= fac - assert np.prod(nprocs) == self.mpi_size + assert xp.prod(nprocs) == self.mpi_size # domain decomposition - breaks = [np.linspace(0.0, 1.0, nproc + 1) for nproc in nprocs] + breaks = [xp.linspace(0.0, 1.0, nproc + 1) for nproc in nprocs] # fill domain array for n in range(self.mpi_size): @@ -1057,14 +1057,14 @@ def _n_mks_load_and_Np_per_clone(self): """Return two arrays: 1) an array of sub_comm.size where the i-th entry corresponds to the number of markers drawn on process i, and 2) an array of size num_clones where the i-th entry corresponds to the number of markers on clone i.""" # number of cells on current process - n_cells_loc = np.prod( + n_cells_loc = xp.prod( self.domain_array[self.mpi_rank, 2::3], dtype=int, ) # array of number of markers on each process at loading stage if self.clone_config is not None: - _n_cells_clone = np.sum(np.prod(self.domain_array[:, 2::3], axis=1, dtype=int)) + _n_cells_clone = xp.sum(xp.prod(self.domain_array[:, 2::3], axis=1, dtype=int)) _n_mks_load_tot = self.clone_config.get_Np_clone(self.Np) _ppc = _n_mks_load_tot / _n_cells_clone else: @@ -1074,14 +1074,14 @@ def _n_mks_load_and_Np_per_clone(self): n_mks_load = self._gather_scalar_in_subcomm_array(int(_ppc * n_cells_loc)) # add deviation from Np to rank 0 - n_mks_load[0] += _n_mks_load_tot - np.sum(n_mks_load) + n_mks_load[0] += _n_mks_load_tot - xp.sum(n_mks_load) # check if all markers are there - assert np.sum(n_mks_load) == _n_mks_load_tot + assert xp.sum(n_mks_load) == _n_mks_load_tot # Np on each clone Np_per_clone = self._gather_scalar_in_intercomm_array(_n_mks_load_tot) - assert np.sum(Np_per_clone) == self.Np + assert xp.sum(Np_per_clone) == self.Np return n_mks_load, Np_per_clone @@ -1092,23 +1092,23 @@ def _allocate_marker_array(self): # number of markers on the local process at loading stage n_mks_load_loc = self.n_mks_load[self._mpi_rank] - bufsize = self.bufsize + 1.0 / np.sqrt(n_mks_load_loc) + bufsize = self.bufsize + 1.0 / xp.sqrt(n_mks_load_loc) # allocate markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer self._n_rows = round(n_mks_load_loc * (1 + bufsize)) - self._markers = np.zeros((self.n_rows, self.n_cols), dtype=float) + self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float) # allocate auxiliary arrays - self._holes = np.zeros(self.n_rows, dtype=bool) - self._ghost_particles = np.zeros(self.n_rows, dtype=bool) - self._valid_mks = np.zeros(self.n_rows, dtype=bool) - self._is_outside_right = np.zeros(self.n_rows, dtype=bool) - self._is_outside_left = np.zeros(self.n_rows, dtype=bool) - self._is_outside = np.zeros(self.n_rows, dtype=bool) + self._holes = xp.zeros(self.n_rows, dtype=bool) + self._ghost_particles = xp.zeros(self.n_rows, dtype=bool) + self._valid_mks = xp.zeros(self.n_rows, dtype=bool) + self._is_outside_right = xp.zeros(self.n_rows, dtype=bool) + self._is_outside_left = xp.zeros(self.n_rows, dtype=bool) + self._is_outside = xp.zeros(self.n_rows, dtype=bool) # create array container (3 x positions, vdim x velocities, weight, s0, w0, ID) for removed markers self._n_lost_markers = 0 - self._lost_markers = np.zeros((int(self.n_rows * 0.5), 10), dtype=float) + self._lost_markers = xp.zeros((int(self.n_rows * 0.5), 10), dtype=float) # arguments for kernels self._args_markers = MarkerArguments( @@ -1223,16 +1223,16 @@ def _generate_sampling_moments(self): # assert len(ns) == len(us) == len(vths) - # ns = np.array(ns) - # us = np.array(us) - # vths = np.array(vths) + # ns = xp.array(ns) + # us = xp.array(us) + # vths = xp.array(vths) # Use the mean of shifts and thermal velocity such that outermost shift+thermal is # new shift + new thermal - # mean_us = np.mean(us, axis=0) - # us_ext = us + vths * np.where(us >= 0, 1, -1) + # mean_us = xp.mean(us, axis=0) + # us_ext = us + vths * xp.where(us >= 0, 1, -1) # us_ext_dist = us_ext - mean_us[None, :] - # new_vths = np.max(np.abs(us_ext_dist), axis=0) + # new_vths = xp.max(xp.abs(us_ext_dist), axis=0) # new_moments = [] @@ -1307,7 +1307,7 @@ def _f_init(*etas, flat_eval=False): out = out0 + out1 if flat_eval: - out = np.squeeze(out) + out = xp.squeeze(out) return out @@ -1316,7 +1316,7 @@ def _f_init(*etas, flat_eval=False): def _load_external( self, n_mks_load_loc: int, - n_mks_load_cum_sum: np.ndarray, + n_mks_load_cum_sum: xp.ndarray, ): """Load markers from external .hdf5 file. @@ -1325,7 +1325,7 @@ def _load_external( n_mks_load_loc: int Number of markers on the local process at loading stage. - n_mks_load_cum_sum: np.ndarray + n_mks_load_cum_sum: xp.ndarray Cumulative sum of number of markers on each process at loading stage. """ if self.mpi_rank == 0: @@ -1349,7 +1349,7 @@ def _load_external( file.close() else: - recvbuf = np.zeros( + recvbuf = xp.zeros( (n_mks_load_loc, self.markers.shape[1]), dtype=float, ) @@ -1491,8 +1491,8 @@ def draw_markers( self.update_ghost_particles() # cumulative sum of number of markers on each process at loading stage. - n_mks_load_cum_sum = np.cumsum(self.n_mks_load) - Np_per_clone_cum_sum = np.cumsum(self.Np_per_clone) + n_mks_load_cum_sum = xp.cumsum(self.n_mks_load) + Np_per_clone_cum_sum = xp.cumsum(self.Np_per_clone) _first_marker_id = (Np_per_clone_cum_sum - self.Np_per_clone)[self.clone_id] + ( n_mks_load_cum_sum - self.n_mks_load )[self._mpi_rank] @@ -1520,9 +1520,9 @@ def draw_markers( self._load_tesselation() if self.type == "sph": self._set_initial_condition() - self.velocities = np.array(self.u_init(self.positions)[0]).T + self.velocities = xp.array(self.u_init(self.positions)[0]).T # set markers ID in last column - self.marker_ids = _first_marker_id + np.arange(n_mks_load_loc, dtype=float) + self.marker_ids = _first_marker_id + xp.arange(n_mks_load_loc, dtype=float) else: if self.mpi_rank == 0 and verbose: print("\nLoading fresh markers:") @@ -1534,7 +1534,7 @@ def draw_markers( # set seed _seed = self.loading_params.seed if _seed is not None: - np.random.seed(_seed) + xp.random.seed(_seed) # counting integers num_loaded_particles_loc = 0 # number of particles alreday loaded (local) @@ -1545,15 +1545,15 @@ def draw_markers( while num_loaded_particles_glob < int(self.Np): # Generate a chunk of random particles num_to_add_glob = min(chunk_size, int(self.Np) - num_loaded_particles_glob) - temp = np.random.rand(num_to_add_glob, 3 + self.vdim) + temp = xp.random.rand(num_to_add_glob, 3 + self.vdim) # check which particles are on the current process domain - is_on_proc_domain = np.logical_and( + is_on_proc_domain = xp.logical_and( temp[:, :3] > self.domain_array[self.mpi_rank, 0::3], temp[:, :3] < self.domain_array[self.mpi_rank, 1::3], ) - valid_idx = np.nonzero(np.all(is_on_proc_domain, axis=1))[0] + valid_idx = xp.nonzero(xp.all(is_on_proc_domain, axis=1))[0] valid_particles = temp[valid_idx] - valid_particles = np.array_split(valid_particles, self.num_clones)[self.clone_id] + valid_particles = xp.array_split(valid_particles, self.num_clones)[self.clone_id] num_valid = valid_particles.shape[0] # Add the valid particles to the phasespace_coords array @@ -1570,7 +1570,7 @@ def draw_markers( # set new n_mks_load self._gather_scalar_in_subcomm_array(num_loaded_particles_loc, out=self.n_mks_load) n_mks_load_loc = self.n_mks_load[self.mpi_rank] - n_mks_load_cum_sum = np.cumsum(self.n_mks_load) + n_mks_load_cum_sum = xp.cumsum(self.n_mks_load) # set new holes in markers array to -1 self._markers[num_loaded_particles_loc:] = -1.0 @@ -1610,11 +1610,11 @@ def draw_markers( # initial velocities - SPH case: v(0) = u(x(0)) for given velocity u(x) if self.type == "sph": self._set_initial_condition() - self.velocities = np.array(self.u_init(self.positions)[0]).T + self.velocities = xp.array(self.u_init(self.positions)[0]).T else: # inverse transform sampling in velocity space - u_mean = np.array(self.loading_params.moments[: self.vdim]) - v_th = np.array(self.loading_params.moments[self.vdim :]) + u_mean = xp.array(self.loading_params.moments[: self.vdim]) + v_th = xp.array(self.loading_params.moments[self.vdim :]) # Particles6D: (1d Maxwellian, 1d Maxwellian, 1d Maxwellian) if self.vdim == 3: @@ -1622,7 +1622,7 @@ def draw_markers( sp.erfinv( 2 * self.velocities - 1, ) - * np.sqrt(2) + * xp.sqrt(2) * v_th + u_mean ) @@ -1632,16 +1632,16 @@ def draw_markers( sp.erfinv( 2 * self.velocities[:, 0] - 1, ) - * np.sqrt(2) + * xp.sqrt(2) * v_th[0] + u_mean[0] ) self._markers[:n_mks_load_loc, 4] = ( - np.sqrt( - -1 * np.log(1 - self.velocities[:, 1]), + xp.sqrt( + -1 * xp.log(1 - self.velocities[:, 1]), ) - * np.sqrt(2) + * xp.sqrt(2) * v_th[1] + u_mean[1] ) @@ -1654,13 +1654,13 @@ def draw_markers( # inversion method for drawing uniformly on the disc if self.spatial == "disc": - self._markers[:n_mks_load_loc, 0] = np.sqrt( + self._markers[:n_mks_load_loc, 0] = xp.sqrt( self._markers[:n_mks_load_loc, 0], ) else: assert self.spatial == "uniform", f'Spatial drawing must be "uniform" or "disc", is {self.spatial}.' - self.marker_ids = _first_marker_id + np.arange(n_mks_load_loc, dtype=float) + self.marker_ids = _first_marker_id + xp.arange(n_mks_load_loc, dtype=float) # set specific initial condition for some particles if self.loading_params.specific_markers is not None: @@ -1681,8 +1681,8 @@ def draw_markers( # check if all particle positions are inside the unit cube [0, 1]^3 n_mks_load_loc = self.n_mks_load[self._mpi_rank] - assert np.all(~self.holes[:n_mks_load_loc]) - assert np.all(self.holes[n_mks_load_loc:]) + assert xp.all(~self.holes[:n_mks_load_loc]) + assert xp.all(self.holes[n_mks_load_loc:]) if self._initialized_sorting and sort: if self.mpi_rank == 0 and verbose: @@ -1755,8 +1755,8 @@ def mpi_sort_markers( # check if all markers are on the right process after sorting if do_test: - all_on_right_proc = np.all( - np.logical_and( + all_on_right_proc = xp.all( + xp.logical_and( self.positions > self.domain_array[self.mpi_rank, 0::3], self.positions < self.domain_array[self.mpi_rank, 1::3], ), @@ -1880,18 +1880,18 @@ def update_weights(self): def reset_marker_ids(self): """Reset the marker ids (last column in marker array) according to the current distribution of particles. The first marker on rank 0 gets the id '0', the last marker on the last rank gets the id 'n_mks_global - 1'.""" - n_mks_proc_cumsum = np.cumsum(self.n_mks_on_each_proc) - n_mks_clone_cumsum = np.cumsum(self.n_mks_on_each_clone) + n_mks_proc_cumsum = xp.cumsum(self.n_mks_on_each_proc) + n_mks_clone_cumsum = xp.cumsum(self.n_mks_on_each_clone) first_marker_id = (n_mks_clone_cumsum - self.n_mks_on_each_clone)[self.clone_id] + ( n_mks_proc_cumsum - self.n_mks_on_each_proc )[self.mpi_rank] - self.marker_ids = first_marker_id + np.arange(self.n_mks_loc, dtype=int) + self.marker_ids = first_marker_id + xp.arange(self.n_mks_loc, dtype=int) @profile def binning( self, components: tuple[bool], - bin_edges: tuple[np.ndarray], + bin_edges: tuple[xp.ndarray], divide_by_jac: bool = True, ): r"""Computes full-f and delta-f distribution functions via marker binning in logical space. @@ -1917,7 +1917,7 @@ def binning( The reconstructed delta-f distribution function. """ - assert np.count_nonzero(components) == len(bin_edges) + assert xp.count_nonzero(components) == len(bin_edges) # volume of a bin bin_vol = 1.0 @@ -1939,13 +1939,13 @@ def binning( _weights0 /= self.domain.jacobian_det(self.positions, remove_outside=False) # _weights0 /= self.velocity_jacobian_det(*self.phasespace_coords.T) - f_slice = np.histogramdd( + f_slice = xp.histogramdd( self.markers_wo_holes_and_ghost[:, slicing], bins=bin_edges, weights=_weights0, )[0] - df_slice = np.histogramdd( + df_slice = xp.histogramdd( self.markers_wo_holes_and_ghost[:, slicing], bins=bin_edges, weights=_weights, @@ -1972,7 +1972,7 @@ def show_distribution_function(self, components, bin_edges): import matplotlib.pyplot as plt - n_dim = np.count_nonzero(components) + n_dim = xp.count_nonzero(components) assert n_dim == 1 or n_dim == 2, f"Distribution function can only be shown in 1D or 2D slices, not {n_dim}." @@ -1988,7 +1988,7 @@ def show_distribution_function(self, components, bin_edges): 4: "$v_2$", 5: "$v_3$", } - indices = np.nonzero(components)[0] + indices = xp.nonzero(components)[0] if n_dim == 1: plt.plot(bin_centers[0], f_slice) @@ -2012,13 +2012,13 @@ def _find_outside_particles(self, axis): self._is_outside_left[self.holes] = False self._is_outside_left[self.ghost_particles] = False - self._is_outside[:] = np.logical_or( + self._is_outside[:] = xp.logical_or( self._is_outside_right, self._is_outside_left, ) # indices or particles that are outside of the logical unit cube - outside_inds = np.nonzero(self._is_outside)[0] + outside_inds = xp.nonzero(self._is_outside)[0] return outside_inds @@ -2045,7 +2045,7 @@ def apply_kinetic_bc(self, newton=False): self.particle_refilling() self._markers[self._is_outside, :-1] = -1.0 - self._n_lost_markers += len(np.nonzero(self._is_outside)[0]) + self._n_lost_markers += len(xp.nonzero(self._is_outside)[0]) for axis in self._periodic_axes: outside_inds = self._find_outside_particles(axis) @@ -2056,8 +2056,8 @@ def apply_kinetic_bc(self, newton=False): self.markers[outside_inds, axis] = self.markers[outside_inds, axis] % 1.0 # set shift for alpha-weighted mid-point computation - outside_right_inds = np.nonzero(self._is_outside_right)[0] - outside_left_inds = np.nonzero(self._is_outside_left)[0] + outside_right_inds = xp.nonzero(self._is_outside_right)[0] + outside_left_inds = xp.nonzero(self._is_outside_left)[0] if newton: self.markers[ outside_right_inds, @@ -2125,12 +2125,12 @@ def particle_refilling(self): for kind in self.bc_refill: # sorting out particles which are out of the domain if kind == "inner": - outside_inds = np.nonzero(self._is_outside_left)[0] + outside_inds = xp.nonzero(self._is_outside_left)[0] self.markers[outside_inds, 0] = 1e-4 r_loss = self.domain.params["a1"] else: - outside_inds = np.nonzero(self._is_outside_right)[0] + outside_inds = xp.nonzero(self._is_outside_right)[0] self.markers[outside_inds, 0] = 1 - 1e-4 r_loss = 1.0 @@ -2179,12 +2179,12 @@ def gyro_transfer(self, outside_inds): Parameters ---------- - outside_inds : np.array (int) + outside_inds : xp.array (int) An array of indices of particles which are outside of the domain. Returns ------- - out : np.array (bool) + out : xp.array (bool) An array of indices of particles where its guiding centers are outside of the domain. """ @@ -2201,18 +2201,18 @@ def gyro_transfer(self, outside_inds): b_cart, xyz = self.equil.b_cart(self.markers[outside_inds, :]) # calculate magnetic field amplitude and normalized magnetic field - absB0 = np.sqrt(b_cart[0] ** 2 + b_cart[1] ** 2 + b_cart[2] ** 2) + absB0 = xp.sqrt(b_cart[0] ** 2 + b_cart[1] ** 2 + b_cart[2] ** 2) norm_b_cart = b_cart / absB0 # calculate parallel and perpendicular velocities - v_parallel = np.einsum("ij,ij->j", v, norm_b_cart) - v_perp = np.cross(norm_b_cart, np.cross(v, norm_b_cart, axis=0), axis=0) - v_perp_square = np.sqrt(v_perp[0] ** 2 + v_perp[1] ** 2 + v_perp[2] ** 2) + v_parallel = xp.einsum("ij,ij->j", v, norm_b_cart) + v_perp = xp.cross(norm_b_cart, xp.cross(v, norm_b_cart, axis=0), axis=0) + v_perp_square = xp.sqrt(v_perp[0] ** 2 + v_perp[1] ** 2 + v_perp[2] ** 2) - assert np.all(np.isclose(v_perp, v - norm_b_cart * v_parallel)) + assert xp.all(xp.isclose(v_perp, v - norm_b_cart * v_parallel)) # calculate Larmor radius - Larmor_r = np.cross(norm_b_cart, v_perp, axis=0) / absB0 * self._epsilon + Larmor_r = xp.cross(norm_b_cart, v_perp, axis=0) / absB0 * self._epsilon # transform cartesian coordinates to logical coordinates # TODO: currently only possible with the geomoetry where its inverse map is defined. @@ -2231,17 +2231,17 @@ def gyro_transfer(self, outside_inds): b_cart = self.equil.b_cart(self.markers[outside_inds, :])[0] # calculate magnetic field amplitude and normalized magnetic field - absB0 = np.sqrt(b_cart[0] ** 2 + b_cart[1] ** 2 + b_cart[2] ** 2) + absB0 = xp.sqrt(b_cart[0] ** 2 + b_cart[1] ** 2 + b_cart[2] ** 2) norm_b_cart = b_cart / absB0 Larmor_r = new_xyz - xyz - Larmor_r /= np.sqrt(Larmor_r[0] ** 2 + Larmor_r[1] ** 2 + Larmor_r[2] ** 2) + Larmor_r /= xp.sqrt(Larmor_r[0] ** 2 + Larmor_r[1] ** 2 + Larmor_r[2] ** 2) - new_v_perp = np.cross(Larmor_r, norm_b_cart, axis=0) * v_perp_square + new_v_perp = xp.cross(Larmor_r, norm_b_cart, axis=0) * v_perp_square self.markers[outside_inds, 3:6] = (norm_b_cart * v_parallel).T + new_v_perp.T - return np.logical_and(1.0 > gc_etas[0], gc_etas[0] > 0.0) + return xp.logical_and(1.0 > gc_etas[0], gc_etas[0] > 0.0) class SortingBoxes: """Boxes used for the sorting of the particles. @@ -2424,26 +2424,26 @@ def _set_boxes(self): n_particles = self._markers_shape[0] n_mkr = int(n_particles / n_box_in) + 1 n_cols = round( - n_mkr * (1 + 1 / np.sqrt(n_mkr) + self._box_bufsize), + n_mkr * (1 + 1 / xp.sqrt(n_mkr) + self._box_bufsize), ) # cartesian boxes - self._boxes = np.zeros((self._n_boxes + 1, n_cols), dtype=int) + self._boxes = xp.zeros((self._n_boxes + 1, n_cols), dtype=int) # TODO: there is still a bug here # the row number in self._boxes should not be n_boxes + 1; this is just a temporary fix to avoid an error that I dont understand. # Must be fixed soon! - self._next_index = np.zeros((self._n_boxes + 1), dtype=int) - self._cumul_next_index = np.zeros((self._n_boxes + 2), dtype=int) - self._neighbours = np.zeros((self._n_boxes, 27), dtype=int) + self._next_index = xp.zeros((self._n_boxes + 1), dtype=int) + self._cumul_next_index = xp.zeros((self._n_boxes + 2), dtype=int) + self._neighbours = xp.zeros((self._n_boxes, 27), dtype=int) # A particle on box i only sees particles in boxes that belong to neighbours[i] initialize_neighbours(self._neighbours, self.nx, self.ny, self.nz) # print(f"{self._rank = }\n{self._neighbours = }") - self._swap_line_1 = np.zeros(self._markers_shape[1]) - self._swap_line_2 = np.zeros(self._markers_shape[1]) + self._swap_line_1 = xp.zeros(self._markers_shape[1]) + self._swap_line_2 = xp.zeros(self._markers_shape[1]) def _set_boundary_boxes(self): """Gather all the boxes that are part of a boundary""" @@ -2604,7 +2604,7 @@ def _sort_boxed_particles_numpy(self): sorting_axis = self._sorting_boxes.box_index if not hasattr(self, "_argsort_array"): - self._argsort_array = np.zeros(self.markers.shape[0], dtype=int) + self._argsort_array = xp.zeros(self.markers.shape[0], dtype=int) self._argsort_array[:] = self._markers[:, sorting_axis].argsort() self._markers[:, :] = self._markers[self._argsort_array] @@ -2633,18 +2633,18 @@ def put_particles_in_boxes(self): self.update_ghost_particles() # if self.verbose: - # valid_box_ids = np.nonzero(self._sorting_boxes._boxes[:, 0] != -1)[0] + # valid_box_ids = xp.nonzero(self._sorting_boxes._boxes[:, 0] != -1)[0] # print(f"Boxes holding at least one particle: {valid_box_ids}") # for i in valid_box_ids: - # n_mks_box = np.count_nonzero(self._sorting_boxes._boxes[i] != -1) + # n_mks_box = xp.count_nonzero(self._sorting_boxes._boxes[i] != -1) # print(f"Number of markers in box {i} is {n_mks_box}") def check_and_assign_particles_to_boxes(self): """Check whether the box array has enough columns (detect load imbalance wrt to sorting boxes), and then assigne the particles to boxes.""" - bcount = np.bincount(np.int64(self.markers_wo_holes[:, -2])) - max_in_box = np.max(bcount) + bcount = xp.bincount(xp.int64(self.markers_wo_holes[:, -2])) + max_in_box = xp.max(bcount) if max_in_box > self._sorting_boxes.boxes.shape[1]: warnings.warn( f'Strong load imbalance detected in sorting boxes: \ @@ -2688,7 +2688,7 @@ def do_sort(self, use_numpy_argsort=False): def remove_ghost_particles(self): self.update_ghost_particles() - new_holes = np.nonzero(self.ghost_particles) + new_holes = xp.nonzero(self.ghost_particles) self._markers[new_holes] = -1.0 self.update_holes() @@ -2980,161 +2980,161 @@ def determine_markers_in_box(self, list_boxes): for i in list_boxes: indices += list(self._sorting_boxes._boxes[i][self._sorting_boxes._boxes[i] != -1]) - indices = np.array(indices, dtype=int) + indices = xp.array(indices, dtype=int) markers_in_box = self.markers[indices] return markers_in_box def get_destinations_box(self): """Find the destination proc for the particles to communicate for the box structure.""" - self._send_info_box = np.zeros(self.mpi_size, dtype=int) - self._send_list_box = [np.zeros((0, self.n_cols))] * self.mpi_size + self._send_info_box = xp.zeros(self.mpi_size, dtype=int) + self._send_list_box = [xp.zeros((0, self.n_cols))] * self.mpi_size # Faces # if self._x_m_proc is not None: self._send_info_box[self._x_m_proc] += len(self._markers_x_m) - self._send_list_box[self._x_m_proc] = np.concatenate((self._send_list_box[self._x_m_proc], self._markers_x_m)) + self._send_list_box[self._x_m_proc] = xp.concatenate((self._send_list_box[self._x_m_proc], self._markers_x_m)) # if self._x_p_proc is not None: self._send_info_box[self._x_p_proc] += len(self._markers_x_p) - self._send_list_box[self._x_p_proc] = np.concatenate((self._send_list_box[self._x_p_proc], self._markers_x_p)) + self._send_list_box[self._x_p_proc] = xp.concatenate((self._send_list_box[self._x_p_proc], self._markers_x_p)) # if self._y_m_proc is not None: self._send_info_box[self._y_m_proc] += len(self._markers_y_m) - self._send_list_box[self._y_m_proc] = np.concatenate((self._send_list_box[self._y_m_proc], self._markers_y_m)) + self._send_list_box[self._y_m_proc] = xp.concatenate((self._send_list_box[self._y_m_proc], self._markers_y_m)) # if self._y_p_proc is not None: self._send_info_box[self._y_p_proc] += len(self._markers_y_p) - self._send_list_box[self._y_p_proc] = np.concatenate((self._send_list_box[self._y_p_proc], self._markers_y_p)) + self._send_list_box[self._y_p_proc] = xp.concatenate((self._send_list_box[self._y_p_proc], self._markers_y_p)) # if self._z_m_proc is not None: self._send_info_box[self._z_m_proc] += len(self._markers_z_m) - self._send_list_box[self._z_m_proc] = np.concatenate((self._send_list_box[self._z_m_proc], self._markers_z_m)) + self._send_list_box[self._z_m_proc] = xp.concatenate((self._send_list_box[self._z_m_proc], self._markers_z_m)) # if self._z_p_proc is not None: self._send_info_box[self._z_p_proc] += len(self._markers_z_p) - self._send_list_box[self._z_p_proc] = np.concatenate((self._send_list_box[self._z_p_proc], self._markers_z_p)) + self._send_list_box[self._z_p_proc] = xp.concatenate((self._send_list_box[self._z_p_proc], self._markers_z_p)) # x-y edges # if self._x_m_y_m_proc is not None: self._send_info_box[self._x_m_y_m_proc] += len(self._markers_x_m_y_m) - self._send_list_box[self._x_m_y_m_proc] = np.concatenate( + self._send_list_box[self._x_m_y_m_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_m_proc], self._markers_x_m_y_m) ) # if self._x_m_y_p_proc is not None: self._send_info_box[self._x_m_y_p_proc] += len(self._markers_x_m_y_p) - self._send_list_box[self._x_m_y_p_proc] = np.concatenate( + self._send_list_box[self._x_m_y_p_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_p_proc], self._markers_x_m_y_p) ) # if self._x_p_y_m_proc is not None: self._send_info_box[self._x_p_y_m_proc] += len(self._markers_x_p_y_m) - self._send_list_box[self._x_p_y_m_proc] = np.concatenate( + self._send_list_box[self._x_p_y_m_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_m_proc], self._markers_x_p_y_m) ) # if self._x_p_y_p_proc is not None: self._send_info_box[self._x_p_y_p_proc] += len(self._markers_x_p_y_p) - self._send_list_box[self._x_p_y_p_proc] = np.concatenate( + self._send_list_box[self._x_p_y_p_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_p_proc], self._markers_x_p_y_p) ) # x-z edges # if self._x_m_z_m_proc is not None: self._send_info_box[self._x_m_z_m_proc] += len(self._markers_x_m_z_m) - self._send_list_box[self._x_m_z_m_proc] = np.concatenate( + self._send_list_box[self._x_m_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_m_z_m_proc], self._markers_x_m_z_m) ) # if self._x_m_z_p_proc is not None: self._send_info_box[self._x_m_z_p_proc] += len(self._markers_x_m_z_p) - self._send_list_box[self._x_m_z_p_proc] = np.concatenate( + self._send_list_box[self._x_m_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_m_z_p_proc], self._markers_x_m_z_p) ) # if self._x_p_z_m_proc is not None: self._send_info_box[self._x_p_z_m_proc] += len(self._markers_x_p_z_m) - self._send_list_box[self._x_p_z_m_proc] = np.concatenate( + self._send_list_box[self._x_p_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_p_z_m_proc], self._markers_x_p_z_m) ) # if self._x_p_z_p_proc is not None: self._send_info_box[self._x_p_z_p_proc] += len(self._markers_x_p_z_p) - self._send_list_box[self._x_p_z_p_proc] = np.concatenate( + self._send_list_box[self._x_p_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_p_z_p_proc], self._markers_x_p_z_p) ) # y-z edges # if self._y_m_z_m_proc is not None: self._send_info_box[self._y_m_z_m_proc] += len(self._markers_y_m_z_m) - self._send_list_box[self._y_m_z_m_proc] = np.concatenate( + self._send_list_box[self._y_m_z_m_proc] = xp.concatenate( (self._send_list_box[self._y_m_z_m_proc], self._markers_y_m_z_m) ) # if self._y_m_z_p_proc is not None: self._send_info_box[self._y_m_z_p_proc] += len(self._markers_y_m_z_p) - self._send_list_box[self._y_m_z_p_proc] = np.concatenate( + self._send_list_box[self._y_m_z_p_proc] = xp.concatenate( (self._send_list_box[self._y_m_z_p_proc], self._markers_y_m_z_p) ) # if self._y_p_z_m_proc is not None: self._send_info_box[self._y_p_z_m_proc] += len(self._markers_y_p_z_m) - self._send_list_box[self._y_p_z_m_proc] = np.concatenate( + self._send_list_box[self._y_p_z_m_proc] = xp.concatenate( (self._send_list_box[self._y_p_z_m_proc], self._markers_y_p_z_m) ) # if self._y_p_z_p_proc is not None: self._send_info_box[self._y_p_z_p_proc] += len(self._markers_y_p_z_p) - self._send_list_box[self._y_p_z_p_proc] = np.concatenate( + self._send_list_box[self._y_p_z_p_proc] = xp.concatenate( (self._send_list_box[self._y_p_z_p_proc], self._markers_y_p_z_p) ) # corners # if self._x_m_y_m_z_m_proc is not None: self._send_info_box[self._x_m_y_m_z_m_proc] += len(self._markers_x_m_y_m_z_m) - self._send_list_box[self._x_m_y_m_z_m_proc] = np.concatenate( + self._send_list_box[self._x_m_y_m_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_m_z_m_proc], self._markers_x_m_y_m_z_m) ) # if self._x_m_y_m_z_p_proc is not None: self._send_info_box[self._x_m_y_m_z_p_proc] += len(self._markers_x_m_y_m_z_p) - self._send_list_box[self._x_m_y_m_z_p_proc] = np.concatenate( + self._send_list_box[self._x_m_y_m_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_m_z_p_proc], self._markers_x_m_y_m_z_p) ) # if self._x_m_y_p_z_m_proc is not None: self._send_info_box[self._x_m_y_p_z_m_proc] += len(self._markers_x_m_y_p_z_m) - self._send_list_box[self._x_m_y_p_z_m_proc] = np.concatenate( + self._send_list_box[self._x_m_y_p_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_p_z_m_proc], self._markers_x_m_y_p_z_m) ) # if self._x_m_y_p_z_p_proc is not None: self._send_info_box[self._x_m_y_p_z_p_proc] += len(self._markers_x_m_y_p_z_p) - self._send_list_box[self._x_m_y_p_z_p_proc] = np.concatenate( + self._send_list_box[self._x_m_y_p_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_p_z_p_proc], self._markers_x_m_y_p_z_p) ) # if self._x_p_y_m_z_m_proc is not None: self._send_info_box[self._x_p_y_m_z_m_proc] += len(self._markers_x_p_y_m_z_m) - self._send_list_box[self._x_p_y_m_z_m_proc] = np.concatenate( + self._send_list_box[self._x_p_y_m_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_m_z_m_proc], self._markers_x_p_y_m_z_m) ) # if self._x_p_y_m_z_p_proc is not None: self._send_info_box[self._x_p_y_m_z_p_proc] += len(self._markers_x_p_y_m_z_p) - self._send_list_box[self._x_p_y_m_z_p_proc] = np.concatenate( + self._send_list_box[self._x_p_y_m_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_m_z_p_proc], self._markers_x_p_y_m_z_p) ) # if self._x_p_y_p_z_m_proc is not None: self._send_info_box[self._x_p_y_p_z_m_proc] += len(self._markers_x_p_y_p_z_m) - self._send_list_box[self._x_p_y_p_z_m_proc] = np.concatenate( + self._send_list_box[self._x_p_y_p_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_p_z_m_proc], self._markers_x_p_y_p_z_m) ) # if self._x_p_y_p_z_p_proc is not None: self._send_info_box[self._x_p_y_p_z_p_proc] += len(self._markers_x_p_y_p_z_p) - self._send_list_box[self._x_p_y_p_z_p_proc] = np.concatenate( + self._send_list_box[self._x_p_y_p_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_p_z_p_proc], self._markers_x_p_y_p_z_p) ) @@ -3144,7 +3144,7 @@ def self_communication_boxes(self): if self._send_info_box[self.mpi_rank] > 0: self.update_holes() - holes_inds = np.nonzero(self.holes)[0] + holes_inds = xp.nonzero(self.holes)[0] if holes_inds.size < self._send_info_box[self.mpi_rank]: warnings.warn( @@ -3166,16 +3166,16 @@ def self_communication_boxes(self): # self.update_holes() # self.update_ghost_particles() # self.update_valid_mks() - # holes_inds = np.nonzero(self.holes)[0] + # holes_inds = xp.nonzero(self.holes)[0] - self.markers[holes_inds[np.arange(self._send_info_box[self.mpi_rank])]] = self._send_list_box[self.mpi_rank] + self.markers[holes_inds[xp.arange(self._send_info_box[self.mpi_rank])]] = self._send_list_box[self.mpi_rank] @profile def communicate_boxes(self, verbose=False): # if verbose: - # n_valid = np.count_nonzero(self.valid_mks) - # n_holes = np.count_nonzero(self.holes) - # n_ghosts = np.count_nonzero(self.ghost_particles) + # n_valid = xp.count_nonzero(self.valid_mks) + # n_holes = xp.count_nonzero(self.holes) + # n_ghosts = xp.count_nonzero(self.ghost_particles) # print(f"before communicate_boxes: {self.mpi_rank = }, {n_valid = } {n_holes = }, {n_ghosts = }") self.prepare_ghost_particles() @@ -3190,9 +3190,9 @@ def communicate_boxes(self, verbose=False): self.update_ghost_particles() # if verbose: - # n_valid = np.count_nonzero(self.valid_mks) - # n_holes = np.count_nonzero(self.holes) - # n_ghosts = np.count_nonzero(self.ghost_particles) + # n_valid = xp.count_nonzero(self.valid_mks) + # n_holes = xp.count_nonzero(self.holes) + # n_ghosts = xp.count_nonzero(self.ghost_particles) # print(f"after communicate_boxes: {self.mpi_rank = }, {n_valid = }, {n_holes = }, {n_ghosts = }") def sendrecv_all_to_all_boxes(self): @@ -3201,7 +3201,7 @@ def sendrecv_all_to_all_boxes(self): for the communication of particles in boundary boxes. """ - self._recv_info_box = np.zeros(self.mpi_comm.Get_size(), dtype=int) + self._recv_info_box = xp.zeros(self.mpi_comm.Get_size(), dtype=int) self.mpi_comm.Alltoall(self._send_info_box, self._recv_info_box) @@ -3212,8 +3212,8 @@ def sendrecv_markers_boxes(self): """ # i-th entry holds the number (not the index) of the first hole to be filled by data from process i - first_hole = np.cumsum(self._recv_info_box) - self._recv_info_box - hole_inds = np.nonzero(self._holes)[0] + first_hole = xp.cumsum(self._recv_info_box) - self._recv_info_box + hole_inds = xp.nonzero(self._holes)[0] # Initialize send and receive commands reqs = [] recvbufs = [] @@ -3224,7 +3224,7 @@ def sendrecv_markers_boxes(self): else: self.mpi_comm.Isend(data, dest=i, tag=self.mpi_comm.Get_rank()) - recvbufs += [np.zeros((N_recv, self._markers.shape[1]), dtype=float)] + recvbufs += [xp.zeros((N_recv, self._markers.shape[1]), dtype=float)] reqs += [self.mpi_comm.Irecv(recvbufs[-1], source=i, tag=i)] # Wait for buffer, then put markers into holes @@ -3247,7 +3247,7 @@ def sendrecv_markers_boxes(self): self.mpi_comm.Abort() # exit() - self._markers[hole_inds[first_hole[i] + np.arange(self._recv_info_box[i])]] = recvbufs[i] + self._markers[hole_inds[first_hole[i] + xp.arange(self._recv_info_box[i])]] = recvbufs[i] test_reqs.pop() reqs[i] = None @@ -3722,11 +3722,11 @@ def eval_density( def eval_sph( self, - eta1: np.ndarray, - eta2: np.ndarray, - eta3: np.ndarray, + eta1: xp.ndarray, + eta2: xp.ndarray, + eta3: xp.ndarray, index: int, - out: np.ndarray = None, + out: xp.ndarray = None, fast: bool = True, kernel_type: str = "gaussian_1d", derivative: int = "0", @@ -3772,12 +3772,12 @@ def eval_sph( h1, h2, h3 : float Radius of the smoothing kernel in each dimension. """ - _shp = np.shape(eta1) - assert _shp == np.shape(eta2) == np.shape(eta3) + _shp = xp.shape(eta1) + assert _shp == xp.shape(eta2) == xp.shape(eta3) if out is not None: - assert _shp == np.shape(out) + assert _shp == xp.shape(out) else: - out = np.zeros_like(eta1) + out = xp.zeros_like(eta1) assert derivative in {0, 1, 2, 3}, f"derivative must be 0, 1, 2 or 3, but is {derivative}." @@ -3860,7 +3860,7 @@ def update_ghost_particles(self): def sendrecv_determine_mtbs( self, - alpha: list | tuple | np.ndarray = (1.0, 1.0, 1.0), + alpha: list | tuple | xp.ndarray = (1.0, 1.0, 1.0), ): """ Determine which markers have to be sent from current process and put them in a new array. @@ -3882,34 +3882,34 @@ def sendrecv_determine_mtbs( Eta-values of shape (n_send, :) according to which the sorting is performed. """ # position that determines the sorting (including periodic shift of boundary conditions) - if not isinstance(alpha, np.ndarray): - alpha = np.array(alpha, dtype=float) + if not isinstance(alpha, xp.ndarray): + alpha = xp.array(alpha, dtype=float) assert alpha.size == 3 - assert np.all(alpha >= 0.0) and np.all(alpha <= 1.0) + assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) bi = self.first_pusher_idx - self._sorting_etas = np.mod( + self._sorting_etas = xp.mod( alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) + (1.0 - alpha) * self.markers[:, bi : bi + 3], 1.0, ) # check which particles are on the current process domain - self._is_on_proc_domain = np.logical_and( + self._is_on_proc_domain = xp.logical_and( self._sorting_etas > self.domain_array[self.mpi_rank, 0::3], self._sorting_etas < self.domain_array[self.mpi_rank, 1::3], ) # to stay on the current process, all three columns must be True - self._can_stay = np.all(self._is_on_proc_domain, axis=1) + self._can_stay = xp.all(self._is_on_proc_domain, axis=1) # holes and ghosts can stay, too self._can_stay[self.holes] = True self._can_stay[self.ghost_particles] = True # True values can stay on the process, False must be sent, already empty rows (-1) cannot be sent - send_inds = np.nonzero(~self._can_stay)[0] + send_inds = xp.nonzero(~self._can_stay)[0] - hole_inds_after_send = np.nonzero(np.logical_or(~self._can_stay, self.holes))[0] + hole_inds_after_send = xp.nonzero(xp.logical_or(~self._can_stay, self.holes))[0] return hole_inds_after_send, send_inds @@ -3928,16 +3928,16 @@ def sendrecv_get_destinations(self, send_inds): """ # One entry for each process - send_info = np.zeros(self.mpi_size, dtype=int) + send_info = xp.zeros(self.mpi_size, dtype=int) # TODO: do not loop over all processes, start with neighbours and work outwards (using while) for i in range(self.mpi_size): - conds = np.logical_and( + conds = xp.logical_and( self._sorting_etas[send_inds] > self.domain_array[i, 0::3], self._sorting_etas[send_inds] < self.domain_array[i, 1::3], ) - self._send_to_i[i] = np.nonzero(np.all(conds, axis=1))[0] + self._send_to_i[i] = xp.nonzero(xp.all(conds, axis=1))[0] send_info[i] = self._send_to_i[i].size self._send_list[i] = self.markers[send_inds][self._send_to_i[i]] @@ -3959,7 +3959,7 @@ def sendrecv_all_to_all(self, send_info): Amount of marticles to be received from i-th process. """ - recv_info = np.zeros(self.mpi_size, dtype=int) + recv_info = xp.zeros(self.mpi_size, dtype=int) self.mpi_comm.Alltoall(send_info, recv_info) @@ -3979,7 +3979,7 @@ def sendrecv_markers(self, recv_info, hole_inds_after_send): """ # i-th entry holds the number (not the index) of the first hole to be filled by data from process i - first_hole = np.cumsum(recv_info) - recv_info + first_hole = xp.cumsum(recv_info) - recv_info # Initialize send and receive commands for i, (data, N_recv) in enumerate(zip(self._send_list, list(recv_info))): @@ -3989,7 +3989,7 @@ def sendrecv_markers(self, recv_info, hole_inds_after_send): else: self.mpi_comm.Isend(data, dest=i, tag=self.mpi_rank) - self._recvbufs[i] = np.zeros((N_recv, self.markers.shape[1]), dtype=float) + self._recvbufs[i] = xp.zeros((N_recv, self.markers.shape[1]), dtype=float) self._reqs[i] = self.mpi_comm.Irecv(self._recvbufs[i], source=i, tag=i) # Wait for buffer, then put markers into holes @@ -4011,12 +4011,12 @@ def sendrecv_markers(self, recv_info, hole_inds_after_send): ) self.mpi_comm.Abort() - self.markers[hole_inds_after_send[first_hole[i] + np.arange(recv_info[i])]] = self._recvbufs[i] + self.markers[hole_inds_after_send[first_hole[i] + xp.arange(recv_info[i])]] = self._recvbufs[i] test_reqs.pop() self._reqs[i] = None - def _gather_scalar_in_subcomm_array(self, scalar: int, out: np.ndarray = None): + def _gather_scalar_in_subcomm_array(self, scalar: int, out: xp.ndarray = None): """Return an array of length sub_comm.size, where the i-th entry corresponds to the value of the scalar on process i. @@ -4025,11 +4025,11 @@ def _gather_scalar_in_subcomm_array(self, scalar: int, out: np.ndarray = None): scalar : int The scalar value on each process. - out : np.ndarray + out : xp.ndarray The returned array (optional). """ if out is None: - _tmp = np.zeros(self.mpi_size, dtype=int) + _tmp = xp.zeros(self.mpi_size, dtype=int) else: assert out.size == self.mpi_size _tmp = out @@ -4044,7 +4044,7 @@ def _gather_scalar_in_subcomm_array(self, scalar: int, out: np.ndarray = None): return _tmp - def _gather_scalar_in_intercomm_array(self, scalar: int, out: np.ndarray = None): + def _gather_scalar_in_intercomm_array(self, scalar: int, out: xp.ndarray = None): """Return an array of length inter_comm.size, where the i-th entry corresponds to the value of the scalar on clone i. @@ -4053,11 +4053,11 @@ def _gather_scalar_in_intercomm_array(self, scalar: int, out: np.ndarray = None) scalar : int The scalar value on each clone. - out : np.ndarray + out : xp.ndarray The returned array (optional). """ if out is None: - _tmp = np.zeros(self.num_clones, dtype=int) + _tmp = xp.zeros(self.num_clones, dtype=int) else: assert out.size == self.num_clones _tmp = out @@ -4086,7 +4086,7 @@ class Tesselation: comm : Intracomm MPI communicator. - domain_array : np.ndarray + domain_array : xp.ndarray A 2d array[float] of shape (comm.Get_size(), 9) holding info on the domain decomposition. sorting_boxes : Particles.SortingBoxes @@ -4098,7 +4098,7 @@ def __init__( tiles_pb: int | float, *, comm: Intracomm = None, - domain_array: np.ndarray = None, + domain_array: xp.ndarray = None, sorting_boxes: Particles.SortingBoxes = None, ): if isinstance(tiles_pb, int): @@ -4116,8 +4116,8 @@ def __init__( assert domain_array is not None if domain_array is None: - self._starts = np.zeros(3) - self._ends = np.ones(3) + self._starts = xp.zeros(3) + self._ends = xp.ones(3) else: self._starts = domain_array[self.rank, 0::3] self._ends = domain_array[self.rank, 1::3] @@ -4140,9 +4140,9 @@ def __init__( if n_boxes == 1: self._dims_mask = [True] * 3 else: - self._dims_mask = np.array(self.boxes_per_dim) > 1 + self._dims_mask = xp.array(self.boxes_per_dim) > 1 - min_tiles = 2 ** np.count_nonzero(self.dims_mask) + min_tiles = 2 ** xp.count_nonzero(self.dims_mask) assert self.tiles_pb >= min_tiles, ( f"At least {min_tiles} tiles per sorting box is enforced, but you have {self.tiles_pb}!" ) @@ -4165,19 +4165,19 @@ def get_tiles(self): # print(f'{self.dims_mask = }') # tiles in one sorting box - self._nt_per_dim = np.array([1, 1, 1]) - _ids = np.nonzero(self._dims_mask)[0] + self._nt_per_dim = xp.array([1, 1, 1]) + _ids = xp.nonzero(self._dims_mask)[0] for fac in factors_vec: _nt = self.nt_per_dim[self._dims_mask] - d = _ids[np.argmin(_nt)] + d = _ids[xp.argmin(_nt)] self._nt_per_dim[d] *= fac # print(f'{_nt = }, {d = }, {self.nt_per_dim = }') - assert np.prod(self.nt_per_dim) == self.tiles_pb + assert xp.prod(self.nt_per_dim) == self.tiles_pb # tiles between [0, box_width] in each direction - self._tile_breaks = [np.linspace(0.0, bw, nt + 1) for bw, nt in zip(self.box_widths, self.nt_per_dim)] - self._tile_midpoints = [(np.roll(tbs, -1)[:-1] + tbs[:-1]) / 2 for tbs in self.tile_breaks] + self._tile_breaks = [xp.linspace(0.0, bw, nt + 1) for bw, nt in zip(self.box_widths, self.nt_per_dim)] + self._tile_midpoints = [(xp.roll(tbs, -1)[:-1] + tbs[:-1]) / 2 for tbs in self.tile_breaks] self._tile_volume = 1.0 for tb in self.tile_breaks: self._tile_volume *= tb[1] @@ -4185,8 +4185,8 @@ def get_tiles(self): def draw_markers(self): """Draw markers on the tile midpoints.""" _, eta1 = self._tile_output_arrays() - eta2 = np.zeros_like(eta1) - eta3 = np.zeros_like(eta1) + eta2 = xp.zeros_like(eta1) + eta3 = xp.zeros_like(eta1) nt_x, nt_y, nt_z = self.nt_per_dim @@ -4197,7 +4197,7 @@ def draw_markers(self): for k in range(self.boxes_per_dim[2]): z_midpoints = self._get_midpoints(k, 2) - xx, yy, zz = np.meshgrid( + xx, yy, zz = xp.meshgrid( x_midpoints, y_midpoints, z_midpoints, @@ -4234,7 +4234,7 @@ def _get_quad_pts(self, n_quad=None): self._tile_quad_pts = [] self._tile_quad_wts = [] for nq, tb in zip(n_quad, self.tile_breaks): - pts_loc, wts_loc = np.polynomial.legendre.leggauss(nq) + pts_loc, wts_loc = xp.polynomial.legendre.leggauss(nq) pts, wts = quadrature_grid(tb[:2], pts_loc, wts_loc) self._tile_quad_pts += [pts[0]] self._tile_quad_wts += [wts[0]] @@ -4261,7 +4261,7 @@ def cell_averages(self, fun, n_quad=None): for k in range(self.boxes_per_dim[2]): z_pts = self._get_box_quad_pts(k, 2) - xx, yy, zz = np.meshgrid( + xx, yy, zz = xp.meshgrid( x_pts.flatten(), y_pts.flatten(), z_pts.flatten(), @@ -4290,9 +4290,9 @@ def _tile_output_arrays(self): * the first with one entry for each tile on one sorting box * the second with one entry for each tile on current process """ - # self._quad_pts = [np.zeros((nt, nq)).flatten() for nt, nq in zip(self.nt_per_dim, self.tile_quad_pts)] - single_box_out = np.zeros(self.nt_per_dim) - out = np.tile(single_box_out, self.boxes_per_dim) + # self._quad_pts = [xp.zeros((nt, nq)).flatten() for nt, nq in zip(self.nt_per_dim, self.tile_quad_pts)] + single_box_out = xp.zeros(self.nt_per_dim) + out = xp.tile(single_box_out, self.boxes_per_dim) return single_box_out, out def _get_midpoints(self, i: int, dim: int): @@ -4313,13 +4313,13 @@ def _get_box_quad_pts(self, i: int, dim: int): Returns ------- - x_pts : np.array + x_pts : xp.array 2d array of shape (n_tiles_pb, n_tile_quad_pts) """ xl = self.starts[dim] + i * self.box_widths[dim] x_tile_breaks = xl + self.tile_breaks[dim][:-1] x_tile_pts = self.tile_quad_pts[dim] - x_pts = np.tile(x_tile_breaks, (x_tile_pts.size, 1)).T + x_tile_pts + x_pts = xp.tile(x_tile_breaks, (x_tile_pts.size, 1)).T + x_tile_pts return x_pts @property diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 39de5e4e6..f662bbac2 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -10,7 +10,7 @@ from struphy.kinetic_background.base import Maxwellian, SumKineticBackground from struphy.pic import utilities_kernels from struphy.pic.base import Particles -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Particles6D(Particles): diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index f0380eb23..53287371c 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -6,7 +6,7 @@ from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments from struphy.pic.base import Particles from struphy.profiling.profiling import ProfileManager -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel @@ -134,7 +134,7 @@ def __init__( comps = ker_args[2] # check marker array column number - assert isinstance(comps, np.ndarray) + assert isinstance(comps, xp.ndarray) assert column_nr + comps.size < particles.n_cols, ( f"{column_nr + comps.size} not smaller than {particles.n_cols = }; not enough columns in marker array !!" ) @@ -146,7 +146,7 @@ def __init__( comps = ker_args[3] # check marker array column number - assert isinstance(comps, np.ndarray) + assert isinstance(comps, xp.ndarray) assert column_nr + comps.size < particles.n_cols, ( f"{column_nr + comps.size} not smaller than {particles.n_cols = }; not enough columns in marker array !!" ) @@ -154,7 +154,7 @@ def __init__( self._init_kernels = init_kernels self._eval_kernels = eval_kernels - self._residuals = np.zeros(self.particles.markers.shape[0]) + self._residuals = xp.zeros(self.particles.markers.shape[0]) self._converged_loc = self._residuals == 1.0 self._not_converged_loc = self._residuals == 0.0 @@ -209,7 +209,7 @@ def __call__(self, dt: float): add_args = ker_args[3] ker( - np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + xp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), column_nr, comps, self.particles.args_markers, @@ -224,7 +224,7 @@ def __call__(self, dt: float): # start stages (e.g. n_stages=4 for RK4) for stage in range(self.n_stages): # start iteration (maxiter=1 for explicit schemes) - n_not_converged = np.empty(1, dtype=int) + n_not_converged = xp.empty(1, dtype=int) n_not_converged[0] = self.particles.n_mks_loc k = 0 @@ -298,12 +298,12 @@ def __call__(self, dt: float): # compute number of non-converged particles (maxiter=1 for explicit schemes) if self.maxiter > 1: self._residuals[:] = markers[:, residual_idx] - max_res = np.max(self._residuals) + max_res = xp.max(self._residuals) if max_res < 0.0: max_res = None self._converged_loc[:] = self._residuals < self._tol self._not_converged_loc[:] = ~self._converged_loc - n_not_converged[0] = np.count_nonzero( + n_not_converged[0] = xp.count_nonzero( self._not_converged_loc, ) diff --git a/src/struphy/pic/pushing/pusher_kernels.py b/src/struphy/pic/pushing/pusher_kernels.py index 429e1c722..cdaa9059b 100644 --- a/src/struphy/pic/pushing/pusher_kernels.py +++ b/src/struphy/pic/pushing/pusher_kernels.py @@ -3035,7 +3035,7 @@ def push_v_sph_pressure( h1, h2, h3 : float Kernel width in respective dimension. - gravity: np.ndarray + gravity: xp.ndarray Constant gravitational force as 3-vector. """ # allocate arrays @@ -3266,7 +3266,7 @@ def push_v_sph_pressure_ideal_gas( h1, h2, h3 : float Kernel width in respective dimension. - gravity: np.ndarray + gravity: xp.ndarray Constant gravitational force as 3-vector. """ # allocate arrays @@ -3498,7 +3498,7 @@ def push_v_viscosity( h1, h2, h3 : float Kernel width in respective dimension. - gravity: np.ndarray + gravity: xp.ndarray Constant gravitational force as 3-vector. """ # allocate arrays diff --git a/src/struphy/pic/sampling_kernels.py b/src/struphy/pic/sampling_kernels.py index 821363a97..ce68d5aff 100644 --- a/src/struphy/pic/sampling_kernels.py +++ b/src/struphy/pic/sampling_kernels.py @@ -93,13 +93,13 @@ def tile_int_kernel( Parameters ---------- - fun: np.ndarray + fun: xp.ndarray The integrand evaluated at the quadrature points (meshgrid). - x_wts, y_wts, z_wts: np.ndarray + x_wts, y_wts, z_wts: xp.ndarray Quadrature weights for tile integral. - out: np.ndarray + out: xp.ndarray The result holding all tile integrals in one sorting box.""" _shp = shape(out) diff --git a/src/struphy/pic/sobol_seq.py b/src/struphy/pic/sobol_seq.py index ff073b1b3..9208d813d 100644 --- a/src/struphy/pic/sobol_seq.py +++ b/src/struphy/pic/sobol_seq.py @@ -19,7 +19,7 @@ from scipy.stats import norm -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp __all__ = ["i4_bit_hi1", "i4_bit_lo0", "i4_sobol_generate", "i4_sobol", "i4_uniform", "prime_ge", "is_prime"] @@ -60,7 +60,7 @@ def i4_bit_hi1(n): Output, integer BIT, the number of bits base 2. """ - i = np.floor(n) + i = xp.floor(n) bit = 0 while i > 0: bit += 1 @@ -105,7 +105,7 @@ def i4_bit_lo0(n): Output, integer BIT, the position of the low 1 bit. """ bit = 1 - i = np.floor(n) + i = xp.floor(n) while i != 2 * (i // 2): bit += 1 i //= 2 @@ -123,7 +123,7 @@ def i4_sobol_generate(dim_num, n, skip=1): Output, real R(M,N), the points. """ - r = np.full((n, dim_num), np.nan) + r = xp.full((n, dim_num), xp.nan) for j in range(n): seed = j + skip r[j, 0:dim_num], next_seed = i4_sobol(dim_num, seed) @@ -222,8 +222,8 @@ def i4_sobol(dim_num, seed): seed_save = -1 # Initialize (part of) V. - v = np.zeros((dim_max, log_max)) - v[0:40, 0] = np.transpose( + v = xp.zeros((dim_max, log_max)) + v[0:40, 0] = xp.transpose( [ 1, 1, @@ -268,7 +268,7 @@ def i4_sobol(dim_num, seed): ] ) - v[2:40, 1] = np.transpose( + v[2:40, 1] = xp.transpose( [ 1, 3, @@ -311,7 +311,7 @@ def i4_sobol(dim_num, seed): ] ) - v[3:40, 2] = np.transpose( + v[3:40, 2] = xp.transpose( [ 7, 5, @@ -353,7 +353,7 @@ def i4_sobol(dim_num, seed): ] ) - v[5:40, 3] = np.transpose( + v[5:40, 3] = xp.transpose( [ 1, 7, @@ -393,7 +393,7 @@ def i4_sobol(dim_num, seed): ] ) - v[7:40, 4] = np.transpose( + v[7:40, 4] = xp.transpose( [ 9, 3, @@ -431,15 +431,15 @@ def i4_sobol(dim_num, seed): ] ) - v[13:40, 5] = np.transpose( + v[13:40, 5] = xp.transpose( [37, 33, 7, 5, 11, 39, 63, 27, 17, 15, 23, 29, 3, 21, 13, 31, 25, 9, 49, 33, 19, 29, 11, 19, 27, 15, 25] ) - v[19:40, 6] = np.transpose( + v[19:40, 6] = xp.transpose( [13, 33, 115, 41, 79, 17, 29, 119, 75, 73, 105, 7, 59, 65, 21, 3, 113, 61, 89, 45, 107] ) - v[37:40, 7] = np.transpose([7, 23, 39]) + v[37:40, 7] = xp.transpose([7, 23, 39]) # Set POLY. poly = [ @@ -518,7 +518,7 @@ def i4_sobol(dim_num, seed): # Expand this bit pattern to separate components of the logical array INCLUD. j = poly[i - 1] - includ = np.zeros(m) + includ = xp.zeros(m) for k in range(m, 0, -1): j2 = j // 2 includ[k - 1] = j != 2 * j2 @@ -532,7 +532,7 @@ def i4_sobol(dim_num, seed): for k in range(1, m + 1): l *= 2 if includ[k - 1]: - newv = np.bitwise_xor(int(newv), int(l * v[i - 1, j - k - 1])) + newv = xp.bitwise_xor(int(newv), int(l * v[i - 1, j - k - 1])) v[i - 1, j - 1] = newv # Multiply columns of V by appropriate power of 2. @@ -543,16 +543,16 @@ def i4_sobol(dim_num, seed): # RECIPD is 1/(common denominator of the elements in V). recipd = 1.0 / (2 * l) - lastq = np.zeros(dim_num) + lastq = xp.zeros(dim_num) - seed = int(np.floor(seed)) + seed = int(xp.floor(seed)) if seed < 0: seed = 0 l = 1 if seed == 0: - lastq = np.zeros(dim_num) + lastq = xp.zeros(dim_num) elif seed == seed_save + 1: # Find the position of the right-hand zero in SEED. @@ -560,12 +560,12 @@ def i4_sobol(dim_num, seed): elif seed <= seed_save: seed_save = 0 - lastq = np.zeros(dim_num) + lastq = xp.zeros(dim_num) for seed_temp in range(int(seed_save), int(seed)): l = i4_bit_lo0(seed_temp) for i in range(1, dim_num + 1): - lastq[i - 1] = np.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) + lastq[i - 1] = xp.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) l = i4_bit_lo0(seed) @@ -573,7 +573,7 @@ def i4_sobol(dim_num, seed): for seed_temp in range(int(seed_save + 1), int(seed)): l = i4_bit_lo0(seed_temp) for i in range(1, dim_num + 1): - lastq[i - 1] = np.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) + lastq[i - 1] = xp.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) l = i4_bit_lo0(seed) @@ -586,10 +586,10 @@ def i4_sobol(dim_num, seed): return # Calculate the new components of QUASI. - quasi = np.zeros(dim_num) + quasi = xp.zeros(dim_num) for i in range(1, dim_num + 1): quasi[i - 1] = lastq[i - 1] * recipd - lastq[i - 1] = np.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) + lastq[i - 1] = xp.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) seed_save = seed seed += 1 @@ -639,11 +639,11 @@ def i4_uniform(a, b, seed): print("I4_UNIFORM - Fatal error!") print(" Input SEED = 0!") - seed = np.floor(seed) + seed = xp.floor(seed) a = round(a) b = round(b) - seed = np.mod(seed, 2147483647) + seed = xp.mod(seed, 2147483647) if seed < 0: seed += 2147483647 @@ -697,7 +697,7 @@ def prime_ge(n): Output, integer P, the smallest prime number that is greater than or equal to N. """ - p = max(np.ceil(n), 2) + p = max(xp.ceil(n), 2) while not is_prime(p): p += 1 @@ -721,7 +721,7 @@ def is_prime(n): return False # All primes >3 are of the form 6n+1 or 6n+5 (6n, 6n+2, 6n+4 are 2-divisible, 6n+3 is 3-divisible) p = 5 - root = int(np.ceil(np.sqrt(n))) + root = int(xp.ceil(xp.sqrt(n))) while p <= root: if n % p == 0 or n % (p + 2) == 0: return False diff --git a/src/struphy/pic/tests/test_accum_vec_H1.py b/src/struphy/pic/tests/test_accum_vec_H1.py index ccec231d8..ab7f3e2aa 100644 --- a/src/struphy/pic/tests/test_accum_vec_H1.py +++ b/src/struphy/pic/tests/test_accum_vec_H1.py @@ -57,7 +57,7 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): from struphy.pic.accumulation.particles_to_grid import AccumulatorVector from struphy.pic.particles import Particles6D from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig if isinstance(MPI.COMM_WORLD, MockComm): @@ -76,7 +76,7 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): params = { "grid": {"Nel": Nel}, - "kinetic": {"test_particles": {"markers": {"Np": Np, "ppc": Np / np.prod(Nel)}}}, + "kinetic": {"test_particles": {"markers": {"Np": Np, "ppc": Np / xp.prod(Nel)}}}, } if mpi_comm is None: clone_config = None @@ -129,12 +129,12 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): _w0 = particles.weights print("Test weights:") - print(f"rank {mpi_rank}:", _w0.shape, np.min(_w0), np.max(_w0)) + print(f"rank {mpi_rank}:", _w0.shape, xp.min(_w0), xp.max(_w0)) _sqrtg = domain.jacobian_det(0.5, 0.5, 0.5) - assert np.isclose(np.min(_w0), _sqrtg) - assert np.isclose(np.max(_w0), _sqrtg) + assert xp.isclose(xp.min(_w0), _sqrtg) + assert xp.isclose(xp.max(_w0), _sqrtg) # mass operators mass_ops = WeightedMassOperators(derham, domain) @@ -151,19 +151,19 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): acc() # sum all MC integrals - _sum_within_clone = np.empty(1, dtype=float) - _sum_within_clone[0] = np.sum(acc.vectors[0].toarray()) + _sum_within_clone = xp.empty(1, dtype=float) + _sum_within_clone[0] = xp.sum(acc.vectors[0].toarray()) if clone_config is not None: clone_config.sub_comm.Allreduce(MPI.IN_PLACE, _sum_within_clone, op=MPI.SUM) print(f"rank {mpi_rank}: {_sum_within_clone = }, {_sqrtg = }") # Check within clone - assert np.isclose(_sum_within_clone, _sqrtg) + assert xp.isclose(_sum_within_clone, _sqrtg) # Check for all clones - _sum_between_clones = np.empty(1, dtype=float) - _sum_between_clones[0] = np.sum(acc.vectors[0].toarray()) + _sum_between_clones = xp.empty(1, dtype=float) + _sum_between_clones[0] = xp.sum(acc.vectors[0].toarray()) if mpi_comm is not None: mpi_comm.Allreduce(MPI.IN_PLACE, _sum_between_clones, op=MPI.SUM) @@ -172,7 +172,7 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): print(f"rank {mpi_rank}: {_sum_between_clones = }, {_sqrtg = }") # Check within clone - assert np.isclose(_sum_between_clones, _sqrtg) + assert xp.isclose(_sum_between_clones, _sqrtg) if __name__ == "__main__": diff --git a/src/struphy/pic/tests/test_accumulation.py b/src/struphy/pic/tests/test_accumulation.py index 599773e83..693df515e 100644 --- a/src/struphy/pic/tests/test_accumulation.py +++ b/src/struphy/pic/tests/test_accumulation.py @@ -61,7 +61,7 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): from struphy.pic.particles import Particles6D from struphy.pic.tests.test_pic_legacy_files.accumulation_kernels_3d import kernel_step_ph_full from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp if isinstance(MPI.COMM_WORLD, MockComm): mpi_comm = None @@ -107,17 +107,17 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): particles.markers[ ~particles.holes, 6, - ] = np.random.rand(particles.n_mks_loc) + ] = xp.random.rand(particles.n_mks_loc) # gather all particles for legacy kernel if mpi_comm is None: - marker_shapes = np.array([particles.markers.shape[0]]) + marker_shapes = xp.array([particles.markers.shape[0]]) else: - marker_shapes = np.zeros(mpi_size, dtype=int) - mpi_comm.Allgather(np.array([particles.markers.shape[0]]), marker_shapes) + marker_shapes = xp.zeros(mpi_size, dtype=int) + mpi_comm.Allgather(xp.array([particles.markers.shape[0]]), marker_shapes) print(rank, marker_shapes) - particles_leg = np.zeros( + particles_leg = xp.zeros( (sum(marker_shapes), particles.markers.shape[1]), dtype=float, ) @@ -128,7 +128,7 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): cumulative_lengths = marker_shapes[0] for i in range(1, mpi_size): - arr_recv = np.zeros( + arr_recv = xp.zeros( (marker_shapes[i], particles.markers.shape[1]), dtype=float, ) @@ -161,10 +161,10 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): for a in range(3): Ni = SPACES.Nbase_1form[a] - vec[a] = np.zeros((Ni[0], Ni[1], Ni[2], 3), dtype=float) + vec[a] = xp.zeros((Ni[0], Ni[1], Ni[2], 3), dtype=float) for b in range(3): - mat[a][b] = np.zeros( + mat[a][b] = xp.zeros( ( Ni[0], Ni[1], @@ -186,21 +186,21 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): SPACES.T[0], SPACES.T[1], SPACES.T[2], - np.array(SPACES.p), - np.array(Nel), - np.array(SPACES.NbaseN), - np.array(SPACES.NbaseD), + xp.array(SPACES.p), + xp.array(Nel), + xp.array(SPACES.NbaseN), + xp.array(SPACES.NbaseD), particles_leg.shape[0], domain.kind_map, domain.params_numpy, domain.T[0], domain.T[1], domain.T[2], - np.array(domain.p), - np.array( + xp.array(domain.p), + xp.array( domain.Nel, ), - np.array(domain.NbaseN), + xp.array(domain.NbaseN), domain.cx, domain.cy, domain.cz, @@ -217,7 +217,7 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): ) end_time = time() - tot_time = np.round(end_time - start_time, 3) + tot_time = xp.round(end_time - start_time, 3) mat[0][0] /= Np mat[0][1] /= Np @@ -250,7 +250,7 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): ACC(1.0, 1.0, 0.0) end_time = time() - tot_time = np.round(end_time - start_time, 3) + tot_time = xp.round(end_time - start_time, 3) if rank == 0 and verbose: print(f"Step ph New took {tot_time} seconds.") diff --git a/src/struphy/pic/tests/test_binning.py b/src/struphy/pic/tests/test_binning.py index 909732cca..e9e69ae73 100644 --- a/src/struphy/pic/tests/test_binning.py +++ b/src/struphy/pic/tests/test_binning.py @@ -47,7 +47,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): LoadingParameters, WeightsParameters, ) - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # Set seed seed = 1234 @@ -79,7 +79,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): # test weights particles.initialize_weights() - v1_bins = np.linspace(-5.0, 5.0, 200, endpoint=True) + v1_bins = xp.linspace(-5.0, 5.0, 200, endpoint=True) dv = v1_bins[1] - v1_bins[0] binned_res, r2 = particles.binning( @@ -89,7 +89,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): v1_plot = v1_bins[:-1] + dv / 2 - ana_res = 1.0 / np.sqrt(2.0 * np.pi) * np.exp(-(v1_plot**2) / 2.0) + ana_res = 1.0 / xp.sqrt(2.0 * xp.pi) * xp.exp(-(v1_plot**2) / 2.0) if show_plot: plt.plot(v1_plot, ana_res, label="Analytical result") @@ -100,7 +100,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - binned_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - binned_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.02, f"Error between binned data and analytical result was {l2_error}" @@ -122,7 +122,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): particles.draw_markers() particles.initialize_weights() - e1_bins = np.linspace(0.0, 1.0, 200, endpoint=True) + e1_bins = xp.linspace(0.0, 1.0, 200, endpoint=True) de = e1_bins[1] - e1_bins[0] binned_res, r2 = particles.binning( @@ -132,7 +132,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): e1_plot = e1_bins[:-1] + de / 2 - ana_res = 1.0 + amp_n * np.cos(2 * np.pi * l_n * e1_plot) + ana_res = 1.0 + amp_n * xp.cos(2 * xp.pi * l_n * e1_plot) if show_plot: plt.plot(e1_plot, ana_res, label="Analytical result") @@ -143,7 +143,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - binned_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - binned_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.02, f"Error between binned data and analytical result was {l2_error}" @@ -182,7 +182,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): particles.draw_markers() particles.initialize_weights() - e1_bins = np.linspace(0.0, 1.0, 200, endpoint=True) + e1_bins = xp.linspace(0.0, 1.0, 200, endpoint=True) de = e1_bins[1] - e1_bins[0] binned_res, r2 = particles.binning( @@ -192,7 +192,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): e1_plot = e1_bins[:-1] + de / 2 - ana_res = n1 + amp_n1 * np.cos(2 * np.pi * l_n1 * e1_plot) + n2 + amp_n2 * np.cos(2 * np.pi * l_n2 * e1_plot) + ana_res = n1 + amp_n1 * xp.cos(2 * xp.pi * l_n1 * e1_plot) + n2 + amp_n2 * xp.cos(2 * xp.pi * l_n2 * e1_plot) # Compare s0 and the sum of two Maxwellians if show_plot: @@ -206,14 +206,14 @@ def test_binning_6D_full_f(mapping, show_plot=False): vth3=(particles.loading_params.moments[5], None), ) - v1 = np.linspace(-10.0, 10.0, 400) - phase_space = np.meshgrid( - np.array([0.0]), - np.array([0.0]), - np.array([0.0]), + v1 = xp.linspace(-10.0, 10.0, 400) + phase_space = xp.meshgrid( + xp.array([0.0]), + xp.array([0.0]), + xp.array([0.0]), v1, - np.array([0.0]), - np.array([0.0]), + xp.array([0.0]), + xp.array([0.0]), ) s0_vals = s0(*phase_space).squeeze() @@ -235,7 +235,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - binned_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - binned_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.04, f"Error between binned data and analytical result was {l2_error}" @@ -280,7 +280,7 @@ def test_binning_6D_delta_f(mapping, show_plot=False): LoadingParameters, WeightsParameters, ) - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # Set seed seed = 1234 @@ -316,7 +316,7 @@ def test_binning_6D_delta_f(mapping, show_plot=False): particles.draw_markers() particles.initialize_weights() - e1_bins = np.linspace(0.0, 1.0, 200, endpoint=True) + e1_bins = xp.linspace(0.0, 1.0, 200, endpoint=True) de = e1_bins[1] - e1_bins[0] binned_res, r2 = particles.binning( @@ -326,7 +326,7 @@ def test_binning_6D_delta_f(mapping, show_plot=False): e1_plot = e1_bins[:-1] + de / 2 - ana_res = amp_n * np.cos(2 * np.pi * l_n * e1_plot) + ana_res = amp_n * xp.cos(2 * xp.pi * l_n * e1_plot) if show_plot: plt.plot(e1_plot, ana_res, label="Analytical result") @@ -337,7 +337,7 @@ def test_binning_6D_delta_f(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - binned_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - binned_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.02, f"Error between binned data and analytical result was {l2_error}" @@ -376,7 +376,7 @@ def test_binning_6D_delta_f(mapping, show_plot=False): particles.draw_markers() particles.initialize_weights() - e1_bins = np.linspace(0.0, 1.0, 200, endpoint=True) + e1_bins = xp.linspace(0.0, 1.0, 200, endpoint=True) de = e1_bins[1] - e1_bins[0] binned_res, r2 = particles.binning( @@ -386,7 +386,7 @@ def test_binning_6D_delta_f(mapping, show_plot=False): e1_plot = e1_bins[:-1] + de / 2 - ana_res = amp_n1 * np.cos(2 * np.pi * l_n1 * e1_plot) + amp_n2 * np.cos(2 * np.pi * l_n2 * e1_plot) + ana_res = amp_n1 * xp.cos(2 * xp.pi * l_n1 * e1_plot) + amp_n2 * xp.cos(2 * xp.pi * l_n2 * e1_plot) # Compare s0 and the sum of two Maxwellians if show_plot: @@ -400,14 +400,14 @@ def test_binning_6D_delta_f(mapping, show_plot=False): vth3=(particles.loading_params.moments[5], None), ) - v1 = np.linspace(-10.0, 10.0, 400) - phase_space = np.meshgrid( - np.array([0.0]), - np.array([0.0]), - np.array([0.0]), + v1 = xp.linspace(-10.0, 10.0, 400) + phase_space = xp.meshgrid( + xp.array([0.0]), + xp.array([0.0]), + xp.array([0.0]), v1, - np.array([0.0]), - np.array([0.0]), + xp.array([0.0]), + xp.array([0.0]), ) s0_vals = s0(*phase_space).squeeze() @@ -429,7 +429,7 @@ def test_binning_6D_delta_f(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - binned_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - binned_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.04, f"Error between binned data and analytical result was {l2_error}" @@ -477,7 +477,7 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): LoadingParameters, WeightsParameters, ) - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # Set seed seed = 1234 @@ -519,7 +519,7 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): # test weights particles.initialize_weights() - v1_bins = np.linspace(-5.0, 5.0, 200, endpoint=True) + v1_bins = xp.linspace(-5.0, 5.0, 200, endpoint=True) dv = v1_bins[1] - v1_bins[0] binned_res, r2 = particles.binning( @@ -531,13 +531,13 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): if comm is None: mpi_res = binned_res else: - mpi_res = np.zeros_like(binned_res) + mpi_res = xp.zeros_like(binned_res) comm.Allreduce(binned_res, mpi_res, op=MPI.SUM) comm.Barrier() v1_plot = v1_bins[:-1] + dv / 2 - ana_res = 1.0 / np.sqrt(2.0 * np.pi) * np.exp(-(v1_plot**2) / 2.0) + ana_res = 1.0 / xp.sqrt(2.0 * xp.pi) * xp.exp(-(v1_plot**2) / 2.0) if show_plot and rank == 0: plt.plot(v1_plot, ana_res, label="Analytical result") @@ -548,7 +548,7 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - mpi_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - mpi_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.03, f"Error between binned data and analytical result was {l2_error}" @@ -571,7 +571,7 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): particles.draw_markers() particles.initialize_weights() - e1_bins = np.linspace(0.0, 1.0, 200, endpoint=True) + e1_bins = xp.linspace(0.0, 1.0, 200, endpoint=True) de = e1_bins[1] - e1_bins[0] binned_res, r2 = particles.binning( @@ -583,13 +583,13 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): if comm is None: mpi_res = binned_res else: - mpi_res = np.zeros_like(binned_res) + mpi_res = xp.zeros_like(binned_res) comm.Allreduce(binned_res, mpi_res, op=MPI.SUM) comm.Barrier() e1_plot = e1_bins[:-1] + de / 2 - ana_res = 1.0 + amp_n * np.cos(2 * np.pi * l_n * e1_plot) + ana_res = 1.0 + amp_n * xp.cos(2 * xp.pi * l_n * e1_plot) if show_plot and rank == 0: plt.plot(e1_plot, ana_res, label="Analytical result") @@ -600,7 +600,7 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - mpi_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - mpi_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.03, f"Error between binned data and analytical result was {l2_error}" @@ -668,7 +668,7 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): particles.draw_markers() particles.initialize_weights() - e1_bins = np.linspace(0.0, 1.0, 200, endpoint=True) + e1_bins = xp.linspace(0.0, 1.0, 200, endpoint=True) de = e1_bins[1] - e1_bins[0] binned_res, r2 = particles.binning( @@ -680,13 +680,13 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): if comm is None: mpi_res = binned_res else: - mpi_res = np.zeros_like(binned_res) + mpi_res = xp.zeros_like(binned_res) comm.Allreduce(binned_res, mpi_res, op=MPI.SUM) comm.Barrier() e1_plot = e1_bins[:-1] + de / 2 - ana_res = n1 + amp_n1 * np.cos(2 * np.pi * l_n1 * e1_plot) + n2 + amp_n2 * np.cos(2 * np.pi * l_n2 * e1_plot) + ana_res = n1 + amp_n1 * xp.cos(2 * xp.pi * l_n1 * e1_plot) + n2 + amp_n2 * xp.cos(2 * xp.pi * l_n2 * e1_plot) # Compare s0 and the sum of two Maxwellians if show_plot and rank == 0: @@ -700,14 +700,14 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): vth3=(particles.loading_params.moments[5], None), ) - v1 = np.linspace(-10.0, 10.0, 400) - phase_space = np.meshgrid( - np.array([0.0]), - np.array([0.0]), - np.array([0.0]), + v1 = xp.linspace(-10.0, 10.0, 400) + phase_space = xp.meshgrid( + xp.array([0.0]), + xp.array([0.0]), + xp.array([0.0]), v1, - np.array([0.0]), - np.array([0.0]), + xp.array([0.0]), + xp.array([0.0]), ) s0_vals = s0(*phase_space).squeeze() @@ -729,7 +729,7 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - mpi_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - mpi_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.04, f"Error between binned data and analytical result was {l2_error}" @@ -774,7 +774,7 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): LoadingParameters, WeightsParameters, ) - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # Set seed seed = 1234 @@ -830,7 +830,7 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): particles.draw_markers() particles.initialize_weights() - e1_bins = np.linspace(0.0, 1.0, 200, endpoint=True) + e1_bins = xp.linspace(0.0, 1.0, 200, endpoint=True) de = e1_bins[1] - e1_bins[0] binned_res, r2 = particles.binning( @@ -842,13 +842,13 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): if comm is None: mpi_res = binned_res else: - mpi_res = np.zeros_like(binned_res) + mpi_res = xp.zeros_like(binned_res) comm.Allreduce(binned_res, mpi_res, op=MPI.SUM) comm.Barrier() e1_plot = e1_bins[:-1] + de / 2 - ana_res = amp_n * np.cos(2 * np.pi * l_n * e1_plot) + ana_res = amp_n * xp.cos(2 * xp.pi * l_n * e1_plot) if show_plot and rank == 0: plt.plot(e1_plot, ana_res, label="Analytical result") @@ -859,7 +859,7 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - mpi_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - mpi_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.02, f"Error between binned data and analytical result was {l2_error}" @@ -929,7 +929,7 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): particles.draw_markers() particles.initialize_weights() - e1_bins = np.linspace(0.0, 1.0, 200, endpoint=True) + e1_bins = xp.linspace(0.0, 1.0, 200, endpoint=True) de = e1_bins[1] - e1_bins[0] binned_res, r2 = particles.binning( @@ -941,13 +941,13 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): if comm is None: mpi_res = binned_res else: - mpi_res = np.zeros_like(binned_res) + mpi_res = xp.zeros_like(binned_res) comm.Allreduce(binned_res, mpi_res, op=MPI.SUM) comm.Barrier() e1_plot = e1_bins[:-1] + de / 2 - ana_res = amp_n1 * np.cos(2 * np.pi * l_n1 * e1_plot) + amp_n2 * np.cos(2 * np.pi * l_n2 * e1_plot) + ana_res = amp_n1 * xp.cos(2 * xp.pi * l_n1 * e1_plot) + amp_n2 * xp.cos(2 * xp.pi * l_n2 * e1_plot) # Compare s0 and the sum of two Maxwellians if show_plot and rank == 0: @@ -961,14 +961,14 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): vth3=(particles.loading_params.moments[5], None), ) - v1 = np.linspace(-10.0, 10.0, 400) - phase_space = np.meshgrid( - np.array([0.0]), - np.array([0.0]), - np.array([0.0]), + v1 = xp.linspace(-10.0, 10.0, 400) + phase_space = xp.meshgrid( + xp.array([0.0]), + xp.array([0.0]), + xp.array([0.0]), v1, - np.array([0.0]), - np.array([0.0]), + xp.array([0.0]), + xp.array([0.0]), ) s0_vals = s0(*phase_space).squeeze() @@ -990,7 +990,7 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): plt.legend() plt.show() - l2_error = np.sqrt(np.sum((ana_res - mpi_res) ** 2)) / np.sqrt(np.sum((ana_res) ** 2)) + l2_error = xp.sqrt(xp.sum((ana_res - mpi_res) ** 2)) / xp.sqrt(xp.sum((ana_res) ** 2)) assert l2_error <= 0.04, f"Error between binned data and analytical result was {l2_error}" diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index 29db9beca..5174c83b5 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -41,7 +41,7 @@ def test_draw(Nel, p, spl_kind, mapping, ppc=10): from struphy.geometry import domains from struphy.pic.particles import Particles6D from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -85,7 +85,7 @@ def test_draw(Nel, p, spl_kind, mapping, ppc=10): particles.initialize_weights() _w0 = particles.weights print("Test weights:") - print(f"rank {rank}:", _w0.shape, np.min(_w0), np.max(_w0)) + print(f"rank {rank}:", _w0.shape, xp.min(_w0), xp.max(_w0)) comm.Barrier() print("Number of particles w/wo holes on each process before sorting : ") @@ -106,17 +106,17 @@ def test_draw(Nel, p, spl_kind, mapping, ppc=10): print("Rank", rank, ":", particles.n_mks_loc, particles.markers.shape[0]) # are all markers in the correct domain? - conds = np.logical_and( + conds = xp.logical_and( particles.markers[:, :3] > derham.domain_array[rank, 0::3], particles.markers[:, :3] < derham.domain_array[rank, 1::3], ) holes = particles.markers[:, 0] == -1.0 - stay = np.all(conds, axis=1) + stay = xp.all(conds, axis=1) - error_mks = particles.markers[np.logical_and(~stay, ~holes)] + error_mks = particles.markers[xp.logical_and(~stay, ~holes)] assert error_mks.size == 0, ( - f"rank {rank} | markers not on correct process: {np.nonzero(np.logical_and(~stay, ~holes))} \n corresponding positions:\n {error_mks[:, :3]}" + f"rank {rank} | markers not on correct process: {xp.nonzero(xp.logical_and(~stay, ~holes))} \n corresponding positions:\n {error_mks[:, :3]}" ) diff --git a/src/struphy/pic/tests/test_mat_vec_filler.py b/src/struphy/pic/tests/test_mat_vec_filler.py index 7edbf7278..c3ad580f4 100644 --- a/src/struphy/pic/tests/test_mat_vec_filler.py +++ b/src/struphy/pic/tests/test_mat_vec_filler.py @@ -1,6 +1,6 @@ import pytest -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @@ -33,12 +33,12 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): print(f"\nNel={Nel}, p={p}, spl_kind={spl_kind}\n") # DR attributes - pn = np.array(DR.p) + pn = xp.array(DR.p) tn1, tn2, tn3 = DR.Vh_fem["0"].knots starts1 = {} - starts1["v0"] = np.array(DR.Vh["0"].starts) + starts1["v0"] = xp.array(DR.Vh["0"].starts) comm.Barrier() sleep(0.02 * (rank + 1)) @@ -94,14 +94,14 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): vec["v2"] += [StencilVector(DR.Vh["2"].spaces[i])._data] # Some filling for testing - fill_mat = np.reshape(np.arange(9, dtype=float), (3, 3)) + 1.0 - fill_vec = np.arange(3, dtype=float) + 1.0 + fill_mat = xp.reshape(xp.arange(9, dtype=float), (3, 3)) + 1.0 + fill_vec = xp.arange(3, dtype=float) + 1.0 # Random points in domain of process (VERY IMPORTANT to be in the right domain, otherwise NON-TRACKED errors occur in filler_kernels !!) dom = DR.domain_array[rank] - eta1s = np.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] - eta2s = np.random.rand(n_markers) * (dom[4] - dom[3]) + dom[3] - eta3s = np.random.rand(n_markers) * (dom[7] - dom[6]) + dom[6] + eta1s = xp.random.rand(n_markers) * (dom[1] - dom[0]) + dom[0] + eta2s = xp.random.rand(n_markers) * (dom[4] - dom[3]) + dom[3] + eta3s = xp.random.rand(n_markers) * (dom[7] - dom[6]) + dom[6] for eta1, eta2, eta3 in zip(eta1s, eta2s, eta3s): comm.Barrier() @@ -118,13 +118,13 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): span3 = bsp.find_span(tn3, DR.p[2], eta3) # non-zero spline values at eta - bn1 = np.empty(DR.p[0] + 1, dtype=float) - bn2 = np.empty(DR.p[1] + 1, dtype=float) - bn3 = np.empty(DR.p[2] + 1, dtype=float) + bn1 = xp.empty(DR.p[0] + 1, dtype=float) + bn2 = xp.empty(DR.p[1] + 1, dtype=float) + bn3 = xp.empty(DR.p[2] + 1, dtype=float) - bd1 = np.empty(DR.p[0], dtype=float) - bd2 = np.empty(DR.p[1], dtype=float) - bd3 = np.empty(DR.p[2], dtype=float) + bd1 = xp.empty(DR.p[0], dtype=float) + bd2 = xp.empty(DR.p[1], dtype=float) + bd3 = xp.empty(DR.p[2], dtype=float) bsp.b_d_splines_slim(tn1, DR.p[0], eta1, span1, bn1, bd1) bsp.b_d_splines_slim(tn2, DR.p[1], eta2, span2, bn2, bd2) @@ -136,9 +136,9 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): ie3 = span3 - pn[2] # global indices of non-vanishing B- and D-splines (no modulo) - glob_n1 = np.arange(ie1, ie1 + pn[0] + 1) - glob_n2 = np.arange(ie2, ie2 + pn[1] + 1) - glob_n3 = np.arange(ie3, ie3 + pn[2] + 1) + glob_n1 = xp.arange(ie1, ie1 + pn[0] + 1) + glob_n2 = xp.arange(ie2, ie2 + pn[1] + 1) + glob_n3 = xp.arange(ie3, ie3 + pn[2] + 1) glob_d1 = glob_n1[:-1] glob_d2 = glob_n2[:-1] @@ -164,10 +164,10 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): # local column indices in _data of non-vanishing B- and D-splines, as sets for comparison cols = [{}, {}, {}] for n in range(3): - cols[n]["NN"] = set(np.arange(2 * pn[n] + 1)) - cols[n]["ND"] = set(np.arange(2 * pn[n])) - cols[n]["DN"] = set(np.arange(1, 2 * pn[n] + 1)) - cols[n]["DD"] = set(np.arange(1, 2 * pn[n])) + cols[n]["NN"] = set(xp.arange(2 * pn[n] + 1)) + cols[n]["ND"] = set(xp.arange(2 * pn[n])) + cols[n]["DN"] = set(xp.arange(1, 2 * pn[n] + 1)) + cols[n]["DD"] = set(xp.arange(1, 2 * pn[n])) # testing vector-valued spaces spaces_vector = ["v1", "v2"] @@ -337,23 +337,23 @@ def assert_mat(mat, rows, cols, row_str, col_str, rank, verbose=False): """ assert len(mat.shape) == 6 # assert non NaN - assert ~np.isnan(mat).any() + assert ~xp.isnan(mat).any() atol = 1e-14 if verbose: print(f"\n({row_str}) ({col_str})") - print(f"rank {rank} | ind_row1: {set(np.where(mat > atol)[0])}") - print(f"rank {rank} | ind_row2: {set(np.where(mat > atol)[1])}") - print(f"rank {rank} | ind_row3: {set(np.where(mat > atol)[2])}") - print(f"rank {rank} | ind_col1: {set(np.where(mat > atol)[3])}") - print(f"rank {rank} | ind_col2: {set(np.where(mat > atol)[4])}") - print(f"rank {rank} | ind_col3: {set(np.where(mat > atol)[5])}") + print(f"rank {rank} | ind_row1: {set(xp.where(mat > atol)[0])}") + print(f"rank {rank} | ind_row2: {set(xp.where(mat > atol)[1])}") + print(f"rank {rank} | ind_row3: {set(xp.where(mat > atol)[2])}") + print(f"rank {rank} | ind_col1: {set(xp.where(mat > atol)[3])}") + print(f"rank {rank} | ind_col2: {set(xp.where(mat > atol)[4])}") + print(f"rank {rank} | ind_col3: {set(xp.where(mat > atol)[5])}") # check if correct indices are non-zero for n, (r, c) in enumerate(zip(row_str, col_str)): - assert set(np.where(mat > atol)[n]) == rows[n][r] - assert set(np.where(mat > atol)[n + 3]) == cols[n][r + c] + assert set(xp.where(mat > atol)[n]) == rows[n][r] + assert set(xp.where(mat > atol)[n + 3]) == cols[n][r + c] # Set matrix back to zero mat[:, :] = 0.0 @@ -384,19 +384,19 @@ def assert_vec(vec, rows, row_str, rank, verbose=False): """ assert len(vec.shape) == 3 # assert non Nan - assert ~np.isnan(vec).any() + assert ~xp.isnan(vec).any() atol = 1e-14 if verbose: print(f"\n({row_str})") - print(f"rank {rank} | ind_row1: {set(np.where(vec > atol)[0])}") - print(f"rank {rank} | ind_row2: {set(np.where(vec > atol)[1])}") - print(f"rank {rank} | ind_row3: {set(np.where(vec > atol)[2])}") + print(f"rank {rank} | ind_row1: {set(xp.where(vec > atol)[0])}") + print(f"rank {rank} | ind_row2: {set(xp.where(vec > atol)[1])}") + print(f"rank {rank} | ind_row3: {set(xp.where(vec > atol)[2])}") # check if correct indices are non-zero for n, r in enumerate(row_str): - assert set(np.where(vec > atol)[n]) == rows[n][r] + assert set(xp.where(vec > atol)[n]) == rows[n][r] # Set vector back to zero vec[:] = 0.0 diff --git a/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py b/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py index 4b0cbc7ae..82c88d7cf 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py @@ -12,7 +12,7 @@ from psydac.ddm.mpi import mpi as MPI import struphy.pic.tests.test_pic_legacy_files.accumulation_kernels_3d as pic_ker_3d -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # import struphy.pic.tests.test_pic_legacy_files.accumulation_kernels_2d as pic_ker_2d @@ -69,22 +69,22 @@ def __init__(self, tensor_space_FEM, domain, basis_u, mpi_comm, use_control, cv_ else: Ni = getattr(self.space, "Nbase_" + str(self.basis_u) + "form")[a] - self.vecs_loc[a] = np.empty((Ni[0], Ni[1], Ni[2]), dtype=float) - self.vecs_glo[a] = np.empty((Ni[0], Ni[1], Ni[2]), dtype=float) + self.vecs_loc[a] = xp.empty((Ni[0], Ni[1], Ni[2]), dtype=float) + self.vecs_glo[a] = xp.empty((Ni[0], Ni[1], Ni[2]), dtype=float) for b in range(3): if self.space.dim == 2: - self.blocks_loc[a][b] = np.empty( + self.blocks_loc[a][b] = xp.empty( (Ni[0], Ni[1], Ni[2], 2 * self.space.p[0] + 1, 2 * self.space.p[1] + 1, self.space.NbaseN[2]), dtype=float, ) - self.blocks_glo[a][b] = np.empty( + self.blocks_glo[a][b] = xp.empty( (Ni[0], Ni[1], Ni[2], 2 * self.space.p[0] + 1, 2 * self.space.p[1] + 1, self.space.NbaseN[2]), dtype=float, ) else: - self.blocks_loc[a][b] = np.empty( + self.blocks_loc[a][b] = xp.empty( ( Ni[0], Ni[1], @@ -95,7 +95,7 @@ def __init__(self, tensor_space_FEM, domain, basis_u, mpi_comm, use_control, cv_ ), dtype=float, ) - self.blocks_glo[a][b] = np.empty( + self.blocks_glo[a][b] = xp.empty( ( Ni[0], Ni[1], @@ -134,16 +134,16 @@ def to_sparse_step1(self): Ni = self.space.Nbase_2form[a] Nj = self.space.Nbase_2form[b] - indices = np.indices(self.blocks_glo[a][b].shape) + indices = xp.indices(self.blocks_glo[a][b].shape) row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() - shift = [np.arange(Ni) - p for Ni, p in zip(Ni[:2], self.space.p[:2])] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni[:2], self.space.p[:2])] if self.space.dim == 2: - shift += [np.zeros(self.space.NbaseN[2], dtype=int)] + shift += [xp.zeros(self.space.NbaseN[2], dtype=int)] else: - shift += [np.arange(Ni[2]) - self.space.p[2]] + shift += [xp.arange(Ni[2]) - self.space.p[2]] col1 = (indices[3] + shift[0][:, None, None, None, None, None]) % Nj[0] col2 = (indices[4] + shift[1][None, :, None, None, None, None]) % Nj[1] @@ -201,16 +201,16 @@ def to_sparse_step3(self): Ni = self.space.Nbase_2form[a] Nj = self.space.Nbase_2form[b] - indices = np.indices(self.blocks_glo[a][b].shape) + indices = xp.indices(self.blocks_glo[a][b].shape) row = (Ni[1] * Ni[2] * indices[0] + Ni[2] * indices[1] + indices[2]).flatten() - shift = [np.arange(Ni) - p for Ni, p in zip(Ni[:2], self.space.p[:2])] + shift = [xp.arange(Ni) - p for Ni, p in zip(Ni[:2], self.space.p[:2])] if self.space.dim == 2: - shift += [np.zeros(self.space.NbaseN[2], dtype=int)] + shift += [xp.zeros(self.space.NbaseN[2], dtype=int)] else: - shift += [np.arange(Ni[2]) - self.space.p[2]] + shift += [xp.arange(Ni[2]) - self.space.p[2]] col1 = (indices[3] + shift[0][:, None, None, None, None, None]) % Nj[0] col2 = (indices[4] + shift[1][None, :, None, None, None, None]) % Nj[1] @@ -528,15 +528,15 @@ def assemble_step3(self, b2_eq, b2): # build global sparse matrix and global vector if self.basis_u == 0: return self.to_sparse_step3(), self.space.Ev_0.dot( - np.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())) + xp.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())) ) elif self.basis_u == 1: return self.to_sparse_step3(), self.space.E1_0.dot( - np.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())) + xp.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())) ) elif self.basis_u == 2: return self.to_sparse_step3(), self.space.E2_0.dot( - np.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())) + xp.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())) ) diff --git a/src/struphy/pic/tests/test_pic_legacy_files/pusher.py b/src/struphy/pic/tests/test_pic_legacy_files/pusher.py index 6bdb74642..b609b69d8 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/pusher.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/pusher.py @@ -1,7 +1,7 @@ import struphy.pic.tests.test_pic_legacy_files.pusher_pos as push_pos import struphy.pic.tests.test_pic_legacy_files.pusher_vel_2d as push_vel_2d import struphy.pic.tests.test_pic_legacy_files.pusher_vel_3d as push_vel_3d -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Pusher: diff --git a/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_2d.py b/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_2d.py index ba32b93bf..fdd4485b5 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_2d.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_2d.py @@ -400,7 +400,7 @@ def evaluate_tensor_product( Returns: -------- - values: double[:, :] values of spline at points from np.meshgrid(eta1, eta2, indexing='ij'). + values: double[:, :] values of spline at points from xp.meshgrid(eta1, eta2, indexing='ij'). """ for i1 in range(len(eta1)): diff --git a/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_3d.py b/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_3d.py index 28e2b5d9c..a40d12fc9 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_3d.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_3d.py @@ -806,7 +806,7 @@ def evaluate_tensor_product( Returns: -------- values: double[:, :, :] values of spline at points from - np.meshgrid(eta1, eta2, eta3, indexing='ij'). + xp.meshgrid(eta1, eta2, eta3, indexing='ij'). """ for i1 in range(len(eta1)): diff --git a/src/struphy/pic/tests/test_pushers.py b/src/struphy/pic/tests/test_pushers.py index f766acf53..4809b7918 100644 --- a/src/struphy/pic/tests/test_pushers.py +++ b/src/struphy/pic/tests/test_pushers.py @@ -34,7 +34,7 @@ def test_push_vxb_analytic(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -125,7 +125,7 @@ def test_push_vxb_analytic(Nel, p, spl_kind, mapping, show_plots=False): ) # compare if markers are the same BEFORE push - assert np.allclose(particles.markers, markers_str.T) + assert xp.allclose(particles.markers, markers_str.T) # push markers dt = 0.1 @@ -135,7 +135,7 @@ def test_push_vxb_analytic(Nel, p, spl_kind, mapping, show_plots=False): pusher_psy(dt) # compare if markers are the same AFTER push - assert np.allclose(particles.markers[:, :6], markers_str.T[:, :6]) + assert xp.allclose(particles.markers[:, :6], markers_str.T[:, :6]) @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @@ -169,7 +169,7 @@ def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -250,8 +250,8 @@ def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): basis_u=2, bc_pos=0, ) - mu0_str = np.zeros(markers_str.shape[1], dtype=float) - pow_str = np.zeros(markers_str.shape[1], dtype=float) + mu0_str = xp.zeros(markers_str.shape[1], dtype=float) + pow_str = xp.zeros(markers_str.shape[1], dtype=float) pusher_psy = Pusher_psy( particles, @@ -271,7 +271,7 @@ def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): ) # compare if markers are the same BEFORE push - assert np.allclose(particles.markers, markers_str.T) + assert xp.allclose(particles.markers, markers_str.T) # push markers dt = 0.1 @@ -281,7 +281,7 @@ def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): pusher_psy(dt) # compare if markers are the same AFTER push - assert np.allclose(particles.markers[:, :6], markers_str.T[:, :6]) + assert xp.allclose(particles.markers[:, :6], markers_str.T[:, :6]) @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @@ -315,7 +315,7 @@ def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -396,8 +396,8 @@ def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): basis_u=1, bc_pos=0, ) - mu0_str = np.zeros(markers_str.shape[1], dtype=float) - pow_str = np.zeros(markers_str.shape[1], dtype=float) + mu0_str = xp.zeros(markers_str.shape[1], dtype=float) + pow_str = xp.zeros(markers_str.shape[1], dtype=float) pusher_psy = Pusher_psy( particles, @@ -417,7 +417,7 @@ def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): ) # compare if markers are the same BEFORE push - assert np.allclose(particles.markers, markers_str.T) + assert xp.allclose(particles.markers, markers_str.T) # push markers dt = 0.1 @@ -427,7 +427,7 @@ def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): pusher_psy(dt) # compare if markers are the same AFTER push - assert np.allclose(particles.markers[:, :6], markers_str.T[:, :6]) + assert xp.allclose(particles.markers[:, :6], markers_str.T[:, :6]) @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @@ -461,7 +461,7 @@ def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -542,8 +542,8 @@ def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): basis_u=0, bc_pos=0, ) - mu0_str = np.zeros(markers_str.shape[1], dtype=float) - pow_str = np.zeros(markers_str.shape[1], dtype=float) + mu0_str = xp.zeros(markers_str.shape[1], dtype=float) + pow_str = xp.zeros(markers_str.shape[1], dtype=float) pusher_psy = Pusher_psy( particles, @@ -563,7 +563,7 @@ def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): ) # compare if markers are the same BEFORE push - assert np.allclose(particles.markers, markers_str.T) + assert xp.allclose(particles.markers, markers_str.T) # push markers dt = 0.1 @@ -573,7 +573,7 @@ def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): pusher_psy(dt) # compare if markers are the same AFTER push - assert np.allclose(particles.markers[:, :6], markers_str.T[:, :6]) + assert xp.allclose(particles.markers[:, :6], markers_str.T[:, :6]) @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @@ -607,7 +607,7 @@ def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -688,8 +688,8 @@ def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): basis_u=2, bc_pos=0, ) - mu0_str = np.random.rand(markers_str.shape[1]) - pow_str = np.zeros(markers_str.shape[1], dtype=float) + mu0_str = xp.random.rand(markers_str.shape[1]) + pow_str = xp.zeros(markers_str.shape[1], dtype=float) pusher_psy = Pusher_psy( particles, @@ -711,7 +711,7 @@ def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): ) # compare if markers are the same BEFORE push - assert np.allclose(particles.markers, markers_str.T) + assert xp.allclose(particles.markers, markers_str.T) # push markers dt = 0.1 @@ -721,7 +721,7 @@ def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): pusher_psy(dt) # compare if markers are the same AFTER push - assert np.allclose(particles.markers[:, :6], markers_str.T[:, :6]) + assert xp.allclose(particles.markers[:, :6], markers_str.T[:, :6]) @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @@ -756,7 +756,7 @@ def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -830,8 +830,8 @@ def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): butcher = ButcherTableau("rk4") # temp fix due to refactoring of ButcherTableau: - butcher._a = np.diag(butcher.a, k=-1) - butcher._a = np.array(list(butcher._a) + [0.0]) + butcher._a = xp.diag(butcher.a, k=-1) + butcher._a = xp.array(list(butcher._a) + [0.0]) pusher_psy = Pusher_psy( particles, @@ -843,7 +843,7 @@ def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): ) # compare if markers are the same BEFORE push - assert np.allclose(particles.markers, markers_str.T) + assert xp.allclose(particles.markers, markers_str.T) # push markers dt = 0.1 @@ -851,12 +851,12 @@ def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): pusher_str.push_step4(markers_str, dt) pusher_psy(dt) - n_mks_load = np.zeros(size, dtype=int) + n_mks_load = xp.zeros(size, dtype=int) - comm.Allgather(np.array(np.shape(particles.markers)[0]), n_mks_load) + comm.Allgather(xp.array(xp.shape(particles.markers)[0]), n_mks_load) - sendcounts = np.zeros(size, dtype=int) - displacements = np.zeros(size, dtype=int) + sendcounts = xp.zeros(size, dtype=int) + displacements = xp.zeros(size, dtype=int) accum_sendcounts = 0.0 for i in range(size): @@ -864,18 +864,18 @@ def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): displacements[i] = accum_sendcounts accum_sendcounts += sendcounts[i] - all_particles_psy = np.zeros((int(accum_sendcounts) * 3,), dtype=float) - all_particles_str = np.zeros((int(accum_sendcounts) * 3,), dtype=float) + all_particles_psy = xp.zeros((int(accum_sendcounts) * 3,), dtype=float) + all_particles_str = xp.zeros((int(accum_sendcounts) * 3,), dtype=float) comm.Barrier() - comm.Allgatherv(np.array(particles.markers[:, :3]), [all_particles_psy, sendcounts, displacements, MPI.DOUBLE]) - comm.Allgatherv(np.array(markers_str.T[:, :3]), [all_particles_str, sendcounts, displacements, MPI.DOUBLE]) + comm.Allgatherv(xp.array(particles.markers[:, :3]), [all_particles_psy, sendcounts, displacements, MPI.DOUBLE]) + comm.Allgatherv(xp.array(markers_str.T[:, :3]), [all_particles_str, sendcounts, displacements, MPI.DOUBLE]) comm.Barrier() - unique_psy = np.unique(all_particles_psy) - unique_str = np.unique(all_particles_str) + unique_psy = xp.unique(all_particles_psy) + unique_str = xp.unique(all_particles_str) - assert np.allclose(unique_psy, unique_str) + assert xp.allclose(unique_psy, unique_str) if __name__ == "__main__": diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index cc98a424c..244427209 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -7,7 +7,7 @@ from struphy.geometry import domains from struphy.pic.particles import Particles6D from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize("nx", [8, 70]) @@ -17,9 +17,9 @@ def test_flattening(nx, ny, nz, algo): from struphy.pic.sorting_kernels import flatten_index, unflatten_index - n1s = np.array(np.random.rand(10) * (nx + 1), dtype=int) - n2s = np.array(np.random.rand(10) * (ny + 1), dtype=int) - n3s = np.array(np.random.rand(10) * (nz + 1), dtype=int) + n1s = xp.array(xp.random.rand(10) * (nx + 1), dtype=int) + n2s = xp.array(xp.random.rand(10) * (ny + 1), dtype=int) + n3s = xp.array(xp.random.rand(10) * (nz + 1), dtype=int) for n1 in n1s: for n2 in n2s: for n3 in n3s: @@ -37,9 +37,9 @@ def test_flattening(nx, ny, nz, algo): def test_flattening(nx, ny, nz, algo): from struphy.pic.sorting_kernels import flatten_index, unflatten_index - n1s = np.array(np.random.rand(10) * (nx + 1), dtype=int) - n2s = np.array(np.random.rand(10) * (ny + 1), dtype=int) - n3s = np.array(np.random.rand(10) * (nz + 1), dtype=int) + n1s = xp.array(xp.random.rand(10) * (nx + 1), dtype=int) + n2s = xp.array(xp.random.rand(10) * (ny + 1), dtype=int) + n3s = xp.array(xp.random.rand(10) * (nz + 1), dtype=int) for n1 in n1s: for n2 in n2s: for n3 in n3s: @@ -57,9 +57,9 @@ def test_flattening(nx, ny, nz, algo): def test_flattening(nx, ny, nz, algo): from struphy.pic.sorting_kernels import flatten_index, unflatten_index - n1s = np.array(np.random.rand(10) * (nx + 1), dtype=int) - n2s = np.array(np.random.rand(10) * (ny + 1), dtype=int) - n3s = np.array(np.random.rand(10) * (nz + 1), dtype=int) + n1s = xp.array(xp.random.rand(10) * (nx + 1), dtype=int) + n2s = xp.array(xp.random.rand(10) * (ny + 1), dtype=int) + n3s = xp.array(xp.random.rand(10) * (nz + 1), dtype=int) for n1 in n1s: for n2 in n2s: for n3 in n3s: diff --git a/src/struphy/pic/tests/test_sph.py b/src/struphy/pic/tests/test_sph.py index 2dcaa3ae8..8e7b5b5a4 100644 --- a/src/struphy/pic/tests/test_sph.py +++ b/src/struphy/pic/tests/test_sph.py @@ -8,7 +8,7 @@ from struphy.initial import perturbations from struphy.pic.particles import ParticlesSPH from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize("boxes_per_dim", [(24, 1, 1)]) @@ -59,9 +59,9 @@ def test_sph_evaluation_1d( pert = {"n": perturbations.ModesCos(ls=(1,), amps=(1e-0,))} if derivative == 0: - fun_exact = lambda e1, e2, e3: 1.5 + np.cos(2 * np.pi * e1) + fun_exact = lambda e1, e2, e3: 1.5 + xp.cos(2 * xp.pi * e1) else: - fun_exact = lambda e1, e2, e3: -2 * np.pi * np.sin(2 * np.pi * e1) + fun_exact = lambda e1, e2, e3: -2 * xp.pi * xp.sin(2 * xp.pi * e1) boundary_params = BoundaryParameters(bc_sph=(bc_x, "periodic", "periodic")) @@ -78,9 +78,9 @@ def test_sph_evaluation_1d( ) # eval points - eta1 = np.linspace(0, 1.0, eval_pts) - eta2 = np.array([0.0]) - eta3 = np.array([0.0]) + eta1 = xp.linspace(0, 1.0, eval_pts) + eta2 = xp.array([0.0]) + eta3 = xp.array([0.0]) particles.draw_markers(sort=False, verbose=False) if comm is not None: @@ -89,7 +89,7 @@ def test_sph_evaluation_1d( h1 = 1 / boxes_per_dim[0] h2 = 1 / boxes_per_dim[1] h3 = 1 / boxes_per_dim[2] - ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + ee1, ee2, ee3 = xp.meshgrid(eta1, eta2, eta3, indexing="ij") test_eval = particles.eval_density( ee1, ee2, @@ -104,11 +104,11 @@ def test_sph_evaluation_1d( if comm is None: all_eval = test_eval else: - all_eval = np.zeros_like(test_eval) + all_eval = xp.zeros_like(test_eval) comm.Allreduce(test_eval, all_eval, op=MPI.SUM) exact_eval = fun_exact(ee1, ee2, ee3) - err_max_norm = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + err_max_norm = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) if rank == 0: print(f"\n{boxes_per_dim = }") @@ -178,19 +178,19 @@ def test_sph_evaluation_2d( pert = {"n": perturbations.ModesCosCos(ls=(1,), ms=(1,), amps=(1e-0,))} if derivative == 0: - fun_exact = lambda e1, e2, e3: 1.5 + np.cos(2 * np.pi * e1) * np.cos(2 * np.pi * e2) + fun_exact = lambda e1, e2, e3: 1.5 + xp.cos(2 * xp.pi * e1) * xp.cos(2 * xp.pi * e2) elif derivative == 1: - fun_exact = lambda e1, e2, e3: -2 * np.pi * np.sin(2 * np.pi * e1) * np.cos(2 * np.pi * e2) + fun_exact = lambda e1, e2, e3: -2 * xp.pi * xp.sin(2 * xp.pi * e1) * xp.cos(2 * xp.pi * e2) else: - fun_exact = lambda e1, e2, e3: -2 * np.pi * np.cos(2 * np.pi * e1) * np.sin(2 * np.pi * e2) + fun_exact = lambda e1, e2, e3: -2 * xp.pi * xp.cos(2 * xp.pi * e1) * xp.sin(2 * xp.pi * e2) # boundary conditions boundary_params = BoundaryParameters(bc_sph=(bc_x, bc_y, "periodic")) # eval points - eta1 = np.linspace(0, 1.0, eval_pts) - eta2 = np.linspace(0, 1.0, eval_pts) - eta3 = np.array([0.0]) + eta1 = xp.linspace(0, 1.0, eval_pts) + eta2 = xp.linspace(0, 1.0, eval_pts) + eta3 = xp.array([0.0]) # particles object particles = ParticlesSPH( @@ -213,7 +213,7 @@ def test_sph_evaluation_2d( h1 = 1 / boxes_per_dim[0] h2 = 1 / boxes_per_dim[1] h3 = 1 / boxes_per_dim[2] - ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + ee1, ee2, ee3 = xp.meshgrid(eta1, eta2, eta3, indexing="ij") test_eval = particles.eval_density( ee1, ee2, @@ -228,11 +228,11 @@ def test_sph_evaluation_2d( if comm is None: all_eval = test_eval else: - all_eval = np.zeros_like(test_eval) + all_eval = xp.zeros_like(test_eval) comm.Allreduce(test_eval, all_eval, op=MPI.SUM) exact_eval = fun_exact(ee1, ee2, ee3) - err_max_norm = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + err_max_norm = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) if rank == 0: print(f"\n{boxes_per_dim = }") @@ -308,9 +308,9 @@ def test_sph_evaluation_3d( boundary_params = BoundaryParameters(bc_sph=(bc_x, bc_y, bc_z)) # eval points - eta1 = np.linspace(0, 1.0, eval_pts) - eta2 = np.linspace(0, 1.0, eval_pts) - eta3 = np.linspace(0, 1.0, eval_pts) + eta1 = xp.linspace(0, 1.0, eval_pts) + eta2 = xp.linspace(0, 1.0, eval_pts) + eta3 = xp.linspace(0, 1.0, eval_pts) # particles object particles = ParticlesSPH( @@ -332,7 +332,7 @@ def test_sph_evaluation_3d( h1 = 1 / boxes_per_dim[0] h2 = 1 / boxes_per_dim[1] h3 = 1 / boxes_per_dim[2] - ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + ee1, ee2, ee3 = xp.meshgrid(eta1, eta2, eta3, indexing="ij") test_eval = particles.eval_density( ee1, ee2, @@ -347,11 +347,11 @@ def test_sph_evaluation_3d( if comm is None: all_eval = test_eval else: - all_eval = np.zeros_like(test_eval) + all_eval = xp.zeros_like(test_eval) comm.Allreduce(test_eval, all_eval, op=MPI.SUM) exact_eval = fun_exact(ee1, ee2, ee3) - err_max_norm = np.max(np.abs(all_eval - exact_eval)) + err_max_norm = xp.max(xp.abs(all_eval - exact_eval)) if rank == 0: print(f"\n{boxes_per_dim = }") @@ -421,17 +421,17 @@ def test_evaluation_SPH_Np_convergence_1d(boxes_per_dim, bc_x, eval_pts, tessela # perturbation]} if bc_x in ("periodic", "fixed"): - fun_exact = lambda e1, e2, e3: 1.5 - np.sin(2 * np.pi * e1) + fun_exact = lambda e1, e2, e3: 1.5 - xp.sin(2 * xp.pi * e1) pert = {"n": perturbations.ModesSin(ls=(1,), amps=(-1e-0,))} elif bc_x == "mirror": - fun_exact = lambda e1, e2, e3: 1.5 - np.cos(2 * np.pi * e1) + fun_exact = lambda e1, e2, e3: 1.5 - xp.cos(2 * xp.pi * e1) pert = {"n": perturbations.ModesCos(ls=(1,), amps=(-1e-0,))} # exact solution - eta1 = np.linspace(0, 1.0, eval_pts) # add offset for non-periodic boundary conditions, TODO: implement Neumann - eta2 = np.array([0.0]) - eta3 = np.array([0.0]) - ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + eta1 = xp.linspace(0, 1.0, eval_pts) # add offset for non-periodic boundary conditions, TODO: implement Neumann + eta2 = xp.array([0.0]) + eta3 = xp.array([0.0]) + ee1, ee2, ee3 = xp.meshgrid(eta1, eta2, eta3, indexing="ij") exact_eval = fun_exact(ee1, ee2, ee3) # boundary conditions @@ -471,7 +471,7 @@ def test_evaluation_SPH_Np_convergence_1d(boxes_per_dim, bc_x, eval_pts, tessela if comm is None: all_eval = test_eval else: - all_eval = np.zeros_like(test_eval) + all_eval = xp.zeros_like(test_eval) comm.Allreduce(test_eval, all_eval, op=MPI.SUM) if show_plot and rank == 0: @@ -481,21 +481,21 @@ def test_evaluation_SPH_Np_convergence_1d(boxes_per_dim, bc_x, eval_pts, tessela plt.title(f"{Np = }, {ppb = }") # plt.savefig(f"fun_{Np}_{ppb}.png") - diff = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + diff = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) err_vec += [diff] print(f"{Np = }, {ppb = }, {diff = }") if tesselation: - fit = np.polyfit(np.log(ppbs), np.log(err_vec), 1) + fit = xp.polyfit(xp.log(ppbs), xp.log(err_vec), 1) xvec = ppbs else: - fit = np.polyfit(np.log(Nps), np.log(err_vec), 1) + fit = xp.polyfit(xp.log(Nps), xp.log(err_vec), 1) xvec = Nps if show_plot and rank == 0: plt.figure(figsize=(12, 8)) plt.loglog(xvec, err_vec, label="Convergence") - plt.loglog(xvec, np.exp(fit[1]) * np.array(xvec) ** (fit[0]), "--", label=f"fit with slope {fit[0]}") + plt.loglog(xvec, xp.exp(fit[1]) * xp.array(xvec) ** (fit[0]), "--", label=f"fit with slope {fit[0]}") plt.legend() plt.show() # plt.savefig(f"Convergence_SPH_{tesselation=}") @@ -506,7 +506,7 @@ def test_evaluation_SPH_Np_convergence_1d(boxes_per_dim, bc_x, eval_pts, tessela if tesselation: assert fit[0] < 2e-3 else: - assert np.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate + assert xp.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate @pytest.mark.parametrize("boxes_per_dim", [(12, 1, 1)]) @@ -542,17 +542,17 @@ def test_evaluation_SPH_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselat # perturbation if bc_x in ("periodic", "fixed"): - fun_exact = lambda e1, e2, e3: 1.5 - np.sin(2 * np.pi * e1) + fun_exact = lambda e1, e2, e3: 1.5 - xp.sin(2 * xp.pi * e1) pert = {"n": perturbations.ModesSin(ls=(1,), amps=(-1e-0,))} elif bc_x == "mirror": - fun_exact = lambda e1, e2, e3: 1.5 - np.cos(2 * np.pi * e1) + fun_exact = lambda e1, e2, e3: 1.5 - xp.cos(2 * xp.pi * e1) pert = {"n": perturbations.ModesCos(ls=(1,), amps=(-1e-0,))} # exact solution - eta1 = np.linspace(0, 1.0, eval_pts) # add offset for non-periodic boundary conditions, TODO: implement Neumann - eta2 = np.array([0.0]) - eta3 = np.array([0.0]) - ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + eta1 = xp.linspace(0, 1.0, eval_pts) # add offset for non-periodic boundary conditions, TODO: implement Neumann + eta2 = xp.array([0.0]) + eta3 = xp.array([0.0]) + ee1, ee2, ee3 = xp.meshgrid(eta1, eta2, eta3, indexing="ij") exact_eval = fun_exact(ee1, ee2, ee3) # boundary conditions @@ -587,7 +587,7 @@ def test_evaluation_SPH_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselat if comm is None: all_eval = test_eval else: - all_eval = np.zeros_like(test_eval) + all_eval = xp.zeros_like(test_eval) comm.Allreduce(test_eval, all_eval, op=MPI.SUM) if show_plot and rank == 0: @@ -598,7 +598,7 @@ def test_evaluation_SPH_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselat # plt.savefig(f"fun_{h1}.png") # error in max-norm - diff = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + diff = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) print(f"{h1 = }, {diff = }") @@ -608,14 +608,14 @@ def test_evaluation_SPH_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselat err_vec += [diff] if tesselation: - fit = np.polyfit(np.log(h_vec[1:5]), np.log(err_vec[1:5]), 1) + fit = xp.polyfit(xp.log(h_vec[1:5]), xp.log(err_vec[1:5]), 1) else: - fit = np.polyfit(np.log(h_vec[:-2]), np.log(err_vec[:-2]), 1) + fit = xp.polyfit(xp.log(h_vec[:-2]), xp.log(err_vec[:-2]), 1) if show_plot and rank == 0: plt.figure(figsize=(12, 8)) plt.loglog(h_vec, err_vec, label="Convergence") - plt.loglog(h_vec, np.exp(fit[1]) * np.array(h_vec) ** (fit[0]), "--", label=f"fit with slope {fit[0]}") + plt.loglog(h_vec, xp.exp(fit[1]) * xp.array(h_vec) ** (fit[0]), "--", label=f"fit with slope {fit[0]}") plt.legend() plt.show() # plt.savefig("Convergence_SPH") @@ -624,7 +624,7 @@ def test_evaluation_SPH_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselat print(f"\n{bc_x = }, {eval_pts = }, {tesselation = }, {fit[0] = }") if not tesselation: - assert np.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate + assert xp.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate @pytest.mark.parametrize("boxes_per_dim", [(12, 1, 1)]) @@ -658,17 +658,17 @@ def test_evaluation_mc_Np_and_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, te # perturbation if bc_x in ("periodic", "fixed"): - fun_exact = lambda e1, e2, e3: 1.5 - np.sin(2 * np.pi * e1) + fun_exact = lambda e1, e2, e3: 1.5 - xp.sin(2 * xp.pi * e1) pert = {"n": perturbations.ModesSin(ls=(1,), amps=(-1e-0,))} elif bc_x == "mirror": - fun_exact = lambda e1, e2, e3: 1.5 - np.cos(2 * np.pi * e1) + fun_exact = lambda e1, e2, e3: 1.5 - xp.cos(2 * xp.pi * e1) pert = {"n": perturbations.ModesCos(ls=(1,), amps=(-1e-0,))} # exact solution - eta1 = np.linspace(0, 1.0, eval_pts) - eta2 = np.array([0.0]) - eta3 = np.array([0.0]) - ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + eta1 = xp.linspace(0, 1.0, eval_pts) + eta2 = xp.array([0.0]) + eta3 = xp.array([0.0]) + ee1, ee2, ee3 = xp.meshgrid(eta1, eta2, eta3, indexing="ij") exact_eval = fun_exact(ee1, ee2, ee3) # boundary conditions @@ -710,11 +710,11 @@ def test_evaluation_mc_Np_and_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, te if comm is None: all_eval = test_eval else: - all_eval = np.zeros_like(test_eval) + all_eval = xp.zeros_like(test_eval) comm.Allreduce(test_eval, all_eval, op=MPI.SUM) # error in max-norm - diff = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + diff = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) err_vec[-1] += [diff] if rank == 0: @@ -726,29 +726,29 @@ def test_evaluation_mc_Np_and_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, te # plt.title(f"{h = }, {Np = }") # # plt.savefig(f"fun_h{h}_N{Np}_ppb{ppb}.png") - err_vec = np.array(err_vec) - err_min = np.min(err_vec) + err_vec = xp.array(err_vec) + err_min = xp.min(err_vec) if show_plot and rank == 0: if tesselation: - h_mesh, n_mesh = np.meshgrid(np.log10(h_arr), np.log10(ppbs), indexing="ij") + h_mesh, n_mesh = xp.meshgrid(xp.log10(h_arr), xp.log10(ppbs), indexing="ij") if not tesselation: - h_mesh, n_mesh = np.meshgrid(np.log10(h_arr), np.log10(Nps), indexing="ij") + h_mesh, n_mesh = xp.meshgrid(xp.log10(h_arr), xp.log10(Nps), indexing="ij") plt.figure(figsize=(6, 6)) - plt.pcolor(h_mesh, n_mesh, np.log10(err_vec), shading="auto") + plt.pcolor(h_mesh, n_mesh, xp.log10(err_vec), shading="auto") plt.title("Error") plt.colorbar(label="log10(error)") plt.xlabel("log10(h)") plt.ylabel("log10(particles)") - min_indices = np.argmin(err_vec, axis=0) + min_indices = xp.argmin(err_vec, axis=0) min_h_values = [] for mi in min_indices: - min_h_values += [np.log10(h_arr[mi])] + min_h_values += [xp.log10(h_arr[mi])] if tesselation: - log_particles = np.log10(ppbs) + log_particles = xp.log10(ppbs) else: - log_particles = np.log10(Nps) + log_particles = xp.log10(Nps) plt.plot(min_h_values, log_particles, "r-", label="Min error h for each Np", linewidth=2) plt.legend() # plt.savefig("SPH_conv_in_h_and_N.png") @@ -760,7 +760,7 @@ def test_evaluation_mc_Np_and_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, te if tesselation: if bc_x == "periodic": - assert np.min(err_vec) < 7.7e-5 + assert xp.min(err_vec) < 7.7e-5 elif bc_x == "fixed": assert err_min < 7.7e-5 else: @@ -807,25 +807,25 @@ def test_evaluation_SPH_Np_convergence_2d(boxes_per_dim, bc_x, bc_y, tesselation # perturbation if bc_x in ("periodic", "fixed"): if bc_y in ("periodic", "fixed"): - fun_exact = lambda x, y, z: 1.5 - np.sin(2 * np.pi / Lx * x) * np.sin(2 * np.pi / Ly * y) + fun_exact = lambda x, y, z: 1.5 - xp.sin(2 * xp.pi / Lx * x) * xp.sin(2 * xp.pi / Ly * y) pert = {"n": perturbations.ModesSinSin(ls=(1,), ms=(1,), amps=(-1e-0,))} elif bc_y == "mirror": - fun_exact = lambda x, y, z: 1.5 - np.sin(2 * np.pi / Lx * x) * np.cos(2 * np.pi / Ly * y) + fun_exact = lambda x, y, z: 1.5 - xp.sin(2 * xp.pi / Lx * x) * xp.cos(2 * xp.pi / Ly * y) pert = {"n": perturbations.ModesSinCos(ls=(1,), ms=(1,), amps=(-1e-0,))} elif bc_x == "mirror": if bc_y in ("periodic", "fixed"): - fun_exact = lambda x, y, z: 1.5 - np.cos(2 * np.pi / Lx * x) * np.sin(2 * np.pi / Ly * y) + fun_exact = lambda x, y, z: 1.5 - xp.cos(2 * xp.pi / Lx * x) * xp.sin(2 * xp.pi / Ly * y) pert = {"n": perturbations.ModesCosSin(ls=(1,), ms=(1,), amps=(-1e-0,))} elif bc_y == "mirror": - fun_exact = lambda x, y, z: 1.5 - np.cos(2 * np.pi / Lx * x) * np.cos(2 * np.pi / Ly * y) + fun_exact = lambda x, y, z: 1.5 - xp.cos(2 * xp.pi / Lx * x) * xp.cos(2 * xp.pi / Ly * y) pert = {"n": perturbations.ModesCosCos(ls=(1,), ms=(1,), amps=(-1e-0,))} # exact solution - eta1 = np.linspace(0, 1.0, 41) - eta2 = np.linspace(0, 1.0, 86) - eta3 = np.array([0.0]) - ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + eta1 = xp.linspace(0, 1.0, 41) + eta2 = xp.linspace(0, 1.0, 86) + eta3 = xp.array([0.0]) + ee1, ee2, ee3 = xp.meshgrid(eta1, eta2, eta3, indexing="ij") x, y, z = domain(eta1, eta2, eta3) exact_eval = fun_exact(x, y, z) @@ -868,11 +868,11 @@ def test_evaluation_SPH_Np_convergence_2d(boxes_per_dim, bc_x, bc_y, tesselation if comm is None: all_eval = test_eval else: - all_eval = np.zeros_like(test_eval) + all_eval = xp.zeros_like(test_eval) comm.Allreduce(test_eval, all_eval, op=MPI.SUM) # error in max-norm - diff = np.max(np.abs(all_eval - exact_eval)) / np.max(np.abs(exact_eval)) + diff = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) err_vec += [diff] if tesselation: @@ -890,16 +890,16 @@ def test_evaluation_SPH_Np_convergence_2d(boxes_per_dim, bc_x, bc_y, tesselation # fig.savefig(f"2d_sph_{Np}_{ppb}.png") if tesselation: - fit = np.polyfit(np.log(ppbs), np.log(err_vec), 1) + fit = xp.polyfit(xp.log(ppbs), xp.log(err_vec), 1) xvec = ppbs else: - fit = np.polyfit(np.log(Nps), np.log(err_vec), 1) + fit = xp.polyfit(xp.log(Nps), xp.log(err_vec), 1) xvec = Nps if show_plot and rank == 0: plt.figure(figsize=(12, 8)) plt.loglog(xvec, err_vec, label="Convergence") - plt.loglog(xvec, np.exp(fit[1]) * np.array(xvec) ** (fit[0]), "--", label=f"fit with slope {fit[0]}") + plt.loglog(xvec, xp.exp(fit[1]) * xp.array(xvec) ** (fit[0]), "--", label=f"fit with slope {fit[0]}") plt.legend() plt.show() # plt.savefig(f"Convergence_SPH_{tesselation=}") @@ -908,7 +908,7 @@ def test_evaluation_SPH_Np_convergence_2d(boxes_per_dim, bc_x, bc_y, tesselation print(f"\n{bc_x = }, {tesselation = }, {fit[0] = }") if not tesselation: - assert np.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate + assert xp.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate if __name__ == "__main__": diff --git a/src/struphy/pic/tests/test_tesselation.py b/src/struphy/pic/tests/test_tesselation.py index 60f3acb4d..19a02cfad 100644 --- a/src/struphy/pic/tests/test_tesselation.py +++ b/src/struphy/pic/tests/test_tesselation.py @@ -10,7 +10,7 @@ from struphy.initial import perturbations from struphy.pic.particles import ParticlesSPH from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @pytest.mark.parametrize("ppb", [8, 12]) @@ -56,20 +56,20 @@ def test_draw(ppb, nx, ny, nz): zl = particles.domain_array[rank, 6] zr = particles.domain_array[rank, 7] - eta1 = np.linspace(xl, xr, tiles_x + 1)[:-1] + (xr - xl) / (2 * tiles_x) - eta2 = np.linspace(yl, yr, tiles_y + 1)[:-1] + (yr - yl) / (2 * tiles_y) - eta3 = np.linspace(zl, zr, tiles_z + 1)[:-1] + (zr - zl) / (2 * tiles_z) + eta1 = xp.linspace(xl, xr, tiles_x + 1)[:-1] + (xr - xl) / (2 * tiles_x) + eta2 = xp.linspace(yl, yr, tiles_y + 1)[:-1] + (yr - yl) / (2 * tiles_y) + eta3 = xp.linspace(zl, zr, tiles_z + 1)[:-1] + (zr - zl) / (2 * tiles_z) - ee1, ee2, ee3 = np.meshgrid(eta1, eta2, eta3, indexing="ij") + ee1, ee2, ee3 = xp.meshgrid(eta1, eta2, eta3, indexing="ij") e1 = ee1.flatten() e2 = ee2.flatten() e3 = ee3.flatten() # print(f'\n{rank = }, {e1 = }') - assert np.allclose(particles.positions[:, 0], e1) - assert np.allclose(particles.positions[:, 1], e2) - assert np.allclose(particles.positions[:, 2], e3) + assert xp.allclose(particles.positions[:, 0], e1) + assert xp.allclose(particles.positions[:, 1], e2) + assert xp.allclose(particles.positions[:, 2], e3) @pytest.mark.parametrize("ppb", [8, 12]) @@ -119,20 +119,20 @@ def test_cell_average(ppb, nx, ny, nz, n_quad, show_plot=False): yl = particles.domain_array[rank, 3] yr = particles.domain_array[rank, 4] - eta1 = np.linspace(xl, xr, tiles_x + 1) - eta2 = np.linspace(yl, yr, tiles_y + 1) + eta1 = xp.linspace(xl, xr, tiles_x + 1) + eta2 = xp.linspace(yl, yr, tiles_y + 1) if ny == nz == 1: plt.figure(figsize=(15, 10)) - plt.plot(particles.positions[:, 0], np.zeros_like(particles.weights), "o", label="markers") + plt.plot(particles.positions[:, 0], xp.zeros_like(particles.weights), "o", label="markers") plt.plot(particles.positions[:, 0], particles.weights, "-o", label="weights") plt.plot( - np.linspace(xl, xr, 100), - particles.f_init(np.linspace(xl, xr, 100), 0.5, 0.5).squeeze(), + xp.linspace(xl, xr, 100), + particles.f_init(xp.linspace(xl, xr, 100), 0.5, 0.5).squeeze(), "--", label="f_init", ) - plt.vlines(np.linspace(xl, xr, nx + 1), 0, 2, label="sorting boxes", color="k") + plt.vlines(xp.linspace(xl, xr, nx + 1), 0, 2, label="sorting boxes", color="k") ax = plt.gca() ax.set_xticks(eta1) ax.set_yticks(eta2) @@ -146,8 +146,8 @@ def test_cell_average(ppb, nx, ny, nz, n_quad, show_plot=False): plt.subplot(1, 2, 1) ax = plt.gca() - ax.set_xticks(np.linspace(0, 1, nx + 1)) - ax.set_yticks(np.linspace(0, 1, ny + 1)) + ax.set_xticks(xp.linspace(0, 1, nx + 1)) + ax.set_yticks(xp.linspace(0, 1, ny + 1)) coloring = particles.weights plt.scatter(particles.positions[:, 0], particles.positions[:, 1], c=coloring, s=40) plt.grid(c="k") @@ -159,12 +159,12 @@ def test_cell_average(ppb, nx, ny, nz, n_quad, show_plot=False): plt.subplot(1, 2, 2) ax = plt.gca() - ax.set_xticks(np.linspace(0, 1, nx + 1)) - ax.set_yticks(np.linspace(0, 1, ny + 1)) + ax.set_xticks(xp.linspace(0, 1, nx + 1)) + ax.set_yticks(xp.linspace(0, 1, ny + 1)) coloring = particles.weights - pos1 = np.linspace(xl, xr, 100) - pos2 = np.linspace(yl, yr, 100) - pp1, pp2 = np.meshgrid(pos1, pos2, indexing="ij") + pos1 = xp.linspace(xl, xr, 100) + pos2 = xp.linspace(yl, yr, 100) + pp1, pp2 = xp.meshgrid(pos1, pos2, indexing="ij") plt.pcolor(pp1, pp2, particles.f_init(pp1, pp2, 0.5).squeeze()) plt.grid(c="k") plt.axis("square") @@ -176,8 +176,8 @@ def test_cell_average(ppb, nx, ny, nz, n_quad, show_plot=False): plt.show() # test - print(f"\n{rank = }, {np.max(np.abs(particles.weights - particles.f_init(particles.positions))) = }") - assert np.max(np.abs(particles.weights - particles.f_init(particles.positions))) < 0.012 + print(f"\n{rank = }, {xp.max(xp.abs(particles.weights - particles.f_init(particles.positions))) = }") + assert xp.max(xp.abs(particles.weights - particles.f_init(particles.positions))) < 0.012 if __name__ == "__main__": diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index 0964ef5e1..fb8fda634 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -5,7 +5,7 @@ OptsRecontructBC, OptsSpatialLoading, ) -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class LoadingParameters: @@ -182,23 +182,23 @@ def __init__( # computations and allocations self._bin_edges = [] for nb, rng in zip(n_bins, ranges): - self._bin_edges += [np.linspace(rng[0], rng[1], nb + 1)] + self._bin_edges += [xp.linspace(rng[0], rng[1], nb + 1)] self._bin_edges = tuple(self.bin_edges) - self._f = np.zeros(n_bins, dtype=float) - self._df = np.zeros(n_bins, dtype=float) + self._f = xp.zeros(n_bins, dtype=float) + self._df = xp.zeros(n_bins, dtype=float) @property def bin_edges(self) -> tuple: return self._bin_edges @property - def f(self) -> np.ndarray: + def f(self) -> xp.ndarray: """The binned distribution function (full-f).""" return self._f @property - def df(self) -> np.ndarray: + def df(self) -> xp.ndarray: """The binned distribution function minus the background (delta-f).""" return self._df @@ -218,19 +218,19 @@ def __init__( pts_e2: int = 16, pts_e3: int = 1, ): - e1 = np.linspace(0.0, 1.0, pts_e1) - e2 = np.linspace(0.0, 1.0, pts_e2) - e3 = np.linspace(0.0, 1.0, pts_e3) - ee1, ee2, ee3 = np.meshgrid(e1, e2, e3, indexing="ij") + e1 = xp.linspace(0.0, 1.0, pts_e1) + e2 = xp.linspace(0.0, 1.0, pts_e2) + e3 = xp.linspace(0.0, 1.0, pts_e3) + ee1, ee2, ee3 = xp.meshgrid(e1, e2, e3, indexing="ij") self._plot_pts = (ee1, ee2, ee3) - self._n_sph = np.zeros(ee1.shape, dtype=float) + self._n_sph = xp.zeros(ee1.shape, dtype=float) @property def plot_pts(self) -> tuple: return self._plot_pts @property - def n_sph(self) -> np.ndarray: + def n_sph(self) -> xp.ndarray: """The evaluated density.""" return self._n_sph @@ -251,15 +251,15 @@ def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): Particles object. """ - res = np.empty(1, dtype=float) + res = xp.empty(1, dtype=float) utils.canonical_kinetic_particles( res, particles.markers, - np.array(derham.p), + xp.array(derham.p), derham.Vh_fem["0"].knots[0], derham.Vh_fem["0"].knots[1], derham.Vh_fem["0"].knots[2], - np.array( + xp.array( derham.V0.coeff_space.starts, ), *domain.args_map, @@ -284,7 +284,7 @@ def get_electron_thermal_energy(density_0_form, derham, domain, nel1, nel2, nel3 Discrete Derham complex. """ - res = np.empty(1, dtype=float) + res = xp.empty(1, dtype=float) utils.thermal_energy( res, density_0_form._operators[0].matrix._data, diff --git a/src/struphy/polar/basic.py b/src/struphy/polar/basic.py index 78b81d4ff..49dc21a52 100644 --- a/src/struphy/polar/basic.py +++ b/src/struphy/polar/basic.py @@ -3,7 +3,7 @@ from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class PolarDerhamSpace(VectorSpace): @@ -209,7 +209,7 @@ class PolarVector(Vector): Element of a PolarDerhamSpace. An instance of a PolarVector consists of two parts: - 1. a list of np.arrays of the polar coeffs (not distributed) + 1. a list of xp.arrays of the polar coeffs (not distributed) 2. a tensor product StencilVector/BlockVector of the parent space with inner rings set to zero (distributed). Parameters @@ -224,7 +224,7 @@ def __init__(self, V): self._dtype = V.dtype # initialize polar coeffs - self._pol = [np.zeros((m, n)) for m, n in zip(V.n_polar, V.n3)] + self._pol = [xp.zeros((m, n)) for m, n in zip(V.n_polar, V.n3)] # full tensor product vector self._tp = V.parent_space.zeros() @@ -241,7 +241,7 @@ def dtype(self): @property def pol(self): - """Polar coefficients as np.array.""" + """Polar coefficients as xp.array.""" return self._pol @pol.setter @@ -327,7 +327,7 @@ def toarray(self, allreduce=False): if self.space.comm is not None and allreduce: self.space.comm.Allreduce(MPI.IN_PLACE, out, op=MPI.SUM) - out = np.concatenate((self.pol[0].flatten(), out)) + out = xp.concatenate((self.pol[0].flatten(), out)) else: out1 = self.tp[0].toarray()[self.space.n_rings[0] * self.space.n[1] * self.space.n3[0] :] @@ -340,7 +340,7 @@ def toarray(self, allreduce=False): self.space.comm.Allreduce(MPI.IN_PLACE, out2, op=MPI.SUM) self.space.comm.Allreduce(MPI.IN_PLACE, out3, op=MPI.SUM) - out = np.concatenate( + out = xp.concatenate( ( self.pol[0].flatten(), out1, @@ -366,7 +366,7 @@ def copy(self, out=None): self._tp.copy(out=w.tp) # copy polar part for n, pl in enumerate(self._pol): - np.copyto(w._pol[n], pl, casting="no") + xp.copyto(w._pol[n], pl, casting="no") return w def __neg__(self): diff --git a/src/struphy/polar/extraction_operators.py b/src/struphy/polar/extraction_operators.py index 9afb9d237..c4588e997 100644 --- a/src/struphy/polar/extraction_operators.py +++ b/src/struphy/polar/extraction_operators.py @@ -1,4 +1,4 @@ -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp # ============================= 2D polar splines (C1) =================================== @@ -47,8 +47,8 @@ def __init__(self, domain, derham): self._pole = (cx[0, 0], cy[0, 0]) - assert np.all(cx[0] == self.pole[0]) - assert np.all(cy[0] == self.pole[1]) + assert xp.all(cx[0] == self.pole[0]) + assert xp.all(cy[0] == self.pole[1]) self._n0 = cx.shape[0] self._n1 = cx.shape[1] @@ -70,14 +70,14 @@ def __init__(self, domain, derham): self._tau = max( [ ((self.cx[1] - self.pole[0]) * (-2)).max(), - ((self.cx[1] - self.pole[0]) - np.sqrt(3) * (self.cy[1] - self.pole[1])).max(), - ((self.cx[1] - self.pole[0]) + np.sqrt(3) * (self.cy[1] - self.pole[1])).max(), + ((self.cx[1] - self.pole[0]) - xp.sqrt(3) * (self.cy[1] - self.pole[1])).max(), + ((self.cx[1] - self.pole[0]) + xp.sqrt(3) * (self.cy[1] - self.pole[1])).max(), ] ) # barycentric coordinates - self._xi_0 = np.zeros((3, self.n1), dtype=float) - self._xi_1 = np.zeros((3, self.n1), dtype=float) + self._xi_0 = xp.zeros((3, self.n1), dtype=float) + self._xi_1 = xp.zeros((3, self.n1), dtype=float) self._xi_0[:, :] = 1 / 3 @@ -85,12 +85,12 @@ def __init__(self, domain, derham): self._xi_1[1, :] = ( 1 / 3 - 1 / (3 * self.tau) * (self.cx[1] - self.pole[0]) - + np.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) + + xp.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) ) self._xi_1[2, :] = ( 1 / 3 - 1 / (3 * self.tau) * (self.cx[1] - self.pole[0]) - - np.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) + - xp.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) ) # remove small values @@ -102,17 +102,17 @@ def __init__(self, domain, derham): # ============= basis extraction operator for discrete 0-forms ================ # first n_rings tp rings --> "polar coeffs" - e0_blocks_ten_to_pol = np.block([self.xi_0, self.xi_1]) + e0_blocks_ten_to_pol = xp.block([self.xi_0, self.xi_1]) self._e_ten_to_pol["0"] = [[csr(e0_blocks_ten_to_pol)]] # ============ basis extraction operator for discrete 1-forms (Hcurl) ========= # first n_rings tp rings --> "polar coeffs" - e1_11_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) - e1_12_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) + e1_11_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) + e1_12_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) - e1_21_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.n1), dtype=float) - e1_22_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) + e1_21_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.n1), dtype=float) + e1_22_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) # 1st component for l in range(2): @@ -135,7 +135,7 @@ def __init__(self, domain, derham): # =============== basis extraction operator for discrete 1-forms (Hdiv) ========= # first n_rings tp rings --> "polar coeffs" - e3_blocks_ten_to_pol = np.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) + e3_blocks_ten_to_pol = xp.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) self._e_ten_to_pol["2"] = [ [csr(e1_22_blocks_ten_to_pol), csr(-e1_21_blocks_ten_to_pol), None], @@ -161,7 +161,7 @@ def __init__(self, domain, derham): self._p_ten_to_ten = {} # first n_rings tp rings --> "polar coeffs" - p0_blocks_ten_to_pol = np.zeros((self.n_polar[0][0], self.n_rings[0][0] * self.n1), dtype=float) + p0_blocks_ten_to_pol = xp.zeros((self.n_polar[0][0], self.n_rings[0][0] * self.n1), dtype=float) # !! NOTE: for odd spline degrees and periodic splines the first Greville point sometimes does NOT start at zero!! if domain.p[1] % 2 != 0 and not (abs(derham.Vh_fem["0"].spaces[1].interpolation_grid[0]) < 1e-14): @@ -176,15 +176,15 @@ def __init__(self, domain, derham): self._p_ten_to_pol["0"] = [[csr(p0_blocks_ten_to_pol)]] # first n_rings + 1 tp rings --> "first tp ring" - p0_blocks_ten_to_ten = np.block([0 * np.identity(self.n1)] * self.n_rings[0][0] + [np.identity(self.n1)]) + p0_blocks_ten_to_ten = xp.block([0 * xp.identity(self.n1)] * self.n_rings[0][0] + [xp.identity(self.n1)]) self._p_ten_to_ten["0"] = [[csr(p0_blocks_ten_to_ten)]] # =========== projection extraction operator for discrete 1-forms (Hcurl) ======== # first n_rings tp rings --> "polar coeffs" - p1_11_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) - p1_22_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) + p1_11_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) + p1_22_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) # !! NOTE: PSYDAC's first integration interval sometimes start at < 0 !! if derham.Vh_fem["3"].spaces[1].histopolation_grid[0] < -1e-14: @@ -196,8 +196,8 @@ def __init__(self, domain, derham): p1_22_blocks_ten_to_pol[1, (self.d1 + 0 * self.d1 // 3) : (self.d1 + 1 * self.d1 // 3)] = 1.0 p1_22_blocks_ten_to_pol[1, (self.d1 + 1 * self.d1 // 3) : (self.d1 + 2 * self.d1 // 3)] = 1.0 - p1_12_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) - p1_21_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.d1), dtype=float) + p1_12_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) + p1_21_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.d1), dtype=float) self._p_ten_to_pol["1"] = [ [csr(p1_11_blocks_ten_to_pol), csr(p1_12_blocks_ten_to_pol), None], @@ -206,26 +206,26 @@ def __init__(self, domain, derham): ] # first n_rings + 1 tp rings --> "first tp ring" - p1_11_blocks_ten_to_ten = np.zeros((self.n1, self.n1), dtype=float) + p1_11_blocks_ten_to_ten = xp.zeros((self.n1, self.n1), dtype=float) # !! NOTE: for odd spline degrees and periodic splines the first Greville point sometimes does NOT start at zero!! if domain.p[1] % 2 != 0 and not (abs(derham.Vh_fem["0"].spaces[1].interpolation_grid[0]) < 1e-14): - p1_11_blocks_ten_to_ten[:, 3 * self.n1 // 3 - 1] = -np.roll(self.xi_1[0], -1) - p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3 - 1] = -np.roll(self.xi_1[1], -1) - p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3 - 1] = -np.roll(self.xi_1[2], -1) + p1_11_blocks_ten_to_ten[:, 3 * self.n1 // 3 - 1] = -xp.roll(self.xi_1[0], -1) + p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3 - 1] = -xp.roll(self.xi_1[1], -1) + p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3 - 1] = -xp.roll(self.xi_1[2], -1) else: p1_11_blocks_ten_to_ten[:, 0 * self.n1 // 3] = -self.xi_1[0] p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3] = -self.xi_1[1] p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3] = -self.xi_1[2] - p1_11_blocks_ten_to_ten += np.identity(self.n1) + p1_11_blocks_ten_to_ten += xp.identity(self.n1) - p1_11_blocks_ten_to_ten = np.block([p1_11_blocks_ten_to_ten, np.identity(self.n1)]) + p1_11_blocks_ten_to_ten = xp.block([p1_11_blocks_ten_to_ten, xp.identity(self.n1)]) - p1_22_blocks_ten_to_ten = np.block([0 * np.identity(self.d1)] * self.n_rings[1][1] + [np.identity(self.d1)]) + p1_22_blocks_ten_to_ten = xp.block([0 * xp.identity(self.d1)] * self.n_rings[1][1] + [xp.identity(self.d1)]) - p1_12_blocks_ten_to_ten = np.zeros((self.d1, (self.n_rings[1][1] + 1) * self.d1), dtype=float) - p1_21_blocks_ten_to_ten = np.zeros((self.n1, (self.n_rings[1][0] + 1) * self.n1), dtype=float) + p1_12_blocks_ten_to_ten = xp.zeros((self.d1, (self.n_rings[1][1] + 1) * self.d1), dtype=float) + p1_21_blocks_ten_to_ten = xp.zeros((self.n1, (self.n_rings[1][0] + 1) * self.n1), dtype=float) self._p_ten_to_ten["1"] = [ [csr(p1_11_blocks_ten_to_ten), csr(p1_12_blocks_ten_to_ten), None], @@ -236,7 +236,7 @@ def __init__(self, domain, derham): # ========== projection extraction operator for discrete 1-forms (Hdiv) ========== # first n_rings tp rings --> "polar coeffs" - p3_blocks_ten_to_pol = np.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) + p3_blocks_ten_to_pol = xp.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) self._p_ten_to_pol["2"] = [ [csr(p1_22_blocks_ten_to_pol), csr(p1_21_blocks_ten_to_pol), None], @@ -245,24 +245,24 @@ def __init__(self, domain, derham): ] # first n_rings + 1 tp rings --> "first tp ring" - p3_blocks_ten_to_ten = np.zeros((self.d1, self.d1), dtype=float) + p3_blocks_ten_to_ten = xp.zeros((self.d1, self.d1), dtype=float) - a0 = np.diff(self.xi_1[1], append=self.xi_1[1, 0]) - a1 = np.diff(self.xi_1[2], append=self.xi_1[2, 0]) + a0 = xp.diff(self.xi_1[1], append=self.xi_1[1, 0]) + a1 = xp.diff(self.xi_1[2], append=self.xi_1[2, 0]) # !! NOTE: PSYDAC's first integration interval sometimes start at < 0 !! if derham.Vh_fem["3"].spaces[1].histopolation_grid[0] < -1e-14: p3_blocks_ten_to_ten[:, (0 * self.n1 // 3 + 1) : (1 * self.n1 // 3 + 1)] = ( - -np.roll(a0, +1)[:, None] - np.roll(a1, +1)[:, None] + -xp.roll(a0, +1)[:, None] - xp.roll(a1, +1)[:, None] ) - p3_blocks_ten_to_ten[:, (1 * self.n1 // 3 + 1) : (2 * self.n1 // 3 + 1)] = -np.roll(a1, +1)[:, None] + p3_blocks_ten_to_ten[:, (1 * self.n1 // 3 + 1) : (2 * self.n1 // 3 + 1)] = -xp.roll(a1, +1)[:, None] else: p3_blocks_ten_to_ten[:, 0 * self.n1 // 3 : 1 * self.n1 // 3] = -a0[:, None] - a1[:, None] p3_blocks_ten_to_ten[:, 1 * self.n1 // 3 : 2 * self.n1 // 3] = -a1[:, None] - p3_blocks_ten_to_ten += np.identity(self.d1) + p3_blocks_ten_to_ten += xp.identity(self.d1) - p3_blocks_ten_to_ten = np.block([p3_blocks_ten_to_ten, np.identity(self.d1)]) + p3_blocks_ten_to_ten = xp.block([p3_blocks_ten_to_ten, xp.identity(self.d1)]) self._p_ten_to_ten["2"] = [ [csr(p1_22_blocks_ten_to_ten), csr(p1_21_blocks_ten_to_ten), None], @@ -295,24 +295,24 @@ def __init__(self, domain, derham): # ======================= discrete gradient ====================================== # "polar coeffs" to "polar coeffs" - grad_pol_to_pol_1 = np.zeros((self.n_polar[1][0], self.n_polar[0][0]), dtype=float) - grad_pol_to_pol_2 = np.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) - grad_pol_to_pol_3 = np.identity(self.n_polar[0][0], dtype=float) + grad_pol_to_pol_1 = xp.zeros((self.n_polar[1][0], self.n_polar[0][0]), dtype=float) + grad_pol_to_pol_2 = xp.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) + grad_pol_to_pol_3 = xp.identity(self.n_polar[0][0], dtype=float) self._grad_pol_to_pol = [[csr(grad_pol_to_pol_1)], [csr(grad_pol_to_pol_2)], [csr(grad_pol_to_pol_3)]] # "polar coeffs" to "first tp ring" - grad_pol_to_ten_1 = np.zeros(((self.n_rings[1][0] + 1) * self.n1, self.n_polar[0][0])) - grad_pol_to_ten_2 = np.zeros(((self.n_rings[1][1] + 1) * self.d1, self.n_polar[0][0])) - grad_pol_to_ten_3 = np.zeros(((self.n_rings[0][0] + 1) * self.n1, self.n_polar[0][0])) + grad_pol_to_ten_1 = xp.zeros(((self.n_rings[1][0] + 1) * self.n1, self.n_polar[0][0])) + grad_pol_to_ten_2 = xp.zeros(((self.n_rings[1][1] + 1) * self.d1, self.n_polar[0][0])) + grad_pol_to_ten_3 = xp.zeros(((self.n_rings[0][0] + 1) * self.n1, self.n_polar[0][0])) grad_pol_to_ten_1[-self.n1 :, :] = -self.xi_1.T self._grad_pol_to_ten = [[csr(grad_pol_to_ten_1)], [csr(grad_pol_to_ten_2)], [csr(grad_pol_to_ten_3)]] # eta_3 direction - grad_e3_1 = np.identity(self.n2, dtype=float) - grad_e3_2 = np.identity(self.n2, dtype=float) + grad_e3_1 = xp.identity(self.n2, dtype=float) + grad_e3_2 = xp.identity(self.n2, dtype=float) grad_e3_3 = grad_1d_matrix(derham.spl_kind[2], self.n2) self._grad_e3 = [[csr(grad_e3_1)], [csr(grad_e3_2)], [csr(grad_e3_3)]] @@ -320,14 +320,14 @@ def __init__(self, domain, derham): # =========================== discrete curl ====================================== # "polar coeffs" to "polar coeffs" - curl_pol_to_pol_12 = np.identity(self.n_polar[1][1], dtype=float) - curl_pol_to_pol_13 = np.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) + curl_pol_to_pol_12 = xp.identity(self.n_polar[1][1], dtype=float) + curl_pol_to_pol_13 = xp.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) - curl_pol_to_pol_21 = np.identity(self.n_polar[1][0], dtype=float) - curl_pol_to_pol_23 = np.zeros((self.n_polar[2][1], self.n_polar[0][0]), dtype=float) + curl_pol_to_pol_21 = xp.identity(self.n_polar[1][0], dtype=float) + curl_pol_to_pol_23 = xp.zeros((self.n_polar[2][1], self.n_polar[0][0]), dtype=float) - curl_pol_to_pol_31 = np.zeros((self.n_polar[3][0], self.n_polar[1][0]), dtype=float) - curl_pol_to_pol_32 = np.zeros((self.n_polar[3][0], self.n_polar[1][1]), dtype=float) + curl_pol_to_pol_31 = xp.zeros((self.n_polar[3][0], self.n_polar[1][0]), dtype=float) + curl_pol_to_pol_32 = xp.zeros((self.n_polar[3][0], self.n_polar[1][1]), dtype=float) self._curl_pol_to_pol = [ [None, csr(-curl_pol_to_pol_12), csr(curl_pol_to_pol_13)], @@ -336,14 +336,14 @@ def __init__(self, domain, derham): ] # "polar coeffs" to "first tp ring" - curl_pol_to_ten_12 = np.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[1][1])) - curl_pol_to_ten_13 = np.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[0][0])) + curl_pol_to_ten_12 = xp.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[1][1])) + curl_pol_to_ten_13 = xp.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[0][0])) - curl_pol_to_ten_21 = np.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[1][0])) - curl_pol_to_ten_23 = np.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[0][0])) + curl_pol_to_ten_21 = xp.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[1][0])) + curl_pol_to_ten_23 = xp.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[0][0])) - curl_pol_to_ten_31 = np.zeros(((self.n_rings[3][0] + 1) * self.n1, self.n_polar[1][0])) - curl_pol_to_ten_32 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[1][1])) + curl_pol_to_ten_31 = xp.zeros(((self.n_rings[3][0] + 1) * self.n1, self.n_polar[1][0])) + curl_pol_to_ten_32 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[1][1])) curl_pol_to_ten_23[-self.n1 :, :] = -self.xi_1.T @@ -361,13 +361,13 @@ def __init__(self, domain, derham): # eta_3 direction curl_e3_12 = grad_1d_matrix(derham.spl_kind[2], self.n2) - curl_e3_13 = np.identity(self.d2) + curl_e3_13 = xp.identity(self.d2) curl_e3_21 = grad_1d_matrix(derham.spl_kind[2], self.n2) - curl_e3_23 = np.identity(self.d2) + curl_e3_23 = xp.identity(self.d2) - curl_e3_31 = np.identity(self.n2) - curl_e3_32 = np.identity(self.n2) + curl_e3_31 = xp.identity(self.n2) + curl_e3_32 = xp.identity(self.n2) self._curl_e3 = [ [None, csr(curl_e3_12), csr(curl_e3_13)], @@ -378,16 +378,16 @@ def __init__(self, domain, derham): # =========================== discrete div ====================================== # "polar coeffs" to "polar coeffs" - div_pol_to_pol_1 = np.zeros((self.n_polar[3][0], self.n_polar[2][0]), dtype=float) - div_pol_to_pol_2 = np.zeros((self.n_polar[3][0], self.n_polar[2][1]), dtype=float) - div_pol_to_pol_3 = np.identity(self.n_polar[3][0], dtype=float) + div_pol_to_pol_1 = xp.zeros((self.n_polar[3][0], self.n_polar[2][0]), dtype=float) + div_pol_to_pol_2 = xp.zeros((self.n_polar[3][0], self.n_polar[2][1]), dtype=float) + div_pol_to_pol_3 = xp.identity(self.n_polar[3][0], dtype=float) self._div_pol_to_pol = [[csr(div_pol_to_pol_1), csr(div_pol_to_pol_2), csr(div_pol_to_pol_3)]] # "polar coeffs" to "first tp ring" - div_pol_to_ten_1 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][0])) - div_pol_to_ten_2 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][1])) - div_pol_to_ten_3 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[3][0])) + div_pol_to_ten_1 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][0])) + div_pol_to_ten_2 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][1])) + div_pol_to_ten_3 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[3][0])) for l in range(2): for j in range(self.d1, 2 * self.d1): @@ -398,8 +398,8 @@ def __init__(self, domain, derham): self._div_pol_to_ten = [[csr(div_pol_to_ten_1), csr(div_pol_to_ten_2), csr(div_pol_to_ten_3)]] # eta_3 direction - div_e3_1 = np.identity(self.d2, dtype=float) - div_e3_2 = np.identity(self.d2, dtype=float) + div_e3_1 = xp.identity(self.d2, dtype=float) + div_e3_2 = xp.identity(self.d2, dtype=float) div_e3_3 = grad_1d_matrix(derham.spl_kind[2], self.n2) self._div_e3 = [[csr(div_e3_1), csr(div_e3_2), csr(div_e3_3)]] @@ -539,13 +539,13 @@ def __init__(self, n0, n1): # =========== extraction operators for discrete 0-forms ================== # extraction operator for basis functions - self.E0_11 = spa.csr_matrix(np.ones((1, n1), dtype=float)) + self.E0_11 = spa.csr_matrix(xp.ones((1, n1), dtype=float)) self.E0_22 = spa.identity((n0 - 1) * n1, format="csr") self.E0 = spa.bmat([[self.E0_11, None], [None, self.E0_22]], format="csr") # global projection extraction operator for interpolation points - self.P0_11 = np.zeros((1, n1), dtype=float) + self.P0_11 = xp.zeros((1, n1), dtype=float) self.P0_11[0, 0] = 1.0 @@ -598,7 +598,7 @@ def __init__(self, n0, n1): # ========= discrete polar gradient matrix =============================== # radial dofs (DN) - G11 = np.zeros(((d0 - 0) * n1, 1), dtype=float) + G11 = xp.zeros(((d0 - 0) * n1, 1), dtype=float) G11[:n1, 0] = -1.0 G12 = spa.kron(grad_1d_1[:, 1:], spa.identity(n1)) @@ -606,7 +606,7 @@ def __init__(self, n0, n1): self.G1 = spa.bmat([[G11, G12]], format="csr") # angular dofs (ND) - G21 = np.zeros(((n0 - 1) * d1, 1), dtype=float) + G21 = xp.zeros(((n0 - 1) * d1, 1), dtype=float) G22 = spa.kron(spa.identity(n0 - 1), grad_1d_2, format="csr") self.G2 = spa.bmat([[G21, G22]], format="csr") @@ -619,13 +619,13 @@ def __init__(self, n0, n1): # 2D vector curl (NN --> ND DN) # angular dofs (ND) - VC11 = np.zeros(((n0 - 1) * d1, 1), dtype=float) + VC11 = xp.zeros(((n0 - 1) * d1, 1), dtype=float) VC12 = spa.kron(spa.identity(n0 - 1), grad_1d_2, format="csr") self.VC1 = spa.bmat([[VC11, VC12]], format="csr") # radial dofs (DN) - VC21 = np.zeros(((d0 - 0) * n1, 1), dtype=float) + VC21 = xp.zeros(((d0 - 0) * n1, 1), dtype=float) VC21[:n1, 0] = 1.0 VC22 = -spa.kron(grad_1d_1[:, 1:], spa.identity(n1)) @@ -687,26 +687,26 @@ def __init__(self, cx, cy): self.Nbase2 = (d0 - 1) * d1 # size of control triangle - self.tau = np.array( + self.tau = xp.array( [ (-2 * (cx[1] - self.x0)).max(), - ((cx[1] - self.x0) - np.sqrt(3) * (cy[1] - self.y0)).max(), - ((cx[1] - self.x0) + np.sqrt(3) * (cy[1] - self.y0)).max(), + ((cx[1] - self.x0) - xp.sqrt(3) * (cy[1] - self.y0)).max(), + ((cx[1] - self.x0) + xp.sqrt(3) * (cy[1] - self.y0)).max(), ] ).max() - self.Xi_0 = np.zeros((3, n1), dtype=float) - self.Xi_1 = np.zeros((3, n1), dtype=float) + self.Xi_0 = xp.zeros((3, n1), dtype=float) + self.Xi_1 = xp.zeros((3, n1), dtype=float) # barycentric coordinates self.Xi_0[:, :] = 1 / 3 self.Xi_1[0, :] = 1 / 3 + 2 / (3 * self.tau) * (cx[1] - self.x0) self.Xi_1[1, :] = ( - 1 / 3 - 1 / (3 * self.tau) * (cx[1] - self.x0) + np.sqrt(3) / (3 * self.tau) * (cy[1] - self.y0) + 1 / 3 - 1 / (3 * self.tau) * (cx[1] - self.x0) + xp.sqrt(3) / (3 * self.tau) * (cy[1] - self.y0) ) self.Xi_1[2, :] = ( - 1 / 3 - 1 / (3 * self.tau) * (cx[1] - self.x0) - np.sqrt(3) / (3 * self.tau) * (cy[1] - self.y0) + 1 / 3 - 1 / (3 * self.tau) * (cx[1] - self.x0) - xp.sqrt(3) / (3 * self.tau) * (cy[1] - self.y0) ) # remove small values @@ -714,13 +714,13 @@ def __init__(self, cx, cy): # =========== extraction operators for discrete 0-forms ================== # extraction operator for basis functions - self.E0_11 = spa.csr_matrix(np.hstack((self.Xi_0, self.Xi_1))) + self.E0_11 = spa.csr_matrix(xp.hstack((self.Xi_0, self.Xi_1))) self.E0_22 = spa.identity((n0 - 2) * n1, format="csr") self.E0 = spa.bmat([[self.E0_11, None], [None, self.E0_22]], format="csr") # global projection extraction operator for interpolation points - self.P0_11 = np.zeros((3, 2 * n1), dtype=float) + self.P0_11 = xp.zeros((3, 2 * n1), dtype=float) self.P0_11[0, n1 + 0 * n1 // 3] = 1.0 self.P0_11[1, n1 + 1 * n1 // 3] = 1.0 @@ -737,8 +737,8 @@ def __init__(self, cx, cy): self.E1C_12 = spa.identity((d0 - 1) * n1) self.E1C_34 = spa.identity((n0 - 2) * d1) - self.E1C_21 = np.zeros((2, 1 * n1), dtype=float) - self.E1C_23 = np.zeros((2, 2 * d1), dtype=float) + self.E1C_21 = xp.zeros((2, 1 * n1), dtype=float) + self.E1C_23 = xp.zeros((2, 2 * d1), dtype=float) # 1st component for s in range(2): @@ -760,22 +760,22 @@ def __init__(self, cx, cy): # extraction operator for interpolation/histopolation in global projector # 1st component - self.P1C_11 = np.zeros((n1, n1), dtype=float) + self.P1C_11 = xp.zeros((n1, n1), dtype=float) self.P1C_12 = spa.identity(n1) self.P1C_23 = spa.identity((d0 - 2) * n1) self.P1C_11[:, 0 * n1 // 3] = -self.Xi_1[0] self.P1C_11[:, 1 * n1 // 3] = -self.Xi_1[1] self.P1C_11[:, 2 * n1 // 3] = -self.Xi_1[2] - self.P1C_11 += np.identity(n1) + self.P1C_11 += xp.identity(n1) # 2nd component - self.P1C_34 = np.zeros((2, 2 * d1), dtype=float) + self.P1C_34 = xp.zeros((2, 2 * d1), dtype=float) self.P1C_45 = spa.identity((n0 - 2) * d1) - self.P1C_34[0, (d1 + 0 * d1 // 3) : (d1 + 1 * d1 // 3)] = np.ones(d1 // 3, dtype=float) - self.P1C_34[1, (d1 + 0 * d1 // 3) : (d1 + 1 * d1 // 3)] = np.ones(d1 // 3, dtype=float) - self.P1C_34[1, (d1 + 1 * d1 // 3) : (d1 + 2 * d1 // 3)] = np.ones(d1 // 3, dtype=float) + self.P1C_34[0, (d1 + 0 * d1 // 3) : (d1 + 1 * d1 // 3)] = xp.ones(d1 // 3, dtype=float) + self.P1C_34[1, (d1 + 0 * d1 // 3) : (d1 + 1 * d1 // 3)] = xp.ones(d1 // 3, dtype=float) + self.P1C_34[1, (d1 + 1 * d1 // 3) : (d1 + 2 * d1 // 3)] = xp.ones(d1 // 3, dtype=float) # combined first and second component self.P1C = spa.bmat( @@ -790,8 +790,8 @@ def __init__(self, cx, cy): # ========================================================================= # ========= extraction operators for discrete 1-forms (H_div) ============= - self.E1D_11 = np.zeros((2, 2 * d1), dtype=float) - self.E1D_13 = np.zeros((2, 1 * n1), dtype=float) + self.E1D_11 = xp.zeros((2, 2 * d1), dtype=float) + self.E1D_13 = xp.zeros((2, 1 * n1), dtype=float) self.E1D_22 = spa.identity((n0 - 2) * d1) self.E1D_34 = spa.identity((d0 - 1) * n1) @@ -834,13 +834,13 @@ def __init__(self, cx, cy): # ========================================================================= # =========== extraction operators for discrete 2-forms =================== - self.E2_1 = np.zeros(((d0 - 1) * d1, d1), dtype=float) + self.E2_1 = xp.zeros(((d0 - 1) * d1, d1), dtype=float) self.E2_2 = spa.identity((d0 - 1) * d1) self.E2 = spa.bmat([[self.E2_1, self.E2_2]], format="csr") # extraction operator for histopolation in global projector - self.P2_11 = np.zeros((d1, d1), dtype=float) + self.P2_11 = xp.zeros((d1, d1), dtype=float) self.P2_12 = spa.identity(d1) self.P2_23 = spa.identity((d0 - 2) * d1) @@ -853,7 +853,7 @@ def __init__(self, cx, cy): # block B self.P2_11[i, 1 * n1 // 3 : 2 * n1 // 3] = -(self.Xi_1[2, (i + 1) % n1] - self.Xi_1[2, i]) - self.P2_11 += np.identity(d1) + self.P2_11 += xp.identity(d1) self.P2 = spa.bmat([[self.P2_11, self.P2_12, None], [None, None, self.P2_23]], format="csr") # ========================================================================= @@ -864,14 +864,14 @@ def __init__(self, cx, cy): # ========================================================================= # ========= discrete polar gradient matrix ================================ - self.G1_1 = np.zeros(((d0 - 1) * n1, 3), dtype=float) + self.G1_1 = xp.zeros(((d0 - 1) * n1, 3), dtype=float) self.G1_1[:n1, :] = -self.Xi_1.T self.G1_2 = spa.kron(grad_1d_1[1:, 2:], spa.identity(n1)) self.G1 = spa.bmat([[self.G1_1, self.G1_2]], format="csr") - self.G2_11 = np.zeros((2, 3), dtype=float) + self.G2_11 = xp.zeros((2, 3), dtype=float) self.G2_11[0, 0] = -1.0 self.G2_11[0, 1] = 1.0 @@ -888,7 +888,7 @@ def __init__(self, cx, cy): # ========= discrete polar curl matrix =================================== # 2D vector curl - self.VC1_11 = np.zeros((2, 3), dtype=float) + self.VC1_11 = xp.zeros((2, 3), dtype=float) self.VC1_11[0, 0] = -1.0 self.VC1_11[0, 1] = 1.0 @@ -900,7 +900,7 @@ def __init__(self, cx, cy): self.VC1 = spa.bmat([[self.VC1_11, None], [None, self.VC1_22]], format="csr") - self.VC2_11 = np.zeros(((d0 - 1) * n1, 3), dtype=float) + self.VC2_11 = xp.zeros(((d0 - 1) * n1, 3), dtype=float) self.VC2_11[:n1, :] = -self.Xi_1.T self.VC2_22 = spa.kron(grad_1d_1[1:, 2:], spa.identity(n1)) @@ -912,7 +912,7 @@ def __init__(self, cx, cy): # 2D scalar curl self.SC1 = -spa.kron(spa.identity(d0 - 1), grad_1d_2) - self.SC2_1 = np.zeros(((d0 - 1) * d1, 2), dtype=float) + self.SC2_1 = xp.zeros(((d0 - 1) * d1, 2), dtype=float) for s in range(2): for j in range(d1): @@ -926,7 +926,7 @@ def __init__(self, cx, cy): # ========================================================================= # ========= discrete polar div matrix ===================================== - self.D1_1 = np.zeros(((d0 - 1) * d1, 2), dtype=float) + self.D1_1 = xp.zeros(((d0 - 1) * d1, 2), dtype=float) for s in range(2): for j in range(d1): @@ -965,24 +965,24 @@ def __init__(self, tensor_space, cx, cy): self.Nbase3_pol = (d0 - 1) * d1 # size of control triangle - self.tau = np.array( - [(-2 * cx[1]).max(), (cx[1] - np.sqrt(3) * cy[1]).max(), (cx[1] + np.sqrt(3) * cy[1]).max()] + self.tau = xp.array( + [(-2 * cx[1]).max(), (cx[1] - xp.sqrt(3) * cy[1]).max(), (cx[1] + xp.sqrt(3) * cy[1]).max()] ).max() - self.Xi_0 = np.zeros((3, n1), dtype=float) - self.Xi_1 = np.zeros((3, n1), dtype=float) + self.Xi_0 = xp.zeros((3, n1), dtype=float) + self.Xi_1 = xp.zeros((3, n1), dtype=float) # barycentric coordinates self.Xi_0[:, :] = 1 / 3 self.Xi_1[0, :] = 1 / 3 + 2 / (3 * self.tau) * cx[1, :, 0] - self.Xi_1[1, :] = 1 / 3 - 1 / (3 * self.tau) * cx[1, :, 0] + np.sqrt(3) / (3 * self.tau) * cy[1, :, 0] - self.Xi_1[2, :] = 1 / 3 - 1 / (3 * self.tau) * cx[1, :, 0] - np.sqrt(3) / (3 * self.tau) * cy[1, :, 0] + self.Xi_1[1, :] = 1 / 3 - 1 / (3 * self.tau) * cx[1, :, 0] + xp.sqrt(3) / (3 * self.tau) * cy[1, :, 0] + self.Xi_1[2, :] = 1 / 3 - 1 / (3 * self.tau) * cx[1, :, 0] - xp.sqrt(3) / (3 * self.tau) * cy[1, :, 0] # =========== extraction operators for discrete 0-forms ================== # extraction operator for basis functions self.E0_pol = spa.bmat( - [[np.hstack((self.Xi_0, self.Xi_1)), None], [None, spa.identity((n0 - 2) * n1)]], format="csr" + [[xp.hstack((self.Xi_0, self.Xi_1)), None], [None, spa.identity((n0 - 2) * n1)]], format="csr" ) self.E0 = spa.kron(self.E0_pol, spa.identity(n2), format="csr") @@ -1005,7 +1005,7 @@ def __init__(self, tensor_space, cx, cy): for j in range(n1): self.E1_1_pol[(d0 - 1) * n1 + s, j] = self.Xi_1[s + 1, j] - self.Xi_0[s + 1, j] - self.E1_1_pol[: (d0 - 1) * n1, n1:] = np.identity((d0 - 1) * n1) + self.E1_1_pol[: (d0 - 1) * n1, n1:] = xp.identity((d0 - 1) * n1) self.E1_1_pol = self.E1_1_pol.tocsr() # 2nd component @@ -1014,7 +1014,7 @@ def __init__(self, tensor_space, cx, cy): self.E1_2_pol[(d0 - 1) * n1 + s, j] = 0.0 self.E1_2_pol[(d0 - 1) * n1 + s, n1 + j] = self.Xi_1[s + 1, (j + 1) % n1] - self.Xi_1[s + 1, j] - self.E1_2_pol[((d0 - 1) * n1 + 2) :, 2 * d1 :] = np.identity((n0 - 2) * d1) + self.E1_2_pol[((d0 - 1) * n1 + 2) :, 2 * d1 :] = xp.identity((n0 - 2) * d1) self.E1_2_pol = self.E1_2_pol.tocsr() # 3rd component @@ -1043,9 +1043,9 @@ def __init__(self, tensor_space, cx, cy): self.P1_1_pol = self.P1_1_pol.tocsr() # 2nd component - self.P1_2_pol[0, (n1 + 0 * n1 // 3) : (n1 + 1 * n1 // 3)] = np.ones((1, n1 // 3), dtype=float) - self.P1_2_pol[1, (n1 + 0 * n1 // 3) : (n1 + 1 * n1 // 3)] = np.ones((1, n1 // 3), dtype=float) - self.P1_2_pol[1, (n1 + 1 * n1 // 3) : (n1 + 2 * n1 // 3)] = np.ones((1, n1 // 3), dtype=float) + self.P1_2_pol[0, (n1 + 0 * n1 // 3) : (n1 + 1 * n1 // 3)] = xp.ones((1, n1 // 3), dtype=float) + self.P1_2_pol[1, (n1 + 0 * n1 // 3) : (n1 + 1 * n1 // 3)] = xp.ones((1, n1 // 3), dtype=float) + self.P1_2_pol[1, (n1 + 1 * n1 // 3) : (n1 + 2 * n1 // 3)] = xp.ones((1, n1 // 3), dtype=float) self.P1_2_pol[2:, 2 * n1 :] = spa.identity((n0 - 2) * d1) self.P1_2_pol = self.P1_2_pol.tocsr() @@ -1073,7 +1073,7 @@ def __init__(self, tensor_space, cx, cy): self.E2_1_pol[s, j] = 0.0 self.E2_1_pol[s, n1 + j] = self.Xi_1[s + 1, (j + 1) % n1] - self.Xi_1[s + 1, j] - self.E2_1_pol[2 : (2 + (n0 - 2) * d1), 2 * n1 :] = np.identity((n0 - 2) * d1) + self.E2_1_pol[2 : (2 + (n0 - 2) * d1), 2 * n1 :] = xp.identity((n0 - 2) * d1) self.E2_1_pol = self.E2_1_pol.tocsr() # 2nd component @@ -1081,11 +1081,11 @@ def __init__(self, tensor_space, cx, cy): for j in range(n1): self.E2_2_pol[s, j] = -(self.Xi_1[s + 1, j] - self.Xi_0[s + 1, j]) - self.E2_2_pol[(2 + (n0 - 2) * d1) :, 1 * n1 :] = np.identity((d0 - 1) * n1) + self.E2_2_pol[(2 + (n0 - 2) * d1) :, 1 * n1 :] = xp.identity((d0 - 1) * n1) self.E2_2_pol = self.E2_2_pol.tocsr() # 3rd component - self.E2_3_pol[:, 1 * d1 :] = np.identity((d0 - 1) * d1) + self.E2_3_pol[:, 1 * d1 :] = xp.identity((d0 - 1) * d1) self.E2_3_pol = self.E2_3_pol.tocsr() # combined first and second component diff --git a/src/struphy/polar/linear_operators.py b/src/struphy/polar/linear_operators.py index 019b7aae0..0412bdba3 100644 --- a/src/struphy/polar/linear_operators.py +++ b/src/struphy/polar/linear_operators.py @@ -6,7 +6,7 @@ from struphy.feec.linear_operators import LinOpWithTransp from struphy.linear_algebra.linalg_kron import kron_matvec_2d from struphy.polar.basic import PolarDerhamSpace, PolarVector -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class PolarExtractionOperator(LinOpWithTransp): @@ -668,7 +668,7 @@ def dot_inner_tp_rings(blocks_e1_e2, blocks_e3, v, out): # loop over codomain components for m, (row_e1_e2, row_e3) in enumerate(zip(blocks_e1_e2, blocks_e3)): - res = np.zeros((n_rows[m], n3_out[m]), dtype=float) + res = xp.zeros((n_rows[m], n3_out[m]), dtype=float) # loop over domain components for n, (block_e1_e2, block_e3) in enumerate(zip(row_e1_e2, row_e3)): @@ -677,7 +677,7 @@ def dot_inner_tp_rings(blocks_e1_e2, blocks_e3, v, out): e1, e2, e3 = in_ends[n] if block_e1_e2 is not None: - tmp = np.zeros((n_rings_in[n], n2, n3_in[n]), dtype=float) + tmp = xp.zeros((n_rings_in[n], n2, n3_in[n]), dtype=float) tmp[:, s2 : e2 + 1, s3 : e3 + 1] = in_vec[n][0 : n_rings_in[n], s2 : e2 + 1, s3 : e3 + 1] res += kron_matvec_2d([block_e1_e2, block_e3], tmp.reshape(n_rings_in[n] * n2, n3_in[n])) @@ -785,7 +785,7 @@ def dot_parts_of_polar(blocks_e1_e2, blocks_e3, v, out): # loop over codomain components for m, (row_e1_e2, row_e3) in enumerate(zip(blocks_e1_e2, blocks_e3)): - res = np.zeros((n_rings_out[m], n2, n3_out[m]), dtype=float) + res = xp.zeros((n_rings_out[m], n2, n3_out[m]), dtype=float) # loop over domain components for n, (block_e1_e2, block_e3) in enumerate(zip(row_e1_e2, row_e3)): @@ -794,7 +794,7 @@ def dot_parts_of_polar(blocks_e1_e2, blocks_e3, v, out): if in_starts[n][0] == 0: s1, s2, s3 = in_starts[n] e1, e2, e3 = in_ends[n] - tmp = np.zeros((n2, n3_in[n]), dtype=float) + tmp = xp.zeros((n2, n3_in[n]), dtype=float) tmp[s2 : e2 + 1, s3 : e3 + 1] = in_tp[n][n_rings_in[n], s2 : e2 + 1, s3 : e3 + 1] res += kron_matvec_2d([block_e1_e2, block_e3], tmp).reshape(n_rings_out[m], n2, n3_out[m]) else: diff --git a/src/struphy/polar/tests/test_legacy_polar_splines.py b/src/struphy/polar/tests/test_legacy_polar_splines.py index b9665ae0b..f4c47060a 100644 --- a/src/struphy/polar/tests/test_legacy_polar_splines.py +++ b/src/struphy/polar/tests/test_legacy_polar_splines.py @@ -12,7 +12,7 @@ def test_polar_splines_2D(plot=False): from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space from struphy.geometry import domains - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp # parameters # number of elements (number of elements in angular direction must be a multiple of 3) @@ -42,8 +42,8 @@ def test_polar_splines_2D(plot=False): fig.set_figheight(10) fig.set_figwidth(10) - el_b_1 = np.linspace(0.0, 1.0, Nel[0] + 1) - el_b_2 = np.linspace(0.0, 1.0, Nel[1] + 1) + el_b_1 = xp.linspace(0.0, 1.0, Nel[0] + 1) + el_b_2 = xp.linspace(0.0, 1.0, Nel[1] + 1) grid_x = domain(el_b_1, el_b_2, 0.0, squeeze_out=True)[0] grid_y = domain(el_b_1, el_b_2, 0.0, squeeze_out=True)[1] @@ -108,7 +108,7 @@ def test_polar_splines_2D(plot=False): ) # plot three new polar splines in V0 - etaplot = [np.linspace(0.0, 1.0, 200), np.linspace(0.0, 1.0, 200)] + etaplot = [xp.linspace(0.0, 1.0, 200), xp.linspace(0.0, 1.0, 200)] xplot = [ domain(etaplot[0], etaplot[1], 0.0, squeeze_out=True)[0], domain(etaplot[0], etaplot[1], 0.0, squeeze_out=True)[1], @@ -123,9 +123,9 @@ def test_polar_splines_2D(plot=False): ax3 = fig.add_subplot(133, projection="3d") # coeffs in polar basis - c0_pol1 = np.zeros(space_2d.E0.shape[0], dtype=float) - c0_pol2 = np.zeros(space_2d.E0.shape[0], dtype=float) - c0_pol3 = np.zeros(space_2d.E0.shape[0], dtype=float) + c0_pol1 = xp.zeros(space_2d.E0.shape[0], dtype=float) + c0_pol2 = xp.zeros(space_2d.E0.shape[0], dtype=float) + c0_pol3 = xp.zeros(space_2d.E0.shape[0], dtype=float) c0_pol1[0] = 1.0 c0_pol2[1] = 1.0 @@ -134,7 +134,7 @@ def test_polar_splines_2D(plot=False): ax1.plot_surface( xplot[0], xplot[1], - space_2d.evaluate_NN(etaplot[0], etaplot[1], np.array([0.0]), c0_pol1, "V0")[:, :, 0], + space_2d.evaluate_NN(etaplot[0], etaplot[1], xp.array([0.0]), c0_pol1, "V0")[:, :, 0], cmap="jet", ) ax1.set_xlabel("R [m]", labelpad=5) @@ -144,7 +144,7 @@ def test_polar_splines_2D(plot=False): ax2.plot_surface( xplot[0], xplot[1], - space_2d.evaluate_NN(etaplot[0], etaplot[1], np.array([0.0]), c0_pol2, "V0")[:, :, 0], + space_2d.evaluate_NN(etaplot[0], etaplot[1], xp.array([0.0]), c0_pol2, "V0")[:, :, 0], cmap="jet", ) ax2.set_xlabel("R [m]", labelpad=5) @@ -154,7 +154,7 @@ def test_polar_splines_2D(plot=False): ax3.plot_surface( xplot[0], xplot[1], - space_2d.evaluate_NN(etaplot[0], etaplot[1], np.array([0.0]), c0_pol3, "V0")[:, :, 0], + space_2d.evaluate_NN(etaplot[0], etaplot[1], xp.array([0.0]), c0_pol3, "V0")[:, :, 0], cmap="jet", ) ax3.set_xlabel("R [m]", labelpad=5) diff --git a/src/struphy/polar/tests/test_polar.py b/src/struphy/polar/tests/test_polar.py index b2b5c8326..a72abca4c 100644 --- a/src/struphy/polar/tests/test_polar.py +++ b/src/struphy/polar/tests/test_polar.py @@ -176,7 +176,7 @@ def test_extraction_ops_and_derivatives(Nel, p, spl_kind): from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -222,11 +222,11 @@ def test_extraction_ops_and_derivatives(Nel, p, spl_kind): b2_pol.tp = b2_tp p3_pol.tp = p3_tp - np.random.seed(1607) - f0_pol.pol = [np.random.rand(f0_pol.pol[0].shape[0], f0_pol.pol[0].shape[1])] - e1_pol.pol = [np.random.rand(e1_pol.pol[n].shape[0], e1_pol.pol[n].shape[1]) for n in range(3)] - b2_pol.pol = [np.random.rand(b2_pol.pol[n].shape[0], b2_pol.pol[n].shape[1]) for n in range(3)] - p3_pol.pol = [np.random.rand(p3_pol.pol[0].shape[0], p3_pol.pol[0].shape[1])] + xp.random.seed(1607) + f0_pol.pol = [xp.random.rand(f0_pol.pol[0].shape[0], f0_pol.pol[0].shape[1])] + e1_pol.pol = [xp.random.rand(e1_pol.pol[n].shape[0], e1_pol.pol[n].shape[1]) for n in range(3)] + b2_pol.pol = [xp.random.rand(b2_pol.pol[n].shape[0], b2_pol.pol[n].shape[1]) for n in range(3)] + p3_pol.pol = [xp.random.rand(p3_pol.pol[0].shape[0], p3_pol.pol[0].shape[1])] f0_pol_leg = f0_pol.toarray(True) e1_pol_leg = e1_pol.toarray(True) @@ -243,10 +243,10 @@ def test_extraction_ops_and_derivatives(Nel, p, spl_kind): r2_pol = derham.extraction_ops["2"].dot(b2_tp) r3_pol = derham.extraction_ops["3"].dot(p3_tp) - assert np.allclose(r0_pol.toarray(True), space.E0.dot(f0_tp_leg)) - assert np.allclose(r1_pol.toarray(True), space.E1.dot(e1_tp_leg)) - assert np.allclose(r2_pol.toarray(True), space.E2.dot(b2_tp_leg)) - assert np.allclose(r3_pol.toarray(True), space.E3.dot(p3_tp_leg)) + assert xp.allclose(r0_pol.toarray(True), space.E0.dot(f0_tp_leg)) + assert xp.allclose(r1_pol.toarray(True), space.E1.dot(e1_tp_leg)) + assert xp.allclose(r2_pol.toarray(True), space.E2.dot(b2_tp_leg)) + assert xp.allclose(r3_pol.toarray(True), space.E3.dot(p3_tp_leg)) # test transposed extraction operators E0T = derham.extraction_ops["0"].transpose() @@ -277,9 +277,9 @@ def test_extraction_ops_and_derivatives(Nel, p, spl_kind): r2_pol = derham.curl.dot(e1_pol) r3_pol = derham.div.dot(b2_pol) - assert np.allclose(r1_pol.toarray(True), space.G.dot(f0_pol_leg)) - assert np.allclose(r2_pol.toarray(True), space.C.dot(e1_pol_leg)) - assert np.allclose(r3_pol.toarray(True), space.D.dot(b2_pol_leg)) + assert xp.allclose(r1_pol.toarray(True), space.G.dot(f0_pol_leg)) + assert xp.allclose(r2_pol.toarray(True), space.C.dot(e1_pol_leg)) + assert xp.allclose(r3_pol.toarray(True), space.D.dot(b2_pol_leg)) # test transposed derivatives GT = derham.grad.transpose() @@ -290,9 +290,9 @@ def test_extraction_ops_and_derivatives(Nel, p, spl_kind): r1_pol = CT.dot(b2_pol) r2_pol = DT.dot(p3_pol) - assert np.allclose(r0_pol.toarray(True), space.G.T.dot(e1_pol_leg)) - assert np.allclose(r1_pol.toarray(True), space.C.T.dot(b2_pol_leg)) - assert np.allclose(r2_pol.toarray(True), space.D.T.dot(p3_pol_leg)) + assert xp.allclose(r0_pol.toarray(True), space.G.T.dot(e1_pol_leg)) + assert xp.allclose(r1_pol.toarray(True), space.C.T.dot(b2_pol_leg)) + assert xp.allclose(r2_pol.toarray(True), space.D.T.dot(p3_pol_leg)) if rank == 0: print("------------- Test passed ---------------------------") @@ -307,7 +307,7 @@ def test_projectors(Nel, p, spl_kind): from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space from struphy.feec.psydac_derham import Derham from struphy.geometry.domains import IGAPolarCylinder - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -338,7 +338,7 @@ def test_projectors(Nel, p, spl_kind): # function to project on physical domain def fun_scalar(x, y, z): - return np.sin(2 * np.pi * (x)) * np.cos(2 * np.pi * y) * np.sin(2 * np.pi * z) + return xp.sin(2 * xp.pi * (x)) * xp.cos(2 * xp.pi * y) * xp.sin(2 * xp.pi * z) fun_vector = [fun_scalar, fun_scalar, fun_scalar] @@ -369,7 +369,7 @@ def fun3(e1, e2, e3): r0_pol_leg = space.projectors.pi_0(fun0) - assert np.allclose(r0_pol.toarray(True), r0_pol_leg) + assert xp.allclose(r0_pol.toarray(True), r0_pol_leg) if rank == 0: print("Test passed for PI_0 polar projector") @@ -385,7 +385,7 @@ def fun3(e1, e2, e3): r1_pol_leg = space.projectors.pi_1(fun1, with_subs=False) - assert np.allclose(r1_pol.toarray(True), r1_pol_leg) + assert xp.allclose(r1_pol.toarray(True), r1_pol_leg) if rank == 0: print("Test passed for PI_1 polar projector") @@ -401,7 +401,7 @@ def fun3(e1, e2, e3): r2_pol_leg = space.projectors.pi_2(fun2, with_subs=False) - assert np.allclose(r2_pol.toarray(True), r2_pol_leg) + assert xp.allclose(r2_pol.toarray(True), r2_pol_leg) if rank == 0: print("Test passed for PI_2 polar projector") @@ -417,7 +417,7 @@ def fun3(e1, e2, e3): r3_pol_leg = space.projectors.pi_3(fun3, with_subs=False) - assert np.allclose(r3_pol.toarray(True), r3_pol_leg) + assert xp.allclose(r3_pol.toarray(True), r3_pol_leg) if rank == 0: print("Test passed for PI_3 polar projector") diff --git a/src/struphy/post_processing/likwid/plot_likwidproject.py b/src/struphy/post_processing/likwid/plot_likwidproject.py index 33b0426af..0e8abb23b 100644 --- a/src/struphy/post_processing/likwid/plot_likwidproject.py +++ b/src/struphy/post_processing/likwid/plot_likwidproject.py @@ -16,7 +16,7 @@ import struphy.post_processing.likwid.likwid_parser as lp import struphy.post_processing.likwid.maxplotlylib as mply import struphy.post_processing.likwid.roofline_plotter as rp -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def clean_string(string_in): @@ -196,16 +196,16 @@ def plot_roofline( fig.update_xaxes( type="log", # Ensure the x-axis is logarithmic - range=[np.log10(xmin), np.log10(xmax)], + range=[xp.log10(xmin), xp.log10(xmax)], title="Operational intensity (FLOP/Byte)", tickvals=xtick_values, # Set where ticks appear ticktext=[str(t) for t in xtick_values], - # ticktext=[f'$10^{{{int(np.log10(t))}}}$' for t in xtick_values] # Set tick labels + # ticktext=[f'$10^{{{int(xp.log10(t))}}}$' for t in xtick_values] # Set tick labels ) fig.update_yaxes( type="log", # Ensure the x-axis is logarithmic - range=[np.log10(ymin), np.log10(ymax)], + range=[xp.log10(ymin), xp.log10(ymax)], title="Performance [GFLOP/s]", tickvals=ytick_values, # Set where ticks appear ticktext=[str(t) for t in ytick_values], diff --git a/src/struphy/post_processing/likwid/plot_time_traces.py b/src/struphy/post_processing/likwid/plot_time_traces.py index 4f3f4eeb8..94b491d03 100644 --- a/src/struphy/post_processing/likwid/plot_time_traces.py +++ b/src/struphy/post_processing/likwid/plot_time_traces.py @@ -7,7 +7,7 @@ # pio.kaleido.scope.mathjax = None import struphy.post_processing.likwid.maxplotlylib as mply -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def glob_to_regex(pat: str) -> str: @@ -121,9 +121,9 @@ def plot_avg_duration_bar_chart( # Compute statistics per region regions = sorted(region_durations.keys()) - avg_durations = [np.mean(region_durations[r]) for r in regions] - min_durations = [np.min(region_durations[r]) for r in regions] - max_durations = [np.max(region_durations[r]) for r in regions] + avg_durations = [xp.mean(region_durations[r]) for r in regions] + min_durations = [xp.min(region_durations[r]) for r in regions] + max_durations = [xp.max(region_durations[r]) for r in regions] yerr = [ [avg - min_ for avg, min_ in zip(avg_durations, min_durations)], [max_ - avg for avg, max_ in zip(avg_durations, max_durations)], @@ -131,7 +131,7 @@ def plot_avg_duration_bar_chart( # Plot bar chart with error bars (min-max spans) plt.figure(figsize=(12, 6)) - x = np.arange(len(regions)) + x = xp.arange(len(regions)) plt.bar(x, avg_durations, yerr=yerr, capsize=5, color="skyblue", edgecolor="k") plt.yscale("log") plt.xticks(x, regions, rotation=45, ha="right") @@ -175,7 +175,7 @@ def plot_gantt_chart_plotly( region_start_times = {} for rank_data in profiling_data["rank_data"].values(): for region_name, info in rank_data.items(): - first_start_time = np.min(info["start_times"]) + first_start_time = xp.min(info["start_times"]) if region_name not in region_start_times or first_start_time < region_start_times[region_name]: region_start_times[region_name] = first_start_time @@ -291,7 +291,7 @@ def plot_gantt_chart( region_start_times = {} for rank_data in profiling_data["rank_data"].values(): for region_name, info in rank_data.items(): - first_start_time = np.min(info["start_times"]) + first_start_time = xp.min(info["start_times"]) if region_name not in region_start_times or first_start_time < region_start_times[region_name]: region_start_times[region_name] = first_start_time diff --git a/src/struphy/post_processing/likwid/roofline_plotter.py b/src/struphy/post_processing/likwid/roofline_plotter.py index 6a49f9c34..abc6bee6f 100644 --- a/src/struphy/post_processing/likwid/roofline_plotter.py +++ b/src/struphy/post_processing/likwid/roofline_plotter.py @@ -4,7 +4,7 @@ import pandas as pd import yaml -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def sort_by_num_threads(bm): @@ -143,14 +143,14 @@ def add_plot_diagonal( bandwidth_GBps, label="", ymax=1e4, - operational_intensity_FLOPpMB=np.arange(0, 1000, 1), + operational_intensity_FLOPpMB=xp.arange(0, 1000, 1), ): max_performance_GFLOP = operational_intensity_FLOPpMB * bandwidth_GBps (line,) = mfig.axs.plot(operational_intensity_FLOPpMB, max_performance_GFLOP) # Specify the y-value where you want to place the text specific_y = ymax # Interpolate to find the corresponding x-value - specific_x = np.interp( + specific_x = xp.interp( specific_y, max_performance_GFLOP, operational_intensity_FLOPpMB, @@ -210,10 +210,10 @@ def get_average_val( xvec.append(x) yvec.append(y) # print('xvec', xvec, 'yvec', yvec) - xvec = np.array(xvec) - yvec = np.array(yvec) + xvec = xp.array(xvec) + yvec = xp.array(yvec) # print('xvec', xvec, 'yvec', yvec) - return np.average(xvec), np.average(yvec), np.std(xvec), np.std(yvec) + return xp.average(xvec), xp.average(yvec), xp.std(xvec), xp.std(yvec) def get_maximum(path, df_index=-1, metric="DP [MFLOP/s] STAT", column_name="Sum"): diff --git a/src/struphy/post_processing/orbits/orbits_tools.py b/src/struphy/post_processing/orbits/orbits_tools.py index fd358b79b..02776b74f 100644 --- a/src/struphy/post_processing/orbits/orbits_tools.py +++ b/src/struphy/post_processing/orbits/orbits_tools.py @@ -6,7 +6,7 @@ from tqdm import tqdm from struphy.post_processing.orbits.orbits_kernels import calculate_guiding_center_from_6d -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def post_process_orbit_guiding_center(path_in, path_kinetics_species, species): @@ -61,7 +61,7 @@ def post_process_orbit_guiding_center(path_in, path_kinetics_species, species): if file.endswith(".npy") ] pproc_nt = len(npy_files_list) - n_markers = np.load(os.path.join(path_orbits, npy_files_list[0])).shape[0] + n_markers = xp.load(os.path.join(path_orbits, npy_files_list[0])).shape[0] # re-ordering npy_files npy_files_list = sorted(npy_files_list) @@ -76,10 +76,10 @@ def post_process_orbit_guiding_center(path_in, path_kinetics_species, species): os.mkdir(path_gc) # temporary marker array - temp = np.empty((n_markers, 7), dtype=float) - etas = np.empty((n_markers, 3), dtype=float) - B_cart = np.empty((n_markers, 3), dtype=float) - lost_particles_mask = np.empty(n_markers, dtype=bool) + temp = xp.empty((n_markers, 7), dtype=float) + etas = xp.empty((n_markers, 3), dtype=float) + B_cart = xp.empty((n_markers, 3), dtype=float) + lost_particles_mask = xp.empty(n_markers, dtype=bool) print("Evaluation of guiding center for " + str(species)) @@ -94,13 +94,13 @@ def post_process_orbit_guiding_center(path_in, path_kinetics_species, species): file_txt = os.path.join(path_gc, npy_files_list[n][:-4] + ".txt") # call .npy file - temp[:, :] = np.load(os.path.join(path_orbits, npy_files_list[n])) + temp[:, :] = xp.load(os.path.join(path_orbits, npy_files_list[n])) # move ids to last column and save - temp = np.roll(temp, -1, axis=1) + temp = xp.roll(temp, -1, axis=1) # sorting out lost particles - lost_particles_mask = np.all(temp[:, :-1] == 0, axis=1) + lost_particles_mask = xp.all(temp[:, :-1] == 0, axis=1) # domain inverse map etas[~lost_particles_mask, :] = domain.inverse_map( @@ -110,7 +110,7 @@ def post_process_orbit_guiding_center(path_in, path_kinetics_species, species): # eval cartesian magnetic filed at marker positions B_cart[~lost_particles_mask, :] = equil.b_cart( - *np.concatenate( + *xp.concatenate( ( etas[:, 0][:, None], etas[:, 1][:, None], @@ -123,10 +123,10 @@ def post_process_orbit_guiding_center(path_in, path_kinetics_species, species): calculate_guiding_center_from_6d(temp, B_cart) # move ids to first column and save - temp = np.roll(temp, 1, axis=1) + temp = xp.roll(temp, 1, axis=1) - np.save(file_npy, temp) - np.savetxt(file_txt, temp[:, :4], fmt="%12.6f", delimiter=", ") + xp.save(file_npy, temp) + xp.savetxt(file_txt, temp[:, :4], fmt="%12.6f", delimiter=", ") def post_process_orbit_classification(path_kinetics_species, species): @@ -168,16 +168,16 @@ def post_process_orbit_classification(path_kinetics_species, species): if file.endswith(".npy") ] pproc_nt = len(npy_files_list) - n_markers = np.load(os.path.join(path_gc, npy_files_list[0])).shape[0] + n_markers = xp.load(os.path.join(path_gc, npy_files_list[0])).shape[0] # re-ordering npy_files npy_files_list = sorted(npy_files_list) # temporary marker array - temp = np.empty((n_markers, 8), dtype=float) - v_parallel = np.empty(n_markers, dtype=float) - trapped_particle_mask = np.empty(n_markers, dtype=bool) - lost_particle_mask = np.empty(n_markers, dtype=bool) + temp = xp.empty((n_markers, 8), dtype=float) + v_parallel = xp.empty(n_markers, dtype=float) + trapped_particle_mask = xp.empty(n_markers, dtype=bool) + lost_particle_mask = xp.empty(n_markers, dtype=bool) print("Classifying guiding center orbits for " + str(species)) @@ -188,16 +188,16 @@ def post_process_orbit_classification(path_kinetics_species, species): # load .npy files file_npy = os.path.join(path_gc, npy_files_list[n]) - temp[:, :-1] = np.load(file_npy) + temp[:, :-1] = xp.load(file_npy) # initial time step if n == 0: v_init = temp[:, 4] - np.save(file_npy, temp) + xp.save(file_npy, temp) continue # synchronizing with former time step - temp[:, -1] = np.load( + temp[:, -1] = xp.load( os.path.join( path_gc, npy_files_list[n - 1], @@ -205,10 +205,10 @@ def post_process_orbit_classification(path_kinetics_species, species): )[:, -1] # call parallel velocity data from .npy file - v_parallel = np.load(os.path.join(path_gc, npy_files_list[n]))[:, 4] + v_parallel = xp.load(os.path.join(path_gc, npy_files_list[n]))[:, 4] # sorting out lost particles - lost_particle_mask = np.all(temp[:, 1:-1] == 0, axis=1) + lost_particle_mask = xp.all(temp[:, 1:-1] == 0, axis=1) # check reverse of parallel velocity trapped_particle_mask[:] = False @@ -221,4 +221,4 @@ def post_process_orbit_classification(path_kinetics_species, species): # assign "-1" at the last index of lost particles temp[lost_particle_mask, -1] = -1 - np.save(file_npy, temp) + xp.save(file_npy, temp) diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index fb663a639..d2071e089 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -19,7 +19,7 @@ from struphy.models.species import ParticleSpecies from struphy.models.variables import PICVariable from struphy.topology.grids import TensorProductGrid -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class ParamsIn: @@ -138,7 +138,7 @@ def create_femfields( fields : dict Nested dictionary holding :class:`~struphy.feec.psydac_derham.SplineFunction`: fields[t][name] contains the Field with the name "name" in the hdf5 file at time t. - t_grid : np.ndarray + t_grid : xp.ndarray Time grid. """ @@ -276,7 +276,7 @@ def eval_femfields( Returns ------- point_data : dict - Nested dictionary holding values of FemFields on the grid as list of 3d np.arrays: + Nested dictionary holding values of FemFields on the grid as list of 3d xp.arrays: point_data[name][t] contains the values of the field with name "name" in fields[t].keys() at time t. If physical is True, physical components of fields are saved. @@ -299,7 +299,7 @@ def eval_femfields( Nel = params_in.grid.Nel - grids_log = [np.linspace(0.0, 1.0, Nel_i * n_i + 1) for Nel_i, n_i in zip(Nel, celldivide)] + grids_log = [xp.linspace(0.0, 1.0, Nel_i * n_i + 1) for Nel_i, n_i in zip(Nel, celldivide)] grids_phy = [ domain(*grids_log)[0], domain(*grids_log)[1], @@ -326,7 +326,7 @@ def eval_femfields( point_data[species][name][t] = [] # scalar spaces - if isinstance(temp_val, np.ndarray): + if isinstance(temp_val, xp.ndarray): if physical: # push-forward if space_id == "H1": @@ -387,7 +387,7 @@ def eval_femfields( def create_vtk( path: str, - t_grid: np.ndarray, + t_grid: xp.ndarray, grids_phy: list, point_data: dict, *, @@ -400,7 +400,7 @@ def create_vtk( path : str Absolute path of where to store the .vts files. Will then be in path/vtk/step_.vts. - t_grid : np.ndarray + t_grid : xp.ndarray Time grid. grids_phy : 3-list @@ -425,7 +425,7 @@ def create_vtk( # time loop nt = len(t_grid) - 1 - log_nt = int(np.log10(nt)) + 1 + log_nt = int(xp.log10(nt)) + 1 print(f"\nCreating vtk in {path} ...") for n, t in enumerate(tqdm(t_grid)): @@ -542,7 +542,7 @@ def post_process_markers( # get number of time steps and markers nt, n_markers, n_cols = files[0]["kinetic/" + species + "/markers"].shape - log_nt = int(np.log10(int(((nt - 1) / step)))) + 1 + log_nt = int(xp.log10(int(((nt - 1) / step)))) + 1 # directory for .txt files and marker index which will be saved path_orbits = os.path.join(path_out, "orbits") @@ -561,8 +561,8 @@ def post_process_markers( os.mkdir(path_orbits) # temporary array - temp = np.empty((n_markers, len(save_index)), order="C") - lost_particles_mask = np.empty(n_markers, dtype=bool) + temp = xp.empty((n_markers, len(save_index)), order="C") + lost_particles_mask = xp.empty(n_markers, dtype=bool) print(f"Evaluation of {n_markers} marker orbits for {species}") @@ -589,28 +589,28 @@ def post_process_markers( # sorting out lost particles ids = temp[:, -1].astype("int") - ids_lost_particles = np.setdiff1d(np.arange(n_markers), ids) - ids_removed_particles = np.nonzero(temp[:, 0] == -1.0)[0] - ids_lost_particles = np.array(list(set(ids_lost_particles) | set(ids_removed_particles)), dtype=int) + ids_lost_particles = xp.setdiff1d(xp.arange(n_markers), ids) + ids_removed_particles = xp.nonzero(temp[:, 0] == -1.0)[0] + ids_lost_particles = xp.array(list(set(ids_lost_particles) | set(ids_removed_particles)), dtype=int) lost_particles_mask[:] = False lost_particles_mask[ids_lost_particles] = True if len(ids_lost_particles) > 0: # lost markers are saved as [0, ..., 0, ids] temp[lost_particles_mask, -1] = ids_lost_particles - ids = np.unique(np.append(ids, ids_lost_particles)) + ids = xp.unique(xp.append(ids, ids_lost_particles)) - assert np.all(sorted(ids) == np.arange(n_markers)) + assert xp.all(sorted(ids) == xp.arange(n_markers)) # compute physical positions (x, y, z) - pos_phys = domain(np.array(temp[~lost_particles_mask, :3]), change_out_order=True) + pos_phys = domain(xp.array(temp[~lost_particles_mask, :3]), change_out_order=True) temp[~lost_particles_mask, :3] = pos_phys # save numpy - np.save(file_npy, temp) + xp.save(file_npy, temp) # move ids to first column and save txt - temp = np.roll(temp, 1, axis=1) - np.savetxt(file_txt, temp[:, (0, 1, 2, 3, -1)], fmt="%12.6f", delimiter=", ") + temp = xp.roll(temp, 1, axis=1) + xp.savetxt(file_txt, temp[:, (0, 1, 2, 3, -1)], fmt="%12.6f", delimiter=", ") # close hdf5 files for file in files: @@ -692,7 +692,7 @@ def post_process_f( path_slice, "grid_" + slice_names[n_gr] + ".npy", ) - np.save(grid_path, grid[:]) + xp.save(grid_path, grid[:]) # compute distribution function for slice_name in tqdm(files[0]["kinetic/" + species + "/f"]): @@ -713,8 +713,8 @@ def post_process_f( data_df += files[rank]["kinetic/" + species + "/df/" + slice_name][::step] # save distribution functions - np.save(os.path.join(path_slice, "f_binned.npy"), data) - np.save(os.path.join(path_slice, "delta_f_binned.npy"), data_df) + xp.save(os.path.join(path_slice, "f_binned.npy"), data) + xp.save(os.path.join(path_slice, "delta_f_binned.npy"), data_df) if compute_bckgr: # bckgr_params = params["kinetic"][species]["background"] @@ -753,11 +753,11 @@ def post_process_f( # check if file exists and is in slice_name if os.path.exists(filename) and current_slice in slice_names: - grid_tot += [np.load(filename)] + grid_tot += [xp.load(filename)] # otherwise evaluate at zero else: - grid_tot += [np.zeros(1)] + grid_tot += [xp.zeros(1)] # v-grid for comp in range(1, f_bckgr.vdim + 1): @@ -769,15 +769,15 @@ def post_process_f( # check if file exists and is in slice_name if os.path.exists(filename) and current_slice in slice_names: - grid_tot += [np.load(filename)] + grid_tot += [xp.load(filename)] # otherwise evaluate at zero else: - grid_tot += [np.zeros(1)] + grid_tot += [xp.zeros(1)] # correct integrating out in v-direction, TODO: check for 5D Maxwellians - factor *= np.sqrt(2 * np.pi) + factor *= xp.sqrt(2 * xp.pi) - grid_eval = np.meshgrid(*grid_tot, indexing="ij") + grid_eval = xp.meshgrid(*grid_tot, indexing="ij") data_bckgr = f_bckgr(*grid_eval).squeeze() @@ -788,9 +788,9 @@ def post_process_f( data_delta_f = data_df # save distribution function - np.save(os.path.join(path_slice, "delta_f_binned.npy"), data_delta_f) + xp.save(os.path.join(path_slice, "delta_f_binned.npy"), data_delta_f) # add extra axis for data_bckgr since data_delta_f has axis for time series - np.save( + xp.save( os.path.join(path_slice, "f_binned.npy"), data_delta_f + data_bckgr[tuple([None])], ) @@ -866,7 +866,7 @@ def post_process_n_sph( eta2 = files[0]["kinetic/" + species + "/n_sph/" + view].attrs["eta2"] eta3 = files[0]["kinetic/" + species + "/n_sph/" + view].attrs["eta3"] - ee1, ee2, ee3 = np.meshgrid( + ee1, ee2, ee3 = xp.meshgrid( eta1, eta2, eta3, @@ -877,7 +877,7 @@ def post_process_n_sph( path_view, "grid_n_sph.npy", ) - np.save(grid_path, (ee1, ee2, ee3)) + xp.save(grid_path, (ee1, ee2, ee3)) # load n_sph data data = files[0]["kinetic/" + species + "/n_sph/" + view][::step].copy() @@ -885,4 +885,4 @@ def post_process_n_sph( data += files[rank]["kinetic/" + species + "/n_sph/" + view][::step] # save distribution functions - np.save(os.path.join(path_view, "n_sph.npy"), data) + xp.save(os.path.join(path_view, "n_sph.npy"), data) diff --git a/src/struphy/post_processing/pproc_struphy.py b/src/struphy/post_processing/pproc_struphy.py index e9fe1b9b4..38dea12ba 100644 --- a/src/struphy/post_processing/pproc_struphy.py +++ b/src/struphy/post_processing/pproc_struphy.py @@ -8,7 +8,7 @@ import struphy.post_processing.orbits.orbits_tools as orbits_pproc import struphy.post_processing.post_processing_tools as pproc from struphy.io.setup import import_parameters_py -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def main( @@ -64,7 +64,7 @@ def main( file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") # save time grid at which post-processing data is created - np.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) + xp.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) if "feec" in file.keys(): exist_fields = True diff --git a/src/struphy/post_processing/profile_struphy.py b/src/struphy/post_processing/profile_struphy.py index 8c738a9a7..578e5c2e4 100644 --- a/src/struphy/post_processing/profile_struphy.py +++ b/src/struphy/post_processing/profile_struphy.py @@ -5,7 +5,7 @@ from matplotlib import pyplot as plt from struphy.post_processing.cprofile_analyser import get_cprofile_data, replace_keys -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp def main(): @@ -150,17 +150,17 @@ def main(): plt.ylabel("time [s]") plt.title("Strong scaling for Nel=" + str(val["Nel"][0]) + " cells") plt.legend(loc="lower left") - plt.loglog(val["mpi_size"], val["time"][0] / 2 ** np.arange(len(val["time"])), "k--", alpha=0.3) + plt.loglog(val["mpi_size"], val["time"][0] / 2 ** xp.arange(len(val["time"])), "k--", alpha=0.3) # weak scaling plot else: plt.plot(val["mpi_size"], val["time"], label=key) plt.xlabel("mpi_size") plt.ylabel("time [s]") plt.title( - "Weak scaling for cells/mpi_size=" + str(np.prod(val["Nel"][0]) / val["mpi_size"][0]) + "=const." + "Weak scaling for cells/mpi_size=" + str(xp.prod(val["Nel"][0]) / val["mpi_size"][0]) + "=const." ) plt.legend(loc="upper left") - # plt.loglog(val['mpi_size'], val['time'][0]*np.ones_like(val['time']), 'k--', alpha=0.3) + # plt.loglog(val['mpi_size'], val['time'][0]*xp.ones_like(val['time']), 'k--', alpha=0.3) plt.xscale("log") plt.show() diff --git a/src/struphy/profiling/profiling.py b/src/struphy/profiling/profiling.py index 568478e40..1ca22fe31 100644 --- a/src/struphy/profiling/profiling.py +++ b/src/struphy/profiling/profiling.py @@ -19,7 +19,7 @@ from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp @lru_cache(maxsize=None) # Cache the import result to avoid repeated imports @@ -171,9 +171,9 @@ def save_to_pickle(cls, file_path): for name, region in cls._regions.items(): local_data[name] = { "ncalls": region.ncalls, - "durations": np.array(region.durations, dtype=np.float64), - "start_times": np.array(region.start_times, dtype=np.float64), - "end_times": np.array(region.end_times, dtype=np.float64), + "durations": xp.array(region.durations, dtype=xp.float64), + "start_times": xp.array(region.start_times, dtype=xp.float64), + "end_times": xp.array(region.end_times, dtype=xp.float64), "config": { "likwid": region.config.likwid, "simulation_label": region.config.simulation_label, @@ -247,7 +247,7 @@ def print_summary(cls): average_duration = total_duration / region.ncalls min_duration = min(region.durations) max_duration = max(region.durations) - std_duration = np.std(region.durations) + std_duration = xp.std(region.durations) else: total_duration = average_duration = min_duration = max_duration = std_duration = 0 @@ -271,16 +271,16 @@ def __init__(self, region_name, time_trace=False): self._region_name = self.config.simulation_label + region_name self._time_trace = time_trace self._ncalls = 0 - self._start_times = np.empty(1, dtype=float) - self._end_times = np.empty(1, dtype=float) - self._durations = np.empty(1, dtype=float) + self._start_times = xp.empty(1, dtype=float) + self._end_times = xp.empty(1, dtype=float) + self._durations = xp.empty(1, dtype=float) self._started = False def __enter__(self): if self._ncalls == len(self._start_times): - self._start_times = np.append(self._start_times, np.zeros_like(self._start_times)) - self._end_times = np.append(self._end_times, np.zeros_like(self._end_times)) - self._durations = np.append(self._durations, np.zeros_like(self._durations)) + self._start_times = xp.append(self._start_times, xp.zeros_like(self._start_times)) + self._end_times = xp.append(self._end_times, xp.zeros_like(self._end_times)) + self._durations = xp.append(self._durations, xp.zeros_like(self._durations)) if self.config.likwid: self._pylikwid().markerstartregion(self.region_name) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 5f92f7373..e90382cb5 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -14,7 +14,7 @@ from struphy.geometry.base import Domain from struphy.io.options import check_option from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class Propagator(metaclass=ABCMeta): @@ -111,7 +111,7 @@ def update_feec_variables(self, **new_coeffs): assert new.space == old.space # calculate maximum of difference abs(new - old) - diffs[var] = np.max(np.abs(new.toarray() - old.toarray())) + diffs[var] = xp.max(xp.abs(new.toarray() - old.toarray())) # copy new coeffs into old new.copy(out=old) @@ -246,9 +246,9 @@ def add_init_kernel( The arguments for the kernel function. """ if comps is None: - comps = np.array([0]) # case for scalar evaluation + comps = xp.array([0]) # case for scalar evaluation else: - comps = np.array(comps, dtype=int) + comps = xp.array(comps, dtype=int) if not hasattr(self, "_init_kernels"): self._init_kernels = [] @@ -297,12 +297,12 @@ def add_eval_kernel( """ if isinstance(alpha, int) or isinstance(alpha, float): alpha = [alpha] * 6 - alpha = np.array(alpha) + alpha = xp.array(alpha) if comps is None: - comps = np.array([0]) # case for scalar evaluation + comps = xp.array([0]) # case for scalar evaluation else: - comps = np.array(comps, dtype=int) + comps = xp.array(comps, dtype=int) if not hasattr(self, "_eval_kernels"): self._eval_kernels = [] diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index 62e8716ce..ba1e15b0b 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -28,7 +28,7 @@ from struphy.pic.pushing.pusher import Pusher from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel @@ -252,14 +252,14 @@ def __call__(self, dt): print("Maxdiff e1 for VlasovMaxwell:", max_de) particles = self.variables.ions.particles buffer_idx = particles.bufferindex - max_diff = np.max( - np.abs( - np.sqrt( + max_diff = xp.max( + xp.abs( + xp.sqrt( particles.markers_wo_holes[:, 3] ** 2 + particles.markers_wo_holes[:, 4] ** 2 + particles.markers_wo_holes[:, 5] ** 2, ) - - np.sqrt( + - xp.sqrt( particles.markers_wo_holes[:, buffer_idx + 3] ** 2 + particles.markers_wo_holes[:, buffer_idx + 4] ** 2 + particles.markers_wo_holes[:, buffer_idx + 5] ** 2, @@ -386,8 +386,8 @@ def __init__( self._e_sum = e.space.zeros() # marker storage - self._f0_values = np.zeros(particles.markers.shape[0], dtype=float) - self._old_weights = np.empty(particles.markers.shape[0], dtype=float) + self._f0_values = xp.zeros(particles.markers.shape[0], dtype=float) + self._old_weights = xp.empty(particles.markers.shape[0], dtype=float) # ================================ # ========= Schur Solver ========= @@ -483,8 +483,8 @@ def __call__(self, dt): print("Status for StepEfieldWeights:", info["success"]) print("Iterations for StepEfieldWeights:", info["niter"]) print("Maxdiff e1 for StepEfieldWeights:", max_de) - max_diff = np.max( - np.abs( + max_diff = xp.max( + xp.abs( self._old_weights[~self.particles[0].holes] - self.particles[0].markers[~self.particles[0].holes, 6], ), @@ -943,14 +943,14 @@ def __init__( # self.particles[0].f0.n, *quad_pts, kind='0', squeeze_out=False, coordinates='logical') # # memory allocation for magnetic field at quadrature points - # self._b_quad1 = np.zeros_like(self._nuh0_at_quad[0]) - # self._b_quad2 = np.zeros_like(self._nuh0_at_quad[0]) - # self._b_quad3 = np.zeros_like(self._nuh0_at_quad[0]) + # self._b_quad1 = xp.zeros_like(self._nuh0_at_quad[0]) + # self._b_quad2 = xp.zeros_like(self._nuh0_at_quad[0]) + # self._b_quad3 = xp.zeros_like(self._nuh0_at_quad[0]) # # memory allocation for (self._b_quad x self._nuh0_at_quad) * self._coupling_vec - # self._vec1 = np.zeros_like(self._nuh0_at_quad[0]) - # self._vec2 = np.zeros_like(self._nuh0_at_quad[0]) - # self._vec3 = np.zeros_like(self._nuh0_at_quad[0]) + # self._vec1 = xp.zeros_like(self._nuh0_at_quad[0]) + # self._vec2 = xp.zeros_like(self._nuh0_at_quad[0]) + # self._vec3 = xp.zeros_like(self._nuh0_at_quad[0]) # FEM spaces and basis extraction operators for u and b u_id = self.derham.space_to_form[u_space] @@ -1578,8 +1578,8 @@ def allocate(self): butcher = self.options.butcher import numpy as np - butcher._a = np.diag(butcher.a, k=-1) - butcher._a = np.array(list(butcher.a) + [0.0]) + butcher._a = xp.diag(butcher.a, k=-1) + butcher._a = xp.array(list(butcher.a) + [0.0]) self._args_pusher_kernel = ( self.domain.args_domain, @@ -1822,10 +1822,10 @@ def __call__(self, dt): # save en_fB_old particles.save_magnetic_energy(PB_b) - en_fB_old = np.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_old = xp.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale en_fB_old /= n_mks_tot - buffer_array = np.array([en_fB_old]) + buffer_array = xp.array([en_fB_old]) if particles.mpi_comm is not None: particles.mpi_comm.Allreduce( @@ -1868,10 +1868,10 @@ def __call__(self, dt): # save en_fB_new particles.save_magnetic_energy(PB_b) - en_fB_new = np.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_new = xp.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale en_fB_new /= n_mks_tot - buffer_array = np.array([en_fB_new]) + buffer_array = xp.array([en_fB_new]) if particles.mpi_comm is not None: particles.mpi_comm.Allreduce( @@ -1915,13 +1915,13 @@ def __call__(self, dt): markers[~holes, first_free_idx : first_free_idx + 3] = markers[~holes, 0:3] # calculate denominator ||z^{n+1, k} - z^n||^2 - sum_u_diff_loc = np.sum((u_diff.toarray() ** 2)) + sum_u_diff_loc = xp.sum((u_diff.toarray() ** 2)) - sum_H_diff_loc = np.sum( + sum_H_diff_loc = xp.sum( (markers[~holes, :3] - markers[~holes, first_init_idx : first_init_idx + 3]) ** 2 ) - buffer_array = np.array([sum_u_diff_loc]) + buffer_array = xp.array([sum_u_diff_loc]) if particles.mpi_comm is not None: particles.mpi_comm.Allreduce( @@ -1932,7 +1932,7 @@ def __call__(self, dt): denominator = buffer_array[0] - buffer_array = np.array([sum_H_diff_loc]) + buffer_array = xp.array([sum_H_diff_loc]) if particles.mpi_comm is not None: particles.mpi_comm.Allreduce( @@ -1959,11 +1959,11 @@ def __call__(self, dt): *self._args_accum_kernel_en_fB_mid, first_free_idx + 3, ) - en_fB_mid = np.sum(markers[~holes, first_free_idx + 3].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_mid = xp.sum(markers[~holes, first_free_idx + 3].dot(markers[~holes, 5])) * self.options.ep_scale en_fB_mid /= n_mks_tot - buffer_array = np.array([en_fB_mid]) + buffer_array = xp.array([en_fB_mid]) if particles.mpi_comm is not None: particles.mpi_comm.Allreduce( @@ -2011,8 +2011,8 @@ def __call__(self, dt): alpha, ) - sum_H_diff_loc = np.sum( - np.abs(markers[~holes, 0:3] - markers[~holes, first_free_idx : first_free_idx + 3]) + sum_H_diff_loc = xp.sum( + xp.abs(markers[~holes, 0:3] - markers[~holes, first_free_idx : first_free_idx + 3]) ) if particles.mpi_comm is not None: @@ -2020,10 +2020,10 @@ def __call__(self, dt): # update en_fB_new particles.save_magnetic_energy(PB_b) - en_fB_new = np.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_new = xp.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale en_fB_new /= n_mks_tot - buffer_array = np.array([en_fB_new]) + buffer_array = xp.array([en_fB_new]) if particles.mpi_comm is not None: particles.mpi_comm.Allreduce( @@ -2042,12 +2042,12 @@ def __call__(self, dt): en_fB_new = buffer_array[0] # calculate total energy difference - e_diff = np.abs(en_U_new + en_fB_new - en_tot_old) + e_diff = xp.abs(en_U_new + en_fB_new - en_tot_old) # calculate ||z^{n+1, k} - z^{n+1, k-1|| - sum_u_diff_loc = np.sum(np.abs(u_new.toarray() - u_old.toarray())) + sum_u_diff_loc = xp.sum(xp.abs(u_new.toarray() - u_old.toarray())) - buffer_array = np.array([sum_u_diff_loc]) + buffer_array = xp.array([sum_u_diff_loc]) if particles.mpi_comm is not None: particles.mpi_comm.Allreduce( @@ -2058,7 +2058,7 @@ def __call__(self, dt): diff = buffer_array[0] - buffer_array = np.array([sum_H_diff_loc]) + buffer_array = xp.array([sum_H_diff_loc]) if particles.mpi_comm is not None: particles.mpi_comm.Allreduce( diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index d02decff9..109638c1a 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -67,7 +67,7 @@ from struphy.pic.particles import Particles5D, Particles6D from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel @@ -1540,13 +1540,13 @@ def __init__(self, a, **params): ] # Initialize Accumulator object for getting density from particles - self._pts_x = 1.0 / (2.0 * self.derham.Nel[0]) * np.polynomial.legendre.leggauss( + self._pts_x = 1.0 / (2.0 * self.derham.Nel[0]) * xp.polynomial.legendre.leggauss( self._nqs[0], )[0] + 1.0 / (2.0 * self.derham.Nel[0]) - self._pts_y = 1.0 / (2.0 * self.derham.Nel[1]) * np.polynomial.legendre.leggauss( + self._pts_y = 1.0 / (2.0 * self.derham.Nel[1]) * xp.polynomial.legendre.leggauss( self._nqs[1], )[0] + 1.0 / (2.0 * self.derham.Nel[1]) - self._pts_z = 1.0 / (2.0 * self.derham.Nel[2]) * np.polynomial.legendre.leggauss( + self._pts_z = 1.0 / (2.0 * self.derham.Nel[2]) * xp.polynomial.legendre.leggauss( self._nqs[2], )[0] + 1.0 / (2.0 * self.derham.Nel[2]) @@ -1586,15 +1586,15 @@ def __call__(self, dt): self._accum_density.accumulate( self._particles, - np.array(self.derham.Nel), - np.array(self._nqs), - np.array( + xp.array(self.derham.Nel), + xp.array(self._nqs), + xp.array( self._pts_x, ), - np.array(self._pts_y), - np.array(self._pts_z), - np.array(self._p_shape), - np.array(self._p_size), + xp.array(self._pts_y), + xp.array(self._pts_z), + xp.array(self._p_shape), + xp.array(self._p_size), ) self._accum_potential.accumulate(self._particles) @@ -1761,18 +1761,18 @@ def __init__( # self._particles.f0.n, *quad_pts, kind='3', squeeze_out=False) # # memory allocation of magnetic field at quadrature points - # self._b_quad1 = np.zeros_like(self._nh0_at_quad) - # self._b_quad2 = np.zeros_like(self._nh0_at_quad) - # self._b_quad3 = np.zeros_like(self._nh0_at_quad) + # self._b_quad1 = xp.zeros_like(self._nh0_at_quad) + # self._b_quad2 = xp.zeros_like(self._nh0_at_quad) + # self._b_quad3 = xp.zeros_like(self._nh0_at_quad) # # memory allocation for self._b_quad x self._nh0_at_quad * self._coupling_const - # self._mat12 = np.zeros_like(self._nh0_at_quad) - # self._mat13 = np.zeros_like(self._nh0_at_quad) - # self._mat23 = np.zeros_like(self._nh0_at_quad) + # self._mat12 = xp.zeros_like(self._nh0_at_quad) + # self._mat13 = xp.zeros_like(self._nh0_at_quad) + # self._mat23 = xp.zeros_like(self._nh0_at_quad) - # self._mat21 = np.zeros_like(self._nh0_at_quad) - # self._mat31 = np.zeros_like(self._nh0_at_quad) - # self._mat32 = np.zeros_like(self._nh0_at_quad) + # self._mat21 = xp.zeros_like(self._nh0_at_quad) + # self._mat31 = xp.zeros_like(self._nh0_at_quad) + # self._mat32 = xp.zeros_like(self._nh0_at_quad) self._type = solver["type"][0] self._tol = solver["tol"] @@ -2581,7 +2581,7 @@ def options(self, new): @profile def allocate(self): # always stabilize - if np.abs(self.options.sigma_1) < 1e-14: + if xp.abs(self.options.sigma_1) < 1e-14: self.options.sigma_1 = 1e-14 if MPI.COMM_WORLD.Get_rank() == 0: print(f"Stabilizing Poisson solve with {self.options.sigma_1 = }") @@ -3003,7 +3003,7 @@ def __call_newton(self, dt): if self._info: print("iteration : ", it, " error : ", err) - if err < tol**2 or np.isnan(err): + if err < tol**2 or xp.isnan(err): break # Newton step @@ -3017,7 +3017,7 @@ def __call_newton(self, dt): un1 -= update mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self.options.nonlin_solver.maxiter - 1 or np.isnan(err): + if it == self.options.nonlin_solver.maxiter - 1 or xp.isnan(err): print( f"!!!WARNING: Maximum iteration in VariationalMomentumAdvection reached - not converged \n {err = } \n {tol**2 = }", ) @@ -3036,7 +3036,7 @@ def __call_picard(self, dt): for it in range(self.options.nonlin_solver.maxiter): # Picard iteration - if err < tol**2 or np.isnan(err): + if err < tol**2 or xp.isnan(err): break # half time step approximation un12 = un.copy(out=self._tmp_un12) @@ -3063,7 +3063,7 @@ def __call_picard(self, dt): # Inverse the mass matrix to get the velocity un1 = self._Mrho.inv.dot(mn1, out=self._tmp_un1) - if it == self.options.nonlin_solver.maxiter - 1 or np.isnan(err): + if it == self.options.nonlin_solver.maxiter - 1 or xp.isnan(err): print( f"!!!WARNING: Maximum iteration in VariationalMomentumAdvection reached - not converged \n {err = } \n {tol**2 = }", ) @@ -3385,7 +3385,7 @@ def __call_newton(self, dt): if self._info: print("iteration : ", it, " error : ", err) - if err < tol**2 or np.isnan(err): + if err < tol**2 or xp.isnan(err): break # Derivative for Newton @@ -3415,7 +3415,7 @@ def __call_newton(self, dt): mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalDensityEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) @@ -3458,7 +3458,7 @@ def _initialize_projectors_and_mass(self): # tmps grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) - self._rhof_values = np.zeros(grid_shape, dtype=float) + self._rhof_values = xp.zeros(grid_shape, dtype=float) # Other mass matrices for newton solve self._M_drho = self.mass_ops.create_weighted_mass("L2", "L2") @@ -3510,20 +3510,20 @@ def _initialize_projectors_and_mass(self): grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) # tmps - self._eval_dl_drho = np.zeros(grid_shape, dtype=float) + self._eval_dl_drho = xp.zeros(grid_shape, dtype=float) - self._uf_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._uf1_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] + self._uf_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._uf1_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] - self._tmp_int_grid = np.zeros(grid_shape, dtype=float) - self._tmp_int_grid2 = np.zeros(grid_shape, dtype=float) - self._rhof_values = np.zeros(grid_shape, dtype=float) - self._rhof1_values = np.zeros(grid_shape, dtype=float) + self._tmp_int_grid = xp.zeros(grid_shape, dtype=float) + self._tmp_int_grid2 = xp.zeros(grid_shape, dtype=float) + self._rhof_values = xp.zeros(grid_shape, dtype=float) + self._rhof1_values = xp.zeros(grid_shape, dtype=float) if self._model == "full": - self._tmp_de_drho = np.zeros(grid_shape, dtype=float) + self._tmp_de_drho = xp.zeros(grid_shape, dtype=float) gam = self._gamma - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -3531,7 +3531,7 @@ def _initialize_projectors_and_mass(self): ) self._proj_rho2_metric_term = deepcopy(metric) - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -3540,7 +3540,7 @@ def _initialize_projectors_and_mass(self): self._proj_drho_metric_term = deepcopy(metric) if self._linearize: - self._init_dener_drho = np.zeros(grid_shape, dtype=float) + self._init_dener_drho = xp.zeros(grid_shape, dtype=float) def _update_Pirho(self, rho): """Update the weights of the `BasisProjectionOperator` Pirho""" @@ -3856,7 +3856,7 @@ def __call_newton(self, dt): if self._info: print("iteration : ", it, " error : ", err) - if err < tol**2 or np.isnan(err): + if err < tol**2 or xp.isnan(err): break # Derivative for Newton @@ -3878,7 +3878,7 @@ def __call_newton(self, dt): # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalEntropyEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) @@ -3965,15 +3965,15 @@ def _initialize_projectors_and_mass(self): ) grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) - self._tmp_int_grid = np.zeros(grid_shape, dtype=float) + self._tmp_int_grid = xp.zeros(grid_shape, dtype=float) if self._model == "full": - self._tmp_de_ds = np.zeros(grid_shape, dtype=float) + self._tmp_de_ds = xp.zeros(grid_shape, dtype=float) if self._linearize: - self._init_dener_ds = np.zeros(grid_shape, dtype=float) + self._init_dener_ds = xp.zeros(grid_shape, dtype=float) gam = self._gamma - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -3981,7 +3981,7 @@ def _initialize_projectors_and_mass(self): ) self._proj_rho2_metric_term = deepcopy(metric) - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -4276,7 +4276,7 @@ def __call_newton(self, dt): if self._info: print("iteration : ", it, " error : ", err) - if err < tol**2 or np.isnan(err): + if err < tol**2 or xp.isnan(err): break # Derivative for Newton @@ -4298,7 +4298,7 @@ def __call_newton(self, dt): # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalMagFieldEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) @@ -4779,7 +4779,7 @@ def __call_picard(self, dt): if self._info: print("iteration : ", it, " error : ", err) - if err < tol**2 or np.isnan(err): + if err < tol**2 or xp.isnan(err): break # Derivative for Newton @@ -4801,7 +4801,7 @@ def __call_picard(self, dt): # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalPBEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) @@ -4850,7 +4850,7 @@ def _initialize_projectors_and_mass(self): grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) - self._tmp_int_grid = np.zeros(grid_shape, dtype=float) + self._tmp_int_grid = xp.zeros(grid_shape, dtype=float) # Inverse mass matrix needed to compute the error self.pc_Mv = preconditioner.MassMatrixDiagonalPreconditioner( @@ -5368,7 +5368,7 @@ def __call_picard(self, dt): if self._info: print("iteration : ", it, " error : ", err) - if err < tol**2 or np.isnan(err): + if err < tol**2 or xp.isnan(err): break # Derivative for Newton @@ -5392,7 +5392,7 @@ def __call_picard(self, dt): # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.massop.dot(un1, out=self._tmp_mn1) - if it == self._nonlin_solver.maxiter - 1 or np.isnan(err): + if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalPBEvolve reached - not converged:\n {err = } \n {tol**2 = }", ) @@ -5441,7 +5441,7 @@ def _initialize_projectors_and_mass(self): grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) - self._tmp_int_grid = np.zeros(grid_shape, dtype=float) + self._tmp_int_grid = xp.zeros(grid_shape, dtype=float) # Inverse mass matrix needed to compute the error self.pc_Mv = preconditioner.MassMatrixDiagonalPreconditioner( @@ -5946,7 +5946,7 @@ def __call_newton(self, dt): if self._info: print("iteration : ", it, " error : ", err) - if (err < tol**2 and it > 0) or np.isnan(err): + if (err < tol**2 and it > 0) or xp.isnan(err): # force at least one iteration break @@ -5984,7 +5984,7 @@ def __call_newton(self, dt): else: sn1 += incr - if it == self._nonlin_solver["maxiter"] - 1 or np.isnan(err): + if it == self._nonlin_solver["maxiter"] - 1 or xp.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalViscosity reached - not converged:\n {err = } \n {tol**2 = }", ) @@ -6120,35 +6120,35 @@ def _initialize_projectors_and_mass(self): grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) - self._guf0_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._guf1_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._guf2_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] + self._guf0_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._guf1_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._guf2_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] - self._guf120_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._guf121_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._guf122_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] + self._guf120_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._guf121_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._guf122_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] - self._uf1_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._uf12_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] + self._uf1_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._uf12_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] - self._gu_sq_values = np.zeros(grid_shape, dtype=float) - self._u_sq_values = np.zeros(grid_shape, dtype=float) - self._gu_init_values = np.zeros(grid_shape, dtype=float) + self._gu_sq_values = xp.zeros(grid_shape, dtype=float) + self._u_sq_values = xp.zeros(grid_shape, dtype=float) + self._gu_init_values = xp.zeros(grid_shape, dtype=float) - self._sf_values = np.zeros(grid_shape, dtype=float) - self._sf1_values = np.zeros(grid_shape, dtype=float) - self._rhof_values = np.zeros(grid_shape, dtype=float) + self._sf_values = xp.zeros(grid_shape, dtype=float) + self._sf1_values = xp.zeros(grid_shape, dtype=float) + self._rhof_values = xp.zeros(grid_shape, dtype=float) - self._e_n1 = np.zeros(grid_shape, dtype=float) - self._e_n = np.zeros(grid_shape, dtype=float) + self._e_n1 = xp.zeros(grid_shape, dtype=float) + self._e_n = xp.zeros(grid_shape, dtype=float) - self._de_s1_values = np.zeros(grid_shape, dtype=float) + self._de_s1_values = xp.zeros(grid_shape, dtype=float) - self._tmp_int_grid = np.zeros(grid_shape, dtype=float) + self._tmp_int_grid = xp.zeros(grid_shape, dtype=float) gam = self._gamma if self._model == "full": - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -6156,7 +6156,7 @@ def _initialize_projectors_and_mass(self): ) self._mass_metric_term = deepcopy(metric) - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -6189,7 +6189,7 @@ def _initialize_projectors_and_mass(self): self.pc_jac.update_mass_operator(self.M_de_ds) elif self._model in ["full_q", "linear_q", "deltaf_q"]: - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -6197,7 +6197,7 @@ def _initialize_projectors_and_mass(self): ) self._mass_metric_term = deepcopy(metric) - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -6205,7 +6205,7 @@ def _initialize_projectors_and_mass(self): ) self._energy_metric = deepcopy(metric) - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -6271,7 +6271,7 @@ def _update_artificial_viscosity(self, un, dt): gu_sq_v += gu1_v[i] gu_sq_v += gu2_v[i] - np.sqrt(gu_sq_v, out=gu_sq_v) + xp.sqrt(gu_sq_v, out=gu_sq_v) gu_sq_v *= dt * self._mu_a # /2 @@ -6712,7 +6712,7 @@ def __call_newton(self, dt): if self._info: print("iteration : ", it, " error : ", err) - if (err < tol**2 and it > 0) or np.isnan(err): + if (err < tol**2 and it > 0) or xp.isnan(err): break if self._model == "full": @@ -6748,7 +6748,7 @@ def __call_newton(self, dt): else: sn1 += incr - if it == self._nonlin_solver["maxiter"] - 1 or np.isnan(err): + if it == self._nonlin_solver["maxiter"] - 1 or xp.isnan(err): print( f"!!!Warning: Maximum iteration in VariationalResistivity reached - not converged:\n {err = } \n {tol**2 = }", ) @@ -6831,7 +6831,7 @@ def __call_newton(self, dt): # if self._info: # print("iteration : ", it, " error : ", err) - # if (err < tol**2 and it > 0) or np.isnan(err): + # if (err < tol**2 and it > 0) or xp.isnan(err): # break # incr = self.inv_jac.dot(self.tot_rhs, out=self._tmp_sn_incr) @@ -6944,26 +6944,26 @@ def _initialize_projectors_and_mass(self): grid_shape = tuple([len(loc_grid) for loc_grid in integration_grid]) - self._cb12_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] - self._cb1_values = [np.zeros(grid_shape, dtype=float) for i in range(3)] + self._cb12_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] + self._cb1_values = [xp.zeros(grid_shape, dtype=float) for i in range(3)] - self._cb_sq_values = np.zeros(grid_shape, dtype=float) - self._cb_sq_values_init = np.zeros(grid_shape, dtype=float) + self._cb_sq_values = xp.zeros(grid_shape, dtype=float) + self._cb_sq_values_init = xp.zeros(grid_shape, dtype=float) - self._sf_values = np.zeros(grid_shape, dtype=float) - self._sf1_values = np.zeros(grid_shape, dtype=float) - self._rhof_values = np.zeros(grid_shape, dtype=float) + self._sf_values = xp.zeros(grid_shape, dtype=float) + self._sf1_values = xp.zeros(grid_shape, dtype=float) + self._rhof_values = xp.zeros(grid_shape, dtype=float) - self._e_n1 = np.zeros(grid_shape, dtype=float) - self._e_n = np.zeros(grid_shape, dtype=float) + self._e_n1 = xp.zeros(grid_shape, dtype=float) + self._e_n = xp.zeros(grid_shape, dtype=float) - self._de_s1_values = np.zeros(grid_shape, dtype=float) + self._de_s1_values = xp.zeros(grid_shape, dtype=float) - self._tmp_int_grid = np.zeros(grid_shape, dtype=float) + self._tmp_int_grid = xp.zeros(grid_shape, dtype=float) gam = self._gamma if self._model == "full": - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -6971,7 +6971,7 @@ def _initialize_projectors_and_mass(self): ) self._mass_metric_term = deepcopy(metric) - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -7004,7 +7004,7 @@ def _initialize_projectors_and_mass(self): self.pc_jac.update_mass_operator(self.M_de_ds) elif self._model in ["full_q", "linear_q", "deltaf_q"]: - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -7012,7 +7012,7 @@ def _initialize_projectors_and_mass(self): ) self._mass_metric_term = deepcopy(metric) - metric = np.power( + metric = xp.power( self.domain.jacobian_det( *integration_grid, ), @@ -7059,7 +7059,7 @@ def _update_artificial_resistivity(self, bn, dt): for j in range(3): cb_sq_v += cb_v[i] * self._sq_term_metric_no_jac[i, j] * cb_v[j] - np.sqrt(cb_sq_v, out=cb_sq_v) + xp.sqrt(cb_sq_v, out=cb_sq_v) cb_sq_v *= dt * self._eta_a @@ -7179,7 +7179,7 @@ class Options: # specific literals OptsTimeSource = Literal["cos", "sin"] # propagator options - omega: float = 2.0 * np.pi + omega: float = 2.0 * xp.pi hfun: OptsTimeSource = "cos" def __post_init__(self): @@ -7206,11 +7206,11 @@ def allocate(self): if self.options.hfun == "cos": def hfun(t): - return np.cos(self.options.omega * t) + return xp.cos(self.options.omega * t) elif self.options.hfun == "sin": def hfun(t): - return np.sin(self.options.omega * t) + return xp.sin(self.options.omega * t) else: raise NotImplementedError(f"{self.options.hfun = } not implemented.") @@ -7555,7 +7555,7 @@ def allocate(self): # get quadrature grid of V0 pts = [grid.flatten() for grid in self.derham.quad_grid_pts["0"]] - mesh_pts = np.meshgrid(*pts, indexing="ij") + mesh_pts = xp.meshgrid(*pts, indexing="ij") # evaluate c(x, y) and metric coeff at local quadrature grid and multiply self._weights = c_fun(*mesh_pts) @@ -7588,13 +7588,13 @@ def allocate(self): for m in range(3): self._M1hw_weights += [[None, None, None]] - self._phi_5d = np.zeros((*self._phi_at_pts.shape, 3, 3), dtype=float) - self._tmp_5d = np.zeros((*self._phi_at_pts.shape, 3, 3), dtype=float) - self._tmp_5dT = np.zeros((3, 3, *self._phi_at_pts.shape), dtype=float) + self._phi_5d = xp.zeros((*self._phi_at_pts.shape, 3, 3), dtype=float) + self._tmp_5d = xp.zeros((*self._phi_at_pts.shape, 3, 3), dtype=float) + self._tmp_5dT = xp.zeros((3, 3, *self._phi_at_pts.shape), dtype=float) self._phi_5d[:, :, :, 0, 1] = self._phi_at_pts * self._jac_det self._phi_5d[:, :, :, 1, 0] = -self._phi_at_pts * self._jac_det self._tmp_5d[:] = self._jac_inv @ self._phi_5d @ self._jac_invT - self._tmp_5dT[:] = np.transpose(self._tmp_5d, axes=(3, 4, 0, 1, 2)) + self._tmp_5dT[:] = xp.transpose(self._tmp_5d, axes=(3, 4, 0, 1, 2)) self._M1hw_weights[0][1] = self._tmp_5dT[0, 1, :, :, :] self._M1hw_weights[1][0] = self._tmp_5dT[1, 0, :, :, :] @@ -7695,7 +7695,7 @@ def __call__(self, dt): self._phi_5d[:, :, :, 0, 1] = self._phi_at_pts * self._jac_det self._phi_5d[:, :, :, 1, 0] = -self._phi_at_pts * self._jac_det self._tmp_5d[:] = self._jac_inv @ self._phi_5d @ self._jac_invT - self._tmp_5dT[:] = np.transpose(self._tmp_5d, axes=(3, 4, 0, 1, 2)) + self._tmp_5dT[:] = xp.transpose(self._tmp_5d, axes=(3, 4, 0, 1, 2)) self._M1hw_weights[0][1] = self._tmp_5dT[0, 1, :, :, :] self._M1hw_weights[1][0] = self._tmp_5dT[1, 0, :, :, :] @@ -8346,9 +8346,9 @@ def allocate(self): A11np = self._M2np + self._A11np_notimedependency if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - A11np += self._stab_sigma * np.identity(A11np.shape[0]) + A11np += self._stab_sigma * xp.identity(A11np.shape[0]) self.A22np = ( - self._stab_sigma * np.identity(A11np.shape[0]) + self._stab_sigma * xp.identity(A11np.shape[0]) + self._nu_e * ( self._Dnp.T @ self._M3np @ self._Dnp @@ -8357,7 +8357,7 @@ def allocate(self): + self._M2Bnp / self._eps_norm ) self._A22prenp = ( - np.identity(self.A22np.shape[0]) * self._stab_sigma + xp.identity(self.A22np.shape[0]) * self._stab_sigma ) # + self._nu_e * (self._Dnp.T @ self._M3np @ self._Dnp) elif self._method_to_solve in ("SparseSolver", "ScipySparse"): A11np += self._stab_sigma * sc.sparse.eye(A11np.shape[0], format="csr") @@ -8540,7 +8540,7 @@ def __call__(self, dt): # Numpy A11np = self._M2np / dt + self._A11np_notimedependency if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): - A11np += self._stab_sigma * np.identity(A11np.shape[0]) + A11np += self._stab_sigma * xp.identity(A11np.shape[0]) _A22prenp = self._A22prenp A22np = self.A22np elif self._method_to_solve in ("SparseSolver", "ScipySparse"): diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 02cccb306..1e9d2658a 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -30,7 +30,7 @@ from struphy.pic.pushing.pusher import Pusher from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel @@ -100,10 +100,10 @@ def allocate(self): # define algorithm butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - butcher._a = np.diag(butcher.a, k=-1) - butcher._a = np.array(list(butcher.a) + [0.0]) + butcher._a = xp.diag(butcher.a, k=-1) + butcher._a = xp.array(list(butcher.a) + [0.0]) args_kernel = ( butcher.a, @@ -842,10 +842,10 @@ def allocate(self): else: butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - butcher._a = np.diag(butcher.a, k=-1) - butcher._a = np.array(list(butcher.a) + [0.0]) + butcher._a = xp.diag(butcher.a, k=-1) + butcher._a = xp.array(list(butcher.a) + [0.0]) kernel = Pyccelkernel(pusher_kernels_gc.push_gc_bxEstar_explicit_multistage) @@ -1290,10 +1290,10 @@ def allocate(self): else: butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - butcher._a = np.diag(butcher.a, k=-1) - butcher._a = np.array(list(butcher.a) + [0.0]) + butcher._a = xp.diag(butcher.a, k=-1) + butcher._a = xp.array(list(butcher.a) + [0.0]) kernel = Pyccelkernel(pusher_kernels_gc.push_gc_Bstar_explicit_multistage) @@ -1432,10 +1432,10 @@ def allocate(self): # choose algorithm self._butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - self._butcher._a = np.diag(self._butcher.a, k=-1) - self._butcher._a = np.array(list(self._butcher.a) + [0.0]) + self._butcher._a = xp.diag(self._butcher.a, k=-1) + self._butcher._a = xp.array(list(self._butcher.a) + [0.0]) particles = self.variables.var.particles @@ -1566,10 +1566,10 @@ def allocate(self): self._butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp - self._butcher._a = np.diag(self._butcher.a, k=-1) - self._butcher._a = np.array(list(self._butcher.a) + [0.0]) + self._butcher._a = xp.diag(self._butcher.a, k=-1) + self._butcher._a = xp.array(list(self._butcher.a) + [0.0]) # instantiate Pusher args_kernel = ( @@ -1728,7 +1728,7 @@ def allocate(self): elif self.options.thermodynamics == "polytropic": kernel = Pyccelkernel(pusher_kernels.push_v_sph_pressure_ideal_gas) - gravity = np.array(self.options.gravity, dtype=float) + gravity = xp.array(self.options.gravity, dtype=float) args_kernel = ( boxes, diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index 67e84decd..b1e3cfe33 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -11,7 +11,7 @@ from struphy.models.variables import FEECVariable from struphy.propagators.base import Propagator from struphy.propagators.propagators_fields import ImplicitDiffusion -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -72,16 +72,16 @@ def test_poisson_M1perp_1d(direction, bc_type, mapping, projected_rhs, show_plot if direction == 0: Nel = [Neli, 1, 1] p = [pi, 1, 1] - e1 = np.linspace(0.0, 1.0, 50) + e1 = xp.linspace(0.0, 1.0, 50) if bc_type == "neumann": spl_kind = [False, True, True] def sol1_xyz(x, y, z): - return np.cos(np.pi / Lx * x) + return xp.cos(xp.pi / Lx * x) def rho1_xyz(x, y, z): - return np.cos(np.pi / Lx * x) * (np.pi / Lx) ** 2 + return xp.cos(xp.pi / Lx * x) * (xp.pi / Lx) ** 2 else: if bc_type == "dirichlet": spl_kind = [False, True, True] @@ -89,24 +89,24 @@ def rho1_xyz(x, y, z): dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): - return np.sin(2 * np.pi / Lx * x) + return xp.sin(2 * xp.pi / Lx * x) def rho1_xyz(x, y, z): - return np.sin(2 * np.pi / Lx * x) * (2 * np.pi / Lx) ** 2 + return xp.sin(2 * xp.pi / Lx * x) * (2 * xp.pi / Lx) ** 2 elif direction == 1: Nel = [1, Neli, 1] p = [1, pi, 1] - e2 = np.linspace(0.0, 1.0, 50) + e2 = xp.linspace(0.0, 1.0, 50) if bc_type == "neumann": spl_kind = [True, False, True] def sol1_xyz(x, y, z): - return np.cos(np.pi / Ly * y) + return xp.cos(xp.pi / Ly * y) def rho1_xyz(x, y, z): - return np.cos(np.pi / Ly * y) * (np.pi / Ly) ** 2 + return xp.cos(xp.pi / Ly * y) * (xp.pi / Ly) ** 2 else: if bc_type == "dirichlet": spl_kind = [True, False, True] @@ -114,10 +114,10 @@ def rho1_xyz(x, y, z): dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): - return np.sin(2 * np.pi / Ly * y) + return xp.sin(2 * xp.pi / Ly * y) def rho1_xyz(x, y, z): - return np.sin(2 * np.pi / Ly * y) * (2 * np.pi / Ly) ** 2 + return xp.sin(2 * xp.pi / Ly * y) * (2 * xp.pi / Ly) ** 2 else: print("Direction should be either 0 or 1") @@ -196,14 +196,14 @@ def rho_pulled(e1, e2, e3): plt.title(f"{Nel = }") plt.legend() - error = np.max(np.abs(analytic_value1 - sol_val1)) + error = xp.max(xp.abs(analytic_value1 - sol_val1)) print(f"{direction = }, {pi = }, {Neli = }, {error=}") errors.append(error) h = 1 / (Neli) h_vec.append(h) - m, _ = np.polyfit(np.log(Nels), np.log(errors), deg=1) + m, _ = xp.polyfit(xp.log(Nels), xp.log(errors), deg=1) print(f"For {pi = }, solution converges in {direction=} with rate {-m = } ") assert -m > (pi + 1 - 0.07) @@ -263,10 +263,10 @@ def test_poisson_M1perp_2d(Nel, p, bc_type, mapping, projected_rhs, show_plot=Fa # manufactured solution in 1D (overwritten for "neumann") def sol1_xyz(x, y, z): - return np.sin(2 * np.pi / Lx * x) + return xp.sin(2 * xp.pi / Lx * x) def rho1_xyz(x, y, z): - return np.sin(2 * np.pi / Lx * x) * (2 * np.pi / Lx) ** 2 + return xp.sin(2 * xp.pi / Lx * x) * (2 * xp.pi / Lx) ** 2 # boundary conditions dirichlet_bc = None @@ -276,11 +276,11 @@ def rho1_xyz(x, y, z): # manufactured solution in 2D def sol2_xyz(x, y, z): - return np.sin(2 * np.pi * x / Lx + 4 * np.pi / Ly * y) + return xp.sin(2 * xp.pi * x / Lx + 4 * xp.pi / Ly * y) def rho2_xyz(x, y, z): - ddx = np.sin(2 * np.pi / Lx * x + 4 * np.pi / Ly * y) * (2 * np.pi / Lx) ** 2 - ddy = np.sin(2 * np.pi / Lx * x + 4 * np.pi / Ly * y) * (4 * np.pi / Ly) ** 2 + ddx = xp.sin(2 * xp.pi / Lx * x + 4 * xp.pi / Ly * y) * (2 * xp.pi / Lx) ** 2 + ddy = xp.sin(2 * xp.pi / Lx * x + 4 * xp.pi / Ly * y) * (4 * xp.pi / Ly) ** 2 return ddx + ddy elif bc_type == "dirichlet": @@ -291,11 +291,11 @@ def rho2_xyz(x, y, z): # manufactured solution in 2D def sol2_xyz(x, y, z): - return np.sin(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) + return xp.sin(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) def rho2_xyz(x, y, z): - ddx = np.sin(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) * (np.pi / Lx) ** 2 - ddy = np.sin(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) * (4 * np.pi / Ly) ** 2 + ddx = xp.sin(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) * (xp.pi / Lx) ** 2 + ddy = xp.sin(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) * (4 * xp.pi / Ly) ** 2 return ddx + ddy elif bc_type == "neumann": @@ -303,19 +303,19 @@ def rho2_xyz(x, y, z): # manufactured solution in 2D def sol2_xyz(x, y, z): - return np.cos(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) + return xp.cos(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) def rho2_xyz(x, y, z): - ddx = np.cos(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) * (np.pi / Lx) ** 2 - ddy = np.cos(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) * (4 * np.pi / Ly) ** 2 + ddx = xp.cos(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) * (xp.pi / Lx) ** 2 + ddy = xp.cos(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) * (4 * xp.pi / Ly) ** 2 return ddx + ddy # manufactured solution in 1D def sol1_xyz(x, y, z): - return np.cos(np.pi / Lx * x) + return xp.cos(xp.pi / Lx * x) def rho1_xyz(x, y, z): - return np.cos(np.pi / Lx * x) * (np.pi / Lx) ** 2 + return xp.cos(xp.pi / Lx * x) * (xp.pi / Lx) ** 2 # create derham object derham = Derham(Nel, p, spl_kind, dirichlet_bc=dirichlet_bc, comm=comm) @@ -328,9 +328,9 @@ def rho1_xyz(x, y, z): Propagator.mass_ops = mass_ops # evaluation grid - e1 = np.linspace(0.0, 1.0, 50) - e2 = np.linspace(0.0, 1.0, 50) - e3 = np.linspace(0.0, 1.0, 1) + e1 = xp.linspace(0.0, 1.0, 50) + e2 = xp.linspace(0.0, 1.0, 50) + e3 = xp.linspace(0.0, 1.0, 1) # pullbacks of right-hand side def rho1_pulled(e1, e2, e3): @@ -415,8 +415,8 @@ def rho2_pulled(e1, e2, e3): analytic_value2 = sol2_xyz(x, y, z) # compute error - error1 = np.max(np.abs(analytic_value1 - sol_val1)) - error2 = np.max(np.abs(analytic_value2 - sol_val2)) + error1 = xp.max(xp.abs(analytic_value1 - sol_val1)) + error2 = xp.max(xp.abs(analytic_value2 - sol_val2)) print(f"{p = }, {bc_type = }, {mapping = }") print(f"{error1 = }") @@ -480,14 +480,14 @@ def test_poisson_M1perp_3d_compare_2p5d(Nel, p, mapping, show_plot=False): dirichlet_bc = ((True, True), (False, False), (False, False)) # evaluation grid - e1 = np.linspace(0.0, 1.0, 50) - e2 = np.linspace(0.0, 1.0, 60) - e3 = np.linspace(0.0, 1.0, 30) + e1 = xp.linspace(0.0, 1.0, 50) + e2 = xp.linspace(0.0, 1.0, 60) + e3 = xp.linspace(0.0, 1.0, 30) # solution and right-hand side on unit cube def rho(e1, e2, e3): - dd1 = np.sin(np.pi * e1) * np.sin(4 * np.pi * e2) * np.cos(2 * np.pi * e3) * (np.pi) ** 2 - dd2 = np.sin(np.pi * e1) * np.sin(4 * np.pi * e2) * np.cos(2 * np.pi * e3) * (4 * np.pi) ** 2 + dd1 = xp.sin(xp.pi * e1) * xp.sin(4 * xp.pi * e2) * xp.cos(2 * xp.pi * e3) * (xp.pi) ** 2 + dd2 = xp.sin(xp.pi * e1) * xp.sin(4 * xp.pi * e2) * xp.cos(2 * xp.pi * e3) * (4 * xp.pi) ** 2 return dd1 + dd2 # create 3d derham object @@ -598,8 +598,8 @@ def rho(e1, e2, e3): sol_val_2p5d = domain.push(_phi_2p5d.spline, e1, e2, e3, kind="0") x, y, z = domain(e1, e2, e3) - print("max diff:", np.max(np.abs(sol_val - sol_val_2p5d))) - assert np.max(np.abs(sol_val - sol_val_2p5d)) < 0.026 + print("max diff:", xp.max(xp.abs(sol_val - sol_val_2p5d))) + assert xp.max(xp.abs(sol_val - sol_val_2p5d)) < 0.026 if show_plot and rank == 0: plt.figure("e1-e2 plane", figsize=(24, 16)) diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 0b1e14de6..659126d14 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -22,7 +22,7 @@ ) from struphy.propagators.base import Propagator from struphy.propagators.propagators_fields import ImplicitDiffusion, Poisson -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel comm = MPI.COMM_WORLD @@ -89,16 +89,16 @@ def test_poisson_1d( if direction == 0: Nel = [Neli, 1, 1] p = [pi, 1, 1] - e1 = np.linspace(0.0, 1.0, 50) + e1 = xp.linspace(0.0, 1.0, 50) if bc_type == "neumann": spl_kind = [False, True, True] def sol1_xyz(x, y, z): - return np.cos(np.pi / Lx * x) + return xp.cos(xp.pi / Lx * x) def rho1_xyz(x, y, z): - return np.cos(np.pi / Lx * x) * (np.pi / Lx) ** 2 + return xp.cos(xp.pi / Lx * x) * (xp.pi / Lx) ** 2 else: if bc_type == "dirichlet": spl_kind = [False, True, True] @@ -106,24 +106,24 @@ def rho1_xyz(x, y, z): dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): - return np.sin(2 * np.pi / Lx * x) + return xp.sin(2 * xp.pi / Lx * x) def rho1_xyz(x, y, z): - return np.sin(2 * np.pi / Lx * x) * (2 * np.pi / Lx) ** 2 + return xp.sin(2 * xp.pi / Lx * x) * (2 * xp.pi / Lx) ** 2 elif direction == 1: Nel = [1, Neli, 1] p = [1, pi, 1] - e2 = np.linspace(0.0, 1.0, 50) + e2 = xp.linspace(0.0, 1.0, 50) if bc_type == "neumann": spl_kind = [True, False, True] def sol1_xyz(x, y, z): - return np.cos(np.pi / Ly * y) + return xp.cos(xp.pi / Ly * y) def rho1_xyz(x, y, z): - return np.cos(np.pi / Ly * y) * (np.pi / Ly) ** 2 + return xp.cos(xp.pi / Ly * y) * (xp.pi / Ly) ** 2 else: if bc_type == "dirichlet": spl_kind = [True, False, True] @@ -131,24 +131,24 @@ def rho1_xyz(x, y, z): dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): - return np.sin(2 * np.pi / Ly * y) + return xp.sin(2 * xp.pi / Ly * y) def rho1_xyz(x, y, z): - return np.sin(2 * np.pi / Ly * y) * (2 * np.pi / Ly) ** 2 + return xp.sin(2 * xp.pi / Ly * y) * (2 * xp.pi / Ly) ** 2 elif direction == 2: Nel = [1, 1, Neli] p = [1, 1, pi] - e3 = np.linspace(0.0, 1.0, 50) + e3 = xp.linspace(0.0, 1.0, 50) if bc_type == "neumann": spl_kind = [True, True, False] def sol1_xyz(x, y, z): - return np.cos(np.pi / Lz * z) + return xp.cos(xp.pi / Lz * z) def rho1_xyz(x, y, z): - return np.cos(np.pi / Lz * z) * (np.pi / Lz) ** 2 + return xp.cos(xp.pi / Lz * z) * (xp.pi / Lz) ** 2 else: if bc_type == "dirichlet": spl_kind = [True, True, False] @@ -156,10 +156,10 @@ def rho1_xyz(x, y, z): dirichlet_bc = tuple(dirichlet_bc) def sol1_xyz(x, y, z): - return np.sin(2 * np.pi / Lz * z) + return xp.sin(2 * xp.pi / Lz * z) def rho1_xyz(x, y, z): - return np.sin(2 * np.pi / Lz * z) * (2 * np.pi / Lz) ** 2 + return xp.sin(2 * xp.pi / Lz * z) * (2 * xp.pi / Lz) ** 2 else: print("Direction should be either 0, 1 or 2") @@ -239,14 +239,14 @@ def rho_pulled(e1, e2, e3): plt.title(f"{Nel = }") plt.legend() - error = np.max(np.abs(analytic_value1 - sol_val1)) + error = xp.max(xp.abs(analytic_value1 - sol_val1)) print(f"{direction = }, {pi = }, {Neli = }, {error=}") errors.append(error) h = 1 / (Neli) h_vec.append(h) - m, _ = np.polyfit(np.log(Nels), np.log(errors), deg=1) + m, _ = xp.polyfit(xp.log(Nels), xp.log(errors), deg=1) print(f"For {pi = }, solution converges in {direction=} with rate {-m = } ") assert -m > (pi + 1 - 0.07) @@ -322,9 +322,9 @@ def test_poisson_accum_1d(mapping, do_plot=False): pert = perturbations.ModesCos(ls=(l,), amps=(amp,)) maxw = Maxwellian3D(n=(1.0, pert)) - pert_exact = lambda x, y, z: amp * np.cos(l * 2 * np.pi / Lx * x) - phi_exact = lambda x, y, z: amp / (l * 2 * np.pi / Lx) ** 2 * np.cos(l * 2 * np.pi / Lx * x) - e_exact = lambda x, y, z: amp / (l * 2 * np.pi / Lx) * np.sin(l * 2 * np.pi / Lx * x) + pert_exact = lambda x, y, z: amp * xp.cos(l * 2 * xp.pi / Lx * x) + phi_exact = lambda x, y, z: amp / (l * 2 * xp.pi / Lx) ** 2 * xp.cos(l * 2 * xp.pi / Lx * x) + e_exact = lambda x, y, z: amp / (l * 2 * xp.pi / Lx) * xp.sin(l * 2 * xp.pi / Lx * x) particles = Particles6D( comm_world=comm, @@ -380,7 +380,7 @@ def test_poisson_accum_1d(mapping, do_plot=False): poisson_solver(dt) # push numerical solution and compare - e1 = np.linspace(0.0, 1.0, 50) + e1 = xp.linspace(0.0, 1.0, 50) e2 = 0.0 e3 = 0.0 @@ -422,7 +422,7 @@ def test_poisson_accum_1d(mapping, do_plot=False): plt.show() - error = np.max(np.abs(num_values_e[0][:, 0, 0] - e_values[:, 0, 0])) / np.max(np.abs(e_values[:, 0, 0])) + error = xp.max(xp.abs(num_values_e[0][:, 0, 0] - e_values[:, 0, 0])) / xp.max(xp.abs(e_values[:, 0, 0])) print(f"{error=}") assert error < 0.0086 @@ -461,10 +461,10 @@ def test_poisson_2d(Nel, p, bc_type, mapping, projected_rhs, show_plot=False): # manufactured solution in 1D (overwritten for "neumann") def sol1_xyz(x, y, z): - return np.sin(2 * np.pi / Lx * x) + return xp.sin(2 * xp.pi / Lx * x) def rho1_xyz(x, y, z): - return np.sin(2 * np.pi / Lx * x) * (2 * np.pi / Lx) ** 2 + return xp.sin(2 * xp.pi / Lx * x) * (2 * xp.pi / Lx) ** 2 # boundary conditions dirichlet_bc = None @@ -474,11 +474,11 @@ def rho1_xyz(x, y, z): # manufactured solution in 2D def sol2_xyz(x, y, z): - return np.sin(2 * np.pi * x / Lx + 4 * np.pi / Ly * y) + return xp.sin(2 * xp.pi * x / Lx + 4 * xp.pi / Ly * y) def rho2_xyz(x, y, z): - ddx = np.sin(2 * np.pi / Lx * x + 4 * np.pi / Ly * y) * (2 * np.pi / Lx) ** 2 - ddy = np.sin(2 * np.pi / Lx * x + 4 * np.pi / Ly * y) * (4 * np.pi / Ly) ** 2 + ddx = xp.sin(2 * xp.pi / Lx * x + 4 * xp.pi / Ly * y) * (2 * xp.pi / Lx) ** 2 + ddy = xp.sin(2 * xp.pi / Lx * x + 4 * xp.pi / Ly * y) * (4 * xp.pi / Ly) ** 2 return ddx + ddy elif bc_type == "dirichlet": @@ -489,11 +489,11 @@ def rho2_xyz(x, y, z): # manufactured solution in 2D def sol2_xyz(x, y, z): - return np.sin(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) + return xp.sin(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) def rho2_xyz(x, y, z): - ddx = np.sin(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) * (np.pi / Lx) ** 2 - ddy = np.sin(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) * (4 * np.pi / Ly) ** 2 + ddx = xp.sin(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) * (xp.pi / Lx) ** 2 + ddy = xp.sin(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) * (4 * xp.pi / Ly) ** 2 return ddx + ddy elif bc_type == "neumann": @@ -501,19 +501,19 @@ def rho2_xyz(x, y, z): # manufactured solution in 2D def sol2_xyz(x, y, z): - return np.cos(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) + return xp.cos(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) def rho2_xyz(x, y, z): - ddx = np.cos(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) * (np.pi / Lx) ** 2 - ddy = np.cos(np.pi * x / Lx) * np.sin(4 * np.pi / Ly * y) * (4 * np.pi / Ly) ** 2 + ddx = xp.cos(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) * (xp.pi / Lx) ** 2 + ddy = xp.cos(xp.pi * x / Lx) * xp.sin(4 * xp.pi / Ly * y) * (4 * xp.pi / Ly) ** 2 return ddx + ddy # manufactured solution in 1D def sol1_xyz(x, y, z): - return np.cos(np.pi / Lx * x) + return xp.cos(xp.pi / Lx * x) def rho1_xyz(x, y, z): - return np.cos(np.pi / Lx * x) * (np.pi / Lx) ** 2 + return xp.cos(xp.pi / Lx * x) * (xp.pi / Lx) ** 2 # create derham object derham = Derham(Nel, p, spl_kind, dirichlet_bc=dirichlet_bc, comm=comm) @@ -526,9 +526,9 @@ def rho1_xyz(x, y, z): Propagator.mass_ops = mass_ops # evaluation grid - e1 = np.linspace(0.0, 1.0, 50) - e2 = np.linspace(0.0, 1.0, 50) - e3 = np.linspace(0.0, 1.0, 1) + e1 = xp.linspace(0.0, 1.0, 50) + e2 = xp.linspace(0.0, 1.0, 50) + e3 = xp.linspace(0.0, 1.0, 1) # pullbacks of right-hand side def rho1_pulled(e1, e2, e3): @@ -625,8 +625,8 @@ def rho2_pulled(e1, e2, e3): analytic_value2 = sol2_xyz(x, y, z) # compute error - error1 = np.max(np.abs(analytic_value1 - sol_val1)) - error2 = np.max(np.abs(analytic_value2 - sol_val2)) + error1 = xp.max(xp.abs(analytic_value1 - sol_val1)) + error2 = xp.max(xp.abs(analytic_value2 - sol_val2)) print(f"{p = }, {bc_type = }, {mapping = }") print(f"{error1 = }") diff --git a/src/struphy/utils/clone_config.py b/src/struphy/utils/clone_config.py index 15e449637..b292331b0 100644 --- a/src/struphy/utils/clone_config.py +++ b/src/struphy/utils/clone_config.py @@ -1,7 +1,7 @@ from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp as np +from struphy.utils.arrays import xp class CloneConfig: @@ -122,10 +122,10 @@ def get_Np_global(self, species_name): if "Np" in markers: return markers["Np"] elif "ppc" in markers: - n_cells = np.prod(self.params["grid"]["Nel"], dtype=int) + n_cells = xp.prod(self.params["grid"]["Nel"], dtype=int) return int(markers["ppc"] * n_cells) elif "ppb" in markers: - n_boxes = np.prod(species["boxes_per_dim"], dtype=int) * self.num_clones + n_boxes = xp.prod(species["boxes_per_dim"], dtype=int) * self.num_clones return int(markers["ppb"] * n_boxes) def print_clone_config(self): @@ -210,7 +210,7 @@ def print_particle_config(self): row = f"{i_clone:6} " # Np = self.params["kinetic"][species_name]["markers"]["Np"] Np = self.get_Np_global(species_name) - n_cells_clone = np.prod(self.params["grid"]["Nel"]) + n_cells_clone = xp.prod(self.params["grid"]["Nel"]) Np_clone = self.get_Np_clone(Np, clone_id=i_clone) ppc_clone = Np_clone / n_cells_clone diff --git a/src/struphy/utils/utils.py b/src/struphy/utils/utils.py index 7b63d14ec..f9898d36f 100644 --- a/src/struphy/utils/utils.py +++ b/src/struphy/utils/utils.py @@ -61,12 +61,12 @@ def save_state(state, libpath=STRUPHY_LIBPATH): def print_all_attr(obj): """Print all object's attributes that do not start with "_" to screen.""" - from struphy.utils.arrays import xp as np + from struphy.utils.arrays import xp for k in dir(obj): if k[0] != "_": v = getattr(obj, k) - if isinstance(v, np.ndarray): + if isinstance(v, xp.ndarray): v = f"{type(getattr(obj, k))} of shape {v.shape}" if "proj_" in k or "quad_grid_" in k: v = "(arrays not displayed)" From 5eac4992971ea90b54cff3edfc011ad6d1696746 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 22 Oct 2025 10:49:54 +0000 Subject: [PATCH 157/292] Resolve "Porting the rest of kinetic models" --- src/struphy/models/kinetic.py | 957 ++++++++---------- src/struphy/models/tests/test_models.py | 10 +- .../propagators/propagators_coupling.py | 147 ++- .../propagators/propagators_markers.py | 4 +- 4 files changed, 539 insertions(+), 579 deletions(-) diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 791477419..c8f6ae452 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -1,6 +1,8 @@ from psydac.ddm.mpi import mpi as MPI +from struphy.feec.projectors import L2Projector from struphy.kinetic_background.base import KineticBackground +from struphy.kinetic_background.maxwellians import Maxwellian3D from struphy.models.base import StruphyModel from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable @@ -321,135 +323,115 @@ class VlasovMaxwellOneSpecies(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species - dct["em_fields"]["e_field"] = "Hcurl" - dct["em_fields"]["b_field"] = "Hdiv" - dct["kinetic"]["species1"] = "Particles6D" - return dct + class EMFields(FieldSpecies): + def __init__(self): + self.e_field = FEECVariable(space="Hcurl") + self.b_field = FEECVariable(space="Hdiv") + self.phi = FEECVariable(space="H1") + self.init_variables() - @staticmethod - def bulk_species(): - return "species1" + class KineticIons(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="Particles6D") + self.init_variables() - @staticmethod - def velocity_scale(): - return "light" + ## propagators - @staticmethod - def propagators_dct(): - return { - propagators_fields.Maxwell: ["e_field", "b_field"], - propagators_markers.PushEta: ["species1"], - propagators_markers.PushVxB: ["species1"], - propagators_coupling.VlasovAmpere: ["e_field", "species1"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - # add special options - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["em_fields"], - option=propagators_fields.ImplicitDiffusion, - dct=dct, - ) - cls.add_option( - species=["kinetic", "species1"], - key="override_eq_params", - option=[False, {"alpha": 1.0, "epsilon": -1.0}], - dct=dct, - ) - return dct + class Propagators: + def __init__(self): + self.maxwell = propagators_fields.Maxwell() + self.push_eta = propagators_markers.PushEta() + self.push_vxb = propagators_markers.PushVxB() + self.coupling_va = propagators_coupling.VlasovAmpere() - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + ## abstract methods - # get species paramaters - species1_params = params["kinetic"]["species1"] + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # equation parameters - if species1_params["options"]["override_eq_params"]: - self._alpha = species1_params["options"]["override_eq_params"]["alpha"] - self._epsilon = species1_params["options"]["override_eq_params"]["epsilon"] - print( - f"\n!!! Override equation parameters: {self._alpha = } and {self._epsilon = }.", - ) - else: - self._alpha = self.equation_params["species1"]["alpha"] - self._epsilon = self.equation_params["species1"]["epsilon"] - - # set background density and mean velocity factors - self.pointer["species1"].f0.moment_factors["u"] = [ - self._epsilon / self._alpha**2, - ] * 3 - - # Initialize background magnetic field from MHD equilibrium - if self.projected_equil: - self._b_background = self.projected_equil.b2 - else: - self._b_background = None - - # propagator parameters - params_maxwell = params["em_fields"]["options"]["Maxwell"]["solver"] - algo_eta = params["kinetic"]["species1"]["options"]["PushEta"]["algo"] - algo_vxb = params["kinetic"]["species1"]["options"]["PushVxB"]["algo"] - params_coupling = params["em_fields"]["options"]["VlasovAmpere"]["solver"] - self._poisson_params = params["em_fields"]["options"]["ImplicitDiffusion"]["solver"] - - # set keyword arguments for propagators - self._kwargs[propagators_fields.Maxwell] = {"solver": params_maxwell} - - self._kwargs[propagators_markers.PushEta] = {"algo": algo_eta} - - self._kwargs[propagators_markers.PushVxB] = { - "algo": algo_vxb, - "kappa": 1.0 / self._epsilon, - "b2": self.pointer["b_field"], - "b2_add": self._b_background, - } - - self._kwargs[propagators_coupling.VlasovAmpere] = { - "c1": self._alpha**2 / self._epsilon, - "c2": 1.0 / self._epsilon, - "solver": params_coupling, - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during the simulation + # 1. instantiate all species + self.em_fields = self.EMFields() + self.kinetic_ions = self.KineticIons() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.maxwell.variables.e = self.em_fields.e_field + self.propagators.maxwell.variables.b = self.em_fields.b_field + self.propagators.push_eta.variables.var = self.kinetic_ions.var + self.propagators.push_vxb.variables.ions = self.kinetic_ions.var + self.propagators.coupling_va.variables.e = self.em_fields.e_field + self.propagators.coupling_va.variables.ions = self.kinetic_ions.var + + # define scalars for update_scalar_quantities self.add_scalar("en_E") self.add_scalar("en_B") - self.add_scalar("en_f", compute="from_particles", species="species1") + self.add_scalar("en_f", compute="from_particles", variable=self.kinetic_ions.var) self.add_scalar("en_tot") - # temporaries + # initial Poisson (not a propagator used in time stepping) + self.initial_poisson = propagators_fields.Poisson() + self.initial_poisson.variables.phi = self.em_fields.phi + + @property + def bulk_species(self): + return self.kinetic_ions + + @property + def velocity_scale(self): + return "light" + + def allocate_helpers(self): self._tmp = xp.empty(1, dtype=float) - def initialize_from_params(self): - """:meta private:""" + def update_scalar_quantities(self): + # e*M1*e/2 + e = self.em_fields.e_field.spline.vector + b = self.em_fields.b_field.spline.vector + + en_E = 0.5 * self.mass_ops.M1.dot_inner(e, e) + self.update_scalar("en_E", en_E) + + en_B = 0.5 * self.mass_ops.M2.dot_inner(b, b) + self.update_scalar("en_B", en_B) + + # alpha^2 / 2 / N * sum_p w_p v_p^2 + particles = self.kinetic_ions.var.particles + alpha = self.kinetic_ions.equation_params.alpha + self._tmp[0] = ( + alpha**2 + / (2 * particles.Np) + * xp.dot( + particles.markers_wo_holes[:, 3] ** 2 + + particles.markers_wo_holes[:, 4] ** 2 + + particles.markers_wo_holes[:, 5] ** 2, + particles.markers_wo_holes[:, 6], + ) + ) + self.update_scalar("en_f", self._tmp[0]) + + # en_tot = en_w + en_e + self.update_scalar("en_tot", en_E + self._tmp[0]) + + def allocate_propagators(self): + """Solve initial Poisson equation. - from struphy.pic.accumulation.particles_to_grid import AccumulatorVector + :meta private: + """ # initialize fields and particles - super().initialize_from_params() + super().allocate_propagators() - if self.rank_world == 0: + if MPI.COMM_WORLD.Get_rank() == 0: print("\nINITIAL POISSON SOLVE:") # use control variate method - self.pointer["species1"].update_weights() + particles = self.kinetic_ions.var.particles + particles.update_weights() # sanity check # self.pointer['species1'].show_distribution_function( @@ -457,61 +439,55 @@ def initialize_from_params(self): # accumulate charge density charge_accum = AccumulatorVector( - self.pointer["species1"], + particles, "H1", Pyccelkernel(accum_kernels.charge_density_0form), self.mass_ops, self.domain.args_domain, ) - charge_accum(self.pointer["species1"].vdim) - # another sanity check: compute FE coeffs of density # charge_accum.show_accumulated_spline_field(self.mass_ops) - # Instantiate Poisson solver - _phi = self.derham.Vh["0"].zeros() - poisson_solver = propagators_fields.ImplicitDiffusion( - _phi, - sigma_1=0.0, - sigma_2=0.0, - sigma_3=1.0, - rho=self._alpha**2 / self._epsilon * charge_accum.vectors[0], - solver=self._poisson_params, - ) + alpha = self.kinetic_ions.equation_params.alpha + epsilon = self.kinetic_ions.equation_params.epsilon + + self.initial_poisson.options.rho = charge_accum + self.initial_poisson.options.rho_coeffs = alpha**2 / epsilon + self.initial_poisson.allocate() # Solve with dt=1. and compute electric field - if self.rank_world == 0: + if MPI.COMM_WORLD.Get_rank() == 0: print("\nSolving initial Poisson problem...") - poisson_solver(1.0) + self.initial_poisson(1.0) - self.derham.grad.dot(-_phi, out=self.pointer["e_field"]) - if self.rank_world == 0: + phi = self.initial_poisson.variables.phi.spline.vector + self.derham.grad.dot(-phi, out=self.em_fields.e_field.spline.vector) + if MPI.COMM_WORLD.Get_rank() == 0: print("Done.") - def update_scalar_quantities(self): - # e*M1*e and b*M2*b - en_E = 0.5 * self.mass_ops.M1.dot_inner(self.pointer["e_field"], self.pointer["e_field"]) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) - self.update_scalar("en_E", en_E) - self.update_scalar("en_B", en_B) - - # alpha^2 / 2 / N * sum_p w_p v_p^2 - self._tmp[0] = ( - self._alpha**2 - / (2 * self.pointer["species1"].Np) - * xp.dot( - self.pointer["species1"].markers_wo_holes[:, 3] ** 2 - + self.pointer["species1"].markers_wo_holes[:, 4] ** 2 - + self.pointer["species1"].markers_wo_holes[:, 5] ** 2, - self.pointer["species1"].markers_wo_holes[:, 6], - ) - ) - - self.update_scalar("en_f", self._tmp[0]) + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "coupling_va.Options" in line: + new_file += [line] + new_file += ["model.initial_poisson.options = model.initial_poisson.Options()\n"] + elif "push_vxb.Options" in line: + new_file += [ + "model.propagators.push_vxb.options = model.propagators.push_vxb.Options(b2_var=model.em_fields.b_field)\n" + ] + elif "set_save_data" in line: + new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] + new_file += ["model.kinetic_ions.set_save_data(binning_plots=(binplot,))\n"] + else: + new_file += [line] - # en_tot = en_w + en_e + en_b - self.update_scalar("en_tot", en_E + en_B + self._tmp[0]) + with open(params_path, "w") as f: + for line in new_file: + f.write(line) class LinearVlasovAmpereOneSpecies(StruphyModel): @@ -581,245 +557,188 @@ class LinearVlasovAmpereOneSpecies(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["em_fields"]["e_field"] = "Hcurl" - dct["kinetic"]["species1"] = "DeltaFParticles6D" - return dct - - @staticmethod - def bulk_species(): - return "species1" + ## species - @staticmethod - def velocity_scale(): - return "light" + class EMFields(FieldSpecies): + def __init__(self): + self.e_field = FEECVariable(space="Hcurl") + self.phi = FEECVariable(space="H1") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_markers.PushEta: ["species1"], - propagators_markers.PushVinEfield: ["species1"], - propagators_coupling.EfieldWeights: ["e_field", "species1"], - propagators_markers.PushVxB: ["species1"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["em_fields"], - option=propagators_fields.ImplicitDiffusion, - dct=dct, - ) - cls.add_option( - species=["kinetic", "species1"], - key="override_eq_params", - option=[False, {"epsilon": -1.0, "alpha": 1.0}], - dct=dct, - ) - return dct + class KineticIons(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="DeltaFParticles6D") + self.init_variables() - def __init__(self, params, comm, clone_config=None, baseclass=False): - """Initializes the model either as the full model or as a baseclass to inherit from. - In case of being a baseclass, the propagators will not be initialized in the __init__ which allows other propagators to be added. + ## propagators - Parameters - ---------- - baseclass : Boolean [optional] - If this model should be used as a baseclass. Default value is False. - """ + class Propagators: + def __init__( + self, + with_B0: bool = True, + with_E0: bool = True, + ): + self.push_eta = propagators_markers.PushEta() + if with_E0: + self.push_vinE = propagators_markers.PushVinEfield() + self.coupling_Eweights = propagators_coupling.EfieldWeights() + if with_B0: + self.push_vxb = propagators_markers.PushVxB() - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + ## abstract methods - from struphy.kinetic_background import maxwellians + def __init__( + self, + with_B0: bool = True, + with_E0: bool = True, + ): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # if model is used as a baseclass - self._baseclass = baseclass + # 1. instantiate all species + self.em_fields = self.EMFields() + self.kinetic_ions = self.KineticIons() - # kinetic parameters - self._species_params = params["kinetic"]["species1"] + # 2. instantiate all propagators + self.propagators = self.Propagators(with_B0=with_B0, with_E0=with_E0) - # Assert Maxwellian background (if list, the first entry is taken) - bckgr_params = self._species_params["background"] - li_bp = list(bckgr_params) - assert li_bp[0] == "Maxwellian3D", "The background distribution function must be a uniform Maxwellian!" - if len(li_bp) > 1: - # overwrite f0 with single Maxwellian - self._f0 = getattr(maxwellians, li_bp[0][:-2])( - maxw_params=bckgr_params[li_bp[0]], - ) - else: - # keep allocated background - self._f0 = self.pointer["species1"].f0 - - # Assert uniformity of the Maxwellian background - assert self._f0.maxw_params["u1"] == 0.0, "The background Maxwellian cannot have shifts in velocity space!" - assert self._f0.maxw_params["u2"] == 0.0, "The background Maxwellian cannot have shifts in velocity space!" - assert self._f0.maxw_params["u3"] == 0.0, "The background Maxwellian cannot have shifts in velocity space!" - assert self._f0.maxw_params["vth1"] == self._f0.maxw_params["vth2"] == self._f0.maxw_params["vth3"], ( - "The background Maxwellian must be isotropic in velocity space!" - ) - self.vth = self._f0.maxw_params["vth1"] - - # Get coupling strength - if self._species_params["options"]["override_eq_params"]: - self.epsilon = self._species_params["options"]["override_eq_params"]["epsilon"] - self.alpha = self._species_params["options"]["override_eq_params"]["alpha"] - if self.rank_world == 0: - print( - f"\n!!! Override equation parameters: {self.epsilon = }, {self.alpha = }.\n", - ) - else: - self.epsilon = self.equation_params["species1"]["epsilon"] - self.alpha = self.equation_params["species1"]["alpha"] - - # allocate memory for evaluating f0 in energy computation - self._f0_values = xp.zeros( - self.pointer["species1"].markers.shape[0], - dtype=float, - ) + # 3. assign variables to propagators + self.propagators.push_eta.variables.var = self.kinetic_ions.var + if with_E0: + self.propagators.push_vinE.variables.var = self.kinetic_ions.var + self.propagators.coupling_Eweights.variables.e = self.em_fields.e_field + self.propagators.coupling_Eweights.variables.ions = self.kinetic_ions.var + if with_B0: + self.propagators.push_vxb.variables.ions = self.kinetic_ions.var - # ==================================================================================== - # Create pointers to background electric potential and field - self._has_background_e = False - if "external_E0" in self.params["em_fields"]["options"].keys(): - e0 = self.params["em_fields"]["options"]["external_E0"] - if e0 != 0.0: - self._has_background_e = True - self._e_background = self.derham.Vh["1"].zeros() - for block in self._e_background._blocks: - block._data[:, :, :] += e0 - - # Get parameters of the background magnetic field - if self.projected_equil: - self._b_background = self.projected_equil.b2 - else: - self._b_background = None - # ==================================================================================== - - # propagator parameters - self._poisson_params = params["em_fields"]["options"]["ImplicitDiffusion"]["solver"] - algo_eta = params["kinetic"]["species1"]["options"]["PushEta"]["algo"] - params_coupling = params["em_fields"]["options"]["EfieldWeights"]["solver"] - - # Initialize propagators/integrators used in splitting substeps - self._kwargs[propagators_markers.PushEta] = { - "algo": algo_eta, - } - - # Only add PushVinEfield if e-field is non-zero, otherwise it is more expensive - if self._has_background_e: - self._kwargs[propagators_markers.PushVinEfield] = { - "e_field": self._e_background, - "kappa": 1.0 / self.epsilon, - } - else: - self._kwargs[propagators_markers.PushVinEfield] = None - - self._kwargs[propagators_coupling.EfieldWeights] = { - "alpha": self.alpha, - "kappa": 1.0 / self.epsilon, - "f0": self._f0, - "solver": params_coupling, - } - - # Only add PushVxB if magnetic field is not zero - self._kwargs[propagators_markers.PushVxB] = None - if self._b_background: - self._kwargs[propagators_markers.PushVxB] = { - "kappa": 1.0 / self.epsilon, - "b2": self._b_background, - } - - # Initialize propagators used in splitting substeps - if not self._baseclass: - self.init_propagators() - - # Scalar variables to be saved during the simulation + # define scalars for update_scalar_quantities self.add_scalar("en_E") - self.add_scalar("en_w", compute="from_particles", species="species1") + self.add_scalar("en_w", compute="from_particles", variable=self.kinetic_ions.var) self.add_scalar("en_tot") - # temporaries + # initial Poisson (not a propagator used in time stepping) + self.initial_poisson = propagators_fields.Poisson() + self.initial_poisson.variables.phi = self.em_fields.phi + + @property + def bulk_species(self): + return self.kinetic_ions + + @property + def velocity_scale(self): + return "light" + + def allocate_helpers(self): self._tmp = xp.empty(1, dtype=float) - self.en_E = 0.0 - def initialize_from_params(self): + def update_scalar_quantities(self): + # e*M1*e/2 + e = self.em_fields.e_field.spline.vector + particles = self.kinetic_ions.var.particles + + en_E = 0.5 * self.mass_ops.M1.dot_inner(e, e) + self.update_scalar("en_E", en_E) + + # evaluate f0 + if not hasattr(self, "_f0"): + backgrounds = self.kinetic_ions.var.backgrounds + if isinstance(backgrounds, list): + self._f0 = backgrounds[0] + else: + self._f0 = backgrounds + self._f0_values = xp.zeros( + self.kinetic_ions.var.particles.markers.shape[0], + dtype=float, + ) + assert isinstance(self._f0, Maxwellian3D) + + self._f0_values[particles.valid_mks] = self._f0(*particles.phasespace_coords.T) + + # alpha^2 * v_th^2 / (2*N) * sum_p s_0 * w_p^2 / f_{0,p} + alpha = self.kinetic_ions.equation_params.alpha + vth = self._f0.maxw_params["vth1"][0] + + self._tmp[0] = ( + alpha**2 + * vth**2 + / (2 * particles.Np) + * xp.dot( + particles.weights**2, # w_p^2 + particles.sampling_density / self._f0_values[particles.valid_mks], # s_{0,p} / f_{0,p} + ) + ) + + self.update_scalar("en_w", self._tmp[0]) + self.update_scalar("en_tot", self._tmp[0] + en_E) + + def allocate_propagators(self): """Solve initial Poisson equation. :meta private: """ - from struphy.pic.accumulation.particles_to_grid import AccumulatorVector - # Initialize fields and particles - super().initialize_from_params() + # initialize fields and particles + super().allocate_propagators() + + if MPI.COMM_WORLD.Get_rank() == 0: + print("\nINITIAL POISSON SOLVE:") + + # use control variate method + particles = self.kinetic_ions.var.particles + particles.update_weights() + + # sanity check + # self.pointer['species1'].show_distribution_function( + # [True] + [False]*5, [xp.linspace(0, 1, 32)]) - # Accumulate charge density + # accumulate charge density charge_accum = AccumulatorVector( - self.pointer["species1"], + particles, "H1", Pyccelkernel(accum_kernels.charge_density_0form), self.mass_ops, self.domain.args_domain, ) - charge_accum(self.pointer["species1"].vdim) - - # Instantiate Poisson solver - _phi = self.derham.Vh["0"].zeros() - poisson_solver = propagators_fields.ImplicitDiffusion( - _phi, - sigma_1=0.0, - sigma_2=0.0, - sigma_3=1.0, - rho=self.alpha**2 / self.epsilon * charge_accum.vectors[0], - solver=self._poisson_params, - ) + # another sanity check: compute FE coeffs of density + # charge_accum.show_accumulated_spline_field(self.mass_ops) - # Solve with dt=1. and compute electric field - if self.rank_world == 0: - print("\nSolving initial Poisson problem...") - poisson_solver(1.0) - self.derham.grad.dot(-_phi, out=self.pointer["e_field"]) - if self.rank_world == 0: - print("Done.") + alpha = self.kinetic_ions.equation_params.alpha + epsilon = self.kinetic_ions.equation_params.epsilon - def update_scalar_quantities(self): - # 0.5 * e^T * M_1 * e - self.en_E = 0.5 * self.mass_ops.M1.dot_inner(self.pointer["e_field"], self.pointer["e_field"]) - self.update_scalar("en_E", self.en_E) + self.initial_poisson.options.rho = charge_accum + self.initial_poisson.options.rho_coeffs = alpha**2 / epsilon + self.initial_poisson.allocate() - # evaluate f0 - self._f0_values[self.pointer["species1"].valid_mks] = self._f0(*self.pointer["species1"].phasespace_coords.T) + # Solve with dt=1. and compute electric field + if MPI.COMM_WORLD.Get_rank() == 0: + print("\nSolving initial Poisson problem...") + self.initial_poisson(1.0) - # alpha^2 * v_th^2 / (2*N) * sum_p s_0 * w_p^2 / f_{0,p} - self._tmp[0] = ( - self.alpha**2 - * self.vth**2 - / (2 * self.pointer["species1"].Np) - * xp.dot( - self.pointer["species1"].weights ** 2, # w_p^2 - self.pointer["species1"].sampling_density - / self._f0_values[self.pointer["species1"].valid_mks], # s_{0,p} / f_{0,p} - ) - ) + phi = self.initial_poisson.variables.phi.spline.vector + self.derham.grad.dot(-phi, out=self.em_fields.e_field.spline.vector) + if MPI.COMM_WORLD.Get_rank() == 0: + print("Done.") - self.update_scalar("en_w", self._tmp[0]) + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "maxwellian_1 + maxwellian_2" in line: + new_file += ["background = maxwellian_1\n"] + elif "maxwellian_1pt =" in line: + new_file += ["maxwellian_1pt = maxwellians.Maxwellian3D(n=(0.0, perturbation))\n"] + elif "set_save_data" in line: + new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] + new_file += ["model.kinetic_ions.set_save_data(binning_plots=(binplot,))\n"] + else: + new_file += [line] - # en_tot = en_w + en_e - if not self._baseclass: - self.update_scalar("en_tot", self._tmp[0] + self.en_E) + with open(params_path, "w") as f: + for line in new_file: + f.write(line) class LinearVlasovMaxwellOneSpecies(LinearVlasovAmpereOneSpecies): @@ -892,80 +811,82 @@ class LinearVlasovMaxwellOneSpecies(LinearVlasovAmpereOneSpecies): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species - dct["em_fields"]["e_field"] = "Hcurl" - dct["em_fields"]["b_field"] = "Hdiv" - dct["kinetic"]["species1"] = "DeltaFParticles6D" - return dct + class EMFields(FieldSpecies): + def __init__(self): + self.e_field = FEECVariable(space="Hcurl") + self.b_field = FEECVariable(space="Hdiv") + self.phi = FEECVariable(space="H1") + self.init_variables() - @staticmethod - def bulk_species(): - return "species1" + class KineticIons(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="DeltaFParticles6D") + self.init_variables() - @staticmethod - def velocity_scale(): - return "light" + ## propagators - @staticmethod - def propagators_dct(): - return { - propagators_markers.PushEta: ["species1"], - propagators_markers.PushVinEfield: ["species1"], - propagators_coupling.EfieldWeights: ["e_field", "species1"], - propagators_markers.PushVxB: ["species1"], - propagators_fields.Maxwell: ["e_field", "b_field"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["em_fields"], - option=propagators_fields.ImplicitDiffusion, - dct=dct, - ) - cls.add_option( - species=["kinetic", "species1"], - key="override_eq_params", - option=[False, {"epsilon": -1.0, "alpha": 1.0}], - dct=dct, - ) - return dct + class Propagators: + def __init__( + self, + with_B0: bool = True, + with_E0: bool = True, + ): + self.push_eta = propagators_markers.PushEta() + if with_E0: + self.push_vinE = propagators_markers.PushVinEfield() + self.coupling_Eweights = propagators_coupling.EfieldWeights() + if with_B0: + self.push_vxb = propagators_markers.PushVxB() + self.maxwell = propagators_fields.Maxwell() - def __init__(self, params, comm, clone_config=None): - super().__init__(params=params, comm=comm, clone_config=clone_config, baseclass=True) + ## abstract methods + + def __init__( + self, + with_B0: bool = True, + with_E0: bool = True, + ): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # propagator parameters - params_maxwell = params["em_fields"]["options"]["Maxwell"]["solver"] + # 1. instantiate all species + self.em_fields = self.EMFields() + self.kinetic_ions = self.KineticIons() - # set keyword arguments for propagators - self._kwargs[propagators_fields.Maxwell] = {"solver": params_maxwell} + # 2. instantiate all propagators + self.propagators = self.Propagators(with_B0=with_B0, with_E0=with_E0) - # Initialize propagators used in splitting substeps - self.init_propagators() + # 3. assign variables to propagators + self.propagators.push_eta.variables.var = self.kinetic_ions.var + if with_E0: + self.propagators.push_vinE.variables.var = self.kinetic_ions.var + self.propagators.coupling_Eweights.variables.e = self.em_fields.e_field + self.propagators.coupling_Eweights.variables.ions = self.kinetic_ions.var + if with_B0: + self.propagators.push_vxb.variables.ions = self.kinetic_ions.var + self.propagators.maxwell.variables.e = self.em_fields.e_field + self.propagators.maxwell.variables.b = self.em_fields.b_field - # magnetic energy - self.add_scalar("en_b") + # define scalars for update_scalar_quantities + self.add_scalar("en_E") + self.add_scalar("en_B") + self.add_scalar("en_w", compute="from_particles", variable=self.kinetic_ions.var) + self.add_scalar("en_tot") - def initialize_from_params(self): - super().initialize_from_params() + # initial Poisson (not a propagator used in time stepping) + self.initial_poisson = propagators_fields.Poisson() + self.initial_poisson.variables.phi = self.em_fields.phi def update_scalar_quantities(self): super().update_scalar_quantities() # 0.5 * b^T * M_2 * b - en_B = 0.5 * self._mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) - self.update_scalar("en_tot", self._tmp[0] + self.en_E + en_B) + b = self.em_fields.b_field.spline.vector + + en_B = 0.5 * self._mass_ops.M2.dot_inner(b, b) + self.update_scalar("en_tot", self.scalar_quantities["en_tot"]["value"][0] + en_B) class DriftKineticElectrostaticAdiabatic(StruphyModel): @@ -1017,153 +938,159 @@ class DriftKineticElectrostaticAdiabatic(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species + + class EMFields(FieldSpecies): + def __init__(self): + self.phi = FEECVariable(space="H1") + self.init_variables() + + class KineticIons(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="Particles5D") + self.init_variables() + + ## propagators + + class Propagators: + def __init__(self): + self.gc_poisson = propagators_fields.ImplicitDiffusion() + self.push_gc_bxe = propagators_markers.PushGuidingCenterBxEstar() + self.push_gc_para = propagators_markers.PushGuidingCenterParallel() + + ## abstract methods - dct["em_fields"]["phi"] = "H1" - dct["kinetic"]["ions"] = "Particles5D" - return dct + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") + + # 1. instantiate all species + self.em_fields = self.EMFields() + self.kinetic_ions = self.KineticIons() + + # 2. instantiate all propagators + self.propagators = self.Propagators() + + # 3. assign variables to propagators + self.propagators.gc_poisson.variables.phi = self.em_fields.phi + self.propagators.push_gc_bxe.variables.ions = self.kinetic_ions.var + self.propagators.push_gc_para.variables.ions = self.kinetic_ions.var + + # define scalars for update_scalar_quantities + self.add_scalar("en_phi") + self.add_scalar("en_particles", compute="from_particles", variable=self.kinetic_ions.var) + self.add_scalar("en_tot") - @staticmethod - def bulk_species(): - return "ions" + @property + def bulk_species(self): + return self.kinetic_ions - @staticmethod - def velocity_scale(): + @property + def velocity_scale(self): return "thermal" - @staticmethod - def propagators_dct(): - return { - propagators_fields.ImplicitDiffusion: ["phi"], - propagators_markers.PushGuidingCenterBxEstar: ["ions"], - propagators_markers.PushGuidingCenterParallel: ["ions"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - # add special options - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["kinetic", "ions"], - key="override_eq_params", - option=[False, {"epsilon": 1.0}], - dct=dct, - ) - return dct + def allocate_helpers(self): + self._tmp3 = xp.empty(1, dtype=float) + self._e_field = self.derham.Vh["1"].zeros() - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + assert self.kinetic_ions.charge_number > 0, "Model written only for positive ions." - from struphy.feec.projectors import L2Projector - from struphy.pic.accumulation.particles_to_grid import AccumulatorVector + def allocate_propagators(self): + """Solve initial Poisson equation. - # prelim - solver_params = params["em_fields"]["options"]["ImplicitDiffusion"]["solver"] - ions_params = params["kinetic"]["ions"] + :meta private: + """ - Z = ions_params["phys_params"]["Z"] - assert Z > 0 # must be positive ions + # initialize fields and particles + super().allocate_propagators() # Poisson right-hand side + particles = self.kinetic_ions.var.particles + Z = self.kinetic_ions.charge_number + epsilon = self.kinetic_ions.equation_params.epsilon + charge_accum = AccumulatorVector( - self.pointer["ions"], + particles, "H1", Pyccelkernel(accum_kernels_gc.gc_density_0form), self.mass_ops, self.domain.args_domain, ) - rho = (charge_accum, self.pointer["ions"]) + rho = charge_accum # get neutralizing background density - if not self.pointer["ions"].control_variate: + if not particles.control_variate: l2_proj = L2Projector("H1", self.mass_ops) - f0e = Z * self.pointer["ions"].f0 + f0e = Z * particles.f0 assert isinstance(f0e, KineticBackground) - rho_eh = l2_proj.get_dofs(f0e.n) + rho_eh = FEECVariable(space="H1") + rho_eh.allocate(derham=self.derham, domain=self.domain) + rho_eh.spline.vector = l2_proj.get_dofs(f0e.n) rho = [rho] rho += [rho_eh] - # Get coupling strength - if ions_params["options"]["override_eq_params"]: - self.epsilon = ions_params["options"]["override_eq_params"]["epsilon"] - print( - f"\n!!! Override equation parameters: {self.epsilon = }.", - ) - else: - self.epsilon = self.equation_params["ions"]["epsilon"] - - # set keyword arguments for propagators - self._kwargs[propagators_fields.ImplicitDiffusion] = { - "sigma_1": 1.0 / self.epsilon**2 / Z, # set to zero for Landau damping test - "sigma_2": 0.0, - "sigma_3": 1.0 / self.epsilon, - "stab_mat": "M0ad", - "diffusion_mat": "M1gyro", - "rho": rho, - "solver": solver_params, - } - - self._kwargs[propagators_markers.PushGuidingCenterBxEstar] = { - "phi": self.pointer["phi"], - "evaluate_e_field": True, - "epsilon": self.epsilon / Z, - "algo": ions_params["options"]["PushGuidingCenterBxEstar"]["algo"], - } - - self._kwargs[propagators_markers.PushGuidingCenterParallel] = { - "phi": self.pointer["phi"], - "evaluate_e_field": True, - "epsilon": self.epsilon / Z, - "algo": ions_params["options"]["PushGuidingCenterParallel"]["algo"], - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # scalar quantities - self.add_scalar("en_phi") - self.add_scalar("en_particles", compute="from_particles", species="ions") - self.add_scalar("en_tot") - - # MPI operations needed for scalar variables - self._tmp3 = xp.empty(1, dtype=float) - self._e_field = self.derham.Vh["1"].zeros() + self.propagators.gc_poisson.options.sigma_1 = 1.0 / epsilon**2 / Z + self.propagators.gc_poisson.options.sigma_2 = 0.0 + self.propagators.gc_poisson.options.sigma_3 = 1.0 / epsilon + self.propagators.gc_poisson.options.stab_mat = "M0ad" + self.propagators.gc_poisson.options.diffusion_mat = "M1perp" + self.propagators.gc_poisson.options.rho = rho + self.propagators.gc_poisson.allocate() def update_scalar_quantities(self): + phi = self.em_fields.phi.spline.vector + particles = self.kinetic_ions.var.particles + epsilon = self.kinetic_ions.equation_params.epsilon + # energy from polarization - e1 = self.derham.grad.dot(-self.pointer["phi"], out=self._e_field) + e1 = self.derham.grad.dot(-phi, out=self._e_field) en_phi1 = 0.5 * self.mass_ops.M1gyro.dot_inner(e1, e1) # energy from adiabatic electrons - en_phi = 0.5 / self.epsilon**2 * self.mass_ops.M0ad.dot_inner(self.pointer["phi"], self.pointer["phi"]) + en_phi = 0.5 / epsilon**2 * self.mass_ops.M0ad.dot_inner(phi, phi) # for Landau damping test # en_phi = 0. # mu_p * |B0(eta_p)| - self.pointer["ions"].save_magnetic_background_energy() + particles.save_magnetic_background_energy() # 1/N sum_p (w_p v_p^2/2 + mu_p |B0|_p) self._tmp3[0] = ( 1 - / self.pointer["ions"].Np + / particles.Np * xp.sum( - self.pointer["ions"].weights * self.pointer["ions"].velocities[:, 0] ** 2 / 2.0 - + self.pointer["ions"].markers_wo_holes_and_ghost[:, 8], + particles.weights * particles.velocities[:, 0] ** 2 / 2.0 + particles.markers_wo_holes_and_ghost[:, 8], ) ) self.update_scalar("en_phi", en_phi + en_phi1) self.update_scalar("en_particles", self._tmp3[0]) self.update_scalar("en_tot", en_phi + en_phi1 + self._tmp3[0]) + + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "BaseUnits(" in line: + new_file += ["base_units = BaseUnits(kBT=1.0)\n"] + elif "push_gc_bxe.Options" in line: + new_file += [ + "model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options(phi=model.em_fields.phi)\n" + ] + elif "push_gc_para.Options" in line: + new_file += [ + "model.propagators.push_gc_para.options = model.propagators.push_gc_para.Options(phi=model.em_fields.phi)\n" + ] + elif "set_save_data" in line: + new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] + new_file += ["model.kinetic_ions.set_save_data(binning_plots=(binplot,))\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index dec9aa532..d684c2707 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -28,12 +28,10 @@ if rank == 0: print(f"\n{fluid_models = }") -kinetic_models = [ - "VlasovAmpereOneSpecies", -] -# for name, obj in inspect.getmembers(kinetic): -# if inspect.isclass(obj) and "models.kinetic" in obj.__module__: -# kinetic_models += [name] +kinetic_models = [] +for name, obj in inspect.getmembers(kinetic): + if inspect.isclass(obj) and "models.kinetic" in obj.__module__: + kinetic_models += [name] if rank == 0: print(f"\n{kinetic_models = }") diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index ba1e15b0b..8a82c084d 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -326,50 +326,85 @@ class EfieldWeights(Propagator): """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - if default: - dct = descend_options_dict(dct, []) + class Variables: + def __init__(self): + self._e: FEECVariable = None + self._ions: PICVariable = None - return dct + @property + def e(self) -> FEECVariable: + return self._e - def __init__( - self, - e: BlockVector, - particles: Particles6D, - *, - alpha: float = 1.0, - kappa: float = 1.0, - f0: Maxwellian = None, - solver=options(default=True)["solver"], - ): - super().__init__(e, particles) + @e.setter + def e(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hcurl" + self._e = new - if f0 is None: - f0 = Maxwellian3D() - assert isinstance(f0, Maxwellian3D) + @property + def ions(self) -> PICVariable: + return self._ions - self._alpha = alpha - self._kappa = kappa - self._f0 = f0 - assert self._f0.maxw_params["vth1"] == self._f0.maxw_params["vth2"] == self._f0.maxw_params["vth3"] - self._vth = self._f0.maxw_params["vth1"] + @ions.setter + def ions(self, new): + assert isinstance(new, PICVariable) + assert new.space in ("Particles6D", "DeltaFParticles6D") + self._ions = new - self._info = solver["info"] + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + alpha: float = 1.0 + kappa: float = 1.0 + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + + def __post_init__(self): + # checks + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._alpha = self.options.alpha + self._kappa = self.options.kappa + + backgrounds = self.variables.ions.backgrounds + # use single Maxwellian + if isinstance(backgrounds, list): + self._f0 = backgrounds[0] + else: + self._f0 = backgrounds + assert isinstance(self._f0, Maxwellian3D), "The background distribution function must be a uniform Maxwellian!" + self._vth = self._f0.maxw_params["vth1"][0] + + self._info = self.options.solver_params.info # Initialize Accumulator object + e = self.variables.e.spline.vector + particles = self.variables.ions.particles + self._accum = Accumulator( particles, "Hcurl", @@ -394,10 +429,10 @@ def __init__( # ================================ # Preconditioner - if solver["type"][1] == None: + if self.options.precond == None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(self.mass_ops.M1) # Define block matrix [[A B], [C I]] (without time step size dt in the diagonals) @@ -408,11 +443,9 @@ def __init__( self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, ) # Instantiate particle pusher @@ -435,14 +468,17 @@ def __init__( ) def __call__(self, dt): + en = self.variables.e.spline.vector + particles = self.variables.ions.particles + # evaluate f0 and accumulate self._f0_values[:] = self._f0( - self.particles[0].markers[:, 0], - self.particles[0].markers[:, 1], - self.particles[0].markers[:, 2], - self.particles[0].markers[:, 3], - self.particles[0].markers[:, 4], - self.particles[0].markers[:, 5], + particles.markers[:, 0], + particles.markers[:, 1], + particles.markers[:, 2], + particles.markers[:, 3], + particles.markers[:, 4], + particles.markers[:, 5], ) self._accum(self._f0_values) @@ -458,25 +494,25 @@ def __call__(self, dt): # new e-field (no tmps created here) self._e_tmp, info = self._schur_solver( - xn=self.feec_vars[0], + xn=en, Byn=self._e_scale, dt=dt, out=self._e_tmp, ) # Store old weights - self._old_weights[~self.particles[0].holes] = self.particles[0].markers_wo_holes[:, 6] + self._old_weights[~particles.holes] = particles.markers_wo_holes[:, 6] # Compute (e^{n+1} + e^n) (no tmps created here) self._e_sum *= 0.0 - self._e_sum += self.feec_vars[0] + self._e_sum += en self._e_sum += self._e_tmp # Update weights self._pusher(dt) # write new coeffs into self.variables - (max_de,) = self.feec_vars_update(self._e_tmp) + max_de = self.update_feec_variables(e=self._e_tmp) # Print out max differences for weights and e-field if self._info: @@ -485,8 +521,7 @@ def __call__(self, dt): print("Maxdiff e1 for StepEfieldWeights:", max_de) max_diff = xp.max( xp.abs( - self._old_weights[~self.particles[0].holes] - - self.particles[0].markers[~self.particles[0].holes, 6], + self._old_weights[~particles.holes] - particles.markers[~particles.holes, 6], ), ) print("Maxdiff weights for StepEfieldWeights:", max_diff) diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 1e9d2658a..fcd091c6d 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -158,7 +158,7 @@ def ions(self) -> PICVariable | SPHVariable: @ions.setter def ions(self, new): assert isinstance(new, PICVariable | SPHVariable) - assert new.space in ("Particles6D", "ParticlesSPH") + assert new.space in ("Particles6D", "DeltaFParticles6D", "ParticlesSPH") self._ions = new def __init__(self): @@ -289,7 +289,7 @@ def var(self) -> PICVariable | SPHVariable: @var.setter def var(self, new): assert isinstance(new, PICVariable | SPHVariable) - assert new.space in ("Particles6D", "ParticlesSPH") + assert new.space in ("Particles6D", "DeltaFParticles6D", "ParticlesSPH") self._var = new def __init__(self): From 6b6163cecb7b31f44662d97c44151659c350f682 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 22 Oct 2025 14:51:42 +0000 Subject: [PATCH 158/292] Use cunumpy --- .gitlab-ci.yml | 3 +- pyproject.toml | 1 + src/struphy/bsplines/bsplines.py | 2 +- .../bsplines/tests/test_bsplines_kernels.py | 3 +- .../bsplines/tests/test_eval_spline_mpi.py | 3 +- src/struphy/console/profile.py | 2 +- src/struphy/diagnostics/console_diagn.py | 2 +- src/struphy/diagnostics/continuous_spectra.py | 6 +- src/struphy/diagnostics/diagn_tools.py | 2 +- .../diagnostics/paraview/mesh_creator.py | 3 +- src/struphy/dispersion_relations/analytic.py | 2 +- src/struphy/dispersion_relations/base.py | 3 +- src/struphy/dispersion_relations/utilities.py | 3 +- src/struphy/eigenvalue_solvers/derivatives.py | 3 +- .../legacy/MHD_eigenvalues_cylinder_1D.py | 2 +- .../control_variates/control_variate.py | 2 +- .../fB_massless_control_variate.py | 2 +- .../fnB_massless_control_variate.py | 3 +- .../massless_control_variate.py | 3 +- .../legacy/emw_operators.py | 2 +- .../legacy/inner_products_1d.py | 3 +- .../legacy/inner_products_2d.py | 2 +- .../legacy/inner_products_3d.py | 2 +- .../eigenvalue_solvers/legacy/l2_error_1d.py | 3 +- .../eigenvalue_solvers/legacy/l2_error_2d.py | 2 +- .../eigenvalue_solvers/legacy/l2_error_3d.py | 2 +- .../legacy/mass_matrices_3d_pre.py | 2 +- .../legacy/massless_operators/fB_arrays.py | 2 +- .../fB_massless_linear_operators.py | 2 +- .../legacy/massless_operators/fB_vv_kernel.py | 2 +- .../legacy/mhd_operators_MF.py | 2 +- .../pro_local/mhd_operators_3d_local.py | 2 +- .../pro_local/projectors_local.py | 2 +- .../shape_function_projectors_L2.py | 2 +- .../shape_function_projectors_local.py | 2 +- .../eigenvalue_solvers/mass_matrices_1d.py | 2 +- .../eigenvalue_solvers/mass_matrices_2d.py | 2 +- .../eigenvalue_solvers/mass_matrices_3d.py | 2 +- .../mhd_axisymmetric_main.py | 2 +- .../mhd_axisymmetric_pproc.py | 3 +- .../eigenvalue_solvers/mhd_operators.py | 2 +- .../eigenvalue_solvers/mhd_operators_core.py | 2 +- .../eigenvalue_solvers/projectors_global.py | 2 +- .../eigenvalue_solvers/spline_space.py | 3 +- src/struphy/examples/_draw_parallel.py | 2 +- .../examples/restelli2018/callables.py | 2 +- src/struphy/feec/basis_projection_ops.py | 2 +- src/struphy/feec/linear_operators.py | 2 +- src/struphy/feec/mass.py | 2 +- src/struphy/feec/preconditioner.py | 2 +- src/struphy/feec/projectors.py | 2 +- src/struphy/feec/psydac_derham.py | 2 +- src/struphy/feec/tests/test_basis_ops.py | 6 +- src/struphy/feec/tests/test_derham.py | 2 +- src/struphy/feec/tests/test_eval_field.py | 3 +- src/struphy/feec/tests/test_field_init.py | 8 +-- src/struphy/feec/tests/test_l2_projectors.py | 2 +- .../feec/tests/test_local_projectors.py | 2 +- .../feec/tests/test_lowdim_nel_is_1.py | 2 +- src/struphy/feec/tests/test_mass_matrices.py | 8 +-- .../feec/tests/test_toarray_struphy.py | 2 +- .../feec/tests/test_tosparse_struphy.py | 2 +- src/struphy/feec/tests/xx_test_preconds.py | 2 +- src/struphy/feec/utilities.py | 2 +- .../feec/utilities_local_projectors.py | 3 +- src/struphy/feec/variational_utilities.py | 2 +- src/struphy/fields_background/base.py | 2 +- .../fields_background/coil_fields/base.py | 3 +- .../coil_fields/coil_fields.py | 3 +- src/struphy/fields_background/equils.py | 2 +- .../tests/test_desc_equil.py | 3 +- .../tests/test_generic_equils.py | 2 +- .../tests/test_mhd_equils.py | 2 +- .../tests/test_numerical_mhd_equil.py | 2 +- src/struphy/geometry/base.py | 2 +- src/struphy/geometry/domains.py | 3 +- src/struphy/geometry/tests/test_domain.py | 17 ++++-- src/struphy/geometry/utilities.py | 2 +- src/struphy/initial/eigenfunctions.py | 2 +- src/struphy/initial/perturbations.py | 2 +- .../initial/tests/test_init_perturbations.py | 2 +- src/struphy/initial/utilities.py | 2 +- src/struphy/io/options.py | 2 +- src/struphy/io/output_handling.py | 3 +- src/struphy/io/setup.py | 2 +- src/struphy/kinetic_background/base.py | 3 +- src/struphy/kinetic_background/maxwellians.py | 3 +- .../kinetic_background/moment_functions.py | 2 +- .../kinetic_background/tests/test_base.py | 2 +- .../tests/test_maxwellians.py | 14 ++--- src/struphy/linear_algebra/linalg_kron.py | 3 +- src/struphy/linear_algebra/saddle_point.py | 2 +- .../tests/test_saddlepoint_massmatrices.py | 5 +- .../tests/test_stencil_dot_kernels.py | 4 +- .../tests/test_stencil_transpose_kernels.py | 4 +- src/struphy/main.py | 2 +- src/struphy/models/base.py | 2 +- src/struphy/models/fluid.py | 2 +- src/struphy/models/hybrid.py | 2 +- src/struphy/models/kinetic.py | 2 +- src/struphy/models/species.py | 2 +- .../models/tests/test_verif_EulerSPH.py | 2 +- .../models/tests/test_verif_LinearMHD.py | 2 +- .../models/tests/test_verif_Maxwell.py | 2 +- .../models/tests/test_verif_Poisson.py | 2 +- .../test_verif_VlasovAmpereOneSpecies.py | 2 +- src/struphy/models/tests/verification.py | 2 +- src/struphy/models/toy.py | 2 +- src/struphy/models/variables.py | 2 +- src/struphy/ode/solvers.py | 2 +- src/struphy/ode/tests/test_ode_feec.py | 2 +- src/struphy/ode/utils.py | 2 +- .../pic/accumulation/particles_to_grid.py | 2 +- src/struphy/pic/base.py | 2 +- src/struphy/pic/particles.py | 3 +- src/struphy/pic/pushing/pusher.py | 2 +- src/struphy/pic/sobol_seq.py | 3 +- src/struphy/pic/tests/test_accum_vec_H1.py | 2 +- src/struphy/pic/tests/test_accumulation.py | 2 +- src/struphy/pic/tests/test_binning.py | 8 +-- src/struphy/pic/tests/test_draw_parallel.py | 2 +- src/struphy/pic/tests/test_mat_vec_filler.py | 3 +- .../test_pic_legacy_files/accumulation.py | 2 +- .../pic/tests/test_pic_legacy_files/pusher.py | 3 +- src/struphy/pic/tests/test_pushers.py | 12 ++-- src/struphy/pic/tests/test_sorting.py | 2 +- src/struphy/pic/tests/test_sph.py | 2 +- src/struphy/pic/tests/test_tesselation.py | 2 +- src/struphy/pic/utilities.py | 3 +- src/struphy/polar/basic.py | 3 +- src/struphy/polar/extraction_operators.py | 2 +- src/struphy/polar/linear_operators.py | 2 +- .../polar/tests/test_legacy_polar_splines.py | 2 +- src/struphy/polar/tests/test_polar.py | 4 +- .../likwid/plot_likwidproject.py | 2 +- .../likwid/plot_time_traces.py | 2 +- .../likwid/roofline_plotter.py | 3 +- .../post_processing/orbits/orbits_tools.py | 2 +- .../post_processing/post_processing_tools.py | 2 +- src/struphy/post_processing/pproc_struphy.py | 2 +- .../post_processing/profile_struphy.py | 2 +- src/struphy/profiling/profiling.py | 3 +- src/struphy/propagators/base.py | 2 +- .../propagators/propagators_coupling.py | 2 +- src/struphy/propagators/propagators_fields.py | 2 +- .../propagators/propagators_markers.py | 12 ++-- .../tests/test_gyrokinetic_poisson.py | 2 +- src/struphy/propagators/tests/test_poisson.py | 2 +- src/struphy/utils/arrays.py | 59 ------------------- src/struphy/utils/clone_config.py | 3 +- src/struphy/utils/cupy_vs_numpy.py | 2 +- src/struphy/utils/utils.py | 2 +- 152 files changed, 202 insertions(+), 272 deletions(-) delete mode 100644 src/struphy/utils/arrays.py diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 3385d72d3..9df6d696b 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -550,12 +550,11 @@ test_cupy: # Install cupy - python3 -m pip install --user cupy-cuda12x - python3 -c "import cupy as cp" + - python3 -m pip install cunumpy # Test numpy backend - - python3 src/struphy/utils/arrays.py - python3 src/struphy/utils/cupy_vs_numpy.py # Test cupy backend - export ARRAY_BACKEND=cupy - - python3 src/struphy/utils/arrays.py - python3 src/struphy/utils/cupy_vs_numpy.py install_tests: diff --git a/pyproject.toml b/pyproject.toml index a2a6e08a1..88e6ce6a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ ] dependencies = [ 'numpy', + 'cunumpy', 'pyccel>=2.0', 'psydac @ git+https://github.com/max-models/psydac-for-struphy.git@devel-tiny', 'scipy', diff --git a/src/struphy/bsplines/bsplines.py b/src/struphy/bsplines/bsplines.py index ab56afd78..a04ee4851 100644 --- a/src/struphy/bsplines/bsplines.py +++ b/src/struphy/bsplines/bsplines.py @@ -16,7 +16,7 @@ """ -from struphy.utils.arrays import xp +import cunumpy as xp __all__ = [ "find_span", diff --git a/src/struphy/bsplines/tests/test_bsplines_kernels.py b/src/struphy/bsplines/tests/test_bsplines_kernels.py index 3800c8653..c1010dd08 100644 --- a/src/struphy/bsplines/tests/test_bsplines_kernels.py +++ b/src/struphy/bsplines/tests/test_bsplines_kernels.py @@ -1,10 +1,9 @@ import time +import cunumpy as xp import pytest from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp - @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @pytest.mark.parametrize("p", [[1, 2, 1], [2, 1, 2], [3, 4, 3]]) diff --git a/src/struphy/bsplines/tests/test_eval_spline_mpi.py b/src/struphy/bsplines/tests/test_eval_spline_mpi.py index 0202963ce..0703fb418 100644 --- a/src/struphy/bsplines/tests/test_eval_spline_mpi.py +++ b/src/struphy/bsplines/tests/test_eval_spline_mpi.py @@ -1,11 +1,10 @@ from sys import int_info from time import sleep +import cunumpy as xp import pytest from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp - @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @pytest.mark.parametrize("p", [[1, 2, 3], [3, 1, 2]]) diff --git a/src/struphy/console/profile.py b/src/struphy/console/profile.py index 11c753a15..9577a00d9 100644 --- a/src/struphy/console/profile.py +++ b/src/struphy/console/profile.py @@ -6,12 +6,12 @@ def struphy_profile(dirs, replace, all, n_lines, print_callers, savefig): import os import pickle + import cunumpy as xp import yaml from matplotlib import pyplot as plt import struphy.utils.utils as utils from struphy.post_processing.cprofile_analyser import get_cprofile_data, replace_keys - from struphy.utils.arrays import xp # Read struphy state file state = utils.read_state() diff --git a/src/struphy/diagnostics/console_diagn.py b/src/struphy/diagnostics/console_diagn.py index 3fa1f11ca..e110d1497 100644 --- a/src/struphy/diagnostics/console_diagn.py +++ b/src/struphy/diagnostics/console_diagn.py @@ -5,13 +5,13 @@ import os import subprocess +import cunumpy as xp import h5py import yaml import struphy import struphy.utils.utils as utils from struphy.diagnostics.diagn_tools import plot_distr_fun, plot_scalars, plots_videos_2d -from struphy.utils.arrays import xp def main(): diff --git a/src/struphy/diagnostics/continuous_spectra.py b/src/struphy/diagnostics/continuous_spectra.py index ce805a45d..5ed069179 100644 --- a/src/struphy/diagnostics/continuous_spectra.py +++ b/src/struphy/diagnostics/continuous_spectra.py @@ -37,8 +37,9 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, the radial location s_spec[m][0], squared eigenfrequencis s_spec[m][1] and global mode index s_spec[m][2] corresponding to slow sound modes for each poloidal mode number m in m_range. """ + import cunumpy as xp + import struphy.bsplines.bsplines as bsp - from struphy.utils.arrays import xp # greville points in radial direction (s) gN_1 = bsp.greville(space.T[0], space.p[0], space.spl_kind[0]) @@ -151,10 +152,9 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, import os import shutil + import cunumpy as xp import yaml - from struphy.utils.arrays import xp - # parse arguments parser = argparse.ArgumentParser( description="Looks for eigenmodes in a given MHD eigenspectrum in a certain poloidal mode number range and plots the continuous shear Alfvén and slow sound spectra (frequency versus radial-like coordinate)." diff --git a/src/struphy/diagnostics/diagn_tools.py b/src/struphy/diagnostics/diagn_tools.py index 88c81409c..d714306f4 100644 --- a/src/struphy/diagnostics/diagn_tools.py +++ b/src/struphy/diagnostics/diagn_tools.py @@ -3,6 +3,7 @@ import shutil import subprocess +import cunumpy as xp import matplotlib.colors as colors import matplotlib.pyplot as plt from scipy.fft import fftfreq, fftn @@ -10,7 +11,6 @@ from tqdm import tqdm from struphy.dispersion_relations import analytic -from struphy.utils.arrays import xp def power_spectrum_2d( diff --git a/src/struphy/diagnostics/paraview/mesh_creator.py b/src/struphy/diagnostics/paraview/mesh_creator.py index 2446d0391..c6c8d89a0 100644 --- a/src/struphy/diagnostics/paraview/mesh_creator.py +++ b/src/struphy/diagnostics/paraview/mesh_creator.py @@ -1,11 +1,10 @@ # from tqdm import tqdm +import cunumpy as xp import vtkmodules.all as vtk from vtkmodules.util.numpy_support import numpy_to_vtk as np2vtk from vtkmodules.util.numpy_support import vtk_to_numpy as vtk2np from vtkmodules.vtkCommonDataModel import vtkUnstructuredGrid -from struphy.utils.arrays import xp - def make_ugrid_and_write_vtu(filename: str, writer, vtk_dir, gvec, s_range, u_range, v_range, periodic): """A helper function to orchestrate operations to run many test cases. diff --git a/src/struphy/dispersion_relations/analytic.py b/src/struphy/dispersion_relations/analytic.py index 756840f9a..fb7a40ccd 100644 --- a/src/struphy/dispersion_relations/analytic.py +++ b/src/struphy/dispersion_relations/analytic.py @@ -1,12 +1,12 @@ "Analytic dispersion relations." +import cunumpy as xp from numpy.polynomial import Polynomial from scipy.optimize import fsolve from struphy.dispersion_relations.base import ContinuousSpectra1D, DispersionRelations1D from struphy.dispersion_relations.utilities import Zplasma from struphy.fields_background.equils import set_defaults -from struphy.utils.arrays import xp class Maxwell1D(DispersionRelations1D): diff --git a/src/struphy/dispersion_relations/base.py b/src/struphy/dispersion_relations/base.py index dd7293867..6994ae2fb 100644 --- a/src/struphy/dispersion_relations/base.py +++ b/src/struphy/dispersion_relations/base.py @@ -2,10 +2,9 @@ from abc import ABCMeta, abstractmethod +import cunumpy as xp from matplotlib import pyplot as plt -from struphy.utils.arrays import xp - class DispersionRelations1D(metaclass=ABCMeta): r""" diff --git a/src/struphy/dispersion_relations/utilities.py b/src/struphy/dispersion_relations/utilities.py index f84aa1acf..b796eb321 100644 --- a/src/struphy/dispersion_relations/utilities.py +++ b/src/struphy/dispersion_relations/utilities.py @@ -1,7 +1,6 @@ +import cunumpy as xp from scipy.special import erfi -from struphy.utils.arrays import xp - def Zplasma(xi, der=0): """ diff --git a/src/struphy/eigenvalue_solvers/derivatives.py b/src/struphy/eigenvalue_solvers/derivatives.py index 738b1537a..0e34cceb1 100644 --- a/src/struphy/eigenvalue_solvers/derivatives.py +++ b/src/struphy/eigenvalue_solvers/derivatives.py @@ -6,10 +6,9 @@ Modules to assemble discrete derivatives. """ +import cunumpy as xp import scipy.sparse as spa -from struphy.utils.arrays import xp - # ================== 1d incident matrix ======================= def grad_1d_matrix(spl_kind, NbaseN): diff --git a/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py b/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py index b7d0d30d7..9fbc38ce8 100644 --- a/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py +++ b/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py @@ -1,3 +1,4 @@ +import cunumpy as xp import scipy as sc import scipy.sparse as spa import scipy.special as sp @@ -9,7 +10,6 @@ import struphy.eigenvalue_solvers.mass_matrices_1d as mass import struphy.eigenvalue_solvers.projectors_global as pro import struphy.eigenvalue_solvers.spline_space as spl -from struphy.utils.arrays import xp # numerical solution of the general ideal MHD eigenvalue problem in a cylinder using 1d B-splines in radial direction diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py index 448298451..a4b95eb06 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/control_variate.py @@ -6,11 +6,11 @@ Class for control variates in delta-f method for current coupling scheme. """ +import cunumpy as xp import scipy.sparse as spa import struphy.feec.basics.kernels_3d as ker import struphy.feec.control_variates.kernels_control_variate as ker_cv -from struphy.utils.arrays import xp class terms_control_variate: diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py index d4713c12a..49be8c3ef 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py @@ -1,9 +1,9 @@ +import cunumpy as xp import scipy.sparse as spa import struphy.feec.basics.kernels_3d as ker import struphy.feec.control_variates.kinetic_extended.fB_massless_kernels_control_variate as ker_cv import struphy.feec.control_variates.kinetic_extended.fnB_massless_cv_kernel_2 as ker_cv2 -from struphy.utils.arrays import xp def bv_right( diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py index b8673c837..d52a3e90e 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py @@ -1,9 +1,8 @@ +import cunumpy as xp import hylife.utilitis_FEEC.basics.kernels_3d as ker import hylife.utilitis_FEEC.control_variates.fnB_massless_kernels_control_variate as ker_cv import scipy.sparse as spa -from struphy.utils.arrays import xp - def bv_pre(tol, n, LO_inv, tensor_space_FEM, p, Nel, idnx, idny, idnz): r""" diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py index 3d09c5bce..4e97422f0 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py @@ -1,9 +1,8 @@ +import cunumpy as xp import hylife.utilitis_FEEC.basics.kernels_3d as ker import hylife.utilitis_FEEC.control_variates.massless_kernels_control_variate as ker_cv import scipy.sparse as spa -from struphy.utils.arrays import xp - def bv_pre(u, uvalue, tensor_space_FEM, p, Nel, idnx, idny, idnz): r""" diff --git a/src/struphy/eigenvalue_solvers/legacy/emw_operators.py b/src/struphy/eigenvalue_solvers/legacy/emw_operators.py index 0a0f16106..e2fd4ed22 100755 --- a/src/struphy/eigenvalue_solvers/legacy/emw_operators.py +++ b/src/struphy/eigenvalue_solvers/legacy/emw_operators.py @@ -6,11 +6,11 @@ Class for 2D/3D linear MHD projection operators. """ +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_3d as ker import struphy.eigenvalue_solvers.legacy.mass_matrices_3d_pre as mass_3d_pre -from struphy.utils.arrays import xp class EMW_operators: diff --git a/src/struphy/eigenvalue_solvers/legacy/inner_products_1d.py b/src/struphy/eigenvalue_solvers/legacy/inner_products_1d.py index 469798bb5..b4f019995 100644 --- a/src/struphy/eigenvalue_solvers/legacy/inner_products_1d.py +++ b/src/struphy/eigenvalue_solvers/legacy/inner_products_1d.py @@ -6,10 +6,9 @@ Modules to compute inner products in 1d. """ +import cunumpy as xp import scipy.sparse as spa -from struphy.utils.arrays import xp - # ======= inner product in V0 ==================== def inner_prod_V0(spline_space, fun, mapping=None): diff --git a/src/struphy/eigenvalue_solvers/legacy/inner_products_2d.py b/src/struphy/eigenvalue_solvers/legacy/inner_products_2d.py index 58ea01f2f..05df4725f 100644 --- a/src/struphy/eigenvalue_solvers/legacy/inner_products_2d.py +++ b/src/struphy/eigenvalue_solvers/legacy/inner_products_2d.py @@ -6,10 +6,10 @@ Modules to compute inner products with given functions in 2D. """ +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_2d as ker -from struphy.utils.arrays import xp # ================ inner product in V0 =========================== diff --git a/src/struphy/eigenvalue_solvers/legacy/inner_products_3d.py b/src/struphy/eigenvalue_solvers/legacy/inner_products_3d.py index 2616857b6..20d95c05c 100644 --- a/src/struphy/eigenvalue_solvers/legacy/inner_products_3d.py +++ b/src/struphy/eigenvalue_solvers/legacy/inner_products_3d.py @@ -6,10 +6,10 @@ Modules to compute inner products with given functions in 3D. """ +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_3d as ker -from struphy.utils.arrays import xp # ================ inner product in V0 =========================== diff --git a/src/struphy/eigenvalue_solvers/legacy/l2_error_1d.py b/src/struphy/eigenvalue_solvers/legacy/l2_error_1d.py index 708d0352e..d568d3207 100644 --- a/src/struphy/eigenvalue_solvers/legacy/l2_error_1d.py +++ b/src/struphy/eigenvalue_solvers/legacy/l2_error_1d.py @@ -6,10 +6,9 @@ Modules to compute L2-errors in 1d. """ +import cunumpy as xp import scipy.sparse as spa -from struphy.utils.arrays import xp - # ======= error in V0 ==================== def l2_error_V0(spline_space, mapping, coeff, fun): diff --git a/src/struphy/eigenvalue_solvers/legacy/l2_error_2d.py b/src/struphy/eigenvalue_solvers/legacy/l2_error_2d.py index ce2142ced..452dd570b 100644 --- a/src/struphy/eigenvalue_solvers/legacy/l2_error_2d.py +++ b/src/struphy/eigenvalue_solvers/legacy/l2_error_2d.py @@ -6,10 +6,10 @@ Modules to compute L2-errors of discrete p-forms with analytical forms in 2D. """ +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_2d as ker -from struphy.utils.arrays import xp # ======= error in V0 ==================== diff --git a/src/struphy/eigenvalue_solvers/legacy/l2_error_3d.py b/src/struphy/eigenvalue_solvers/legacy/l2_error_3d.py index 93dcc4053..7553e3a83 100644 --- a/src/struphy/eigenvalue_solvers/legacy/l2_error_3d.py +++ b/src/struphy/eigenvalue_solvers/legacy/l2_error_3d.py @@ -6,10 +6,10 @@ Modules to compute L2-errors of discrete p-forms with analytical forms in 3D. """ +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_3d as ker -from struphy.utils.arrays import xp # ======= error in V0 ==================== diff --git a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py index fedff3c90..eecb68293 100644 --- a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py +++ b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py @@ -6,11 +6,11 @@ Modules to obtain preconditioners for mass matrices in 3D. """ +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.spline_space as spl import struphy.linear_algebra.linalg_kron as linkron -from struphy.utils.arrays import xp # ================ inverse mass matrix in V0 =========================== diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py index 193d51a11..e74302878 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py @@ -1,13 +1,13 @@ import time import timeit +import cunumpy as xp import scipy.sparse as spa from psydac.ddm.mpi import mpi as MPI import struphy.geometry.mappings_3d as mapping3d import struphy.geometry.mappings_3d_fast as mapping_fast import struphy.linear_algebra.linalg_kernels as linalg -from struphy.utils.arrays import xp class Temp_arrays: diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py index c8b04f877..2ddddc64a 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py @@ -1,11 +1,11 @@ import time +import cunumpy as xp import scipy.sparse as spa import struphy.feec.massless_operators.fB_bb_kernel as bb_kernel import struphy.feec.massless_operators.fB_bv_kernel as bv_kernel import struphy.feec.massless_operators.fB_vv_kernel as vv_kernel -from struphy.utils.arrays import xp class Massless_linear_operators: diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py index 449e31e50..4a5a2dbe0 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py @@ -1,9 +1,9 @@ +import cunumpy as xp from numpy import empty, exp, floor, zeros import struphy.bsplines.bsplines_kernels as bsp import struphy.geometry.mappings_kernels as mapping_fast import struphy.linear_algebra.linalg_kernels as linalg -from struphy.utils.arrays import xp # ========================================================================================== diff --git a/src/struphy/eigenvalue_solvers/legacy/mhd_operators_MF.py b/src/struphy/eigenvalue_solvers/legacy/mhd_operators_MF.py index 7f3775389..560314458 100644 --- a/src/struphy/eigenvalue_solvers/legacy/mhd_operators_MF.py +++ b/src/struphy/eigenvalue_solvers/legacy/mhd_operators_MF.py @@ -1,9 +1,9 @@ +import cunumpy as xp import scipy.sparse as spa from struphy.eigenvalue_solvers.projectors_global import Projectors_tensor_3d from struphy.eigenvalue_solvers.spline_space import Tensor_spline_space from struphy.linear_algebra.linalg_kron import kron_matvec_3d, kron_solve_3d -from struphy.utils.arrays import xp # ================================================================================================= diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py index 5d6c133e2..0c96616de 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py @@ -8,13 +8,13 @@ import sys +import cunumpy as xp import scipy.sparse as spa import source_run.kernels_projectors_evaluation as ker_eva import struphy.feec.basics.kernels_3d as ker_loc_3d import struphy.feec.bsplines as bsp import struphy.feec.projectors.pro_local.kernels_projectors_local_mhd as ker_loc -from struphy.utils.arrays import xp class projectors_local_mhd: diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py index f7d366c99..dacc4f243 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py @@ -6,11 +6,11 @@ Classes for local projectors in 1D and 3D based on quasi-spline interpolation and histopolation. """ +import cunumpy as xp import scipy.sparse as spa import struphy.feec.bsplines as bsp import struphy.feec.projectors.pro_local.kernels_projectors_local as ker_loc -from struphy.utils.arrays import xp # ======================= 1d ==================================== diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py index 04705a6a5..20814c8ac 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py @@ -5,12 +5,12 @@ Classes for local projectors in 1D and 3D based on quasi-spline interpolation and histopolation. """ +import cunumpy as xp import scipy.sparse as spa from psydac.ddm.mpi import mpi as MPI import struphy.feec.bsplines as bsp import struphy.feec.projectors.shape_pro_local.shape_L2_projector_kernel as ker_loc -from struphy.utils.arrays import xp # ======================= 3d ==================================== diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py index dd5a313f2..7c9425e47 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py @@ -5,12 +5,12 @@ Classes for local projectors in 1D and 3D based on quasi-spline interpolation and histopolation. """ +import cunumpy as xp import scipy.sparse as spa from psydac.ddm.mpi import mpi as MPI import struphy.feec.bsplines as bsp import struphy.feec.projectors.shape_pro_local.shape_local_projector_kernel as ker_loc -from struphy.utils.arrays import xp # ======================= 3d ==================================== diff --git a/src/struphy/eigenvalue_solvers/mass_matrices_1d.py b/src/struphy/eigenvalue_solvers/mass_matrices_1d.py index 9c41ca85d..b5013c088 100644 --- a/src/struphy/eigenvalue_solvers/mass_matrices_1d.py +++ b/src/struphy/eigenvalue_solvers/mass_matrices_1d.py @@ -2,10 +2,10 @@ # # Copyright 2020 Florian Holderied +import cunumpy as xp import scipy.sparse as spa import struphy.bsplines.bsplines as bsp -from struphy.utils.arrays import xp # ======= mass matrices in 1D ==================== diff --git a/src/struphy/eigenvalue_solvers/mass_matrices_2d.py b/src/struphy/eigenvalue_solvers/mass_matrices_2d.py index e15d5e4c7..c2c6c3ae9 100644 --- a/src/struphy/eigenvalue_solvers/mass_matrices_2d.py +++ b/src/struphy/eigenvalue_solvers/mass_matrices_2d.py @@ -2,10 +2,10 @@ # # Copyright 2020 Florian Holderied +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_2d as ker -from struphy.utils.arrays import xp # ================ mass matrix in V0 =========================== diff --git a/src/struphy/eigenvalue_solvers/mass_matrices_3d.py b/src/struphy/eigenvalue_solvers/mass_matrices_3d.py index f7846f0b2..05f019e13 100644 --- a/src/struphy/eigenvalue_solvers/mass_matrices_3d.py +++ b/src/struphy/eigenvalue_solvers/mass_matrices_3d.py @@ -2,10 +2,10 @@ # # Copyright 2020 Florian Holderied (florian.holderied@ipp.mpg.de) +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_3d as ker -from struphy.utils.arrays import xp # ================ mass matrix in V0 =========================== diff --git a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py index e0ee63cc5..c44c97335 100644 --- a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py +++ b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py @@ -32,11 +32,11 @@ def solve_mhd_ev_problem_2d(num_params, eq_mhd, n_tor, basis_tor="i", path_out=N import os import time + import cunumpy as xp import scipy.sparse as spa from struphy.eigenvalue_solvers.mhd_operators import MHDOperators from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space - from struphy.utils.arrays import xp print("\nStart of eigenspectrum calculation for toroidal mode number", n_tor) print("") diff --git a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py index 779b9463e..10e9e274e 100644 --- a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py +++ b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py @@ -3,10 +3,9 @@ def main(): import argparse import os + import cunumpy as xp import yaml - from struphy.utils.arrays import xp - # parse arguments parser = argparse.ArgumentParser(description="Restrict a full .npy eigenspectrum to a range of eigenfrequencies.") diff --git a/src/struphy/eigenvalue_solvers/mhd_operators.py b/src/struphy/eigenvalue_solvers/mhd_operators.py index c0ffd0658..5f1462c24 100644 --- a/src/struphy/eigenvalue_solvers/mhd_operators.py +++ b/src/struphy/eigenvalue_solvers/mhd_operators.py @@ -3,11 +3,11 @@ # Copyright 2021 Florian Holderied (florian.holderied@ipp.mpg.de) +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.legacy.mass_matrices_3d_pre as mass_3d_pre from struphy.eigenvalue_solvers.mhd_operators_core import MHDOperatorsCore -from struphy.utils.arrays import xp class MHDOperators: diff --git a/src/struphy/eigenvalue_solvers/mhd_operators_core.py b/src/struphy/eigenvalue_solvers/mhd_operators_core.py index 45712afae..76752ccc3 100644 --- a/src/struphy/eigenvalue_solvers/mhd_operators_core.py +++ b/src/struphy/eigenvalue_solvers/mhd_operators_core.py @@ -3,12 +3,12 @@ # Copyright 2021 Florian Holderied (florian.holderied@ipp.mpg.de) +import cunumpy as xp import scipy.sparse as spa import struphy.eigenvalue_solvers.kernels_projectors_global_mhd as ker import struphy.eigenvalue_solvers.mass_matrices_2d as mass_2d import struphy.eigenvalue_solvers.mass_matrices_3d as mass_3d -from struphy.utils.arrays import xp class MHDOperatorsCore: diff --git a/src/struphy/eigenvalue_solvers/projectors_global.py b/src/struphy/eigenvalue_solvers/projectors_global.py index 7c66c3397..fa7b69958 100644 --- a/src/struphy/eigenvalue_solvers/projectors_global.py +++ b/src/struphy/eigenvalue_solvers/projectors_global.py @@ -6,11 +6,11 @@ Classes for commuting projectors in 1D, 2D and 3D based on global spline interpolation and histopolation. """ +import cunumpy as xp import scipy.sparse as spa import struphy.bsplines.bsplines as bsp from struphy.linear_algebra.linalg_kron import kron_lusolve_2d, kron_lusolve_3d, kron_matvec_2d, kron_matvec_3d -from struphy.utils.arrays import xp # ======================= 1d ==================================== diff --git a/src/struphy/eigenvalue_solvers/spline_space.py b/src/struphy/eigenvalue_solvers/spline_space.py index b71bc867b..b36ad73db 100644 --- a/src/struphy/eigenvalue_solvers/spline_space.py +++ b/src/struphy/eigenvalue_solvers/spline_space.py @@ -6,11 +6,10 @@ Basic modules to create tensor-product finite element spaces of univariate B-splines. """ +import cunumpy as xp import matplotlib import scipy.sparse as spa -from struphy.utils.arrays import xp - matplotlib.rcParams.update({"font.size": 16}) import matplotlib.pyplot as plt diff --git a/src/struphy/examples/_draw_parallel.py b/src/struphy/examples/_draw_parallel.py index 62fae8cd2..e95598b3c 100644 --- a/src/struphy/examples/_draw_parallel.py +++ b/src/struphy/examples/_draw_parallel.py @@ -1,9 +1,9 @@ +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.feec.psydac_derham import Derham from struphy.geometry import domains from struphy.pic.particles import Particles6D -from struphy.utils.arrays import xp def main(): diff --git a/src/struphy/examples/restelli2018/callables.py b/src/struphy/examples/restelli2018/callables.py index 54404ab27..505a60f90 100644 --- a/src/struphy/examples/restelli2018/callables.py +++ b/src/struphy/examples/restelli2018/callables.py @@ -1,6 +1,6 @@ "Analytical callables needed for the simulation of the Two-Fluid Quasi-Neutral Model by Restelli." -from struphy.utils.arrays import xp +import cunumpy as xp class RestelliForcingTerm: diff --git a/src/struphy/feec/basis_projection_ops.py b/src/struphy/feec/basis_projection_ops.py index 4a8b804b0..3bce4701f 100644 --- a/src/struphy/feec/basis_projection_ops.py +++ b/src/struphy/feec/basis_projection_ops.py @@ -1,3 +1,4 @@ +import cunumpy as xp from psydac.api.settings import PSYDAC_BACKEND_GPYCCEL from psydac.ddm.mpi import mpi as MPI from psydac.fem.basic import FemSpace @@ -14,7 +15,6 @@ from struphy.feec.utilities import RotationMatrix from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.linear_operators import PolarExtractionOperator -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 61a08e90e..3a29ea821 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -1,6 +1,7 @@ import itertools from abc import abstractmethod +import cunumpy as xp from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI from psydac.linalg.basic import LinearOperator, Vector, VectorSpace @@ -10,7 +11,6 @@ from struphy.feec.utilities import apply_essential_bc_to_array from struphy.polar.basic import PolarDerhamSpace -from struphy.utils.arrays import xp class LinOpWithTransp(LinearOperator): diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 1facbc2a0..76a26bd4d 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -1,6 +1,7 @@ import inspect from copy import deepcopy +import cunumpy as xp from psydac.api.settings import PSYDAC_BACKEND_GPYCCEL from psydac.ddm.mpi import mpi as MPI from psydac.fem.tensor import TensorFemSpace @@ -16,7 +17,6 @@ from struphy.feec.utilities import RotationMatrix from struphy.geometry.base import Domain from struphy.polar.linear_operators import PolarExtractionOperator -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 783eb062d..f3016b532 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -1,3 +1,4 @@ +import cunumpy as xp from psydac.api.essential_bc import apply_essential_bc_stencil from psydac.ddm.cart import CartDecomposition, DomainDecomposition from psydac.fem.tensor import TensorFemSpace @@ -11,7 +12,6 @@ from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.mass import WeightedMassOperator -from struphy.utils.arrays import xp class MassMatrixPreconditioner(LinearOperator): diff --git a/src/struphy/feec/projectors.py b/src/struphy/feec/projectors.py index 5a43cfc74..a291d18bd 100644 --- a/src/struphy/feec/projectors.py +++ b/src/struphy/feec/projectors.py @@ -1,3 +1,4 @@ +import cunumpy as xp from psydac.api.settings import PSYDAC_BACKEND_GPYCCEL from psydac.ddm.mpi import mpi as MPI from psydac.feec.global_projectors import GlobalProjector @@ -37,7 +38,6 @@ from struphy.kernel_arguments.local_projectors_args_kernels import LocalProjectorsArguments from struphy.polar.basic import PolarVector from struphy.polar.linear_operators import PolarExtractionOperator -from struphy.utils.arrays import xp class CommutingProjector: diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index fbbddae28..f296d95a9 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import importlib.metadata +import cunumpy as xp import psydac.core.bsplines as bsp from psydac.ddm.cart import DomainDecomposition from psydac.ddm.mpi import MockComm, MockMPI @@ -33,7 +34,6 @@ from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator -from struphy.utils.arrays import xp class Derham: diff --git a/src/struphy/feec/tests/test_basis_ops.py b/src/struphy/feec/tests/test_basis_ops.py index e46522592..7b06e09e1 100644 --- a/src/struphy/feec/tests/test_basis_ops.py +++ b/src/struphy/feec/tests/test_basis_ops.py @@ -14,6 +14,7 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): """ from time import time + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -24,7 +25,6 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): from struphy.feec.psydac_derham import Derham from struphy.fields_background.equils import HomogenSlab from struphy.geometry import domains - from struphy.utils.arrays import xp # mpi communicator MPI_COMM = MPI.COMM_WORLD @@ -464,6 +464,7 @@ def test_some_basis_ops(Nel, p, spl_kind, mapping): ) @pytest.mark.parametrize("mapping", [["IGAPolarCylinder", {"a": 1.0, "Lz": 3.0}]]) def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.mhd_operators import MHDOperators @@ -474,7 +475,6 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal from struphy.fields_background.equils import ScrewPinch from struphy.geometry import domains from struphy.polar.basic import PolarVector - from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -726,7 +726,7 @@ def assert_ops(mpi_rank, res_PSY, res_STR, verbose=False, MPI_COMM=None): TODO """ - from struphy.utils.arrays import xp + import cunumpy as xp if verbose: if MPI_COMM is not None: diff --git a/src/struphy/feec/tests/test_derham.py b/src/struphy/feec/tests/test_derham.py index fff60bbfc..c5c0e57ea 100644 --- a/src/struphy/feec/tests/test_derham.py +++ b/src/struphy/feec/tests/test_derham.py @@ -7,6 +7,7 @@ def test_psydac_derham(Nel, p, spl_kind): """Remark: p=even projectors yield slightly different results, pass with atol=1e-3.""" + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -14,7 +15,6 @@ def test_psydac_derham(Nel, p, spl_kind): from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import compare_arrays - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/feec/tests/test_eval_field.py b/src/struphy/feec/tests/test_eval_field.py index 8213853b1..6148fa41e 100644 --- a/src/struphy/feec/tests/test_eval_field.py +++ b/src/struphy/feec/tests/test_eval_field.py @@ -1,9 +1,8 @@ +import cunumpy as xp import pytest from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp - @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @pytest.mark.parametrize("p", [[3, 2, 4]]) diff --git a/src/struphy/feec/tests/test_field_init.py b/src/struphy/feec/tests/test_field_init.py index 248c32d58..292edf76a 100644 --- a/src/struphy/feec/tests/test_field_init.py +++ b/src/struphy/feec/tests/test_field_init.py @@ -9,11 +9,11 @@ def test_bckgr_init_const(Nel, p, spl_kind, spaces, vec_comps): """Test field background initialization of "LogicalConst" with multiple fields in params.""" + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.feec.psydac_derham import Derham from struphy.io.options import FieldsBackground - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -64,6 +64,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show import inspect + import cunumpy as xp from matplotlib import pyplot as plt from psydac.ddm.mpi import mpi as MPI @@ -72,7 +73,6 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB from struphy.geometry import domains from struphy.io.options import FieldsBackground - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -1083,13 +1083,13 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): """Test field perturbation with ModesSin + ModesCos on top of of "LogicalConst" with multiple fields in params.""" + import cunumpy as xp from matplotlib import pyplot as plt from psydac.ddm.mpi import mpi as MPI from struphy.feec.psydac_derham import Derham from struphy.initial.perturbations import ModesCos, ModesSin from struphy.io.options import FieldsBackground - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -1312,12 +1312,12 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): def test_noise_init(Nel, p, spl_kind, space, direction): """Only tests 1d noise ('e1', 'e2', 'e3') !!""" + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import compare_arrays from struphy.initial.perturbations import Noise - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/feec/tests/test_l2_projectors.py b/src/struphy/feec/tests/test_l2_projectors.py index 500190075..6d3695caa 100644 --- a/src/struphy/feec/tests/test_l2_projectors.py +++ b/src/struphy/feec/tests/test_l2_projectors.py @@ -1,5 +1,6 @@ import inspect +import cunumpy as xp import matplotlib.pyplot as plt import pytest from psydac.ddm.mpi import mpi as MPI @@ -8,7 +9,6 @@ from struphy.feec.projectors import L2Projector from struphy.feec.psydac_derham import Derham from struphy.geometry import domains -from struphy.utils.arrays import xp @pytest.mark.parametrize("Nel", [[16, 32, 1]]) diff --git a/src/struphy/feec/tests/test_local_projectors.py b/src/struphy/feec/tests/test_local_projectors.py index a7549600c..ce2abda01 100644 --- a/src/struphy/feec/tests/test_local_projectors.py +++ b/src/struphy/feec/tests/test_local_projectors.py @@ -1,6 +1,7 @@ import inspect import time +import cunumpy as xp import matplotlib.pyplot as plt import pytest from psydac.ddm.mpi import MockComm @@ -12,7 +13,6 @@ from struphy.feec.local_projectors_kernels import fill_matrix_column from struphy.feec.psydac_derham import Derham from struphy.feec.utilities_local_projectors import get_one_spline, get_span_and_basis, get_values_and_indices_splines -from struphy.utils.arrays import xp def get_span_and_basis(pts, space): diff --git a/src/struphy/feec/tests/test_lowdim_nel_is_1.py b/src/struphy/feec/tests/test_lowdim_nel_is_1.py index 476c8f0c4..fdbe6be3d 100644 --- a/src/struphy/feec/tests/test_lowdim_nel_is_1.py +++ b/src/struphy/feec/tests/test_lowdim_nel_is_1.py @@ -7,13 +7,13 @@ def test_lowdim_derham(Nel, p, spl_kind, do_plot=False): """Test Nel=1 in various directions.""" + import cunumpy as xp from matplotlib import pyplot as plt from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector from struphy.feec.psydac_derham import Derham - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/feec/tests/test_mass_matrices.py b/src/struphy/feec/tests/test_mass_matrices.py index c106abc0b..a57d67cdc 100644 --- a/src/struphy/feec/tests/test_mass_matrices.py +++ b/src/struphy/feec/tests/test_mass_matrices.py @@ -12,6 +12,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): """Compare Struphy mass matrices to Struphy-legacy mass matrices.""" + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.mhd_operators import MHDOperators @@ -21,7 +22,6 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): from struphy.feec.utilities import RotationMatrix, compare_arrays, create_equal_random_arrays from struphy.fields_background.equils import ScrewPinch, ShearedSlab from struphy.geometry import domains - from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -376,6 +376,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): """Compare Struphy polar mass matrices to Struphy-legacy polar mass matrices.""" + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.mhd_operators import MHDOperators @@ -386,7 +387,6 @@ def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): from struphy.fields_background.equils import ScrewPinch from struphy.geometry import domains from struphy.polar.basic import PolarVector - from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -575,6 +575,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots import time + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from psydac.linalg.solvers import inverse @@ -584,7 +585,6 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots from struphy.feec.utilities import create_equal_random_arrays from struphy.fields_background.equils import ScrewPinch, ShearedSlab from struphy.geometry import domains - from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -882,6 +882,7 @@ def test_mass_preconditioner_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show import time + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from psydac.linalg.solvers import inverse @@ -892,7 +893,6 @@ def test_mass_preconditioner_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show from struphy.fields_background.equils import ScrewPinch from struphy.geometry import domains from struphy.polar.basic import PolarVector - from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() diff --git a/src/struphy/feec/tests/test_toarray_struphy.py b/src/struphy/feec/tests/test_toarray_struphy.py index db266d108..0279293ea 100644 --- a/src/struphy/feec/tests/test_toarray_struphy.py +++ b/src/struphy/feec/tests/test_toarray_struphy.py @@ -12,13 +12,13 @@ def test_toarray_struphy(Nel, p, spl_kind, mapping): TODO """ + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.feec.mass import WeightedMassOperators from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import compare_arrays, create_equal_random_arrays from struphy.geometry import domains - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/feec/tests/test_tosparse_struphy.py b/src/struphy/feec/tests/test_tosparse_struphy.py index d9b96aee0..020ab9e29 100644 --- a/src/struphy/feec/tests/test_tosparse_struphy.py +++ b/src/struphy/feec/tests/test_tosparse_struphy.py @@ -14,6 +14,7 @@ def test_tosparse_struphy(Nel, p, spl_kind, mapping): TODO """ + import cunumpy as xp from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI @@ -21,7 +22,6 @@ def test_tosparse_struphy(Nel, p, spl_kind, mapping): from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import create_equal_random_arrays from struphy.geometry import domains - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/feec/tests/xx_test_preconds.py b/src/struphy/feec/tests/xx_test_preconds.py index 1bd145fc6..267e0279a 100644 --- a/src/struphy/feec/tests/xx_test_preconds.py +++ b/src/struphy/feec/tests/xx_test_preconds.py @@ -12,6 +12,7 @@ ], ) def test_mass_preconditioner(Nel, p, spl_kind, mapping): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -21,7 +22,6 @@ def test_mass_preconditioner(Nel, p, spl_kind, mapping): from struphy.feec.preconditioner import MassMatrixPreconditioner from struphy.feec.psydac_derham import Derham from struphy.geometry import domains - from struphy.utils.arrays import xp MPI_COMM = MPI.COMM_WORLD diff --git a/src/struphy/feec/utilities.py b/src/struphy/feec/utilities.py index e2952572f..aa3912857 100644 --- a/src/struphy/feec/utilities.py +++ b/src/struphy/feec/utilities.py @@ -1,3 +1,4 @@ +import cunumpy as xp from psydac.api.essential_bc import apply_essential_bc_stencil from psydac.fem.tensor import TensorFemSpace from psydac.fem.vector import VectorFemSpace @@ -8,7 +9,6 @@ import struphy.feec.utilities_kernels as kernels from struphy.feec import banded_to_stencil_kernels as bts from struphy.polar.basic import PolarVector -from struphy.utils.arrays import xp class RotationMatrix: diff --git a/src/struphy/feec/utilities_local_projectors.py b/src/struphy/feec/utilities_local_projectors.py index 09fed1e54..9ff91a120 100644 --- a/src/struphy/feec/utilities_local_projectors.py +++ b/src/struphy/feec/utilities_local_projectors.py @@ -1,5 +1,6 @@ +import cunumpy as xp + from struphy.feec.local_projectors_kernels import are_quadrature_points_zero, get_rows, select_quasi_points -from struphy.utils.arrays import xp def split_points( diff --git a/src/struphy/feec/variational_utilities.py b/src/struphy/feec/variational_utilities.py index b2549d42b..12695bfa2 100644 --- a/src/struphy/feec/variational_utilities.py +++ b/src/struphy/feec/variational_utilities.py @@ -1,5 +1,6 @@ from copy import deepcopy +import cunumpy as xp from psydac.linalg.basic import IdentityOperator, Vector from psydac.linalg.block import BlockVector from psydac.linalg.solvers import inverse @@ -12,7 +13,6 @@ ) from struphy.feec.linear_operators import LinOpWithTransp from struphy.feec.psydac_derham import Derham -from struphy.utils.arrays import xp class BracketOperator(LinOpWithTransp): diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index c6178d579..ccc3db90e 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -2,11 +2,11 @@ from abc import ABCMeta, abstractmethod +import cunumpy as xp from matplotlib import pyplot as plt from pyevtk.hl import gridToVTK from struphy.geometry.base import Domain -from struphy.utils.arrays import xp class FluidEquilibrium(metaclass=ABCMeta): diff --git a/src/struphy/fields_background/coil_fields/base.py b/src/struphy/fields_background/coil_fields/base.py index 55b0720e8..3b068c62d 100644 --- a/src/struphy/fields_background/coil_fields/base.py +++ b/src/struphy/fields_background/coil_fields/base.py @@ -1,10 +1,9 @@ from abc import ABCMeta, abstractmethod +import cunumpy as xp from matplotlib import pyplot as plt from pyevtk.hl import gridToVTK -from struphy.utils.arrays import xp - class CoilMagneticField(metaclass=ABCMeta): """ diff --git a/src/struphy/fields_background/coil_fields/coil_fields.py b/src/struphy/fields_background/coil_fields/coil_fields.py index 2629cc707..e1f66c71d 100644 --- a/src/struphy/fields_background/coil_fields/coil_fields.py +++ b/src/struphy/fields_background/coil_fields/coil_fields.py @@ -1,6 +1,7 @@ +import cunumpy as xp + from struphy.feec.psydac_derham import Derham from struphy.fields_background.coil_fields.base import CoilMagneticField, load_csv_data -from struphy.utils.arrays import xp class RatGUI(CoilMagneticField): diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 25d3d00ad..93688b55f 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -7,6 +7,7 @@ import warnings from time import time +import cunumpy as xp from scipy.integrate import odeint, quad from scipy.interpolate import RectBivariateSpline, UnivariateSpline from scipy.optimize import fsolve, minimize @@ -28,7 +29,6 @@ NumericalMHDequilibrium, ) from struphy.fields_background.mhd_equil.eqdsk import readeqdsk -from struphy.utils.arrays import xp from struphy.utils.utils import read_state, subp_run diff --git a/src/struphy/fields_background/tests/test_desc_equil.py b/src/struphy/fields_background/tests/test_desc_equil.py index e7ae0872c..5aca31b8d 100644 --- a/src/struphy/fields_background/tests/test_desc_equil.py +++ b/src/struphy/fields_background/tests/test_desc_equil.py @@ -1,10 +1,9 @@ import importlib.util +import cunumpy as xp import pytest from matplotlib import pyplot as plt -from struphy.utils.arrays import xp - desc_spec = importlib.util.find_spec("desc") diff --git a/src/struphy/fields_background/tests/test_generic_equils.py b/src/struphy/fields_background/tests/test_generic_equils.py index 6ed2c6178..77ca8baaa 100644 --- a/src/struphy/fields_background/tests/test_generic_equils.py +++ b/src/struphy/fields_background/tests/test_generic_equils.py @@ -1,3 +1,4 @@ +import cunumpy as xp import pytest from matplotlib import pyplot as plt @@ -5,7 +6,6 @@ GenericCartesianFluidEquilibrium, GenericCartesianFluidEquilibriumWithB, ) -from struphy.utils.arrays import xp def test_generic_equils(show=False): diff --git a/src/struphy/fields_background/tests/test_mhd_equils.py b/src/struphy/fields_background/tests/test_mhd_equils.py index 076891914..494d707b3 100644 --- a/src/struphy/fields_background/tests/test_mhd_equils.py +++ b/src/struphy/fields_background/tests/test_mhd_equils.py @@ -1,7 +1,7 @@ +import cunumpy as xp import pytest from struphy.fields_background import equils -from struphy.utils.arrays import xp @pytest.mark.parametrize( diff --git a/src/struphy/fields_background/tests/test_numerical_mhd_equil.py b/src/struphy/fields_background/tests/test_numerical_mhd_equil.py index 1002c37c8..aa1278d5d 100644 --- a/src/struphy/fields_background/tests/test_numerical_mhd_equil.py +++ b/src/struphy/fields_background/tests/test_numerical_mhd_equil.py @@ -1,7 +1,7 @@ +import cunumpy as xp import pytest from struphy.fields_background.base import FluidEquilibrium, LogicalMHDequilibrium -from struphy.utils.arrays import xp @pytest.mark.parametrize( diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index 9de4acede..a789913e7 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -3,6 +3,7 @@ from abc import ABCMeta, abstractmethod +import cunumpy as xp import h5py from scipy.sparse import csc_matrix, kron from scipy.sparse.linalg import splu, spsolve @@ -11,7 +12,6 @@ from struphy.geometry import evaluation_kernels, transform_kernels from struphy.kernel_arguments.pusher_args_kernels import DomainArguments from struphy.linear_algebra import linalg_kron -from struphy.utils.arrays import xp class Domain(metaclass=ABCMeta): diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index 5e1649dc7..024b980c9 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -2,6 +2,8 @@ import copy +import cunumpy as xp + from struphy.fields_background.base import AxisymmMHDequilibrium from struphy.fields_background.equils import EQDSKequilibrium from struphy.geometry.base import ( @@ -12,7 +14,6 @@ interp_mapping, ) from struphy.geometry.utilities import field_line_tracing -from struphy.utils.arrays import xp class Tokamak(PoloidalSplineTorus): diff --git a/src/struphy/geometry/tests/test_domain.py b/src/struphy/geometry/tests/test_domain.py index 76db82315..c9a489331 100644 --- a/src/struphy/geometry/tests/test_domain.py +++ b/src/struphy/geometry/tests/test_domain.py @@ -4,8 +4,9 @@ def test_prepare_arg(): """Tests prepare_arg static method in domain base class.""" + import cunumpy as xp + from struphy.geometry.base import Domain - from struphy.utils.arrays import xp def a1(e1, e2, e3): return e1 * e2 @@ -158,9 +159,10 @@ def a_vec(e1, e2, e3): def test_evaluation_mappings(mapping): """Tests domain object creation with default parameters and evaluation of metric coefficients.""" + import cunumpy as xp + from struphy.geometry import domains from struphy.geometry.base import Domain - from struphy.utils.arrays import xp # arrays: arr1 = xp.linspace(0.0, 1.0, 4) @@ -317,9 +319,10 @@ def test_evaluation_mappings(mapping): def test_pullback(): """Tests pullbacks to p-forms.""" + import cunumpy as xp + from struphy.geometry import domains from struphy.geometry.base import Domain - from struphy.utils.arrays import xp # arrays: arr1 = xp.linspace(0.0, 1.0, 4) @@ -476,9 +479,10 @@ def fun(x, y, z): def test_pushforward(): """Tests pushforward of p-forms.""" + import cunumpy as xp + from struphy.geometry import domains from struphy.geometry.base import Domain - from struphy.utils.arrays import xp # arrays: arr1 = xp.linspace(0.0, 1.0, 4) @@ -635,9 +639,10 @@ def fun(e1, e2, e3): def test_transform(): """Tests transformation of p-forms.""" + import cunumpy as xp + from struphy.geometry import domains from struphy.geometry.base import Domain - from struphy.utils.arrays import xp # arrays: arr1 = xp.linspace(0.0, 1.0, 4) @@ -817,7 +822,7 @@ def fun(e1, e2, e3): # """ # # from struphy.geometry import domains -# from struphy.utils.arrays import xp +# import cunumpy as xp # # # arrays: # arr1 = xp.linspace(0., 1., 4) diff --git a/src/struphy/geometry/utilities.py b/src/struphy/geometry/utilities.py index 0d5c07d93..ccc159692 100644 --- a/src/struphy/geometry/utilities.py +++ b/src/struphy/geometry/utilities.py @@ -3,6 +3,7 @@ from typing import Callable +import cunumpy as xp import numpy as np # from typing import TYPE_CHECKING @@ -17,7 +18,6 @@ from struphy.geometry.utilities_kernels import weighted_arc_lengths_flux_surface from struphy.io.options import GivenInBasis from struphy.linear_algebra.linalg_kron import kron_lusolve_2d -from struphy.utils.arrays import xp def field_line_tracing( diff --git a/src/struphy/initial/eigenfunctions.py b/src/struphy/initial/eigenfunctions.py index 093fe21ac..2f66fcacb 100644 --- a/src/struphy/initial/eigenfunctions.py +++ b/src/struphy/initial/eigenfunctions.py @@ -1,11 +1,11 @@ import os +import cunumpy as xp import yaml from psydac.api.discretization import discretize from sympde.topology import Derham, Line from struphy.fields_background.equils import set_defaults -from struphy.utils.arrays import xp class InitialMHDAxisymHdivEigFun: diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index a062c3f61..767f899c9 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -3,12 +3,12 @@ from dataclasses import dataclass +import cunumpy as xp import scipy import scipy.special from struphy.initial.base import Perturbation from struphy.io.options import GivenInBasis, NoiseDirections, check_option -from struphy.utils.arrays import xp @dataclass diff --git a/src/struphy/initial/tests/test_init_perturbations.py b/src/struphy/initial/tests/test_init_perturbations.py index ee71c22c8..2f5fa3176 100644 --- a/src/struphy/initial/tests/test_init_perturbations.py +++ b/src/struphy/initial/tests/test_init_perturbations.py @@ -20,6 +20,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False): """Test the initialization Field.initialize_coeffs with all "Modes" classes in perturbations.py.""" + import cunumpy as xp from matplotlib import pyplot as plt from psydac.ddm.mpi import mpi as MPI @@ -29,7 +30,6 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False from struphy.initial import perturbations from struphy.initial.base import Perturbation from struphy.models.variables import FEECVariable - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/initial/utilities.py b/src/struphy/initial/utilities.py index 274b73ba6..7af042249 100644 --- a/src/struphy/initial/utilities.py +++ b/src/struphy/initial/utilities.py @@ -1,10 +1,10 @@ import os +import cunumpy as xp import h5py from struphy.fields_background.equils import set_defaults from struphy.io.output_handling import DataContainer -from struphy.utils.arrays import xp class InitFromOutput: diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index fff19cb48..cecd4b72c 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -2,10 +2,10 @@ from dataclasses import dataclass from typing import Literal, get_args +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.physics.physics import ConstantsOfNature -from struphy.utils.arrays import xp ## Literal options diff --git a/src/struphy/io/output_handling.py b/src/struphy/io/output_handling.py index ed9c563d4..734a2d1a3 100644 --- a/src/struphy/io/output_handling.py +++ b/src/struphy/io/output_handling.py @@ -1,10 +1,9 @@ import ctypes import os +import cunumpy as xp import h5py -from struphy.utils.arrays import xp - class DataContainer: """ diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index 58bcdccf0..78a295b35 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -5,12 +5,12 @@ import sys from types import ModuleType +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.geometry.base import Domain from struphy.io.options import DerhamOptions from struphy.topology.grids import TensorProductGrid -from struphy.utils.arrays import xp def import_parameters_py(params_path: str) -> ModuleType: diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 5da0a1568..41c7b6575 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -3,9 +3,10 @@ from abc import ABCMeta, abstractmethod from typing import Callable +import cunumpy as xp + from struphy.fields_background.base import FluidEquilibriumWithB from struphy.initial.base import Perturbation -from struphy.utils.arrays import xp class KineticBackground(metaclass=ABCMeta): diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index ea541eb70..c7dea067a 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -2,11 +2,12 @@ from typing import Callable +import cunumpy as xp + from struphy.fields_background.base import FluidEquilibriumWithB from struphy.fields_background.equils import set_defaults from struphy.initial.base import Perturbation from struphy.kinetic_background.base import Maxwellian -from struphy.utils.arrays import xp class Maxwellian3D(Maxwellian): diff --git a/src/struphy/kinetic_background/moment_functions.py b/src/struphy/kinetic_background/moment_functions.py index b573e657a..39b22c6e1 100644 --- a/src/struphy/kinetic_background/moment_functions.py +++ b/src/struphy/kinetic_background/moment_functions.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 "Analytical moment functions." -from struphy.utils.arrays import xp +import cunumpy as xp class ITPA_density: diff --git a/src/struphy/kinetic_background/tests/test_base.py b/src/struphy/kinetic_background/tests/test_base.py index 5dc077c24..8a2e89d28 100644 --- a/src/struphy/kinetic_background/tests/test_base.py +++ b/src/struphy/kinetic_background/tests/test_base.py @@ -1,10 +1,10 @@ def test_kinetic_background_magics(show_plot=False): """Test the magic commands __sum__, __mul__ and __sub__ of the Maxwellian base class.""" + import cunumpy as xp import matplotlib.pyplot as plt from struphy.kinetic_background.maxwellians import Maxwellian3D - from struphy.utils.arrays import xp Nel = [32, 1, 1] e1 = xp.linspace(0.0, 1.0, Nel[0]) diff --git a/src/struphy/kinetic_background/tests/test_maxwellians.py b/src/struphy/kinetic_background/tests/test_maxwellians.py index 44e3679d4..edf24af4c 100644 --- a/src/struphy/kinetic_background/tests/test_maxwellians.py +++ b/src/struphy/kinetic_background/tests/test_maxwellians.py @@ -8,10 +8,10 @@ def test_maxwellian_3d_uniform(Nel, show_plot=False): Asserts that the results over the domain and velocity space correspond to the analytical computation. """ + import cunumpy as xp import matplotlib.pyplot as plt from struphy.kinetic_background.maxwellians import Maxwellian3D - from struphy.utils.arrays import xp e1 = xp.linspace(0.0, 1.0, Nel[0]) e2 = xp.linspace(0.0, 1.0, Nel[1]) @@ -93,11 +93,11 @@ def test_maxwellian_3d_uniform(Nel, show_plot=False): def test_maxwellian_3d_perturbed(Nel, show_plot=False): """Tests the Maxwellian3D class for perturbations.""" + import cunumpy as xp import matplotlib.pyplot as plt from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import Maxwellian3D - from struphy.utils.arrays import xp e1 = xp.linspace(0.0, 1.0, Nel[0]) v1 = xp.linspace(-5.0, 5.0, 128) @@ -255,6 +255,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): import inspect + import cunumpy as xp import matplotlib.pyplot as plt from struphy.fields_background import equils @@ -263,7 +264,6 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): from struphy.initial import perturbations from struphy.initial.base import Perturbation from struphy.kinetic_background.maxwellians import Maxwellian3D - from struphy.utils.arrays import xp e1 = xp.linspace(0.0, 1.0, Nel[0]) e2 = xp.linspace(0.0, 1.0, Nel[1]) @@ -666,10 +666,10 @@ def test_maxwellian_2d_uniform(Nel, show_plot=False): Asserts that the results over the domain and velocity space correspond to the analytical computation. """ + import cunumpy as xp import matplotlib.pyplot as plt from struphy.kinetic_background.maxwellians import GyroMaxwellian2D - from struphy.utils.arrays import xp e1 = xp.linspace(0.0, 1.0, Nel[0]) e2 = xp.linspace(0.0, 1.0, Nel[1]) @@ -759,11 +759,11 @@ def test_maxwellian_2d_uniform(Nel, show_plot=False): def test_maxwellian_2d_perturbed(Nel, show_plot=False): """Tests the GyroMaxwellian2D class for perturbations.""" + import cunumpy as xp import matplotlib.pyplot as plt from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import GyroMaxwellian2D - from struphy.utils.arrays import xp e1 = xp.linspace(0.0, 1.0, Nel[0]) v1 = xp.linspace(-5.0, 5.0, 128) @@ -1017,6 +1017,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): import inspect + import cunumpy as xp import matplotlib.pyplot as plt from struphy.fields_background import equils @@ -1025,7 +1026,6 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): from struphy.initial import perturbations from struphy.initial.base import Perturbation from struphy.kinetic_background.maxwellians import GyroMaxwellian2D - from struphy.utils.arrays import xp e1 = xp.linspace(0.0, 1.0, Nel[0]) e2 = xp.linspace(0.0, 1.0, Nel[1]) @@ -1417,13 +1417,13 @@ def test_canonical_maxwellian_uniform(Nel, show_plot=False): Asserts that the results over the domain and velocity space correspond to the analytical computation. """ + import cunumpy as xp import matplotlib.pyplot as plt from struphy.fields_background import equils from struphy.geometry import domains from struphy.initial import perturbations from struphy.kinetic_background.maxwellians import CanonicalMaxwellian - from struphy.utils.arrays import xp e1 = xp.linspace(0.0, 1.0, Nel[0]) e2 = xp.linspace(0.0, 1.0, Nel[1]) diff --git a/src/struphy/linear_algebra/linalg_kron.py b/src/struphy/linear_algebra/linalg_kron.py index 29be15d60..2e0dd57dc 100644 --- a/src/struphy/linear_algebra/linalg_kron.py +++ b/src/struphy/linear_algebra/linalg_kron.py @@ -13,11 +13,10 @@ [r_M11, rM12, ... , r_MNO]] """ +import cunumpy as xp from scipy.linalg import solve_circulant from scipy.sparse.linalg import splu -from struphy.utils.arrays import xp - def kron_matvec_2d(kmat, vec2d): """ diff --git a/src/struphy/linear_algebra/saddle_point.py b/src/struphy/linear_algebra/saddle_point.py index efa4dbbbe..47dd36190 100644 --- a/src/struphy/linear_algebra/saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point.py @@ -1,5 +1,6 @@ from typing import Union +import cunumpy as xp import scipy as sc from psydac.linalg.basic import LinearOperator, Vector from psydac.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace @@ -7,7 +8,6 @@ from psydac.linalg.solvers import inverse from struphy.linear_algebra.tests.test_saddlepoint_massmatrices import _plot_residual_norms -from struphy.utils.arrays import xp class SaddlePointSolver: diff --git a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py index 9f196411e..ed5ff4d8a 100644 --- a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py +++ b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py @@ -13,6 +13,7 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m import time + import cunumpy as xp import scipy as sc from psydac.ddm.mpi import mpi as MPI from psydac.linalg.basic import IdentityOperator @@ -29,7 +30,6 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m from struphy.geometry import domains from struphy.initial import perturbations from struphy.linear_algebra.saddle_point import SaddlePointSolver - from struphy.utils.arrays import xp mpi_comm = MPI.COMM_WORLD mpi_rank = mpi_comm.Get_rank() @@ -372,11 +372,10 @@ def _plot_residual_norms(residual_norms): def _plot_velocity(data_reshaped): + import cunumpy as xp import matplotlib import matplotlib.pyplot as plt - from struphy.utils.arrays import xp - matplotlib.use("Agg") x = xp.linspace(0, 1, 30) diff --git a/src/struphy/linear_algebra/tests/test_stencil_dot_kernels.py b/src/struphy/linear_algebra/tests/test_stencil_dot_kernels.py index 0475b418a..d2c2238ff 100644 --- a/src/struphy/linear_algebra/tests/test_stencil_dot_kernels.py +++ b/src/struphy/linear_algebra/tests/test_stencil_dot_kernels.py @@ -13,13 +13,13 @@ def test_1d(Nel, p, spl_kind, domain_ind, codomain_ind): a) the result from kernel in struphy.linear_algebra.stencil_dot_kernels.matvec_1d_kernel b) the result from Stencil .dot with precompiled=True""" + import cunumpy as xp from psydac.api.settings import PSYDAC_BACKEND_GPYCCEL from psydac.ddm.mpi import mpi as MPI from psydac.linalg.stencil import StencilMatrix, StencilVector from struphy.feec.psydac_derham import Derham from struphy.linear_algebra.stencil_dot_kernels import matvec_1d_kernel - from struphy.utils.arrays import xp # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" @@ -134,13 +134,13 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): a) the result from kernel in struphy.linear_algebra.stencil_dot_kernels.matvec_1d_kernel b) the result from Stencil .dot with precompiled=True""" + import cunumpy as xp from psydac.api.settings import PSYDAC_BACKEND_GPYCCEL from psydac.ddm.mpi import mpi as MPI from psydac.linalg.stencil import StencilMatrix, StencilVector from struphy.feec.psydac_derham import Derham from struphy.linear_algebra.stencil_dot_kernels import matvec_3d_kernel - from struphy.utils.arrays import xp # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" diff --git a/src/struphy/linear_algebra/tests/test_stencil_transpose_kernels.py b/src/struphy/linear_algebra/tests/test_stencil_transpose_kernels.py index 486e302f5..1125a980c 100644 --- a/src/struphy/linear_algebra/tests/test_stencil_transpose_kernels.py +++ b/src/struphy/linear_algebra/tests/test_stencil_transpose_kernels.py @@ -13,13 +13,13 @@ def test_1d(Nel, p, spl_kind, domain_ind, codomain_ind): a) the result from kernel in struphy.linear_algebra.stencil_transpose_kernels.transpose_1d_kernel b) the result from Stencil .transpose with precompiled=True""" + import cunumpy as xp from psydac.api.settings import PSYDAC_BACKEND_GPYCCEL from psydac.ddm.mpi import mpi as MPI from psydac.linalg.stencil import StencilMatrix from struphy.feec.psydac_derham import Derham from struphy.linear_algebra.stencil_transpose_kernels import transpose_1d_kernel - from struphy.utils.arrays import xp # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" @@ -128,13 +128,13 @@ def test_3d(Nel, p, spl_kind, domain_ind, codomain_ind): a) the result from kernel in struphy.linear_algebra.stencil_transpose_kernels.transpose_3d_kernel b) the result from Stencil .transpose with precompiled=True""" + import cunumpy as xp from psydac.api.settings import PSYDAC_BACKEND_GPYCCEL from psydac.ddm.mpi import mpi as MPI from psydac.linalg.stencil import StencilMatrix from struphy.feec.psydac_derham import Derham from struphy.linear_algebra.stencil_transpose_kernels import transpose_3d_kernel - from struphy.utils.arrays import xp # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" diff --git a/src/struphy/main.py b/src/struphy/main.py index 3b06b671e..524e19eb7 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -8,6 +8,7 @@ import time from typing import Optional, TypedDict +import cunumpy as xp import h5py from line_profiler import profile from psydac.ddm.mpi import MockMPI @@ -38,7 +39,6 @@ from struphy.profiling.profiling import ProfileManager from struphy.topology import grids from struphy.topology.grids import TensorProductGrid -from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index d910cc41b..d18942154 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -5,6 +5,7 @@ from functools import reduce from textwrap import indent +import cunumpy as xp import yaml from line_profiler import profile from psydac.ddm.mpi import MockMPI @@ -35,7 +36,6 @@ from struphy.profiling.profiling import ProfileManager from struphy.propagators.base import Propagator from struphy.topology.grids import TensorProductGrid -from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig from struphy.utils.utils import dict_to_yaml, read_state diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 98983da4b..2b832ba00 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -1,3 +1,4 @@ +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -9,7 +10,6 @@ from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.polar.basic import PolarVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.utils.arrays import xp rank = MPI.COMM_WORLD.Get_rank() diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index 0d765981d..b7801a050 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -1,3 +1,4 @@ +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.models.base import StruphyModel @@ -6,7 +7,6 @@ from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.polar.basic import PolarVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel rank = MPI.COMM_WORLD.Get_rank() diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index c8f6ae452..9a01e38fd 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -1,3 +1,4 @@ +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.feec.projectors import L2Projector @@ -9,7 +10,6 @@ from struphy.pic.accumulation import accum_kernels, accum_kernels_gc from struphy.pic.accumulation.particles_to_grid import AccumulatorVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel rank = MPI.COMM_WORLD.Get_rank() diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 4ad4c14d6..56a338ca9 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,6 +1,7 @@ import warnings from abc import ABCMeta, abstractmethod +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.io.options import Units @@ -13,7 +14,6 @@ LoadingParameters, WeightsParameters, ) -from struphy.utils.arrays import xp class Species(metaclass=ABCMeta): diff --git a/src/struphy/models/tests/test_verif_EulerSPH.py b/src/struphy/models/tests/test_verif_EulerSPH.py index fa9b2bcfa..79a248ac9 100644 --- a/src/struphy/models/tests/test_verif_EulerSPH.py +++ b/src/struphy/models/tests/test_verif_EulerSPH.py @@ -1,5 +1,6 @@ import os +import cunumpy as xp import pytest from matplotlib import pyplot as plt from matplotlib.ticker import FormatStrFormatter @@ -19,7 +20,6 @@ WeightsParameters, ) from struphy.topology import grids -from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") diff --git a/src/struphy/models/tests/test_verif_LinearMHD.py b/src/struphy/models/tests/test_verif_LinearMHD.py index 3ed92aca0..5cbbbc9fd 100644 --- a/src/struphy/models/tests/test_verif_LinearMHD.py +++ b/src/struphy/models/tests/test_verif_LinearMHD.py @@ -1,5 +1,6 @@ import os +import cunumpy as xp import pytest from psydac.ddm.mpi import mpi as MPI @@ -11,7 +12,6 @@ from struphy.io.options import BaseUnits, DerhamOptions, EnvironmentOptions, FieldsBackground, Time from struphy.kinetic_background import maxwellians from struphy.topology import grids -from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "verification_tests") diff --git a/src/struphy/models/tests/test_verif_Maxwell.py b/src/struphy/models/tests/test_verif_Maxwell.py index dc1a4794c..e97675e7b 100644 --- a/src/struphy/models/tests/test_verif_Maxwell.py +++ b/src/struphy/models/tests/test_verif_Maxwell.py @@ -1,5 +1,6 @@ import os +import cunumpy as xp import pytest from matplotlib import pyplot as plt from psydac.ddm.mpi import mpi as MPI @@ -14,7 +15,6 @@ from struphy.kinetic_background import maxwellians from struphy.models.toy import Maxwell from struphy.topology import grids -from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") diff --git a/src/struphy/models/tests/test_verif_Poisson.py b/src/struphy/models/tests/test_verif_Poisson.py index 289886cb3..5b62d61ab 100644 --- a/src/struphy/models/tests/test_verif_Poisson.py +++ b/src/struphy/models/tests/test_verif_Poisson.py @@ -1,5 +1,6 @@ import os +import cunumpy as xp from matplotlib import pyplot as plt from psydac.ddm.mpi import mpi as MPI @@ -18,7 +19,6 @@ WeightsParameters, ) from struphy.topology import grids -from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") diff --git a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py index 980352a04..a8c0b8b13 100644 --- a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py @@ -1,5 +1,6 @@ import os +import cunumpy as xp import h5py from matplotlib import pyplot as plt from matplotlib.ticker import FormatStrFormatter @@ -19,7 +20,6 @@ WeightsParameters, ) from struphy.topology import grids -from struphy.utils.arrays import xp test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") diff --git a/src/struphy/models/tests/verification.py b/src/struphy/models/tests/verification.py index d54263ee5..75b0a1047 100644 --- a/src/struphy/models/tests/verification.py +++ b/src/struphy/models/tests/verification.py @@ -2,6 +2,7 @@ import pickle from pathlib import Path +import cunumpy as xp import h5py import yaml from matplotlib import pyplot as plt @@ -11,7 +12,6 @@ import struphy from struphy.post_processing import pproc_struphy -from struphy.utils.arrays import xp def VlasovAmpereOneSpecies_weakLandau( diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index 666e87649..bad2b7916 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -1,3 +1,4 @@ +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.feec.projectors import L2Projector @@ -6,7 +7,6 @@ from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers -from struphy.utils.arrays import xp rank = MPI.COMM_WORLD.Get_rank() diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 3985f4b35..d71f4644d 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -4,6 +4,7 @@ from abc import ABCMeta, abstractmethod from typing import TYPE_CHECKING +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.feec.psydac_derham import Derham, SplineFunction @@ -21,7 +22,6 @@ from struphy.pic import particles from struphy.pic.base import Particles from struphy.pic.particles import ParticlesSPH -from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig if TYPE_CHECKING: diff --git a/src/struphy/ode/solvers.py b/src/struphy/ode/solvers.py index 4216198ff..c6d6366b9 100644 --- a/src/struphy/ode/solvers.py +++ b/src/struphy/ode/solvers.py @@ -1,10 +1,10 @@ from inspect import signature +import cunumpy as xp from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector from struphy.ode.utils import ButcherTableau -from struphy.utils.arrays import xp class ODEsolverFEEC: diff --git a/src/struphy/ode/tests/test_ode_feec.py b/src/struphy/ode/tests/test_ode_feec.py index 9d1a49be5..dadd3aad3 100644 --- a/src/struphy/ode/tests/test_ode_feec.py +++ b/src/struphy/ode/tests/test_ode_feec.py @@ -20,6 +20,7 @@ def test_exp_growth(spaces, algo, show_plots=False): """Solve dy/dt = omega*y for different feec variables y and with all available solvers from the ButcherTableau.""" + import cunumpy as xp from matplotlib import pyplot as plt from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector @@ -28,7 +29,6 @@ def test_exp_growth(spaces, algo, show_plots=False): from struphy.feec.psydac_derham import Derham from struphy.ode.solvers import ODEsolverFEEC from struphy.ode.utils import ButcherTableau - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/ode/utils.py b/src/struphy/ode/utils.py index 997225d42..6748d07f1 100644 --- a/src/struphy/ode/utils.py +++ b/src/struphy/ode/utils.py @@ -1,7 +1,7 @@ from dataclasses import dataclass from typing import Literal, get_args -from struphy.utils.arrays import xp +import cunumpy as xp OptsButcher = Literal[ "rk4", diff --git a/src/struphy/pic/accumulation/particles_to_grid.py b/src/struphy/pic/accumulation/particles_to_grid.py index d694c017b..06d67a6df 100644 --- a/src/struphy/pic/accumulation/particles_to_grid.py +++ b/src/struphy/pic/accumulation/particles_to_grid.py @@ -1,5 +1,6 @@ "Base classes for particle deposition (accumulation) on the grid." +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilMatrix, StencilVector @@ -12,7 +13,6 @@ from struphy.pic.accumulation.filter import AccumFilter, FilterParameters from struphy.pic.base import Particles from struphy.profiling.profiling import ProfileManager -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 824c51c0a..0f92323ae 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -14,6 +14,7 @@ class Intracomm: x = None +import cunumpy as xp from line_profiler import profile from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI @@ -53,7 +54,6 @@ class Intracomm: WeightsParameters, ) from struphy.utils import utils -from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig from struphy.utils.pyccel import Pyccelkernel diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index f662bbac2..7258b3cd5 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -1,5 +1,7 @@ import copy +import cunumpy as xp + from struphy.fields_background import equils from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB @@ -10,7 +12,6 @@ from struphy.kinetic_background.base import Maxwellian, SumKineticBackground from struphy.pic import utilities_kernels from struphy.pic.base import Particles -from struphy.utils.arrays import xp class Particles6D(Particles): diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 53287371c..190525de9 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -1,12 +1,12 @@ "Accelerated particle pushing." +import cunumpy as xp from line_profiler import profile from psydac.ddm.mpi import mpi as MPI from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments from struphy.pic.base import Particles from struphy.profiling.profiling import ProfileManager -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel diff --git a/src/struphy/pic/sobol_seq.py b/src/struphy/pic/sobol_seq.py index 9208d813d..ce965cc8f 100644 --- a/src/struphy/pic/sobol_seq.py +++ b/src/struphy/pic/sobol_seq.py @@ -17,10 +17,9 @@ from __future__ import division +import cunumpy as xp from scipy.stats import norm -from struphy.utils.arrays import xp - __all__ = ["i4_bit_hi1", "i4_bit_lo0", "i4_sobol_generate", "i4_sobol", "i4_uniform", "prime_ge", "is_prime"] diff --git a/src/struphy/pic/tests/test_accum_vec_H1.py b/src/struphy/pic/tests/test_accum_vec_H1.py index ab7f3e2aa..7ed52b153 100644 --- a/src/struphy/pic/tests/test_accum_vec_H1.py +++ b/src/struphy/pic/tests/test_accum_vec_H1.py @@ -47,6 +47,7 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): import copy + import cunumpy as xp from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI @@ -57,7 +58,6 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): from struphy.pic.accumulation.particles_to_grid import AccumulatorVector from struphy.pic.particles import Particles6D from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp from struphy.utils.clone_config import CloneConfig if isinstance(MPI.COMM_WORLD, MockComm): diff --git a/src/struphy/pic/tests/test_accumulation.py b/src/struphy/pic/tests/test_accumulation.py index 693df515e..012c73e33 100644 --- a/src/struphy/pic/tests/test_accumulation.py +++ b/src/struphy/pic/tests/test_accumulation.py @@ -48,6 +48,7 @@ def test_accumulation(Nel, p, spl_kind, mapping, Np=40, verbose=False): def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): from time import time + import cunumpy as xp from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI @@ -61,7 +62,6 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): from struphy.pic.particles import Particles6D from struphy.pic.tests.test_pic_legacy_files.accumulation_kernels_3d import kernel_step_ph_full from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp if isinstance(MPI.COMM_WORLD, MockComm): mpi_comm = None diff --git a/src/struphy/pic/tests/test_binning.py b/src/struphy/pic/tests/test_binning.py index e9e69ae73..a6a1dde6e 100644 --- a/src/struphy/pic/tests/test_binning.py +++ b/src/struphy/pic/tests/test_binning.py @@ -35,6 +35,7 @@ def test_binning_6D_full_f(mapping, show_plot=False): name and specification of the mapping """ + import cunumpy as xp import matplotlib.pyplot as plt from psydac.ddm.mpi import mpi as MPI @@ -47,7 +48,6 @@ def test_binning_6D_full_f(mapping, show_plot=False): LoadingParameters, WeightsParameters, ) - from struphy.utils.arrays import xp # Set seed seed = 1234 @@ -268,6 +268,7 @@ def test_binning_6D_delta_f(mapping, show_plot=False): name and specification of the mapping """ + import cunumpy as xp import matplotlib.pyplot as plt from psydac.ddm.mpi import mpi as MPI @@ -280,7 +281,6 @@ def test_binning_6D_delta_f(mapping, show_plot=False): LoadingParameters, WeightsParameters, ) - from struphy.utils.arrays import xp # Set seed seed = 1234 @@ -464,6 +464,7 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): name and specification of the mapping """ + import cunumpy as xp import matplotlib.pyplot as plt from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI @@ -477,7 +478,6 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): LoadingParameters, WeightsParameters, ) - from struphy.utils.arrays import xp # Set seed seed = 1234 @@ -761,6 +761,7 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): name and specification of the mapping """ + import cunumpy as xp import matplotlib.pyplot as plt from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI @@ -774,7 +775,6 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): LoadingParameters, WeightsParameters, ) - from struphy.utils.arrays import xp # Set seed seed = 1234 diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index 5174c83b5..cf95f4dc7 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -35,13 +35,13 @@ def test_draw(Nel, p, spl_kind, mapping, ppc=10): """Asserts whether all particles are on the correct process after `particles.mpi_sort_markers()`.""" + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.feec.psydac_derham import Derham from struphy.geometry import domains from struphy.pic.particles import Particles6D from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/pic/tests/test_mat_vec_filler.py b/src/struphy/pic/tests/test_mat_vec_filler.py index c3ad580f4..491e1f20e 100644 --- a/src/struphy/pic/tests/test_mat_vec_filler.py +++ b/src/struphy/pic/tests/test_mat_vec_filler.py @@ -1,7 +1,6 @@ +import cunumpy as xp import pytest -from struphy.utils.arrays import xp - @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @pytest.mark.parametrize("p", [[1, 2, 3]]) diff --git a/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py b/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py index 82c88d7cf..2ee5ba7b2 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py @@ -8,11 +8,11 @@ import time +import cunumpy as xp import scipy.sparse as spa from psydac.ddm.mpi import mpi as MPI import struphy.pic.tests.test_pic_legacy_files.accumulation_kernels_3d as pic_ker_3d -from struphy.utils.arrays import xp # import struphy.pic.tests.test_pic_legacy_files.accumulation_kernels_2d as pic_ker_2d diff --git a/src/struphy/pic/tests/test_pic_legacy_files/pusher.py b/src/struphy/pic/tests/test_pic_legacy_files/pusher.py index b609b69d8..518e19ee0 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/pusher.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/pusher.py @@ -1,7 +1,8 @@ +import cunumpy as xp + import struphy.pic.tests.test_pic_legacy_files.pusher_pos as push_pos import struphy.pic.tests.test_pic_legacy_files.pusher_vel_2d as push_vel_2d import struphy.pic.tests.test_pic_legacy_files.pusher_vel_3d as push_vel_3d -from struphy.utils.arrays import xp class Pusher: diff --git a/src/struphy/pic/tests/test_pushers.py b/src/struphy/pic/tests/test_pushers.py index 4809b7918..d64076cd1 100644 --- a/src/struphy/pic/tests/test_pushers.py +++ b/src/struphy/pic/tests/test_pushers.py @@ -23,6 +23,7 @@ ], ) def test_push_vxb_analytic(Nel, p, spl_kind, mapping, show_plots=False): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space @@ -34,7 +35,6 @@ def test_push_vxb_analytic(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -158,6 +158,7 @@ def test_push_vxb_analytic(Nel, p, spl_kind, mapping, show_plots=False): ], ) def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space @@ -169,7 +170,6 @@ def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -304,6 +304,7 @@ def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): ], ) def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space @@ -315,7 +316,6 @@ def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -450,6 +450,7 @@ def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): ], ) def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space @@ -461,7 +462,6 @@ def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -596,6 +596,7 @@ def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): ], ) def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space @@ -607,7 +608,6 @@ def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -744,6 +744,7 @@ def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): ], ) def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space @@ -756,7 +757,6 @@ def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): from struphy.pic.pushing.pusher import Pusher as Pusher_psy from struphy.pic.tests.test_pic_legacy_files.pusher import Pusher as Pusher_str from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index 244427209..337dfaede 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -1,5 +1,6 @@ from time import time +import cunumpy as xp import pytest from psydac.ddm.mpi import mpi as MPI @@ -7,7 +8,6 @@ from struphy.geometry import domains from struphy.pic.particles import Particles6D from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters -from struphy.utils.arrays import xp @pytest.mark.parametrize("nx", [8, 70]) diff --git a/src/struphy/pic/tests/test_sph.py b/src/struphy/pic/tests/test_sph.py index 8e7b5b5a4..64b7b0cc3 100644 --- a/src/struphy/pic/tests/test_sph.py +++ b/src/struphy/pic/tests/test_sph.py @@ -1,3 +1,4 @@ +import cunumpy as xp import pytest from matplotlib import pyplot as plt from psydac.ddm.mpi import MockComm @@ -8,7 +9,6 @@ from struphy.initial import perturbations from struphy.pic.particles import ParticlesSPH from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters -from struphy.utils.arrays import xp @pytest.mark.parametrize("boxes_per_dim", [(24, 1, 1)]) diff --git a/src/struphy/pic/tests/test_tesselation.py b/src/struphy/pic/tests/test_tesselation.py index 19a02cfad..cf6ed922e 100644 --- a/src/struphy/pic/tests/test_tesselation.py +++ b/src/struphy/pic/tests/test_tesselation.py @@ -1,5 +1,6 @@ from time import time +import cunumpy as xp import pytest from matplotlib import pyplot as plt from psydac.ddm.mpi import mpi as MPI @@ -10,7 +11,6 @@ from struphy.initial import perturbations from struphy.pic.particles import ParticlesSPH from struphy.pic.utilities import BoundaryParameters, LoadingParameters, WeightsParameters -from struphy.utils.arrays import xp @pytest.mark.parametrize("ppb", [8, 12]) diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index fb8fda634..3ae645557 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -1,3 +1,5 @@ +import cunumpy as xp + import struphy.pic.utilities_kernels as utils from struphy.io.options import ( OptsLoading, @@ -5,7 +7,6 @@ OptsRecontructBC, OptsSpatialLoading, ) -from struphy.utils.arrays import xp class LoadingParameters: diff --git a/src/struphy/polar/basic.py b/src/struphy/polar/basic.py index 49dc21a52..6a1c2c2f2 100644 --- a/src/struphy/polar/basic.py +++ b/src/struphy/polar/basic.py @@ -1,10 +1,9 @@ +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from psydac.linalg.basic import Vector, VectorSpace from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector -from struphy.utils.arrays import xp - class PolarDerhamSpace(VectorSpace): """ diff --git a/src/struphy/polar/extraction_operators.py b/src/struphy/polar/extraction_operators.py index c4588e997..f58c88f59 100644 --- a/src/struphy/polar/extraction_operators.py +++ b/src/struphy/polar/extraction_operators.py @@ -1,4 +1,4 @@ -from struphy.utils.arrays import xp +import cunumpy as xp # ============================= 2D polar splines (C1) =================================== diff --git a/src/struphy/polar/linear_operators.py b/src/struphy/polar/linear_operators.py index 0412bdba3..03a3e504a 100644 --- a/src/struphy/polar/linear_operators.py +++ b/src/struphy/polar/linear_operators.py @@ -1,3 +1,4 @@ +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector, BlockVectorSpace from psydac.linalg.stencil import StencilVector, StencilVectorSpace @@ -6,7 +7,6 @@ from struphy.feec.linear_operators import LinOpWithTransp from struphy.linear_algebra.linalg_kron import kron_matvec_2d from struphy.polar.basic import PolarDerhamSpace, PolarVector -from struphy.utils.arrays import xp class PolarExtractionOperator(LinOpWithTransp): diff --git a/src/struphy/polar/tests/test_legacy_polar_splines.py b/src/struphy/polar/tests/test_legacy_polar_splines.py index f4c47060a..be2bfb654 100644 --- a/src/struphy/polar/tests/test_legacy_polar_splines.py +++ b/src/struphy/polar/tests/test_legacy_polar_splines.py @@ -7,12 +7,12 @@ def test_polar_splines_2D(plot=False): sys.path.append("..") + import cunumpy as xp import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space from struphy.geometry import domains - from struphy.utils.arrays import xp # parameters # number of elements (number of elements in angular direction must be a multiple of 3) diff --git a/src/struphy/polar/tests/test_polar.py b/src/struphy/polar/tests/test_polar.py index a72abca4c..ac0113c4f 100644 --- a/src/struphy/polar/tests/test_polar.py +++ b/src/struphy/polar/tests/test_polar.py @@ -167,6 +167,7 @@ def test_spaces(Nel, p, spl_kind): @pytest.mark.parametrize("p", [[3, 2, 2]]) @pytest.mark.parametrize("spl_kind", [[False, True, True], [False, True, False]]) def test_extraction_ops_and_derivatives(Nel, p, spl_kind): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space @@ -176,7 +177,6 @@ def test_extraction_ops_and_derivatives(Nel, p, spl_kind): from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() @@ -302,12 +302,12 @@ def test_extraction_ops_and_derivatives(Nel, p, spl_kind): @pytest.mark.parametrize("p", [[4, 3, 2]]) @pytest.mark.parametrize("spl_kind", [[False, True, True], [False, True, False]]) def test_projectors(Nel, p, spl_kind): + import cunumpy as xp from psydac.ddm.mpi import mpi as MPI from struphy.eigenvalue_solvers.spline_space import Spline_space_1d, Tensor_spline_space from struphy.feec.psydac_derham import Derham from struphy.geometry.domains import IGAPolarCylinder - from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/post_processing/likwid/plot_likwidproject.py b/src/struphy/post_processing/likwid/plot_likwidproject.py index 0e8abb23b..feda5d3b6 100644 --- a/src/struphy/post_processing/likwid/plot_likwidproject.py +++ b/src/struphy/post_processing/likwid/plot_likwidproject.py @@ -7,6 +7,7 @@ import re import sys +import cunumpy as xp import matplotlib.pyplot as plt import pandas as pd import plotly.express as px @@ -16,7 +17,6 @@ import struphy.post_processing.likwid.likwid_parser as lp import struphy.post_processing.likwid.maxplotlylib as mply import struphy.post_processing.likwid.roofline_plotter as rp -from struphy.utils.arrays import xp def clean_string(string_in): diff --git a/src/struphy/post_processing/likwid/plot_time_traces.py b/src/struphy/post_processing/likwid/plot_time_traces.py index 94b491d03..ed0a34010 100644 --- a/src/struphy/post_processing/likwid/plot_time_traces.py +++ b/src/struphy/post_processing/likwid/plot_time_traces.py @@ -2,12 +2,12 @@ import pickle import re +import cunumpy as xp import matplotlib.pyplot as plt import plotly.io as pio # pio.kaleido.scope.mathjax = None import struphy.post_processing.likwid.maxplotlylib as mply -from struphy.utils.arrays import xp def glob_to_regex(pat: str) -> str: diff --git a/src/struphy/post_processing/likwid/roofline_plotter.py b/src/struphy/post_processing/likwid/roofline_plotter.py index abc6bee6f..3a4808bdc 100644 --- a/src/struphy/post_processing/likwid/roofline_plotter.py +++ b/src/struphy/post_processing/likwid/roofline_plotter.py @@ -1,11 +1,10 @@ import glob import pickle +import cunumpy as xp import pandas as pd import yaml -from struphy.utils.arrays import xp - def sort_by_num_threads(bm): sorted_arrays = {} diff --git a/src/struphy/post_processing/orbits/orbits_tools.py b/src/struphy/post_processing/orbits/orbits_tools.py index 02776b74f..97eee89af 100644 --- a/src/struphy/post_processing/orbits/orbits_tools.py +++ b/src/struphy/post_processing/orbits/orbits_tools.py @@ -1,12 +1,12 @@ import os import shutil +import cunumpy as xp import h5py import yaml from tqdm import tqdm from struphy.post_processing.orbits.orbits_kernels import calculate_guiding_center_from_6d -from struphy.utils.arrays import xp def post_process_orbit_guiding_center(path_in, path_kinetics_species, species): diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index d2071e089..74a6288f6 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -2,6 +2,7 @@ import pickle import shutil +import cunumpy as xp import h5py import yaml from tqdm import tqdm @@ -19,7 +20,6 @@ from struphy.models.species import ParticleSpecies from struphy.models.variables import PICVariable from struphy.topology.grids import TensorProductGrid -from struphy.utils.arrays import xp class ParamsIn: diff --git a/src/struphy/post_processing/pproc_struphy.py b/src/struphy/post_processing/pproc_struphy.py index 38dea12ba..044ec0de8 100644 --- a/src/struphy/post_processing/pproc_struphy.py +++ b/src/struphy/post_processing/pproc_struphy.py @@ -2,13 +2,13 @@ import pickle import shutil +import cunumpy as xp import h5py import yaml import struphy.post_processing.orbits.orbits_tools as orbits_pproc import struphy.post_processing.post_processing_tools as pproc from struphy.io.setup import import_parameters_py -from struphy.utils.arrays import xp def main( diff --git a/src/struphy/post_processing/profile_struphy.py b/src/struphy/post_processing/profile_struphy.py index 578e5c2e4..da4632555 100644 --- a/src/struphy/post_processing/profile_struphy.py +++ b/src/struphy/post_processing/profile_struphy.py @@ -1,11 +1,11 @@ import pickle import sys +import cunumpy as xp import yaml from matplotlib import pyplot as plt from struphy.post_processing.cprofile_analyser import get_cprofile_data, replace_keys -from struphy.utils.arrays import xp def main(): diff --git a/src/struphy/profiling/profiling.py b/src/struphy/profiling/profiling.py index 1ca22fe31..e96749614 100644 --- a/src/struphy/profiling/profiling.py +++ b/src/struphy/profiling/profiling.py @@ -17,10 +17,9 @@ # Import the profiling configuration class and context manager from functools import lru_cache +import cunumpy as xp from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp - @lru_cache(maxsize=None) # Cache the import result to avoid repeated imports def _import_pylikwid(): diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index e90382cb5..f69d4c7fe 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Literal +import cunumpy as xp from psydac.linalg.block import BlockVector from psydac.linalg.stencil import StencilVector @@ -14,7 +15,6 @@ from struphy.geometry.base import Domain from struphy.io.options import check_option from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable -from struphy.utils.arrays import xp class Propagator(metaclass=ABCMeta): diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index 8a82c084d..9d6e2fd4e 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from typing import Literal +import cunumpy as xp from line_profiler import profile from psydac.ddm.mpi import mpi as MPI from psydac.linalg.block import BlockVector @@ -28,7 +29,6 @@ from struphy.pic.pushing.pusher import Pusher from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 109638c1a..750917a95 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -5,6 +5,7 @@ from dataclasses import dataclass from typing import Callable, Literal, get_args +import cunumpy as xp import scipy as sc from line_profiler import profile from matplotlib import pyplot as plt @@ -67,7 +68,6 @@ from struphy.pic.particles import Particles5D, Particles6D from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index fcd091c6d..69aa86159 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Callable, Literal, get_args +import cunumpy as xp from line_profiler import profile from numpy import array, polynomial, random from psydac.ddm.mpi import mpi as MPI @@ -30,7 +31,6 @@ from struphy.pic.pushing.pusher import Pusher from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel @@ -100,7 +100,7 @@ def allocate(self): # define algorithm butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp + import cunumpy as xp butcher._a = xp.diag(butcher.a, k=-1) butcher._a = xp.array(list(butcher.a) + [0.0]) @@ -842,7 +842,7 @@ def allocate(self): else: butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp + import cunumpy as xp butcher._a = xp.diag(butcher.a, k=-1) butcher._a = xp.array(list(butcher.a) + [0.0]) @@ -1290,7 +1290,7 @@ def allocate(self): else: butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp + import cunumpy as xp butcher._a = xp.diag(butcher.a, k=-1) butcher._a = xp.array(list(butcher.a) + [0.0]) @@ -1432,7 +1432,7 @@ def allocate(self): # choose algorithm self._butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp + import cunumpy as xp self._butcher._a = xp.diag(self._butcher.a, k=-1) self._butcher._a = xp.array(list(self._butcher.a) + [0.0]) @@ -1566,7 +1566,7 @@ def allocate(self): self._butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: - from struphy.utils.arrays import xp + import cunumpy as xp self._butcher._a = xp.diag(self._butcher.a, k=-1) self._butcher._a = xp.array(list(self._butcher.a) + [0.0]) diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index b1e3cfe33..ce29c8141 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -1,3 +1,4 @@ +import cunumpy as xp import matplotlib.pyplot as plt import pytest from psydac.ddm.mpi import mpi as MPI @@ -11,7 +12,6 @@ from struphy.models.variables import FEECVariable from struphy.propagators.base import Propagator from struphy.propagators.propagators_fields import ImplicitDiffusion -from struphy.utils.arrays import xp comm = MPI.COMM_WORLD rank = comm.Get_rank() diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 659126d14..78567f8d1 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -1,3 +1,4 @@ +import cunumpy as xp import matplotlib.pyplot as plt import pytest from psydac.ddm.mpi import mpi as MPI @@ -22,7 +23,6 @@ ) from struphy.propagators.base import Propagator from struphy.propagators.propagators_fields import ImplicitDiffusion, Poisson -from struphy.utils.arrays import xp from struphy.utils.pyccel import Pyccelkernel comm = MPI.COMM_WORLD diff --git a/src/struphy/utils/arrays.py b/src/struphy/utils/arrays.py deleted file mode 100644 index ae2b9ef82..000000000 --- a/src/struphy/utils/arrays.py +++ /dev/null @@ -1,59 +0,0 @@ -import os -from types import ModuleType -from typing import TYPE_CHECKING, Literal - -BackendType = Literal["numpy", "cupy"] - - -class ArrayBackend: - def __init__( - self, - backend: BackendType = "numpy", - verbose: bool = False, - ) -> None: - assert backend.lower() in ["numpy", "cupy"], "Array backend must be either 'numpy' or 'cupy'." - - self._backend: BackendType = "cupy" if backend.lower() == "cupy" else "numpy" - - # Import numpy/cupy - if self.backend == "cupy": - try: - import cupy as cp - - self._xp = cp - except ImportError: - if verbose: - print("CuPy not available.") - self._backend = "numpy" - - if self.backend == "numpy": - import numpy as np - - self._xp = np - - assert isinstance(self.xp, ModuleType) - - if verbose: - print(f"Using {self.xp.__name__} backend.") - - @property - def backend(self) -> BackendType: - return self._backend - - @property - def xp(self) -> ModuleType: - return self._xp - - -# TODO: Make this configurable via environment variable or config file. -array_backend = ArrayBackend( - backend="cupy" if os.getenv("ARRAY_BACKEND", "numpy").lower() == "cupy" else "numpy", - verbose=True, -) - -# TYPE_CHECKING is True when type checking (e.g., mypy), but False at runtime. -# This allows us to use autocompletion for xp (i.e., numpy/cupy) as if numpy was imported. -if TYPE_CHECKING: - import numpy as xp -else: - xp = array_backend.xp diff --git a/src/struphy/utils/clone_config.py b/src/struphy/utils/clone_config.py index b292331b0..d23bee47c 100644 --- a/src/struphy/utils/clone_config.py +++ b/src/struphy/utils/clone_config.py @@ -1,8 +1,7 @@ +import cunumpy as xp from psydac.ddm.mpi import MockComm from psydac.ddm.mpi import mpi as MPI -from struphy.utils.arrays import xp - class CloneConfig: """ diff --git a/src/struphy/utils/cupy_vs_numpy.py b/src/struphy/utils/cupy_vs_numpy.py index d32044a58..43f214d73 100644 --- a/src/struphy/utils/cupy_vs_numpy.py +++ b/src/struphy/utils/cupy_vs_numpy.py @@ -1,6 +1,6 @@ import time -from arrays import xp +import cunumpy as xp def main(N=8192): diff --git a/src/struphy/utils/utils.py b/src/struphy/utils/utils.py index f9898d36f..1707f0c07 100644 --- a/src/struphy/utils/utils.py +++ b/src/struphy/utils/utils.py @@ -61,7 +61,7 @@ def save_state(state, libpath=STRUPHY_LIBPATH): def print_all_attr(obj): """Print all object's attributes that do not start with "_" to screen.""" - from struphy.utils.arrays import xp + import cunumpy as xp for k in dir(obj): if k[0] != "_": From db8623c1332f1d66c0323469abdc4f11660e4e55 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 22 Oct 2025 17:29:02 +0000 Subject: [PATCH 159/292] Update tutorials --- src/struphy/fields_background/base.py | 42 +-- .../test_verif_VlasovAmpereOneSpecies.py | 2 +- src/struphy/pic/particles.py | 2 +- tutorials/tutorial_01_parameter_files.ipynb | 4 +- tutorials/tutorial_02_test_particles.ipynb | 4 +- ...l_03_smoothed_particle_hydrodynamics.ipynb | 116 +++---- ...ipynb => tutorial_04_vlasov_maxwell.ipynb} | 13 +- .../tutorial_05_mapped_domains.ipynb | 9 +- .../tutorial_06_mhd_equilibria.ipynb | 2 +- .../tutorial_07_data_structures.ipynb | 321 ++---------------- 10 files changed, 108 insertions(+), 407 deletions(-) rename tutorials/{tutorial_05_vlasov_maxwell.ipynb => tutorial_04_vlasov_maxwell.ipynb} (97%) rename tutorials_old/tutorial_04_mapped_domains.ipynb => tutorials/tutorial_05_mapped_domains.ipynb (98%) rename tutorials_old/tutorial_05_mhd_equilibria.ipynb => tutorials/tutorial_06_mhd_equilibria.ipynb (99%) rename tutorials_old/tutorial_11_data_structures.ipynb => tutorials/tutorial_07_data_structures.ipynb (76%) diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index c6178d579..cda9ef42d 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -960,17 +960,17 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): # poloidal plane grid fig = plt.figure(figsize=(13, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): - xp = x[:, :, int(n * jump)].squeeze() + xpp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = xp.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xpp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" else: - pc1 = xp + pc1 = xpp pc2 = yp l1 = "x" l2 = "y" @@ -1013,17 +1013,17 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): fig = plt.figure(figsize=(13, 2 * 6.5)) ax = fig.add_subplot() for m in range(2): - xp = xt[:, m, :].squeeze() + xpp = xt[:, m, :].squeeze() yp = yt[:, m, :].squeeze() zp = zt[:, m, :].squeeze() if self.domain.__class__.__name__ in torus_mappings: - tc1 = xp + tc1 = xpp tc2 = yp l1 = "x" l2 = "y" else: - tc1 = xp + tc1 = xpp tc2 = zp l1 = "x" l2 = "z" @@ -1060,17 +1060,17 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): # Jacobian determinant fig = plt.figure(figsize=(13, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): - xp = x[:, :, int(n * jump)].squeeze() + xpp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = xp.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xpp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" else: - pc1 = xp + pc1 = xpp pc2 = yp l1 = "x" l2 = "y" @@ -1090,17 +1090,17 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): # pressure fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): - xp = x[:, :, int(n * jump)].squeeze() + xpp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = xp.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xpp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" else: - pc1 = xp + pc1 = xpp pc2 = yp l1 = "x" l2 = "y" @@ -1120,17 +1120,17 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): # density fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): - xp = x[:, :, int(n * jump)].squeeze() + xpp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = xp.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xpp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" else: - pc1 = xp + pc1 = xpp pc2 = yp l1 = "x" l2 = "y" @@ -1150,17 +1150,17 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): # magnetic field strength fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): - xp = x[:, :, int(n * jump)].squeeze() + xpp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = xp.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xpp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" else: - pc1 = xp + pc1 = xpp pc2 = yp l1 = "x" l2 = "y" @@ -1180,17 +1180,17 @@ def show(self, n1=16, n2=33, n3=21, n_planes=5): # current density fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5)) for n in range(n_planes): - xp = x[:, :, int(n * jump)].squeeze() + xpp = x[:, :, int(n * jump)].squeeze() yp = y[:, :, int(n * jump)].squeeze() zp = z[:, :, int(n * jump)].squeeze() if self.domain.__class__.__name__ in torus_mappings: - pc1 = xp.sqrt(xp**2 + yp**2) + pc1 = xp.sqrt(xpp**2 + yp**2) pc2 = zp l1 = "R" l2 = "Z" else: - pc1 = xp + pc1 = xpp pc2 = yp l1 = "x" l2 = "y" diff --git a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py index 980352a04..89f99cc8e 100644 --- a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py @@ -116,7 +116,7 @@ def E_exact(t): r = 0.3677 omega = 1.4156 phi = 0.5362 - return 2 * eps**2 * xp.pi / k**2 * r**2 * xp.exp(2 * gamma * t) * xp.cos(omega * t - phi) ** 2 + return 16 * eps**2 * r**2 * xp.exp(2 * gamma * t) * 2 * xp.pi * xp.cos(omega * t - phi) ** 2 / 2 # get parameters dt = time_opts.dt diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index f662bbac2..aef8ffb1a 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -44,7 +44,7 @@ def __init__( # default number of diagnostics and auxiliary columns self._n_cols_diagnostics = kwargs.pop("n_cols_diagn", 0) self._n_cols_aux = kwargs.pop("n_cols_aux", 5) - print(kwargs.keys()) + super().__init__(**kwargs) # call projected mhd equilibrium in case of CanonicalMaxwellian diff --git a/tutorials/tutorial_01_parameter_files.ipynb b/tutorials/tutorial_01_parameter_files.ipynb index f05cbb734..7621ed6b1 100644 --- a/tutorials/tutorial_01_parameter_files.ipynb +++ b/tutorials/tutorial_01_parameter_files.ipynb @@ -121,10 +121,10 @@ "equil = None\n", "\n", "# grid\n", - "grid = None\n", + "grid = grids.TensorProductGrid()\n", "\n", "# derham options\n", - "derham_opts = None" + "derham_opts = DerhamOptions()" ] }, { diff --git a/tutorials/tutorial_02_test_particles.ipynb b/tutorials/tutorial_02_test_particles.ipynb index c65d01cbd..e95b6df81 100644 --- a/tutorials/tutorial_02_test_particles.ipynb +++ b/tutorials/tutorial_02_test_particles.ipynb @@ -142,10 +142,10 @@ "equil = None\n", "\n", "# grid\n", - "grid = None\n", + "grid = grids.TensorProductGrid()\n", "\n", "# derham options\n", - "derham_opts = None" + "derham_opts = DerhamOptions()" ] }, { diff --git a/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb b/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb index ffe72d95c..922282da1 100644 --- a/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb +++ b/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "d44f9387", + "id": "0", "metadata": {}, "source": [ "# 3 - Smoothed-particle hydrodynamics\n", @@ -22,7 +22,7 @@ }, { "cell_type": "markdown", - "id": "29c3fc2b", + "id": "1", "metadata": {}, "source": [ "## Pressure-less fluid flow in Beltrami force field\n", @@ -46,7 +46,7 @@ { "cell_type": "code", "execution_count": null, - "id": "bb4344bf", + "id": "2", "metadata": {}, "outputs": [], "source": [ @@ -72,7 +72,7 @@ }, { "cell_type": "markdown", - "id": "0f9a7269", + "id": "3", "metadata": {}, "source": [ "We shall use the Strang splitting algorithm for the propagators, and simulate in Cartesian geometry (Cuboid):" @@ -81,7 +81,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9e7a378a", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -106,7 +106,7 @@ }, { "cell_type": "markdown", - "id": "3b88acc5", + "id": "5", "metadata": {}, "source": [ "The Beltrami flow can be specified as a lambda function and then passed to `GenericCartesianFluidEquilibrium`. This fluid equilibirum will be set as initial condition below." @@ -115,7 +115,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11c45bd8", + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -138,7 +138,7 @@ }, { "cell_type": "markdown", - "id": "45d672b4", + "id": "7", "metadata": {}, "source": [ "We will also need a grid and Derham complex for projecting the fluid equilibirum onto a spline basis:" @@ -147,7 +147,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d166e75d", + "id": "8", "metadata": {}, "outputs": [], "source": [ @@ -163,7 +163,7 @@ }, { "cell_type": "markdown", - "id": "9e7ecca1", + "id": "9", "metadata": {}, "source": [ "In the next step, the light-weight instance of the model is created. We set the parameter `epsilon` to 1.0, appearing in the propagator [PushVinEfield](https://struphy.pages.mpcdf.de/struphy/sections/subsections/propagators_markers.html#struphy.propagators.propagators_markers.PushVinEfield). Moreover, we launch with 1000 particles and save 100 % percent of them through `n_markers=1.0`." @@ -172,7 +172,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0d22885b", + "id": "10", "metadata": {}, "outputs": [], "source": [ @@ -195,7 +195,7 @@ }, { "cell_type": "markdown", - "id": "3f577b80", + "id": "11", "metadata": {}, "source": [ "We now set the propagator options. Here, it is important to pass the pressure from the Beltrami flow as auxiliary field to `push_v`:" @@ -204,7 +204,7 @@ { "cell_type": "code", "execution_count": null, - "id": "263332e0", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -219,7 +219,7 @@ }, { "cell_type": "markdown", - "id": "a2fa6f8b", + "id": "13", "metadata": {}, "source": [ "The initial condition of the species is defined by a background function, plus possible perturbations, which we ignore here. The background of the speceis is set to the Beltrami flow:" @@ -228,7 +228,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d658ceb0", + "id": "14", "metadata": {}, "outputs": [], "source": [ @@ -238,7 +238,7 @@ }, { "cell_type": "markdown", - "id": "d6504849", + "id": "15", "metadata": {}, "source": [ "Let us start the run, post-process the raw data, load the post-processed data and plot the particle trajectories:" @@ -247,7 +247,7 @@ { "cell_type": "code", "execution_count": null, - "id": "be28bb9a", + "id": "16", "metadata": {}, "outputs": [], "source": [ @@ -269,7 +269,7 @@ { "cell_type": "code", "execution_count": null, - "id": "0eadf863", + "id": "17", "metadata": {}, "outputs": [], "source": [ @@ -282,7 +282,7 @@ { "cell_type": "code", "execution_count": null, - "id": "de74bc5d", + "id": "18", "metadata": {}, "outputs": [], "source": [ @@ -292,7 +292,7 @@ { "cell_type": "code", "execution_count": null, - "id": "1f18caa7", + "id": "19", "metadata": {}, "outputs": [], "source": [ @@ -307,7 +307,7 @@ " [-1.0, 0.0, +1.0])\n", "\n", "dt = time_opts.dt\n", - "Nt = simdata.time_grid_size - 1\n", + "Nt = simdata.t_grid.size - 1\n", "interval = Nt/20\n", "plot_ct = 0\n", "for i in range(Nt):\n", @@ -329,7 +329,7 @@ }, { "cell_type": "markdown", - "id": "d3eafa60", + "id": "20", "metadata": {}, "source": [ "Let us perform another simulation, similar to the previous one. We will save the results in the folder `sim_2`:" @@ -338,7 +338,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9561df94", + "id": "21", "metadata": {}, "outputs": [], "source": [ @@ -348,7 +348,7 @@ }, { "cell_type": "markdown", - "id": "34422982", + "id": "22", "metadata": {}, "source": [ "This time, we shall draw the markers on a regular grid obtained from a tesselation of the domain: " @@ -357,7 +357,7 @@ { "cell_type": "code", "execution_count": null, - "id": "fe9a414c", + "id": "23", "metadata": {}, "outputs": [], "source": [ @@ -382,7 +382,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d3354fe7", + "id": "24", "metadata": {}, "outputs": [], "source": [ @@ -398,7 +398,7 @@ { "cell_type": "code", "execution_count": null, - "id": "69ada6ed", + "id": "25", "metadata": {}, "outputs": [], "source": [ @@ -409,7 +409,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b2bd98c9", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -429,7 +429,7 @@ { "cell_type": "code", "execution_count": null, - "id": "98659a03", + "id": "27", "metadata": {}, "outputs": [], "source": [ @@ -442,7 +442,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b78574d6", + "id": "28", "metadata": {}, "outputs": [], "source": [ @@ -452,7 +452,7 @@ { "cell_type": "code", "execution_count": null, - "id": "05d19960", + "id": "29", "metadata": {}, "outputs": [], "source": [ @@ -467,7 +467,7 @@ " [-1.0, 0.0, +1.0])\n", "\n", "dt = time_opts.dt\n", - "Nt = simdata.time_grid_size - 1\n", + "Nt = simdata.t_grid.size - 1\n", "interval = Nt/20\n", "plot_ct = 0\n", "for i in range(Nt):\n", @@ -489,7 +489,7 @@ }, { "cell_type": "markdown", - "id": "99087247", + "id": "30", "metadata": {}, "source": [ "## Gas expansion\n", @@ -537,7 +537,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a6f19621", + "id": "31", "metadata": {}, "outputs": [], "source": [ @@ -561,7 +561,7 @@ }, { "cell_type": "markdown", - "id": "8245c936", + "id": "32", "metadata": {}, "source": [ "We start with the generic imports and also import the model `EulerSPH`:" @@ -570,7 +570,7 @@ { "cell_type": "code", "execution_count": null, - "id": "77c93c55", + "id": "33", "metadata": {}, "outputs": [], "source": [ @@ -596,7 +596,7 @@ }, { "cell_type": "markdown", - "id": "8fa5cb3c", + "id": "34", "metadata": {}, "source": [ "Here, it is important to set the base unit `kBT` in order to derive the velocity unit:" @@ -605,7 +605,7 @@ { "cell_type": "code", "execution_count": null, - "id": "405a42f0", + "id": "35", "metadata": {}, "outputs": [], "source": [ @@ -630,7 +630,7 @@ }, { "cell_type": "markdown", - "id": "a2a23116", + "id": "36", "metadata": {}, "source": [ "As background, which goes into the initial condition below, we define a Gaussian blob in the xy-plane:" @@ -639,7 +639,7 @@ { "cell_type": "code", "execution_count": null, - "id": "c56515fb", + "id": "37", "metadata": {}, "outputs": [], "source": [ @@ -655,7 +655,7 @@ }, { "cell_type": "markdown", - "id": "814c12cb", + "id": "38", "metadata": {}, "source": [ "We also need a grid and Derham complex for projecting the fluid background on a spline basis:" @@ -664,7 +664,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8fffd408", + "id": "39", "metadata": {}, "outputs": [], "source": [ @@ -680,7 +680,7 @@ }, { "cell_type": "markdown", - "id": "35bc5741", + "id": "40", "metadata": {}, "source": [ "Now we create the light-weight instance of `EulerSPH`, without the optional propagator for particles in a magnetic background field. Note as well that we shall reject particles whose weight is below a certain threshold in order to save computing time:" @@ -689,7 +689,7 @@ { "cell_type": "code", "execution_count": null, - "id": "fd88cb8c", + "id": "41", "metadata": {}, "outputs": [], "source": [ @@ -713,7 +713,7 @@ }, { "cell_type": "markdown", - "id": "6f5818f5", + "id": "42", "metadata": {}, "source": [ "For visualization of the result, we want to save a binning plot and a kernel density plot (sph evaluation of the fluid density):" @@ -722,7 +722,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3be20f57", + "id": "43", "metadata": {}, "outputs": [], "source": [ @@ -742,7 +742,7 @@ }, { "cell_type": "markdown", - "id": "f1f36fc7", + "id": "44", "metadata": {}, "source": [ "We choose `gaussian_2d` as the smoothing kernel for sph evaluations during the pressure step:" @@ -751,7 +751,7 @@ { "cell_type": "code", "execution_count": null, - "id": "57b7a966", + "id": "45", "metadata": {}, "outputs": [], "source": [ @@ -765,7 +765,7 @@ }, { "cell_type": "markdown", - "id": "cab5de8f", + "id": "46", "metadata": {}, "source": [ "Now we set the initial condition and run the simulation. The time steps take long at the beginning, but get faster towards the end, when particles are spread out over the domain. The simulation launched in the console will be a lot faster than in the notebook, especially when using MPI, or compiling with GPU." @@ -774,7 +774,7 @@ { "cell_type": "code", "execution_count": null, - "id": "33f8717c", + "id": "47", "metadata": {}, "outputs": [], "source": [ @@ -785,7 +785,7 @@ { "cell_type": "code", "execution_count": null, - "id": "18ed3d41", + "id": "48", "metadata": {}, "outputs": [], "source": [ @@ -807,7 +807,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ed420df4", + "id": "49", "metadata": {}, "outputs": [], "source": [ @@ -820,7 +820,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3d1e2242", + "id": "50", "metadata": {}, "outputs": [], "source": [ @@ -829,7 +829,7 @@ }, { "cell_type": "markdown", - "id": "db6498c1", + "id": "51", "metadata": {}, "source": [ "The above output of the `simdata` object tells us where to find the post-processed simulation data:" @@ -838,7 +838,7 @@ { "cell_type": "code", "execution_count": null, - "id": "14d59054", + "id": "52", "metadata": {}, "outputs": [], "source": [ @@ -869,7 +869,7 @@ { "cell_type": "code", "execution_count": null, - "id": "d3b3befa", + "id": "53", "metadata": {}, "outputs": [], "source": [ @@ -927,12 +927,12 @@ { "cell_type": "code", "execution_count": null, - "id": "d86a7691", + "id": "54", "metadata": {}, "outputs": [], "source": [ "dt = time_opts.dt\n", - "Nt = simdata.time_grid_size - 1\n", + "Nt = simdata.t_grid.size - 1\n", "\n", "positions = orbits[:, :, :3]\n", "\n", diff --git a/tutorials/tutorial_05_vlasov_maxwell.ipynb b/tutorials/tutorial_04_vlasov_maxwell.ipynb similarity index 97% rename from tutorials/tutorial_05_vlasov_maxwell.ipynb rename to tutorials/tutorial_04_vlasov_maxwell.ipynb index fd8bcac5b..f35f0443f 100644 --- a/tutorials/tutorial_05_vlasov_maxwell.ipynb +++ b/tutorials/tutorial_04_vlasov_maxwell.ipynb @@ -4,7 +4,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# 9 - Vlasov-Ampère equations\n", + "# 4 - Vlasov-Ampère equations\n", "\n", "The equations we will solve are described in the model [VlasovAmpereOneSpecies](https://struphy.pages.mpcdf.de/struphy/sections/subsections/models_kinetic.html#struphy.models.kinetic.VlasovAmpereOneSpecies).\n", "To create the default parameter file from the console:\n", @@ -344,17 +344,6 @@ "plt.title('Initial electric field')\n", "plt.legend();" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# electric energy\n", - "\n", - "# plt.plot(time_vec, np.log(energy_E))" - ] } ], "metadata": { diff --git a/tutorials_old/tutorial_04_mapped_domains.ipynb b/tutorials/tutorial_05_mapped_domains.ipynb similarity index 98% rename from tutorials_old/tutorial_04_mapped_domains.ipynb rename to tutorials/tutorial_05_mapped_domains.ipynb index 4ef0a0ef0..f26724fdf 100644 --- a/tutorials_old/tutorial_04_mapped_domains.ipynb +++ b/tutorials/tutorial_05_mapped_domains.ipynb @@ -4,7 +4,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# 4 - Mapped domains with polar singularity\n", + "# 5 - Mapped domains with polar singularity\n", "\n", "This tutorial provides access to [Struphy domains](https://struphy.pages.mpcdf.de/struphy/sections/domains.html) which can be defined from analytical formulas or through third-party software (VMEC, GVEC, DESC, etc.).\n", "\n", @@ -408,13 +408,6 @@ "source": [ "domain.show()" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { diff --git a/tutorials_old/tutorial_05_mhd_equilibria.ipynb b/tutorials/tutorial_06_mhd_equilibria.ipynb similarity index 99% rename from tutorials_old/tutorial_05_mhd_equilibria.ipynb rename to tutorials/tutorial_06_mhd_equilibria.ipynb index ad7d89358..176374284 100644 --- a/tutorials_old/tutorial_05_mhd_equilibria.ipynb +++ b/tutorials/tutorial_06_mhd_equilibria.ipynb @@ -4,7 +4,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# 5 - MHD equilibria\n", + "# 6 - MHD equilibria\n", "\n", "This tutorial provides acces to the [Struphy MHD equlibrium interface](https://struphy.pages.mpcdf.de/struphy/sections/subsections/mhd_equils.html). We shall plot some available equilibria, in particular:\n", "\n", diff --git a/tutorials_old/tutorial_11_data_structures.ipynb b/tutorials/tutorial_07_data_structures.ipynb similarity index 76% rename from tutorials_old/tutorial_11_data_structures.ipynb rename to tutorials/tutorial_07_data_structures.ipynb index d59e3da86..62727c733 100644 --- a/tutorials_old/tutorial_11_data_structures.ipynb +++ b/tutorials/tutorial_07_data_structures.ipynb @@ -4,7 +4,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# 11 - Struphy data structures\n", + "# 7 - Struphy data structures\n", "\n", "In this tutorial we will learn about the data structures used in Struphy to store FEEC and particle variables. \n", "\n", @@ -874,7 +874,7 @@ "from struphy.pic import particles\n", "\n", "for name, obj in inspect.getmembers(particles):\n", - " if inspect.isclass(obj):\n", + " if inspect.isclass(obj) and obj.__module__ == particles.__name__:\n", " print(obj)" ] }, @@ -882,7 +882,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Next, we want to instantiate the model [Vlasov](https://struphy.pages.mpcdf.de/struphy/sections/models.html#struphy.models.toy.Vlasov) and then access its PIC data structures. For this we set the I/O paths to default and create the standard parameter files. We shall also consider the model [GuidingCenter](https://struphy.pages.mpcdf.de/struphy/sections/models.html#struphy.models.toy.GuidingCenter) as an example of a 5D PIC model:" + "We now create an instance of `Particles6D` with default parameters:" ] }, { @@ -891,108 +891,13 @@ "metadata": {}, "outputs": [], "source": [ - "!{'struphy --set-i d'}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!{'struphy params Vlasov -y'}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "!{'struphy params GuidingCenter -y'}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Let us read-in the generated parameter files:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "import struphy\n", - "import yaml\n", + "from struphy.pic.particles import Particles6D\n", + "from struphy.pic.utilities import LoadingParameters\n", "\n", - "vl_name = os.path.join(struphy.__path__[0], 'io/inp', 'params_Vlasov.yml')\n", - "dk_name = os.path.join(struphy.__path__[0], 'io/inp', 'params_GuidingCenter.yml')\n", - "\n", - "with open(vl_name) as file:\n", - " vl_params = yaml.load(file, Loader=yaml.FullLoader)\n", - " \n", - "with open(dk_name) as file:\n", - " dk_params = yaml.load(file, Loader=yaml.FullLoader)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Even before instantiating a model, we can look at its `species`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from struphy.models.toy import Vlasov, GuidingCenter\n", - "from psydac.ddm.mpi import mpi as MPI\n", + "loading_params = LoadingParameters(Np=120)\n", + "particles = Particles6D(loading_params=loading_params)\n", "\n", - "print(Vlasov.species()['kinetic'])\n", - "print(GuidingCenter.species()['kinetic'])" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now create an instance of `Vlasov` and initialize the variables (here just `ions`) from the default parameter file:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "comm = MPI.COMM_WORLD\n", - "\n", - "vlasov = Vlasov(vl_params, comm)\n", - "vlasov.initialize_from_params()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The screen output shows some information on the current instance of `Vlasov`. The corresponding instance of the `Particles6D` class can be accessed via the `pointer` attribute:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "p6d = vlasov.pointer['ions']\n", - "print(type(p6d))" + "particles.draw_markers()" ] }, { @@ -1008,17 +913,15 @@ "metadata": {}, "outputs": [], "source": [ - "print(f'{p6d.Np = }')\n", - "print(f'{p6d.markers.shape = }')\n", - "print(f'{p6d.markers_wo_holes.shape = }')" + "print(f'{particles.Np = }')\n", + "print(f'{particles.markers.shape = }')\n", + "print(f'{particles.markers_wo_holes.shape = }')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "In the class `Particles6D`, there can be at most 16 attributes attached ed to each marker, corresponding to the number of columns. Moreover, we see that the `markers` row number is larger than the amount of particles, in order not to run out of space during communication.\n", - "\n", "Let us look at some parameters/coordinates of the first five markers:" ] }, @@ -1028,203 +931,19 @@ "metadata": {}, "outputs": [], "source": [ - "print(f'p6d.positions[:5] = \\n{p6d.positions[:5]}\\n')\n", - "print(f'p6d.velocities[:5] = \\n{p6d.velocities[:5]}\\n')\n", - "print(f'p6d.phasespace_coords[:5] = \\n{p6d.phasespace_coords[:5]}\\n')\n", - "print(f'{p6d.weights[:5] = }\\n')\n", - "print(f'{p6d.sampling_density[:5] = }\\n')\n", - "print(f'{p6d.weights0[:5] = }\\n')\n", - "print(f'{p6d.marker_ids[:5] = }')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Remark that the marker positions in Struphy are always in the unit cube `[0, 1]^3`. The above attributes of the `Particles` class have a setter method. The setter method only works for setting, for instance the `positions`, **for all markers** on the current process: " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "p6d.positions = np.zeros(p6d.positions.shape)\n", - "print(f'p6d.positions[:5] = \\n{p6d.positions[:5]}\\n')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we want to set the positions of a selected group of markers, we have to do this by hand in the `markers` array. However, the column indices for `positions` can be accessed from the `index` dictionary:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "p6d.markers[2, p6d.index['pos']] = 1.\n", - "print(f'p6d.positions[:5] = \\n{p6d.positions[:5]}\\n')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we want to set just one position coordinate of a specific particle, we can check the corresponding `index` and write directly into the markers array:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(p6d.index['pos'])\n", - "\n", - "p6d.markers[2, 1] = 2.\n", - "print(f'p6d.positions[:5] = \\n{p6d.positions[:5]}\\n')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### MPI sort markers\n", - "\n", - "Let look at the size of `Particles.markers` in a parallel run:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def show_shapes():\n", - "\n", - " import os\n", - " import struphy\n", - " import yaml\n", - " \n", - " from struphy.models.toy import Vlasov\n", - " from psydac.ddm.mpi import mpi as MPI\n", - "\n", - " vl_name = os.path.join(struphy.__path__[0], 'io/inp', 'params_Vlasov.yml')\n", - "\n", - " with open(vl_name) as file:\n", - " vl_params = yaml.load(file, Loader=yaml.FullLoader)\n", - " \n", - " comm = MPI.COMM_WORLD\n", - " rank = comm.Get_rank()\n", - "\n", - " vlasov = Vlasov(vl_params, comm)\n", - " vlasov.initialize_from_params()\n", - " \n", - " p6d = vlasov.pointer['ions']\n", - " \n", - " out = f'{rank = }, {p6d.Np = }, {p6d.markers.shape = }, {p6d.markers_wo_holes.shape = }, {p6d.n_mks_loc = }'\n", - "\n", - " #out = f'{rank = }, before update: {x0[s[0], s[1], :] = }:'\n", - "\n", - " return out\n", - "\n", - "with ipp.Cluster(engines='mpi', n=4) as rc:\n", - " view = rc.broadcast_view()\n", - " r = view.apply_sync(show_shapes)\n", - " print(\"\\n\".join(r))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here, the total number of markers drawn (on all processes) `n_mks` is 6720, the row-size of the markers array on each process is 2141, and the number of drawn markers on each process `n_mks_loc` is \n", - "\n", - "* 1691 on rank 0\n", - "* 1709 on rank 1\n", - "* 1646 on rank 2\n", - "* 1674 on rank 3\n", - "\n", - "The differences arise because of the random draw in the decomposed domain.\n", - "\n", - "Let us check how the `mpi_sort_markers` algorithm works:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def show_sorting():\n", - "\n", - " import os\n", - " import struphy\n", - " import yaml\n", - " \n", - " from struphy.models.toy import Vlasov\n", - " from psydac.ddm.mpi import mpi as MPI\n", - "\n", - " vl_name = os.path.join(struphy.__path__[0], 'io/inp', 'params_Vlasov.yml')\n", - "\n", - " with open(vl_name) as file:\n", - " vl_params = yaml.load(file, Loader=yaml.FullLoader)\n", - " \n", - " comm = MPI.COMM_WORLD\n", - " rank = comm.Get_rank()\n", - "\n", - " vlasov = Vlasov(vl_params, comm)\n", - " vlasov.initialize_from_params()\n", - " \n", - " p6d = vlasov.pointer['ions']\n", - " \n", - " out = f'\\n{rank = }, {p6d.domain_array[rank] = }, \\n\\ninitial array: p6d.markers[:5, :3] = \\n{p6d.markers[:5, :3]}'\n", - "\n", - " if rank == 0:\n", - " p6d.markers[1:3, 1] = 0.6\n", - " elif rank == 1:\n", - " p6d.markers[3, 1] = 0.1\n", - "\n", - " out += f'\\nafter update, before send: p6d.markers[:5, :3] = \\n{p6d.markers[:5, :3]}'\n", - "\n", - " p6d.mpi_sort_markers()\n", - " \n", - " out += f'\\nafter send: p6d.markers[:5, :3] = \\n{p6d.markers[:5, :3]}'\n", - "\n", - " return out\n", - "\n", - "with ipp.Cluster(engines='mpi', n=2) as rc:\n", - " view = rc.broadcast_view()\n", - " r = view.apply_sync(show_sorting)\n", - " print(\"\\n\".join(r))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here, the `domain_array` has to be read as follows: \n", - "\n", - "* In the first direction, the domain is `[0., 1.]` with 12 elements on both processes\n", - "* In the second direction, the domain is `[0, .5]` on rank 0 and `[.5, 1.]` on rank 1, both with 7 elements\n", - "* In the third direction, the domain is `[0., 1.]` with 4 elements on both processes\n", - "\n", - "Therefore, after the initial drawing, rank 0 holds only markers with $\\eta_2 < 0.5$ and rank 1 holds only markers with $\\eta_2 > 0.5$, which is confirmed by looking at the second column of the `initial_arrays`.\n", - "\n", - "We now update $\\eta_2 = 0.6$ for the 2nd and the 3rd marker on rank 0, and $\\eta_2 = 0.1$ for the 4th marker on rank 1. These markers are thus not on the correct process anymore, according to `domain_array`. \n", - "\n", - "After `mpi_sort_markers`, the 2nd and 3rd marker are sent from ranks `0 -> 1` and the 4th marker is sent from ranks `1 -> 0`. The latter fills the hole in the 2nd row on rank 0, however the hole in the 3rd row (designated with `-1.`) remains open, because there is no incoming marker to fill it. " + "print(f'{particles.positions[:5] = }\\n')\n", + "print(f'{particles.velocities[:5] = }\\n')\n", + "print(f'{particles.phasespace_coords[:5] = }\\n')\n", + "print(f'{particles.weights[:5] = }\\n')\n", + "print(f'{particles.sampling_density[:5] = }\\n')\n", + "print(f'{particles.weights0[:5] = }\\n')\n", + "print(f'{particles.marker_ids[:5] = }')" ] } ], "metadata": { "kernelspec": { - "display_name": "tenv", + "display_name": "env", "language": "python", "name": "python3" }, @@ -1238,7 +957,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.12.3" } }, "nbformat": 4, From 645207835ca031b6dc3b29c9588de68eec2d8718 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 23 Oct 2025 08:48:39 +0200 Subject: [PATCH 160/292] Added struphy[mpi] to [dev] dependencies --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 88e6ce6a7..d099ae936 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ phys = [ 'desc-opt', ] dev = [ + "struphy[mpi]", "notebook", "autopep8", "isort", From b2afd19746e8950321edfd42029882cb5a88254e Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 23 Oct 2025 09:32:26 +0000 Subject: [PATCH 161/292] Port two hybrid --- src/struphy/console/test.py | 2 +- src/struphy/models/hybrid.py | 567 ++++++++---------- src/struphy/models/tests/test_models.py | 6 +- src/struphy/pic/accumulation/accum_kernels.py | 51 -- .../propagators/propagators_coupling.py | 165 ++--- src/struphy/propagators/propagators_fields.py | 152 ++--- 6 files changed, 428 insertions(+), 515 deletions(-) diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 9e82ee146..ebcff34d4 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -19,7 +19,7 @@ def struphy_test( Test identifier: "unit", "models", "fluid", "kinetic", "hybrid", "toy", "verification" or a model name. mpi : int - Number of MPI processes used in tests (must be >1, default=2). + Number of MPI processes used in tests (default=1). with_desc : bool Whether to include DESC equilibrium in unit tests (mem consuming). diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index b7801a050..b5e7c2547 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -5,6 +5,7 @@ from struphy.models.species import FieldSpecies, FluidSpecies, ParticleSpecies from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.pic.accumulation import accum_kernels, accum_kernels_gc +from struphy.pic.accumulation.particles_to_grid import AccumulatorVector from struphy.polar.basic import PolarVector from struphy.propagators import propagators_coupling, propagators_fields, propagators_markers from struphy.utils.pyccel import Pyccelkernel @@ -68,182 +69,109 @@ class LinearMHDVlasovCC(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species - dct["em_fields"]["b_field"] = "Hdiv" - dct["fluid"]["mhd"] = {"density": "L2", "velocity": "Hdiv", "pressure": "L2"} - dct["kinetic"]["energetic_ions"] = "Particles6D" - return dct + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - @staticmethod - def bulk_species(): - return "mhd" + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.velocity = FEECVariable(space="Hdiv") + self.pressure = FEECVariable(space="L2") + self.init_variables() - @staticmethod - def velocity_scale(): - return "alfvén" + class EnergeticIons(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="Particles6D") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_fields.CurrentCoupling6DDensity: ["mhd_velocity"], - propagators_fields.ShearAlfven: ["mhd_velocity", "b_field"], - propagators_coupling.CurrentCoupling6DCurrent: ["energetic_ions", "mhd_velocity"], - propagators_markers.PushEta: ["energetic_ions"], - propagators_markers.PushVxB: ["energetic_ions"], - propagators_fields.Magnetosonic: ["mhd_density", "mhd_velocity", "mhd_pressure"], - } + ## propagators - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + class Propagators: + def __init__(self): + self.couple_dens = propagators_fields.CurrentCoupling6DDensity() + self.shear_alf = propagators_fields.ShearAlfven() + self.couple_curr = propagators_coupling.CurrentCoupling6DCurrent() + self.push_eta = propagators_markers.PushEta() + self.push_vxb = propagators_markers.PushVxB() + self.mag_sonic = propagators_fields.Magnetosonic() - # add special options - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["fluid", "mhd"], - key="u_space", - option="Hdiv", - dct=dct, - ) - return dct + ## abstract methods - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - from struphy.polar.basic import PolarVector + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + self.energetic_ions = self.EnergeticIons() - # prelim - e_ions_params = self.kinetic["energetic_ions"]["params"] + # 2. instantiate all propagators + self.propagators = self.Propagators() - # extract necessary parameters - u_space = params["fluid"]["mhd"]["options"]["u_space"] - params_alfven = params["fluid"]["mhd"]["options"]["ShearAlfven"] - params_sonic = params["fluid"]["mhd"]["options"]["Magnetosonic"] - params_eta = params["kinetic"]["energetic_ions"]["options"]["PushEta"] - params_vxb = params["kinetic"]["energetic_ions"]["options"]["PushVxB"] - params_density = params["fluid"]["mhd"]["options"]["CurrentCoupling6DDensity"] - params_current = params["kinetic"]["energetic_ions"]["options"]["CurrentCoupling6DCurrent"] + # 3. assign variables to propagators + self.propagators.couple_dens.variables.u = self.mhd.velocity - # compute coupling parameters - Ab = params["fluid"]["mhd"]["phys_params"]["A"] - Ah = params["kinetic"]["energetic_ions"]["phys_params"]["A"] - epsilon = self.equation_params["energetic_ions"]["epsilon"] + self.propagators.shear_alf.variables.u = self.mhd.velocity + self.propagators.shear_alf.variables.b = self.em_fields.b_field - if abs(epsilon - 1) < 1e-6: - epsilon = 1.0 + self.propagators.couple_curr.variables.ions = self.energetic_ions.var + self.propagators.couple_curr.variables.u = self.mhd.velocity - self._Ab = Ab - self._Ah = Ah + self.propagators.push_eta.variables.var = self.energetic_ions.var + self.propagators.push_vxb.variables.ions = self.energetic_ions.var - # add control variate to mass_ops object - if self.pointer["energetic_ions"].control_variate: - self.mass_ops.weights["f0"] = self.pointer["energetic_ions"].f0 + self.propagators.mag_sonic.variables.n = self.mhd.density + self.propagators.mag_sonic.variables.u = self.mhd.velocity + self.propagators.mag_sonic.variables.p = self.mhd.pressure - # project background magnetic field (2-form) and background pressure (3-form) - self._b_eq = self.derham.P["2"]( - [ - self.equil.b2_1, - self.equil.b2_2, - self.equil.b2_3, - ] - ) - self._p_eq = self.derham.P["3"](self.equil.p3) - self._ones = self._p_eq.space.zeros() + # define scalars for update_scalar_quantities + self.add_scalar("en_U", compute="from_field") + self.add_scalar("en_p", compute="from_field") + self.add_scalar("en_B", compute="from_field") + self.add_scalar("en_f", compute="from_particles", variable=self.energetic_ions.var) + self.add_scalar("en_tot", summands=["en_U", "en_p", "en_B", "en_f"]) + self.add_scalar("n_lost_particles", compute="from_particles", variable=self.energetic_ions.var) + + @property + def bulk_species(self): + return self.mhd + + @property + def velocity_scale(self): + return "alfvén" + def allocate_helpers(self): + self._ones = self.projected_equil.p3.space.zeros() if isinstance(self._ones, PolarVector): self._ones.tp[:] = 1.0 else: self._ones[:] = 1.0 - # set keyword arguments for propagators - if params_density["turn_off"]: - self._kwargs[propagators_fields.CurrentCoupling6DDensity] = None - else: - self._kwargs[propagators_fields.CurrentCoupling6DDensity] = { - "particles": self.pointer["energetic_ions"], - "u_space": u_space, - "b_eq": self._b_eq, - "b_tilde": self.pointer["b_field"], - "Ab": Ab, - "Ah": Ah, - "epsilon": epsilon, - "solver": params_density["solver"], - "filter": params_density["filter"], - "boundary_cut": params_density["boundary_cut"], - } - - if params_alfven["turn_off"]: - self._kwargs[propagators_fields.ShearAlfven] = None - else: - self._kwargs[propagators_fields.ShearAlfven] = { - "u_space": u_space, - "solver": params_alfven["solver"], - } - - if params_current["turn_off"]: - self._kwargs[propagators_coupling.CurrentCoupling6DCurrent] = None - else: - self._kwargs[propagators_coupling.CurrentCoupling6DCurrent] = { - "u_space": u_space, - "b_eq": self._b_eq, - "b_tilde": self.pointer["b_field"], - "Ab": Ab, - "Ah": Ah, - "epsilon": epsilon, - "solver": params_current["solver"], - "filter": params_current["filter"], - "boundary_cut": params_current["boundary_cut"], - } - - self._kwargs[propagators_markers.PushEta] = { - "algo": params_eta["algo"], - } - - self._kwargs[propagators_markers.PushVxB] = { - "algo": params_vxb["algo"], - "kappa": 1.0 / epsilon, - "b2": self.pointer["b_field"], - "b2_add": self._b_eq, - } - - if params_sonic["turn_off"]: - self._kwargs[propagators_fields.Magnetosonic] = None - else: - self._kwargs[propagators_fields.Magnetosonic] = { - "u_space": u_space, - "b": self.pointer["b_field"], - "solver": params_sonic["solver"], - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation: - self.add_scalar("en_U", compute="from_field") - self.add_scalar("en_p", compute="from_field") - self.add_scalar("en_B", compute="from_field") - self.add_scalar("en_f", compute="from_particles", species="energetic_ions") - self.add_scalar("en_tot", summands=["en_U", "en_p", "en_B", "en_f"]) - self.add_scalar("n_lost_particles", compute="from_particles", species="energetic_ions") - - # temporary vectors for scalar quantities: self._tmp = xp.empty(1, dtype=float) self._n_lost_particles = xp.empty(1, dtype=float) + # add control variate to mass_ops object + if self.energetic_ions.var.particles.control_variate: + self.mass_ops.weights["f0"] = self.energetic_ions.var.particles.f0 + + self._Ah = self.energetic_ions.mass_number + self._Ab = self.mhd.mass_number + def update_scalar_quantities(self): # perturbed fields - en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.pointer["mhd_velocity"], self.pointer["mhd_velocity"]) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) - en_p = self.pointer["mhd_pressure"].inner(self._ones) / (5 / 3 - 1) + u = self.mhd.velocity.spline.vector + p = self.mhd.pressure.spline.vector + b = self.em_fields.b_field.spline.vector + particles = self.energetic_ions.var.particles + + en_U = 0.5 * self.mass_ops.M2n.dot_inner(u, u) + en_B = 0.5 * self.mass_ops.M2.dot_inner(b, b) + en_p = p.inner(self._ones) / (5 / 3 - 1) self.update_scalar("en_U", en_U) self.update_scalar("en_B", en_B) @@ -253,12 +181,10 @@ def update_scalar_quantities(self): self._tmp[0] = ( self._Ah / self._Ab - * self.pointer["energetic_ions"] - .markers_wo_holes[:, 6] - .dot( - self.pointer["energetic_ions"].markers_wo_holes[:, 3] ** 2 - + self.pointer["energetic_ions"].markers_wo_holes[:, 4] ** 2 - + self.pointer["energetic_ions"].markers_wo_holes[:, 5] ** 2, + * particles.markers_wo_holes[:, 6].dot( + particles.markers_wo_holes[:, 3] ** 2 + + particles.markers_wo_holes[:, 4] ** 2 + + particles.markers_wo_holes[:, 5] ** 2, ) / (2) ) @@ -267,16 +193,47 @@ def update_scalar_quantities(self): self.update_scalar("en_tot", en_U + en_B + en_p + self._tmp[0]) # Print number of lost ions - self._n_lost_particles[0] = self.pointer["energetic_ions"].n_lost_markers + self._n_lost_particles[0] = particles.n_lost_markers self.update_scalar("n_lost_particles", self._n_lost_particles[0]) - if self.rank_world == 0: + if rank == 0: print( "ratio of lost particles: ", - self._n_lost_particles[0] / self.pointer["energetic_ions"].Np * 100, + self._n_lost_particles[0] / particles.Np * 100, "%", ) + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "mag_sonic.Options" in line: + new_file += [ + "model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n" + ] + elif "couple_dens.Options" in line: + new_file += [ + "model.propagators.couple_dens.options = model.propagators.couple_dens.Options(energetic_ions=model.energetic_ions.var,\n" + ] + new_file += [ + " b_tilde=model.em_fields.b_field)\n" + ] + elif "couple_curr.Options" in line: + new_file += [ + "model.propagators.couple_curr.options = model.propagators.couple_curr.Options(b_tilde=model.em_fields.b_field)\n" + ] + elif "set_save_data" in line: + new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] + new_file += ["model.energetic_ions.set_save_data(binning_plots=(binplot,))\n"] + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class LinearMHDVlasovPC(StruphyModel): r""" @@ -898,199 +855,175 @@ class ColdPlasmaVlasov(StruphyModel): 6. :class:`~struphy.propagators.propagators_coupling.VlasovAmpere` """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} + ## species - dct["em_fields"]["e_field"] = "Hcurl" - dct["em_fields"]["b_field"] = "Hdiv" - dct["fluid"]["cold_electrons"] = {"j": "Hcurl"} - dct["kinetic"]["hot_electrons"] = "Particles6D" - return dct + class EMFields(FieldSpecies): + def __init__(self): + self.e_field = FEECVariable(space="Hcurl") + self.b_field = FEECVariable(space="Hdiv") + self.phi = FEECVariable(space="H1") + self.init_variables() - @staticmethod - def bulk_species(): - return "cold_electrons" + class ThermalElectrons(FluidSpecies): + def __init__(self): + self.current = FEECVariable(space="Hcurl") + self.init_variables() - @staticmethod - def velocity_scale(): - return "light" + class HotElectrons(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="Particles6D") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_fields.Maxwell: ["e_field", "b_field"], - propagators_fields.OhmCold: ["cold_electrons_j", "e_field"], - propagators_fields.JxBCold: ["cold_electrons_j"], - propagators_markers.PushEta: ["hot_electrons"], - propagators_markers.PushVxB: ["hot_electrons"], - propagators_coupling.VlasovAmpere: ["e_field", "hot_electrons"], - } + ## propagators - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] + class Propagators: + def __init__(self): + self.maxwell = propagators_fields.Maxwell() + self.ohm = propagators_fields.OhmCold() + self.jxb = propagators_fields.JxBCold() + self.push_eta = propagators_markers.PushEta() + self.push_vxb = propagators_markers.PushVxB() + self.coupling_va = propagators_coupling.VlasovAmpere() - # add special options - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["em_fields"], - option=propagators_fields.ImplicitDiffusion, - dct=dct, - ) - return dct + ## abstract methods - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + def __init__(self): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # Get rank and size - self._rank = self.rank_world + # 1. instantiate all species + self.em_fields = self.EMFields() + self.thermal_elec = self.ThermalElectrons() + self.hot_elec = self.HotElectrons() - # prelim - hot_params = params["kinetic"]["hot_electrons"] + # 2. instantiate all propagators + self.propagators = self.Propagators() - # model parameters - self._alpha = xp.abs( - self.equation_params["cold_electrons"]["alpha"], - ) - self._epsilon_cold = self.equation_params["cold_electrons"]["epsilon"] - self._epsilon_hot = self.equation_params["hot_electrons"]["epsilon"] + # 3. assign variables to propagators + self.propagators.maxwell.variables.e = self.em_fields.e_field + self.propagators.maxwell.variables.b = self.em_fields.b_field - self._nu = hot_params["phys_params"]["Z"] / params["fluid"]["cold_electrons"]["phys_params"]["Z"] + self.propagators.ohm.variables.j = self.thermal_elec.current + self.propagators.ohm.variables.e = self.em_fields.e_field - # Initialize background magnetic field from MHD equilibrium - self._b_background = self.derham.P["2"]( - [ - self.equil.b2_1, - self.equil.b2_2, - self.equil.b2_3, - ] - ) + self.propagators.jxb.variables.j = self.thermal_elec.current - # propagator parameters - params_maxwell = params["em_fields"]["options"]["Maxwell"]["solver"] - params_ohmcold = params["fluid"]["cold_electrons"]["options"]["OhmCold"]["solver"] - params_jxbcold = params["fluid"]["cold_electrons"]["options"]["JxBCold"]["solver"] - algo_eta = params["kinetic"]["hot_electrons"]["options"]["PushEta"]["algo"] - algo_vxb = params["kinetic"]["hot_electrons"]["options"]["PushVxB"]["algo"] - params_coupling = params["em_fields"]["options"]["VlasovAmpere"]["solver"] - self._poisson_params = params["em_fields"]["options"]["ImplicitDiffusion"]["solver"] + self.propagators.push_eta.variables.var = self.hot_elec.var + self.propagators.push_vxb.variables.ions = self.hot_elec.var - # set keyword arguments for propagators - self._kwargs[propagators_fields.Maxwell] = {"solver": params_maxwell} + self.propagators.coupling_va.variables.e = self.em_fields.e_field + self.propagators.coupling_va.variables.ions = self.hot_elec.var - self._kwargs[propagators_fields.OhmCold] = { - "alpha": self._alpha, - "epsilon": self._epsilon_cold, - "solver": params_ohmcold, - } + # define scalars for update_scalar_quantities + self.add_scalar("en_E") + self.add_scalar("en_B") + self.add_scalar("en_J") + self.add_scalar("en_f", compute="from_particles", variable=self.hot_elec.var) + self.add_scalar("en_tot") - self._kwargs[propagators_fields.JxBCold] = { - "epsilon": self._epsilon_cold, - "solver": params_jxbcold, - } + # initial Poisson (not a propagator used in time stepping) + self.initial_poisson = propagators_fields.Poisson() + self.initial_poisson.variables.phi = self.em_fields.phi - self._kwargs[propagators_markers.PushEta] = {"algo": algo_eta} + @property + def bulk_species(self): + return self.thermal_elec - self._kwargs[propagators_markers.PushVxB] = { - "algo": algo_vxb, - "kappa": 1.0 / self._epsilon_hot, - "b2": self.pointer["b_field"], - "b2_add": self._b_background, - } + @property + def velocity_scale(self): + return "light" - self._kwargs[propagators_coupling.VlasovAmpere] = { - "c1": self._alpha**2 / self._epsilon_hot, - "c2": 1.0 / self._epsilon_hot, - "solver": params_coupling, - } + def allocate_helpers(self): + self._tmp = xp.empty(1, dtype=float) - # Initialize propagators used in splitting substeps - self.init_propagators() + def update_scalar_quantities(self): + # e*M1*e/2 + e = self.em_fields.e_field.spline.vector + en_E = 0.5 * self.mass_ops.M1.dot_inner(e, e) + self.update_scalar("en_E", en_E) - # Scalar variables to be saved during simulation - self.add_scalar("en_E") - self.add_scalar("en_B") - self.add_scalar("en_J") - self.add_scalar("en_f", compute="from_particles", species="hot_electrons") - self.add_scalar("en_tot") + # alpha^2 / 2 / N * sum_p w_p v_p^2 + particles = self.hot_elec.var.particles + alpha = self.hot_elec.equation_params.alpha + self._tmp[0] = ( + alpha**2 + / (2 * particles.Np) + * xp.dot( + particles.markers_wo_holes[:, 3] ** 2 + + particles.markers_wo_holes[:, 4] ** 2 + + particles.markers_wo_holes[:, 5] ** 2, + particles.markers_wo_holes[:, 6], + ) + ) + self.update_scalar("en_f", self._tmp[0]) - # temporaries - self._tmp = xp.empty(1, dtype=float) + # en_tot = en_w + en_e + self.update_scalar("en_tot", en_E + self._tmp[0]) + + def allocate_propagators(self): + """Solve initial Poisson equation. - def initialize_from_params(self): - """:meta private:""" - from psydac.linalg.stencil import StencilVector + :meta private: + """ - from struphy.pic.accumulation.particles_to_grid import AccumulatorVector + # initialize fields and particles + super().allocate_propagators() - # Initialize fields and particles - super().initialize_from_params() + if MPI.COMM_WORLD.Get_rank() == 0: + print("\nINITIAL POISSON SOLVE:") - # Accumulate charge density + # use control variate method + particles = self.hot_elec.var.particles + particles.update_weights() + + # sanity check + # self.pointer['species1'].show_distribution_function( + # [True] + [False]*5, [xp.linspace(0, 1, 32)]) + + # accumulate charge density charge_accum = AccumulatorVector( - self.pointer["hot_electrons"], + particles, "H1", - Pyccelkernel(accum_kernels.vlasov_maxwell_poisson), + Pyccelkernel(accum_kernels.charge_density_0form), self.mass_ops, self.domain.args_domain, ) - charge_accum() - # Locally subtract mean charge for solvability with periodic bc - if xp.all(charge_accum.vectors[0].space.periods): - charge_accum._vectors[0][:] -= xp.mean( - charge_accum.vectors[0].toarray()[charge_accum.vectors[0].toarray() != 0], - ) + # another sanity check: compute FE coeffs of density + # charge_accum.show_accumulated_spline_field(self.mass_ops) - # Instantiate Poisson solver - _phi = StencilVector(self.derham.Vh["0"]) - poisson_solver = propagators_fields.ImplicitDiffusion( - _phi, - sigma_1=0, - rho=self._nu * self._alpha**2 / self._epsilon_cold * charge_accum.vectors[0], - x0=self._nu * self._alpha**2 / self._epsilon_cold * charge_accum.vectors[0], - solver=self._poisson_params, - ) + alpha = self.hot_elec.equation_params.alpha + epsilon = self.hot_elec.equation_params.epsilon - # Solve with dt=1. and compute electric field - poisson_solver(1.0) - self.derham.grad.dot(-_phi, out=self.pointer["e_field"]) + self.initial_poisson.options.rho = charge_accum + self.initial_poisson.options.rho_coeffs = alpha**2 / epsilon + self.initial_poisson.allocate() - def update_scalar_quantities(self): - en_E = 0.5 * self.mass_ops.M1.dot_inner(self.pointer["e_field"], self.pointer["e_field"]) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) - en_J = ( - 0.5 - * self._alpha**2 - * self.mass_ops.M1ninv.dot_inner(self.pointer["cold_electrons_j"], self.pointer["cold_electrons_j"]) - ) - self.update_scalar("en_E", en_E) - self.update_scalar("en_B", en_B) - self.update_scalar("en_J", en_J) + # Solve with dt=1. and compute electric field + if MPI.COMM_WORLD.Get_rank() == 0: + print("\nSolving initial Poisson problem...") + self.initial_poisson(1.0) - # nu alpha^2 eps_h / eps_c / 2 / N * sum_p w_p v_p^2 - self._tmp[0] = ( - self._nu - * self._alpha**2 - * self._epsilon_hot - / self._epsilon_cold - / (2 * self.pointer["hot_electrons"].Np) - * xp.dot( - self.pointer["hot_electrons"].markers_wo_holes[:, 3] ** 2 - + self.pointer["hot_electrons"].markers_wo_holes[:, 4] ** 2 - + self.pointer["hot_electrons"].markers_wo_holes[:, 5] ** 2, - self.pointer["hot_electrons"].markers_wo_holes[:, 6], - ) - ) + phi = self.initial_poisson.variables.phi.spline.vector + self.derham.grad.dot(-phi, out=self.em_fields.e_field.spline.vector) + if MPI.COMM_WORLD.Get_rank() == 0: + print("Done.") - self.update_scalar("en_f", self._tmp[0]) + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "coupling_va.Options" in line: + new_file += [line] + new_file += ["model.initial_poisson.options = model.initial_poisson.Options()\n"] + elif "set_save_data" in line: + new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] + new_file += ["model.hot_elec.set_save_data(binning_plots=(binplot,))\n"] + else: + new_file += [line] - # en_tot = en_E + en_B + en_J + en_w - self.update_scalar("en_tot", en_E + en_B + en_J + self._tmp[0]) + with open(params_path, "w") as f: + for line in new_file: + f.write(line) diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index d684c2707..aa2af4013 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -35,7 +35,11 @@ if rank == 0: print(f"\n{kinetic_models = }") -hybrid_models = ["LinearMHDDriftkineticCC"] +hybrid_models = [ + "LinearMHDDriftkineticCC", + "LinearMHDVlasovCC", + "ColdPlasmaVlasov", +] # for name, obj in inspect.getmembers(hybrid): # if inspect.isclass(obj) and "models.hybrid" in obj.__module__: # hybrid_models += [name] diff --git a/src/struphy/pic/accumulation/accum_kernels.py b/src/struphy/pic/accumulation/accum_kernels.py index 20d807a53..ab0a49469 100644 --- a/src/struphy/pic/accumulation/accum_kernels.py +++ b/src/struphy/pic/accumulation/accum_kernels.py @@ -487,57 +487,6 @@ def linear_vlasov_ampere( # -- removed omp: #$ omp end parallel -def vlasov_maxwell_poisson( - args_markers: "MarkerArguments", - args_derham: "DerhamArguments", - args_domain: "DomainArguments", - vec: "float[:,:,:]", -): - r""" - Accumulates the charge density in V0 - - .. math:: - - \rho_p^\mu = w_p \,. - - Parameters - ---------- - - Note - ---- - The above parameter list contains only the model specific input arguments. - """ - - markers = args_markers.markers - Np = args_markers.Np - - # -- removed omp: #$ omp parallel private (ip, eta1, eta2, eta3, filling) - # -- removed omp: #$ omp for reduction ( + :vec) - for ip in range(shape(markers)[0]): - # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: - continue - - # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - - # filling = w_p - filling = markers[ip, 6] / Np - - particle_to_mat_kernels.vec_fill_b_v0( - args_derham, - eta1, - eta2, - eta3, - vec, - filling, - ) - - # -- removed omp: #$ omp end parallel - - @stack_array("dfm", "df_inv", "df_inv_t", "g_inv", "v", "df_inv_times_v", "filling_m", "filling_v") def vlasov_maxwell( args_markers: "MarkerArguments", diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index 9d6e2fd4e..76c2396ef 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -868,88 +868,107 @@ class CurrentCoupling6DCurrent(Propagator): :ref:`time_discret`: Crank-Nicolson (implicit mid-point). System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`. """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["filter"] = { - "use_filter": None, - "modes": (1), - "repeat": 1, - "alpha": 0.5, - } - dct["boundary_cut"] = { - "e1": 0.0, - "e2": 0.0, - "e3": 0.0, - } - dct["turn_off"] = False + class Variables: + def __init__(self): + self._ions: PICVariable = None + self._u: FEECVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def ions(self) -> PICVariable: + return self._ions - return dct + @ions.setter + def ions(self, new): + assert isinstance(new, PICVariable) + assert new.space in ("Particles6D") + self._ions = new - def __init__( - self, - particles: Particles6D, - u: BlockVector, - *, - u_space: str, - b_eq: BlockVector | PolarVector, - b_tilde: BlockVector | PolarVector, - Ab: int = 1, - Ah: int = 1, - epsilon: float = 1.0, - solver: dict = options(default=True)["solver"], - filter: dict = options(default=True)["filter"], - boundary_cut: dict = options(default=True)["boundary_cut"], - ): - super().__init__(particles, u) + @property + def u(self) -> FEECVariable: + return self._u - if u_space == "H1vec": - self._space_key_int = 0 - else: - self._space_key_int = int( - self.derham.space_to_form[u_space], - ) + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new - self._b_eq = b_eq - self._b_tilde = b_tilde + def __init__(self): + self.variables = self.Variables() - self._info = solver["info"] + @dataclass + class Options: + # propagator options + b_tilde: FEECVariable = None + u_space: OptsVecSpace = "Hdiv" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + filter_params: FilterParameters = None + boundary_cut: tuple = (0.0, 0.0, 0.0) + + def __post_init__(self): + # checks + check_option(self.u_space, OptsVecSpace) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + assert self.b_tilde.space == "Hdiv" + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._space_key_int = int(self.derham.space_to_form[self.options.u_space]) + + particles = self.variables.ions.particles + u = self.variables.u.spline.vector + self._b_eq = self.projected_equil.b2 + self._b_tilde = self.options.b_tilde.spline.vector + + self._info = self.options.solver_params.info if self.derham.comm is None: self._rank = 0 else: self._rank = self.derham.comm.Get_rank() + Ah = self.variables.ions.species.mass_number + Ab = self.variables.u.species.mass_number + epsilon = self.variables.ions.species.equation_params.epsilon + self._coupling_mat = Ah / Ab / epsilon**2 self._coupling_vec = Ah / Ab / epsilon self._scale_push = 1.0 / epsilon - self._boundary_cut_e1 = boundary_cut["e1"] + self._boundary_cut_e1 = self.options.boundary_cut[0] # load accumulator self._accumulator = Accumulator( particles, - u_space, + self.options.u_space, Pyccelkernel(accum_kernels.cc_lin_mhd_6d_2), self.mass_ops, self.domain.args_domain, add_vector=True, symmetry="symm", - filter_params=filter, + filter_params=self.options.filter_params, ) # if self.particles[0].control_variate: @@ -988,7 +1007,7 @@ def __init__( # self._vec3 = xp.zeros_like(self._nuh0_at_quad[0]) # FEM spaces and basis extraction operators for u and b - u_id = self.derham.space_to_form[u_space] + u_id = self.derham.space_to_form[self.options.u_space] self._EuT = self.derham.extraction_ops[u_id].transpose() self._EbT = self.derham.extraction_ops["2"].transpose() @@ -1002,15 +1021,15 @@ def __init__( self._u_avg2 = self._EuT.codomain.zeros() # load particle pusher kernel - if u_space == "Hcurl": + if self.options.u_space == "Hcurl": kernel = Pyccelkernel(pusher_kernels.push_bxu_Hcurl) - elif u_space == "Hdiv": + elif self.options.u_space == "Hdiv": kernel = Pyccelkernel(pusher_kernels.push_bxu_Hdiv) - elif u_space == "H1vec": + elif self.options.u_space == "H1vec": kernel = Pyccelkernel(pusher_kernels.push_bxu_H1vec) else: raise ValueError( - f'{u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', + f'{self.options.u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', ) # instantiate Pusher @@ -1037,10 +1056,10 @@ def __init__( _A = getattr(self.mass_ops, "M" + u_id + "n") # preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(_A) _BC = -1 / 4 * self._accumulator.operators[0] @@ -1048,17 +1067,15 @@ def __init__( self._schur_solver = SchurSolver( _A, _BC, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, ) def __call__(self, dt): # pointer to old coefficients - un = self.feec_vars[0] + particles = self.variables.ions.particles + un = self.variables.u.spline.vector # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) self._b_eq.copy(out=self._b_full1) @@ -1128,11 +1145,11 @@ def __call__(self, dt): self._pusher(self._scale_push * dt) # write new coeffs into Propagator.variables - max_du = self.feec_vars_update(un1) + max_du = self.update_feec_variables(u=un1) # update weights in case of control variate - if self.particles[0].control_variate: - self.particles[0].update_weights() + if particles.control_variate: + particles.update_weights() if self._info and self._rank == 0: print("Status for CurrentCoupling6DCurrent:", info["success"]) diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index 750917a95..f7ac3a3c3 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -1687,65 +1687,70 @@ class CurrentCoupling6DDensity(Propagator): :ref:`time_discret`: Crank-Nicolson (implicit mid-point). """ - @staticmethod - def options(default=False): - dct = {} - dct["solver"] = { - "type": [ - ("pbicgstab", "MassMatrixPreconditioner"), - ("bicgstab", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["filter"] = { - "use_filter": None, - "modes": (1), - "repeat": 1, - "alpha": 0.5, - } - dct["boundary_cut"] = { - "e1": 0.0, - "e2": 0.0, - "e3": 0.0, - } - dct["turn_off"] = False - if default: - dct = descend_options_dict(dct, []) + class Variables: + def __init__(self): + self._u: FEECVariable = None - return dct + @property + def u(self) -> FEECVariable: + return self._u - def __init__( - self, - u: BlockVector, - *, - particles: Particles6D, - u_space: str, - b_eq: BlockVector | PolarVector, - b_tilde: BlockVector | PolarVector, - Ab: int = 1, - Ah: int = 1, - epsilon: float = 1.0, - solver: dict = options(default=True)["solver"], - filter: dict = options(default=True)["filter"], - boundary_cut: dict = options(default=True)["boundary_cut"], - ): - super().__init__(u) + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new - # assert parameters and expose some quantities to self - if u_space == "H1vec": - self._space_key_int = 0 - else: - self._space_key_int = int( - self.derham.space_to_form[u_space], - ) + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # propagator options + energetic_ions: PICVariable = None + b_tilde: FEECVariable = None + u_space: OptsVecSpace = "Hdiv" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + filter_params: FilterParameters = None + boundary_cut: tuple = (0.0, 0.0, 0.0) + + def __post_init__(self): + # checks + check_option(self.u_space, OptsVecSpace) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + assert self.energetic_ions.space == "Particles6D" + assert self.b_tilde.space == "Hdiv" + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new - self._particles = particles - self._b_eq = b_eq - self._b_tilde = b_tilde + @profile + def allocate(self): + self._space_key_int = int(self.derham.space_to_form[self.options.u_space]) + + particles = self.options.energetic_ions.particles + u = self.variables.u.spline.vector + self._b_eq = self.projected_equil.b2 + self._b_tilde = self.options.b_tilde.spline.vector # if self._particles.control_variate: @@ -1774,52 +1779,57 @@ def __init__( # self._mat31 = xp.zeros_like(self._nh0_at_quad) # self._mat32 = xp.zeros_like(self._nh0_at_quad) - self._type = solver["type"][0] - self._tol = solver["tol"] - self._maxiter = solver["maxiter"] - self._info = solver["info"] - self._verbose = solver["verbose"] + self._type = self.options.solver + self._tol = self.options.solver_params.tol + self._maxiter = self.options.solver_params.maxiter + self._info = self.options.solver_params.info + self._verbose = self.options.solver_params.verbose + self._recycle = self.options.solver_params.recycle + + Ah = self.options.energetic_ions.species.mass_number + Ab = self.variables.u.species.mass_number + epsilon = self.options.energetic_ions.species.equation_params.epsilon self._coupling_const = Ah / Ab / epsilon - self._boundary_cut_e1 = boundary_cut["e1"] + self._boundary_cut_e1 = self.options.boundary_cut[0] # load accumulator self._accumulator = Accumulator( particles, - u_space, + self.options.u_space, Pyccelkernel(accum_kernels.cc_lin_mhd_6d_1), self.mass_ops, self.domain.args_domain, add_vector=False, symmetry="asym", - filter_params=filter, + filter_params=self.options.filter_params, ) # transposed extraction operator PolarVector --> BlockVector (identity map in case of no polar splines) self._E2T = self.derham.extraction_ops["2"].transpose() # mass matrix in system (M - dt/2 * A)*u^(n + 1) = (M + dt/2 * A)*u^n - u_id = self.derham.space_to_form[u_space] + u_id = self.derham.space_to_form[self.options.u_space] self._M = getattr(self.mass_ops, "M" + u_id + "n") # preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) + pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(self._M) # linear solver self._solver = inverse( self._M, - solver["type"][0], + self.options.solver, pc=pc, - x0=self.feec_vars[0], + x0=self.variables.u.spline.vector, tol=self._tol, maxiter=self._maxiter, verbose=self._verbose, - recycle=solver["recycle"], + recycle=self._recycle, ) # temporary vectors to avoid memory allocation @@ -1831,7 +1841,7 @@ def __init__( def __call__(self, dt): # pointer to old coefficients - un = self.feec_vars[0] + un = self.variables.u.spline.vector # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) self._b_eq.copy(out=self._b_full1) @@ -1894,7 +1904,7 @@ def __call__(self, dt): info = self._solver._info # write new coeffs into Propagator.variables - max_du = self.feec_vars_update(un1) + max_du = self.update_feec_variables(u=un1) if self._info and MPI.COMM_WORLD.Get_rank() == 0: print("Status for CurrentCoupling6DDensity:", info["success"]) From 94068bc3154812cfe694db062c026d31ae1081ab Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 23 Oct 2025 14:33:06 +0200 Subject: [PATCH 162/292] Added a basic template for the Github actions --- .github/actions/compile/action.yml | 19 +++ .../action.yml | 23 +++ .../clone-and-install-struphy/action.yml | 23 +++ .../actions/install/macos-latest/action.yml | 26 ++++ .../actions/install/ubuntu-latest/action.yml | 21 +++ .github/actions/tests/models/action.yml | 16 ++ .github/actions/tests/quickstart/action.yml | 17 +++ .github/actions/tests/unit/action.yml | 9 ++ .github/workflows/docs.yml | 85 +++++++++++ .github/workflows/publish.yml | 33 +++++ .github/workflows/static_analysis.yml | 137 ++++++++++++++++++ .github/workflows/testing.yml | 91 ++++++++++++ 12 files changed, 500 insertions(+) create mode 100644 .github/actions/compile/action.yml create mode 100644 .github/actions/install/clone-and-install-struphy-editable/action.yml create mode 100644 .github/actions/install/clone-and-install-struphy/action.yml create mode 100644 .github/actions/install/macos-latest/action.yml create mode 100644 .github/actions/install/ubuntu-latest/action.yml create mode 100644 .github/actions/tests/models/action.yml create mode 100644 .github/actions/tests/quickstart/action.yml create mode 100644 .github/actions/tests/unit/action.yml create mode 100644 .github/workflows/docs.yml create mode 100644 .github/workflows/publish.yml create mode 100644 .github/workflows/static_analysis.yml create mode 100644 .github/workflows/testing.yml diff --git a/.github/actions/compile/action.yml b/.github/actions/compile/action.yml new file mode 100644 index 000000000..46657900f --- /dev/null +++ b/.github/actions/compile/action.yml @@ -0,0 +1,19 @@ +name: "Compile kernels with pyccel" + +runs: + using: composite + steps: + - name: Checkout pyccel version + shell: bash + run: | + pyccel --version + + - name: Compile kernels + shell: bash + run: | + struphy compile --status + struphy compile -y --language ${{ matrix.compile-language }} || ( + echo "Initial compile failed. Removing compiled kernels and trying again..." && + struphy compile -d -y && + struphy compile -y --language ${{ matrix.compile-language }} + ) diff --git a/.github/actions/install/clone-and-install-struphy-editable/action.yml b/.github/actions/install/clone-and-install-struphy-editable/action.yml new file mode 100644 index 000000000..d69df67be --- /dev/null +++ b/.github/actions/install/clone-and-install-struphy-editable/action.yml @@ -0,0 +1,23 @@ +name: "Clone and install struphy" + +runs: + using: composite + steps: + # TODO: Remove cloning of struphy repo + # Temporary until we moved Struphy to Github + - name: Clone struphy + shell: bash + run: | + git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git + cd struphy/ + git checkout run-318-tests-with-oversubscribe + + - name: Install struphy + shell: bash + run: | + cd struphy/ + pip install --upgrade pip + pip uninstall -y gvec + pip install -e ".[dev]" + pip list + struphy -h diff --git a/.github/actions/install/clone-and-install-struphy/action.yml b/.github/actions/install/clone-and-install-struphy/action.yml new file mode 100644 index 000000000..11baef1a6 --- /dev/null +++ b/.github/actions/install/clone-and-install-struphy/action.yml @@ -0,0 +1,23 @@ +name: "Clone and install struphy" + +runs: + using: composite + steps: + # TODO: Remove cloning of struphy repo + # Temporary until we moved Struphy to Github + - name: Clone struphy + shell: bash + run: | + git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git + cd struphy/ + git checkout run-318-tests-with-oversubscribe + + - name: Install struphy + shell: bash + run: | + cd struphy/ + pip install --upgrade pip + pip uninstall -y gvec + pip install ".[phys]" + pip list + struphy -h diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml new file mode 100644 index 000000000..c14682f3a --- /dev/null +++ b/.github/actions/install/macos-latest/action.yml @@ -0,0 +1,26 @@ +name: "Install MacOS prereqs" + +runs: + using: composite + steps: + - name: Install dependencies + shell: bash + run: | + brew update + brew install python3 + brew install gcc + brew install openblas + brew install lapack + brew install open-mpi + brew install libomp + brew install git + brew install pandoc + brew install cmake + brew link --overwrite cmake + cmake --version + make -v + system_profiler SPHardwareDataType + export FC=`which gfortran` # for gvec + export CC=`which gcc` # for gvec + export CXX=`which g++` # for gvec + brew install netcdf-fortran diff --git a/.github/actions/install/ubuntu-latest/action.yml b/.github/actions/install/ubuntu-latest/action.yml new file mode 100644 index 000000000..abcaba465 --- /dev/null +++ b/.github/actions/install/ubuntu-latest/action.yml @@ -0,0 +1,21 @@ +name: "Install ubuntu prereqs" + +runs: + using: composite + steps: + - name: Install dependencies + shell: bash + run: | + sudo apt install -y software-properties-common + sudo add-apt-repository -y ppa:deadsnakes/ppa + sudo apt update -y + sudo apt install -y python3-pip + sudo apt install -y python3-venv + sudo apt install -y gfortran gcc + sudo apt install -y liblapack-dev libopenmpi-dev + sudo apt install -y libblas-dev openmpi-bin + sudo apt install -y libomp-dev libomp5 + sudo apt install -y git + sudo apt install -y pandoc + sudo apt install -y libnetcdf-dev + sudo apt install -y g++ liblapack3 cmake cmake-curses-gui zlib1g-dev libnetcdf-dev libnetcdff-dev diff --git a/.github/actions/tests/models/action.yml b/.github/actions/tests/models/action.yml new file mode 100644 index 000000000..c3fcee1bd --- /dev/null +++ b/.github/actions/tests/models/action.yml @@ -0,0 +1,16 @@ +name: "Run model tests" + +runs: + using: composite + steps: + - name: Install dependencies + shell: bash + run: | + struphy compile --status + struphy test LinearMHD + struphy test toy + struphy test models --mpi 1 + struphy test models --mpi 2 + struphy test models --mpi 3 + struphy test models --mpi 4 + struphy test verification --mpi 4 diff --git a/.github/actions/tests/quickstart/action.yml b/.github/actions/tests/quickstart/action.yml new file mode 100644 index 000000000..3845360d3 --- /dev/null +++ b/.github/actions/tests/quickstart/action.yml @@ -0,0 +1,17 @@ +name: "Run quickstart" + +runs: + using: composite + steps: + - name: Run quickstart + shell: bash + run: | + struphy -p + struphy -h + struphy params VlasovAmpereOneSpecies + ls -1a + mv params_VlasovAmpereOneSpecies.py test.py + python3 test.py + ls sim_1/ + mpirun -n 2 python3 test.py + LINE_PROFILE=1 mpirun -n 2 python3 test.py diff --git a/.github/actions/tests/unit/action.yml b/.github/actions/tests/unit/action.yml new file mode 100644 index 000000000..2d0e70a2d --- /dev/null +++ b/.github/actions/tests/unit/action.yml @@ -0,0 +1,9 @@ +name: "Run unit tests" + +runs: + using: composite + steps: + - name: Run unit tests + shell: bash + run: | + struphy test unit diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 000000000..79caa16e7 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,85 @@ +name: Deploy docs to GitHub Pages + +on: + push: + branches: ["devel", "main"] # TODO: Set to main only after release + workflow_dispatch: + +defaults: + run: + shell: bash + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: "pages" + cancel-in-progress: false + +jobs: + build-and-deploy: + runs-on: ubuntu-latest + container: + image: texlive/texlive:latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + + - name: Cache pip + uses: actions/cache@v3 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('**/pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-pip- + + - name: Install pandoc + run: | + apt-get update + apt-get install -y pandoc + + # - name: Install TeX Live with pdflatex + # run: | + # sudo apt-get update + # sudo apt-get install -y texlive-full + # sudo apt-get install -y texlive texlive-latex-base texlive-latex-extra texlive-fonts-recommended texlive-science + + - name: Check if pdflatex is available + run: | + which pdflatex + pdflatex --version + + - name: Set up virtual environment and install project + run: | + python -m venv env + source env/bin/activate + pip install --upgrade pip + pip install ".[test, dev, docs]" + + - name: Build Sphinx docs + run: | + source env/bin/activate + cd docs + make html + + - name: Setup Pages + uses: actions/configure-pages@v5 + + - name: Upload built docs + uses: actions/upload-pages-artifact@v3 + with: + path: docs/build/html/ + + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 000000000..02f9fc5ea --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,33 @@ +name: Publish Python Package + +on: + push: + branches: + - main + +jobs: + build-and-publish: + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v3 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + + - name: Install build tools + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build the package + run: python -m build + + - name: Publish to PyPI + env: + TWINE_USERNAME: __token__ # Use API token + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: twine upload dist/* diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml new file mode 100644 index 000000000..e9b165132 --- /dev/null +++ b/.github/workflows/static_analysis.yml @@ -0,0 +1,137 @@ +name: Static analysis + +on: + push: + branches: + - main + - devel + pull_request: + branches: + - main + - devel + +defaults: + run: + shell: bash + +jobs: + cloc: + runs-on: ubuntu-latest + steps: + - name: Checkout the code + uses: actions/checkout@v4 + + - name: Download and run cloc + run: | + curl -s https://raw.githubusercontent.com/AlDanial/cloc/master/cloc > cloc + chmod +x cloc + ./cloc --version + ./cloc $(git ls-files) + + check-formatting: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v5 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + # You can test your matrix by printing the current Python version + - name: Display Python version + run: python -c "import sys; print(sys.version)" + + # Cache pip dependencies + - name: Cache pip + uses: actions/cache@v3 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('**/pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-pip- + + - name: Install prerequisites (Ubuntu) + uses: ./.github/actions/install/ubuntu-latest + + - name: Install struphy + uses: ./.github/actions/install/clone-and-install-struphy-editable + + - name: Lint the repo + run: | + struphy lint all --output-format plain --verbose + + # black: + # runs-on: ubuntu-latest + # continue-on-error: true + # steps: + # - name: Checkout the code + # uses: actions/checkout@v4 + + # - name: Code formatting with black + # run: | + # pip install black "black[jupyter]" + # # black --check src/ + # # black --check tutorials/ + + # isort: + # runs-on: ubuntu-latest + # continue-on-error: true + # steps: + # - name: Checkout the code + # uses: actions/checkout@v4 + + # - name: Code formatting with isort + # run: | + # pip install isort + # # isort --check src/ + # # isort --check tutorials/ + + # mypy: + # runs-on: ubuntu-latest + # continue-on-error: true + # steps: + # - name: Checkout the code + # uses: actions/checkout@v4 + + # - name: Type checking with mypy + # run: | + # pip install mypy + # mypy src/ || true + + # prospector: + # runs-on: ubuntu-latest + # continue-on-error: true + # steps: + # - name: Checkout the code + # uses: actions/checkout@v4 + + # - name: Code analysis with prospector + # run: | + # pip install prospector + # prospector src/ || true + + # ruff: + # runs-on: ubuntu-latest + # continue-on-error: true + # steps: + # - name: Checkout the code + # uses: actions/checkout@v4 + + # - name: Linting with ruff + # run: | + # pip install ruff + # ruff check src/ || true + + # pylint: + # runs-on: ubuntu-latest + # continue-on-error: true + # steps: + # - name: Checkout the code + # uses: actions/checkout@v4 + + # - name: Linting with pylint + # run: | + # pip install pylint + # pylint src/ || true diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml new file mode 100644 index 000000000..892ee7e9c --- /dev/null +++ b/.github/workflows/testing.yml @@ -0,0 +1,91 @@ +name: Testing + +on: + push: + branches: + - main + - devel + pull_request: + branches: + - main + - devel + +jobs: + test: + runs-on: ${{ matrix.os }} + env: + OMPI_MCA_rmaps_base_oversubscribe: 1 # Linux + PRRTE_MCA_rmaps_base_oversubscribe: 1 # MacOS + strategy: + fail-fast: false + matrix: + python-version: ["3.12"] + os: [ubuntu-latest, "macos-latest"] + compile-language: ["fortran", "c"] + test-type: ["unit", "model", "quickstart"] + + steps: + # Checkout the repository + - name: Checkout code + uses: actions/checkout@v5 + + # https://docs.github.com/en/actions/tutorials/build-and-test-code/python + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + # You can test your matrix by printing the current Python version + - name: Display Python version + run: python -c "import sys; print(sys.version)" + + # Cache pip dependencies + - name: Cache pip + uses: actions/cache@v3 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('**/pyproject.toml') }} + restore-keys: | + ${{ runner.os }}-pip- + + # Install prereqs + # I don't think it's possible to use a single action for this because + # we can't use ${matrix.os} in an if statement, so we have to use two different actions. + - name: Install prerequisites (Ubuntu) + if: matrix.os == 'ubuntu-latest' + uses: ./.github/actions/install/ubuntu-latest + + - name: Install prerequisites (macOS) + if: matrix.os == 'macos-latest' + uses: ./.github/actions/install/macos-latest + + # Check that mpirun oversubscribing works, doesn't work unless OMPI_MCA_rmaps_base_oversubscribe==1 + - name: Test mpirun + run: | + echo $OMPI_MCA_rmaps_base_oversubscribe + echo $PRRTE_MCA_rmaps_base_oversubscribe + pip install mpi4py -U + which mpirun + mpirun --version + mpirun --oversubscribe --report-bindings -n 4 python -c "from mpi4py import MPI; comm=MPI.COMM_WORLD; print(f'Hello from rank {comm.Get_rank()} of {comm.Get_size()}'); assert comm.Get_size()==4" + + # Clone struphy-ci-testing + - name: Install struphy + uses: ./.github/actions/install/clone-and-install-struphy + + # Compile + - name: Compile kernels + uses: ./.github/actions/compile + + # Run tests + - name: Run unit tests + if: matrix.test-type == 'unit' + uses: ./.github/actions/tests/unit + + - name: Run model tests + if: matrix.test-type == 'model' + uses: ./.github/actions/tests/models + + - name: Run quickstart tests + if: matrix.test-type == 'quickstart' + uses: ./.github/actions/tests/quickstart From ef5cdf6d6d2e6a539465a5fe6450e064cec47d88 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 23 Oct 2025 14:39:15 +0200 Subject: [PATCH 163/292] Removed the cloning of struphy --- .../clone-and-install-struphy-editable/action.yml | 9 --------- .../actions/install/clone-and-install-struphy/action.yml | 9 --------- 2 files changed, 18 deletions(-) diff --git a/.github/actions/install/clone-and-install-struphy-editable/action.yml b/.github/actions/install/clone-and-install-struphy-editable/action.yml index d69df67be..1e815d487 100644 --- a/.github/actions/install/clone-and-install-struphy-editable/action.yml +++ b/.github/actions/install/clone-and-install-struphy-editable/action.yml @@ -3,19 +3,10 @@ name: "Clone and install struphy" runs: using: composite steps: - # TODO: Remove cloning of struphy repo - # Temporary until we moved Struphy to Github - - name: Clone struphy - shell: bash - run: | - git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git - cd struphy/ - git checkout run-318-tests-with-oversubscribe - name: Install struphy shell: bash run: | - cd struphy/ pip install --upgrade pip pip uninstall -y gvec pip install -e ".[dev]" diff --git a/.github/actions/install/clone-and-install-struphy/action.yml b/.github/actions/install/clone-and-install-struphy/action.yml index 11baef1a6..f86a066ba 100644 --- a/.github/actions/install/clone-and-install-struphy/action.yml +++ b/.github/actions/install/clone-and-install-struphy/action.yml @@ -3,19 +3,10 @@ name: "Clone and install struphy" runs: using: composite steps: - # TODO: Remove cloning of struphy repo - # Temporary until we moved Struphy to Github - - name: Clone struphy - shell: bash - run: | - git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git - cd struphy/ - git checkout run-318-tests-with-oversubscribe - name: Install struphy shell: bash run: | - cd struphy/ pip install --upgrade pip pip uninstall -y gvec pip install ".[phys]" From 23ab98c6a69814c981b355c1b72797170c496458 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 20:25:33 +0200 Subject: [PATCH 164/292] Add '-y' option to struphy params command --- .github/actions/tests/quickstart/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/tests/quickstart/action.yml b/.github/actions/tests/quickstart/action.yml index 3845360d3..b78ade5e7 100644 --- a/.github/actions/tests/quickstart/action.yml +++ b/.github/actions/tests/quickstart/action.yml @@ -8,7 +8,7 @@ runs: run: | struphy -p struphy -h - struphy params VlasovAmpereOneSpecies + struphy params VlasovAmpereOneSpecies -y ls -1a mv params_VlasovAmpereOneSpecies.py test.py python3 test.py From 574c8738943071e9bc7fbf5cf1f40b7f5fd44760 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 21:33:53 +0200 Subject: [PATCH 165/292] Added netcdf and hdf5 install --- .github/actions/install/macos-latest/action.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index c14682f3a..39ae3996c 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -15,6 +15,8 @@ runs: brew install libomp brew install git brew install pandoc + brew install netcdf + brew install hdf5 brew install cmake brew link --overwrite cmake cmake --version From d6fe949b80336381a3f15eb7cb394083ea98c276 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 21:34:33 +0200 Subject: [PATCH 166/292] Commented out the quickstart test, we can add it back after merging 318 --- .github/workflows/testing.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 892ee7e9c..aed1183b3 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -86,6 +86,6 @@ jobs: if: matrix.test-type == 'model' uses: ./.github/actions/tests/models - - name: Run quickstart tests - if: matrix.test-type == 'quickstart' - uses: ./.github/actions/tests/quickstart + #- name: Run quickstart tests + # if: matrix.test-type == 'quickstart' + # uses: ./.github/actions/tests/quickstart From 5287118ce18ab6ccab348d1770617bec5b991f63 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 21:51:41 +0200 Subject: [PATCH 167/292] Removed netcdf --- .github/actions/install/macos-latest/action.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index 39ae3996c..9b13f421f 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -15,7 +15,6 @@ runs: brew install libomp brew install git brew install pandoc - brew install netcdf brew install hdf5 brew install cmake brew link --overwrite cmake From 485a1f2ff0b215a8cf767981c902908e6d46d834 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 21:54:24 +0200 Subject: [PATCH 168/292] Install with phys,mpi --- .../install/clone-and-install-struphy-editable/action.yml | 2 +- .github/actions/install/clone-and-install-struphy/action.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/actions/install/clone-and-install-struphy-editable/action.yml b/.github/actions/install/clone-and-install-struphy-editable/action.yml index 1e815d487..49d871e15 100644 --- a/.github/actions/install/clone-and-install-struphy-editable/action.yml +++ b/.github/actions/install/clone-and-install-struphy-editable/action.yml @@ -9,6 +9,6 @@ runs: run: | pip install --upgrade pip pip uninstall -y gvec - pip install -e ".[dev]" + pip install -e ".[dev,mpi]" pip list struphy -h diff --git a/.github/actions/install/clone-and-install-struphy/action.yml b/.github/actions/install/clone-and-install-struphy/action.yml index f86a066ba..c28e6e688 100644 --- a/.github/actions/install/clone-and-install-struphy/action.yml +++ b/.github/actions/install/clone-and-install-struphy/action.yml @@ -9,6 +9,6 @@ runs: run: | pip install --upgrade pip pip uninstall -y gvec - pip install ".[phys]" + pip install ".[phys,mpi]" pip list struphy -h From e1aa86ddb058f4146f750fa2401fa336c40195e8 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 21:56:02 +0200 Subject: [PATCH 169/292] Removed quickstart from matrix --- .github/workflows/testing.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index aed1183b3..f410539f1 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -22,7 +22,7 @@ jobs: python-version: ["3.12"] os: [ubuntu-latest, "macos-latest"] compile-language: ["fortran", "c"] - test-type: ["unit", "model", "quickstart"] + test-type: ["unit", "model"] #, "quickstart"] steps: # Checkout the repository From cd71fccd69408ba44bb49d767c766d7e564013c5 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:13:11 +0200 Subject: [PATCH 170/292] Updated model test --- .github/actions/tests/models/action.yml | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/.github/actions/tests/models/action.yml b/.github/actions/tests/models/action.yml index c3fcee1bd..0a9fafc16 100644 --- a/.github/actions/tests/models/action.yml +++ b/.github/actions/tests/models/action.yml @@ -7,10 +7,10 @@ runs: shell: bash run: | struphy compile --status - struphy test LinearMHD - struphy test toy - struphy test models --mpi 1 - struphy test models --mpi 2 - struphy test models --mpi 3 - struphy test models --mpi 4 - struphy test verification --mpi 4 + struphy compile --status + struphy test models --fast + struphy test models --fast --mpi 2 + struphy test models --fast --verification --mpi 1 + struphy test models --fast --verification --mpi 4 + struphy test models --fast --verification --mpi 4 --nclones 2 + struphy test DriftKineticElectrostaticAdiabatic --mpi 2 --nclones 2 \ No newline at end of file From 513b446784dde345ac892a6fd2252f08db2a790d Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:39:58 +0200 Subject: [PATCH 171/292] Added brew install pkgconf --- .github/actions/install/macos-latest/action.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index 9b13f421f..15c5e503e 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -12,6 +12,7 @@ runs: brew install openblas brew install lapack brew install open-mpi + brew install pkgconf brew install libomp brew install git brew install pandoc From 1a5a828110d7f6b128acdf611eb1b55377c7e441 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:45:45 +0200 Subject: [PATCH 172/292] Run ruff check --select I src/**/*.py --- .github/workflows/static_analysis.yml | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index e9b165132..e879e04e0 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -112,17 +112,16 @@ jobs: # pip install prospector # prospector src/ || true - # ruff: - # runs-on: ubuntu-latest - # continue-on-error: true - # steps: - # - name: Checkout the code - # uses: actions/checkout@v4 + ruff: + runs-on: ubuntu-latest + steps: + - name: Checkout the code + uses: actions/checkout@v4 - # - name: Linting with ruff - # run: | - # pip install ruff - # ruff check src/ || true + - name: Linting with ruff + run: | + pip install ruff + ruff check --select I src/**/*.py # pylint: # runs-on: ubuntu-latest From 50734b689eaaf815428c6e317d0b1828c934e6e3 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:47:33 +0200 Subject: [PATCH 173/292] Removed pylint quote style --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 988284b70..e17c88325 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -145,7 +145,6 @@ combine_as_imports = true [tool.pylint] max-line-length = 120 -quote-style = "single" [tool.ruff] line-length = 120 From 47002ac32e411e8d6f55bc55a0806318bec86289 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:51:11 +0200 Subject: [PATCH 174/292] Formatting --- .../kinetic_extended/massless_kernels_control_variate.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_kernels_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_kernels_control_variate.py index c56b67711..28ce7bbb2 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_kernels_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_kernels_control_variate.py @@ -27,6 +27,7 @@ def uvpre( bn3: "float[:,:,:,:]", ): from numpy import empty, exp, zeros + # -- removed omp: #$ omp parallel # -- removed omp: #$ omp do private (ie1, ie2, ie3, q1, q2, q3, il1, il2, il3, value) From ede6e771ea79f13a67ef4ff2aa71525dbd320c22 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:53:12 +0200 Subject: [PATCH 175/292] Updated isort settings --- pyproject.toml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e17c88325..6141afb69 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -138,10 +138,10 @@ ignore = [ [tool.isort] profile = "black" line_length = 120 -multi_line_output = 3 -include_trailing_comma = true -force_grid_wrap = 0 -combine_as_imports = true +# multi_line_output = 3 +# include_trailing_comma = true +# force_grid_wrap = 0 +# combine_as_imports = true [tool.pylint] max-line-length = 120 From bf6f5e5c49c436d63f66b8b4035f06d2b3dc6e68 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:55:30 +0200 Subject: [PATCH 176/292] Moved src/struphy/utils/set_release_dependencies.py to utils/ --- {src/struphy/utils => utils}/set_release_dependencies.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename {src/struphy/utils => utils}/set_release_dependencies.py (100%) diff --git a/src/struphy/utils/set_release_dependencies.py b/utils/set_release_dependencies.py similarity index 100% rename from src/struphy/utils/set_release_dependencies.py rename to utils/set_release_dependencies.py From 436f042da3d5cbac724070ebdd29d91084df5dff Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:56:19 +0200 Subject: [PATCH 177/292] Formatting --- utils/set_release_dependencies.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/set_release_dependencies.py b/utils/set_release_dependencies.py index a08717060..96e29343a 100644 --- a/utils/set_release_dependencies.py +++ b/utils/set_release_dependencies.py @@ -1,8 +1,8 @@ import importlib.metadata import re +import tomllib import tomli_w -import tomllib def get_min_bound(entry): From b1bd0093c47de76f8e55692c7a362e65f7ca3632 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:57:40 +0200 Subject: [PATCH 178/292] Added isort to check tutorials --- .github/workflows/static_analysis.yml | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index e879e04e0..45b4422bc 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -75,18 +75,18 @@ jobs: # # black --check src/ # # black --check tutorials/ - # isort: - # runs-on: ubuntu-latest - # continue-on-error: true - # steps: - # - name: Checkout the code - # uses: actions/checkout@v4 + isort: + runs-on: ubuntu-latest + continue-on-error: true + steps: + - name: Checkout the code + uses: actions/checkout@v4 - # - name: Code formatting with isort - # run: | - # pip install isort - # # isort --check src/ - # # isort --check tutorials/ + - name: Code formatting with isort + run: | + pip install isort + # isort --check src/ + # isort --check doc/tutorials/ # mypy: # runs-on: ubuntu-latest From a66d0181ed936105907ac6b7802f4caae3ca0952 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 22:59:18 +0200 Subject: [PATCH 179/292] Renamed check-formatting to struphy-lint-all --- .github/workflows/static_analysis.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index 45b4422bc..c2f485b78 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -28,7 +28,7 @@ jobs: ./cloc --version ./cloc $(git ls-files) - check-formatting: + struphy lint all: runs-on: ubuntu-latest steps: - name: Checkout code @@ -85,8 +85,8 @@ jobs: - name: Code formatting with isort run: | pip install isort - # isort --check src/ - # isort --check doc/tutorials/ + isort --check src/ + isort --check doc/tutorials/ # mypy: # runs-on: ubuntu-latest From 285719a25e9d96db1304de4f40b2f10306035454 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 23 Oct 2025 23:39:05 +0200 Subject: [PATCH 180/292] Added trailing commas --- src/struphy/bsplines/bsplines_kernels.py | 16 +- src/struphy/bsplines/evaluation_kernels_1d.py | 7 +- src/struphy/bsplines/evaluation_kernels_3d.py | 289 +++- .../bsplines/tests/test_eval_spline_mpi.py | 11 +- src/struphy/console/compile.py | 18 +- src/struphy/console/format.py | 22 +- src/struphy/console/params.py | 2 +- src/struphy/console/profile.py | 4 +- src/struphy/console/run.py | 2 +- src/struphy/console/tests/test_console.py | 26 +- src/struphy/diagnostics/continuous_spectra.py | 9 +- src/struphy/diagnostics/diagn_tools.py | 4 +- .../diagnostics/paraview/mesh_creator.py | 7 +- src/struphy/dispersion_relations/analytic.py | 14 +- .../kernels_projectors_global.py | 62 +- .../legacy/MHD_eigenvalues_cylinder_1D.py | 45 +- .../fB_massless_control_variate.py | 2 +- .../fB_massless_kernels_control_variate.py | 42 +- .../fnB_massless_control_variate.py | 4 +- .../fnB_massless_kernels_control_variate.py | 56 +- .../massless_control_variate.py | 4 +- .../massless_kernels_control_variate.py | 128 +- .../legacy/emw_operators.py | 6 +- .../legacy/mass_matrices_3d_pre.py | 18 +- .../legacy/massless_operators/fB_bv_kernel.py | 24 +- .../fB_massless_linear_operators.py | 33 +- .../legacy/massless_operators/fB_vv_kernel.py | 6 +- .../pro_local/mhd_operators_3d_local.py | 116 +- .../pro_local/projectors_local.py | 48 +- .../shape_L2_projector_kernel.py | 54 +- .../shape_function_projectors_L2.py | 37 +- .../shape_function_projectors_local.py | 41 +- .../shape_local_projector_kernel.py | 141 +- .../eigenvalue_solvers/mass_matrices_3d.py | 9 +- .../mhd_axisymmetric_main.py | 14 +- .../mhd_axisymmetric_pproc.py | 5 +- .../eigenvalue_solvers/mhd_operators.py | 45 +- .../eigenvalue_solvers/mhd_operators_core.py | 284 ++-- .../eigenvalue_solvers/projectors_global.py | 91 +- .../eigenvalue_solvers/spline_space.py | 109 +- src/struphy/examples/_draw_parallel.py | 2 +- src/struphy/feec/basis_projection_ops.py | 14 +- src/struphy/feec/linear_operators.py | 4 +- src/struphy/feec/local_projectors_kernels.py | 20 +- src/struphy/feec/mass.py | 56 +- src/struphy/feec/mass_kernels.py | 12 +- src/struphy/feec/preconditioner.py | 2 +- src/struphy/feec/projectors.py | 27 +- src/struphy/feec/psydac_derham.py | 16 +- src/struphy/feec/tests/test_basis_ops.py | 13 +- src/struphy/feec/tests/test_derham.py | 32 +- src/struphy/feec/tests/test_eval_field.py | 18 +- src/struphy/feec/tests/test_field_init.py | 89 +- src/struphy/feec/tests/test_l2_projectors.py | 10 +- .../feec/tests/test_local_projectors.py | 46 +- .../feec/tests/test_lowdim_nel_is_1.py | 14 +- src/struphy/feec/tests/test_mass_matrices.py | 62 +- .../feec/tests/test_toarray_struphy.py | 3 +- .../feec/tests/test_tosparse_struphy.py | 23 +- src/struphy/feec/utilities.py | 12 +- .../feec/utilities_local_projectors.py | 23 +- src/struphy/feec/variational_utilities.py | 23 +- src/struphy/fields_background/base.py | 15 +- .../fields_background/coil_fields/base.py | 6 +- .../coil_fields/coil_fields.py | 20 +- src/struphy/fields_background/equils.py | 36 +- .../mhd_equil/eqdsk/readeqdsk.py | 6 +- .../tests/test_desc_equil.py | 12 +- src/struphy/geometry/base.py | 10 +- src/struphy/geometry/domains.py | 2 +- src/struphy/geometry/mappings_kernels.py | 192 ++- src/struphy/geometry/transform_kernels.py | 24 +- src/struphy/initial/eigenfunctions.py | 11 +- src/struphy/initial/perturbations.py | 2 +- .../initial/tests/test_init_perturbations.py | 18 +- src/struphy/io/output_handling.py | 8 +- src/struphy/io/setup.py | 36 +- src/struphy/kinetic_background/base.py | 6 +- src/struphy/kinetic_background/maxwellians.py | 6 +- .../tests/test_maxwellians.py | 142 +- src/struphy/linear_algebra/linalg_kron.py | 15 +- src/struphy/linear_algebra/saddle_point.py | 15 +- .../tests/test_saddle_point_propagator.py | 2 +- .../tests/test_saddlepoint_massmatrices.py | 2 +- src/struphy/main.py | 34 +- src/struphy/models/base.py | 12 +- src/struphy/models/fluid.py | 136 +- src/struphy/models/hybrid.py | 24 +- src/struphy/models/kinetic.py | 6 +- src/struphy/models/species.py | 6 +- src/struphy/models/tests/test_models.py | 10 +- .../models/tests/test_verif_EulerSPH.py | 6 +- .../models/tests/test_verif_LinearMHD.py | 6 +- .../models/tests/test_verif_Maxwell.py | 22 +- .../models/tests/test_verif_Poisson.py | 6 +- .../test_verif_VlasovAmpereOneSpecies.py | 4 +- src/struphy/models/tests/test_xxpproc.py | 2 +- src/struphy/models/tests/verification.py | 40 +- src/struphy/models/toy.py | 25 +- src/struphy/models/variables.py | 6 +- src/struphy/ode/tests/test_ode_feec.py | 12 +- .../pic/accumulation/accum_kernels_gc.py | 110 +- .../accumulation/particle_to_mat_kernels.py | 42 +- src/struphy/pic/base.py | 144 +- src/struphy/pic/particles.py | 3 +- src/struphy/pic/pushing/pusher.py | 18 +- src/struphy/pic/sobol_seq.py | 14 +- src/struphy/pic/sph_eval_kernels.py | 14 +- src/struphy/pic/tests/test_accum_vec_H1.py | 7 +- src/struphy/pic/tests/test_accumulation.py | 3 +- src/struphy/pic/tests/test_binning.py | 16 +- src/struphy/pic/tests/test_mat_vec_filler.py | 30 +- .../test_pic_legacy_files/accumulation.py | 12 +- .../accumulation_kernels_3d.py | 192 ++- .../test_pic_legacy_files/mappings_3d.py | 212 ++- .../test_pic_legacy_files/mappings_3d_fast.py | 305 +++- .../tests/test_pic_legacy_files/pusher_pos.py | 1257 +++++++++++++++-- .../test_pic_legacy_files/pusher_vel_2d.py | 200 ++- .../test_pic_legacy_files/pusher_vel_3d.py | 294 +++- .../spline_evaluation_3d.py | 266 +++- src/struphy/pic/tests/test_pushers.py | 24 +- src/struphy/pic/tests/test_sorting.py | 3 +- src/struphy/pic/tests/test_sph.py | 66 +- src/struphy/pic/tests/test_tesselation.py | 2 +- src/struphy/polar/basic.py | 2 +- src/struphy/polar/extraction_operators.py | 11 +- src/struphy/polar/linear_operators.py | 9 +- .../likwid/plot_likwidproject.py | 2 +- .../likwid/plot_time_traces.py | 6 +- src/struphy/post_processing/pproc_struphy.py | 26 +- .../post_processing/profile_struphy.py | 4 +- src/struphy/propagators/base.py | 4 +- .../propagators/propagators_coupling.py | 12 +- src/struphy/propagators/propagators_fields.py | 156 +- .../propagators/propagators_markers.py | 6 +- .../tests/test_gyrokinetic_poisson.py | 29 +- src/struphy/propagators/tests/test_poisson.py | 25 +- src/struphy/utils/clone_config.py | 2 +- src/struphy/utils/test_clone_config.py | 6 +- src/struphy/utils/utils.py | 6 +- 140 files changed, 5385 insertions(+), 1474 deletions(-) diff --git a/src/struphy/bsplines/bsplines_kernels.py b/src/struphy/bsplines/bsplines_kernels.py index 17374f178..9fa1a9521 100644 --- a/src/struphy/bsplines/bsplines_kernels.py +++ b/src/struphy/bsplines/bsplines_kernels.py @@ -83,7 +83,13 @@ def find_span(t: "Final[float[:]]", p: "int", eta: "float") -> "int": @pure def basis_funs( - t: "Final[float[:]]", p: "int", eta: "float", span: "int", left: "float[:]", right: "float[:]", values: "float[:]" + t: "Final[float[:]]", + p: "int", + eta: "float", + span: "int", + left: "float[:]", + right: "float[:]", + values: "float[:]", ): """ Parameters @@ -595,7 +601,13 @@ def basis_funs_and_der( @pure @stack_array("values_b") def basis_funs_1st_der( - t: "Final[float[:]]", p: "int", eta: "float", span: "int", left: "float[:]", right: "float[:]", values: "float[:]" + t: "Final[float[:]]", + p: "int", + eta: "float", + span: "int", + left: "float[:]", + right: "float[:]", + values: "float[:]", ): """ Parameters diff --git a/src/struphy/bsplines/evaluation_kernels_1d.py b/src/struphy/bsplines/evaluation_kernels_1d.py index a6ec8b7a5..6510eafff 100644 --- a/src/struphy/bsplines/evaluation_kernels_1d.py +++ b/src/struphy/bsplines/evaluation_kernels_1d.py @@ -61,7 +61,12 @@ def evaluation_kernel_1d(p1: int, basis1: "Final[float[:]]", ind1: "Final[int[:] @pure @stack_array("tmp1", "tmp2") def evaluate( - kind1: int, t1: "Final[float[:]]", p1: int, ind1: "Final[int[:,:]]", coeff: "Final[float[:]]", eta1: float + kind1: int, + t1: "Final[float[:]]", + p1: int, + ind1: "Final[int[:,:]]", + coeff: "Final[float[:]]", + eta1: float, ) -> float: """ Point-wise evaluation of a spline. diff --git a/src/struphy/bsplines/evaluation_kernels_3d.py b/src/struphy/bsplines/evaluation_kernels_3d.py index 8ccaa252b..a6c900616 100644 --- a/src/struphy/bsplines/evaluation_kernels_3d.py +++ b/src/struphy/bsplines/evaluation_kernels_3d.py @@ -246,47 +246,212 @@ def evaluate_tensor_product( for i3 in range(len(eta3)): if kind == 0: spline_values[i1, i2, i3] = evaluate_3d( - 1, 1, 1, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 1, + 1, + 1, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 11: spline_values[i1, i2, i3] = evaluate_3d( - 2, 1, 1, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 2, + 1, + 1, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 12: spline_values[i1, i2, i3] = evaluate_3d( - 1, 2, 1, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 1, + 2, + 1, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 13: spline_values[i1, i2, i3] = evaluate_3d( - 1, 1, 2, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 1, + 1, + 2, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 21: spline_values[i1, i2, i3] = evaluate_3d( - 1, 2, 2, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 1, + 2, + 2, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 22: spline_values[i1, i2, i3] = evaluate_3d( - 2, 1, 2, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 2, + 1, + 2, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 23: spline_values[i1, i2, i3] = evaluate_3d( - 2, 2, 1, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 2, + 2, + 1, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 3: spline_values[i1, i2, i3] = evaluate_3d( - 2, 2, 2, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 2, + 2, + 2, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 41: spline_values[i1, i2, i3] = evaluate_3d( - 3, 1, 1, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 3, + 1, + 1, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 42: spline_values[i1, i2, i3] = evaluate_3d( - 1, 3, 1, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 1, + 3, + 1, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 43: spline_values[i1, i2, i3] = evaluate_3d( - 1, 1, 3, t1, t2, t3, p1, p2, p3, ind1, ind2, ind3, coeff, eta1[i1], eta2[i2], eta3[i3] + 1, + 1, + 3, + t1, + t2, + t3, + p1, + p2, + p3, + ind1, + ind2, + ind3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) @@ -1051,7 +1216,17 @@ def eval_spline_mpi( b3 = bd3 value = eval_spline_mpi_kernel( - pn[0] - kind[0], pn[1] - kind[1], pn[2] - kind[2], b1, b2, b3, span1, span2, span3, _data, starts + pn[0] - kind[0], + pn[1] - kind[1], + pn[2] - kind[2], + b1, + b2, + b3, + span1, + span2, + span3, + _data, + starts, ) return value @@ -1196,7 +1371,17 @@ def eval_spline_mpi_tensor_product_fast( b3 = bd3 values[i, j, k] = eval_spline_mpi_kernel( - pn[0] - kind[0], pn[1] - kind[1], pn[2] - kind[2], b1, b2, b3, span1, span2, span3, _data, starts + pn[0] - kind[0], + pn[1] - kind[1], + pn[2] - kind[2], + b1, + b2, + b3, + span1, + span2, + span3, + _data, + starts, ) @@ -1262,7 +1447,17 @@ def eval_spline_mpi_tensor_product_fixed( b3[:] = b3s[k, :] values[i, j, k] = eval_spline_mpi_kernel( - pn[0] - kind[0], pn[1] - kind[1], pn[2] - kind[2], b1, b2, b3, span1, span2, span3, _data, starts + pn[0] - kind[0], + pn[1] - kind[1], + pn[2] - kind[2], + b1, + b2, + b3, + span1, + span2, + span3, + _data, + starts, ) @@ -1320,7 +1515,16 @@ def eval_spline_mpi_matrix( continue # point not in process domain values[i, j, k] = eval_spline_mpi( - eta1[i, j, k], eta2[i, j, k], eta3[i, j, k], _data, kind, pn, tn1, tn2, tn3, starts + eta1[i, j, k], + eta2[i, j, k], + eta3[i, j, k], + _data, + kind, + pn, + tn1, + tn2, + tn3, + starts, ) @@ -1382,7 +1586,16 @@ def eval_spline_mpi_sparse_meshgrid( continue # point not in process domain values[i, j, k] = eval_spline_mpi( - eta1[i, 0, 0], eta2[0, j, 0], eta3[0, 0, k], _data, kind, pn, tn1, tn2, tn3, starts + eta1[i, 0, 0], + eta2[0, j, 0], + eta3[0, 0, k], + _data, + kind, + pn, + tn1, + tn2, + tn3, + starts, ) @@ -1432,7 +1645,16 @@ def eval_spline_mpi_markers( continue # point not in process domain values[ip] = eval_spline_mpi( - markers[ip, 0], markers[ip, 1], markers[ip, 2], _data, kind, pn, tn1, tn2, tn3, starts + markers[ip, 0], + markers[ip, 1], + markers[ip, 2], + _data, + kind, + pn, + tn1, + tn2, + tn3, + starts, ) @@ -1453,20 +1675,39 @@ def get_spans(eta1: float, eta2: float, eta3: float, args_derham: "DerhamArgumen # get spline values at eta bsplines_kernels.b_d_splines_slim( - args_derham.tn1, args_derham.pn[0], eta1, int(span1), args_derham.bn1, args_derham.bd1 + args_derham.tn1, + args_derham.pn[0], + eta1, + int(span1), + args_derham.bn1, + args_derham.bd1, ) bsplines_kernels.b_d_splines_slim( - args_derham.tn2, args_derham.pn[1], eta2, int(span2), args_derham.bn2, args_derham.bd2 + args_derham.tn2, + args_derham.pn[1], + eta2, + int(span2), + args_derham.bn2, + args_derham.bd2, ) bsplines_kernels.b_d_splines_slim( - args_derham.tn3, args_derham.pn[2], eta3, int(span3), args_derham.bn3, args_derham.bd3 + args_derham.tn3, + args_derham.pn[2], + eta3, + int(span3), + args_derham.bn3, + args_derham.bd3, ) return span1, span2, span3 def eval_0form_spline_mpi( - span1: int, span2: int, span3: int, args_derham: "DerhamArguments", form_coeffs: "float[:,:,:]" + span1: int, + span2: int, + span3: int, + args_derham: "DerhamArguments", + form_coeffs: "float[:,:,:]", ) -> float: """Single-point evaluation of Derham 0-form spline defined by form_coeffs, given N-spline values (in bn) and knot span indices span.""" @@ -1602,7 +1843,11 @@ def eval_2form_spline_mpi( def eval_3form_spline_mpi( - span1: int, span2: int, span3: int, args_derham: "DerhamArguments", form_coeffs: "float[:,:,:]" + span1: int, + span2: int, + span3: int, + args_derham: "DerhamArguments", + form_coeffs: "float[:,:,:]", ) -> float: """Single-point evaluation of Derham 0-form spline defined by form_coeffs, given D-spline values (in bd) and knot span indices span.""" diff --git a/src/struphy/bsplines/tests/test_eval_spline_mpi.py b/src/struphy/bsplines/tests/test_eval_spline_mpi.py index 0703fb418..923fc8ea6 100644 --- a/src/struphy/bsplines/tests/test_eval_spline_mpi.py +++ b/src/struphy/bsplines/tests/test_eval_spline_mpi.py @@ -700,7 +700,10 @@ def test_eval_tensor_product_grid(Nel, p, spl_kind, n_markers=10): # Histopolation grids spaces = derham.Vh_fem["3"].spaces ptsG, wtsG, spans, bases, subs = prepare_projection_of_basis( - spaces, spaces, derham.Vh["3"].starts, derham.Vh["3"].ends + spaces, + spaces, + derham.Vh["3"].starts, + derham.Vh["3"].ends, ) eta1s = ptsG[0].flatten() eta2s = ptsG[1].flatten() @@ -715,9 +718,9 @@ def test_eval_tensor_product_grid(Nel, p, spl_kind, n_markers=10): comm.Barrier() sleep(0.02 * (rank + 1)) - print(f"rank {rank} | {eta1s = }") - print(f"rank {rank} | {eta2s = }") - print(f"rank {rank} | {eta3s = }\n") + print(f"rank {rank} | {eta1s =}") + print(f"rank {rank} | {eta2s =}") + print(f"rank {rank} | {eta3s =}\n") comm.Barrier() # compare spline evaluation routines diff --git a/src/struphy/console/compile.py b/src/struphy/console/compile.py index 0cb628b18..432e4fa1f 100644 --- a/src/struphy/console/compile.py +++ b/src/struphy/console/compile.py @@ -4,7 +4,17 @@ def struphy_compile( - language, compiler, compiler_config, omp_pic, omp_feec, delete, status, verbose, dependencies, time_execution, yes + language, + compiler, + compiler_config, + omp_pic, + omp_feec, + delete, + status, + verbose, + dependencies, + time_execution, + yes, ): """Compile Struphy kernels. All files that contain "kernels" are detected automatically and saved to state.yml. @@ -187,9 +197,9 @@ def struphy_compile( deps = depmod.get_dependencies(ker.replace(".py", so_suffix)) deps_li = deps.split(" ") print("-" * 28) - print(f"{ker = }") + print(f"{ker =}") for dep in deps_li: - print(f"{dep = }") + print(f"{dep =}") else: # struphy and psydac (change dir not to be in source path) @@ -258,7 +268,7 @@ def struphy_compile( # only install (from .whl) if psydac not up-to-date if psydac_ver < struphy_ver: print( - f"You have psydac version {psydac_ver}, but version {struphy_ver} is available. Please re-install struphy (e.g. pip install .)\n" + f"You have psydac version {psydac_ver}, but version {struphy_ver} is available. Please re-install struphy (e.g. pip install .)\n", ) sys.exit(1) else: diff --git a/src/struphy/console/format.py b/src/struphy/console/format.py index 747e2d0c1..eec22f784 100644 --- a/src/struphy/console/format.py +++ b/src/struphy/console/format.py @@ -576,7 +576,7 @@ def parse_json_file_to_html(json_file_path, html_output_path): "", "", "Code Analysis Report", - ] + ], ) # Include external CSS and JS libraries @@ -586,7 +586,7 @@ def parse_json_file_to_html(json_file_path, html_output_path): "", "", "", - ] + ], ) # Custom CSS for light mode and code prettification @@ -700,7 +700,7 @@ def parse_json_file_to_html(json_file_path, html_output_path): text-align: center; color: #999999; } -""" +""", ) html_content.append("") @@ -715,7 +715,7 @@ def parse_json_file_to_html(json_file_path, html_output_path): }); }); -""" +""", ) html_content.extend(["", "", "

Code Issues Report

"]) @@ -729,7 +729,7 @@ def parse_json_file_to_html(json_file_path, html_output_path):

Total Issues: {total_issues}

Number of files: {total_files}

-""" +""", ) # Navigation menu @@ -753,7 +753,7 @@ def parse_json_file_to_html(json_file_path, html_output_path): f"""
File: {display_name} -""" +""", ) issue_data = {} @@ -803,7 +803,7 @@ def parse_json_file_to_html(json_file_path, html_output_path): f"{code} - " f"{message}
" f"Location: " - f"{display_name}:{row}:{column}
" + f"{display_name}:{row}:{column}
", ) html_content.append("

") @@ -846,7 +846,7 @@ def parse_json_file_to_html(json_file_path, html_output_path): html_content.append( # f"
" # f"{line_number}{line_content}
" - f"{line_number}: {line_content}" + f"{line_number}: {line_content}", ) html_content.append("") # Include fix details if available @@ -854,12 +854,12 @@ def parse_json_file_to_html(json_file_path, html_output_path): html_content.append("
") html_content.append( f"

Fix Available ({fix.get('applicability', 'Unknown')}): " - f"ruff check --select ALL --fix {display_name}

" + f"ruff check --select ALL --fix {display_name}

", ) html_content.append("
") else: html_content.append( - f"

Cannot read file {filename} or invalid row {row}.

" + f"

Cannot read file {filename} or invalid row {row}.

", ) html_content.append("") @@ -873,7 +873,7 @@ def parse_json_file_to_html(json_file_path, html_output_path):

Generated by on {time.strftime("%Y-%m-%d %H:%M:%S")}

-""" +""", ) html_content.extend(["", ""]) diff --git a/src/struphy/console/params.py b/src/struphy/console/params.py index ade6c5ea5..4916cd1c0 100644 --- a/src/struphy/console/params.py +++ b/src/struphy/console/params.py @@ -29,7 +29,7 @@ def struphy_params(model_name: str, params_path: str, yes: bool = False, check_f except AttributeError: pass - print(f"{model_name = }") + print(f"{model_name =}") # print units if check_file: diff --git a/src/struphy/console/profile.py b/src/struphy/console/profile.py index 9577a00d9..a2f10ce66 100644 --- a/src/struphy/console/profile.py +++ b/src/struphy/console/profile.py @@ -106,7 +106,7 @@ def struphy_profile(dirs, replace, all, n_lines, print_callers, savefig): + "ncalls".ljust(15) + "tottime".ljust(15) + "percall".ljust(15) - + "cumtime".ljust(15) + + "cumtime".ljust(15), ) print("-" * 154) for position, key in enumerate(dicts[0].keys()): @@ -207,7 +207,7 @@ def struphy_profile(dirs, replace, all, n_lines, print_callers, savefig): ax.set( title="Weak scaling for cells/mpi_size=" + str(xp.prod(val["Nel"][0]) / val["mpi_size"][0]) - + "=const." + + "=const.", ) ax.legend(loc="upper left") # ax.loglog(val['mpi_size'], val['time'][0]*xp.ones_like(val['time']), 'k--', alpha=0.3) diff --git a/src/struphy/console/run.py b/src/struphy/console/run.py index 1845230e9..3fa8ee56d 100644 --- a/src/struphy/console/run.py +++ b/src/struphy/console/run.py @@ -210,7 +210,7 @@ def struphy_run( if likwid: command = likwid_command + command + ["--likwid"] - print(f"Running with likwid with {likwid_repetitions = }") + print(f"Running with likwid with {likwid_repetitions =}") f.write(f"# Launching likwid {likwid_repetitions} times with likwid-mpirun\n") for i in range(likwid_repetitions): f.write(f"\n\n# Run number {i + 1:03}\n") diff --git a/src/struphy/console/tests/test_console.py b/src/struphy/console/tests/test_console.py index abdd51384..5855e7cc3 100644 --- a/src/struphy/console/tests/test_console.py +++ b/src/struphy/console/tests/test_console.py @@ -290,7 +290,7 @@ def mock_remove(path): # Otherwise, we will not remove all the *_tmp.py files # We can not use the real os.remove becuase then # the state and all compiled files will be removed - print(f"{path = }") + print(f"{path =}") if "_tmp.py" in path: print("Not mock remove") os_remove(path) @@ -318,19 +318,19 @@ def mock_remove(path): time_execution=time_execution, yes=yes, ) - print(f"{language = }") - print(f"{compiler = }") - print(f"{omp_pic = }") - print(f"{omp_feec = }") - print(f"{delete = }") + print(f"{language =}") + print(f"{compiler =}") + print(f"{omp_pic =}") + print(f"{omp_feec =}") + print(f"{delete =}") print(f"{status} = ") - print(f"{verbose = }") - print(f"{dependencies = }") - print(f"{time_execution = }") - print(f"{yes = }") - print(f"{mock_save_state.call_count = }") - print(f"{mock_subprocess_run.call_count = }") - print(f"{mock_os_remove.call_count = }") + print(f"{verbose =}") + print(f"{dependencies =}") + print(f"{time_execution =}") + print(f"{yes =}") + print(f"{mock_save_state.call_count =}") + print(f"{mock_subprocess_run.call_count =}") + print(f"{mock_os_remove.call_count =}") if delete: print("if delete") diff --git a/src/struphy/diagnostics/continuous_spectra.py b/src/struphy/diagnostics/continuous_spectra.py index 5ed069179..c31d789c7 100644 --- a/src/struphy/diagnostics/continuous_spectra.py +++ b/src/struphy/diagnostics/continuous_spectra.py @@ -157,7 +157,7 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, # parse arguments parser = argparse.ArgumentParser( - description="Looks for eigenmodes in a given MHD eigenspectrum in a certain poloidal mode number range and plots the continuous shear Alfvén and slow sound spectra (frequency versus radial-like coordinate)." + description="Looks for eigenmodes in a given MHD eigenspectrum in a certain poloidal mode number range and plots the continuous shear Alfvén and slow sound spectra (frequency versus radial-like coordinate).", ) parser.add_argument("m_l_alfvén", type=int, help="lower bound of poloidal mode number range for Alfvénic modes") @@ -252,7 +252,12 @@ def get_mhd_continua_2d(space, domain, omega2, U_eig, m_range, omega_A, div_tol, fem_1d_2 = Spline_space_1d(Nel[1], p[1], spl_kind[1], nq_el[1], dirichlet_bc[1]) fem_2d = Tensor_spline_space( - [fem_1d_1, fem_1d_2], polar_ck, domain.cx[:, :, 0], domain.cy[:, :, 0], n_tor=n_tor, basis_tor="i" + [fem_1d_1, fem_1d_2], + polar_ck, + domain.cx[:, :, 0], + domain.cy[:, :, 0], + n_tor=n_tor, + basis_tor="i", ) # load and analyze spectrum diff --git a/src/struphy/diagnostics/diagn_tools.py b/src/struphy/diagnostics/diagn_tools.py index d714306f4..b9e66dbb6 100644 --- a/src/struphy/diagnostics/diagn_tools.py +++ b/src/struphy/diagnostics/diagn_tools.py @@ -173,7 +173,7 @@ def power_spectrum_2d( # print(f"{intersec = }") # print(f"{[omega[intersec[n]] for n in range(fit_branches)]}") assert len(intersec) == fit_branches, ( - f"Number of found branches {len(intersec)} is not {fit_branches = }! \ + f"Number of found branches {len(intersec)} is not {fit_branches =}! \ Try to lower 'noise_level' or increase 'extr_order'." ) k_fit += [k] @@ -184,7 +184,7 @@ def power_spectrum_2d( coeffs = [] for m, om in omega_fit.items(): coeffs += [xp.polyfit(k_fit, om, deg=fit_degree[n])] - print(f"\nFitted {coeffs = }") + print(f"\nFitted {coeffs =}") if do_plot: _, ax = plt.subplots(1, 1, figsize=(10, 10)) diff --git a/src/struphy/diagnostics/paraview/mesh_creator.py b/src/struphy/diagnostics/paraview/mesh_creator.py index c6c8d89a0..4bf83211c 100644 --- a/src/struphy/diagnostics/paraview/mesh_creator.py +++ b/src/struphy/diagnostics/paraview/mesh_creator.py @@ -37,7 +37,12 @@ def make_ugrid_and_write_vtu(filename: str, writer, vtk_dir, gvec, s_range, u_ra point_data = {} cell_data = {} vtk_points, suv_points, xyz_points, point_indices = gen_vtk_points( - gvec, s_range, u_range, v_range, point_data, cell_data + gvec, + s_range, + u_range, + v_range, + point_data, + cell_data, ) print("vtk_points.GetNumberOfPoints()", vtk_points.GetNumberOfPoints(), flush=True) diff --git a/src/struphy/dispersion_relations/analytic.py b/src/struphy/dispersion_relations/analytic.py index fb7a40ccd..a54355428 100644 --- a/src/struphy/dispersion_relations/analytic.py +++ b/src/struphy/dispersion_relations/analytic.py @@ -261,7 +261,15 @@ class FluidSlabITG(DispersionRelations1D): def __init__(self, vstar=10.0, vi=1.0, Z=1.0, kz=1.0, gamma=5 / 3): super().__init__( - "wave 1", "wave 2", "wave 3", velocity_scale="thermal", vstar=vstar, vi=vi, Z=Z, kz=kz, gamma=gamma + "wave 1", + "wave 2", + "wave 3", + velocity_scale="thermal", + vstar=vstar, + vi=vi, + Z=Z, + kz=kz, + gamma=gamma, ) def __call__(self, k): @@ -1018,7 +1026,7 @@ def __call__(self, x, m, n): # slow sound continuum specs["slow_sound"] = xp.sqrt( - gamma * p(x) * F(x, m, n) ** 2 / (rho(x) * (gamma * p(x) + By(x) ** 2 + Bz(x) ** 2)) + gamma * p(x) * F(x, m, n) ** 2 / (rho(x) * (gamma * p(x) + By(x) ** 2 + Bz(x) ** 2)), ) return specs @@ -1125,7 +1133,7 @@ def __call__(self, r, m, n): # slow sound continuum specs["slow_sound"] = xp.sqrt( - gamma * p(r) * F(r, m, n) ** 2 / (rho(r) * (gamma * p(r) + Bt(r) ** 2 + Bz(r) ** 2)) + gamma * p(r) * F(r, m, n) ** 2 / (rho(r) * (gamma * p(r) + Bt(r) ** 2 + Bz(r) ** 2)), ) return specs diff --git a/src/struphy/eigenvalue_solvers/kernels_projectors_global.py b/src/struphy/eigenvalue_solvers/kernels_projectors_global.py index 4f6c01392..f01cefc1c 100644 --- a/src/struphy/eigenvalue_solvers/kernels_projectors_global.py +++ b/src/struphy/eigenvalue_solvers/kernels_projectors_global.py @@ -24,7 +24,13 @@ def kernel_int_2d(nq1: "int", nq2: "int", w1: "float[:]", w2: "float[:]", mat_f: # ========= kernel for integration in 3d ================== def kernel_int_3d( - nq1: "int", nq2: "int", nq3: "int", w1: "float[:]", w2: "float[:]", w3: "float[:]", mat_f: "float[:,:,:]" + nq1: "int", + nq2: "int", + nq3: "int", + w1: "float[:]", + w2: "float[:]", + w3: "float[:]", + mat_f: "float[:,:,:]", ) -> "float": f_loc = 0.0 @@ -47,7 +53,11 @@ def kernel_int_3d( # ========= kernel for integration along eta1 direction, reducing to a 2d array ============================ def kernel_int_2d_eta1( - subs1: "int[:]", subs_cum1: "int[:]", w1: "float[:,:]", mat_f: "float[:,:,:]", f_int: "float[:,:]" + subs1: "int[:]", + subs_cum1: "int[:]", + w1: "float[:,:]", + mat_f: "float[:,:,:]", + f_int: "float[:,:]", ): n1, n2 = shape(f_int) @@ -66,7 +76,11 @@ def kernel_int_2d_eta1( # ========= kernel for integration along eta2 direction, reducing to a 2d array ============================ def kernel_int_2d_eta2( - subs2: "int[:]", subs_cum2: "int[:]", w2: "float[:,:]", mat_f: "float[:,:,:]", f_int: "float[:,:]" + subs2: "int[:]", + subs_cum2: "int[:]", + w2: "float[:,:]", + mat_f: "float[:,:,:]", + f_int: "float[:,:]", ): n1, n2 = shape(f_int) @@ -167,7 +181,11 @@ def kernel_int_2d_eta1_eta2_old(w1: "float[:,:]", w2: "float[:,:]", mat_f: "floa # ========= kernel for integration along eta1 direction, reducing to a 3d array ============================ def kernel_int_3d_eta1( - subs1: "int[:]", subs_cum1: "int[:]", w1: "float[:,:]", mat_f: "float[:,:,:,:]", f_int: "float[:,:,:]" + subs1: "int[:]", + subs_cum1: "int[:]", + w1: "float[:,:]", + mat_f: "float[:,:,:,:]", + f_int: "float[:,:,:]", ): n1, n2, n3 = shape(f_int) @@ -186,7 +204,11 @@ def kernel_int_3d_eta1( def kernel_int_3d_eta1_transpose( - subs1: "int[:]", subs_cum1: "int[:]", w1: "float[:,:]", f_int: "float[:,:,:]", mat_f: "float[:,:,:,:]" + subs1: "int[:]", + subs_cum1: "int[:]", + w1: "float[:,:]", + f_int: "float[:,:,:]", + mat_f: "float[:,:,:,:]", ): n1, n2, n3 = shape(f_int) @@ -205,7 +227,11 @@ def kernel_int_3d_eta1_transpose( # ========= kernel for integration along eta2 direction, reducing to a 3d array ============================ def kernel_int_3d_eta2( - subs2: "int[:]", subs_cum2: "int[:]", w2: "float[:,:]", mat_f: "float[:,:,:,:]", f_int: "float[:,:,:]" + subs2: "int[:]", + subs_cum2: "int[:]", + w2: "float[:,:]", + mat_f: "float[:,:,:,:]", + f_int: "float[:,:,:]", ): n1, n2, n3 = shape(f_int) @@ -224,7 +250,11 @@ def kernel_int_3d_eta2( def kernel_int_3d_eta2_transpose( - subs2: "int[:]", subs_cum2: "int[:]", w2: "float[:,:]", f_int: "float[:,:,:]", mat_f: "float[:,:,:,:]" + subs2: "int[:]", + subs_cum2: "int[:]", + w2: "float[:,:]", + f_int: "float[:,:,:]", + mat_f: "float[:,:,:,:]", ): n1, n2, n3 = shape(f_int) @@ -243,7 +273,11 @@ def kernel_int_3d_eta2_transpose( # ========= kernel for integration along eta3 direction, reducing to a 3d array ============================ def kernel_int_3d_eta3( - subs3: "int[:]", subs_cum3: "int[:]", w3: "float[:,:]", mat_f: "float[:,:,:,:]", f_int: "float[:,:,:]" + subs3: "int[:]", + subs_cum3: "int[:]", + w3: "float[:,:]", + mat_f: "float[:,:,:,:]", + f_int: "float[:,:,:]", ): n1, n2, n3 = shape(f_int) @@ -262,7 +296,11 @@ def kernel_int_3d_eta3( def kernel_int_3d_eta3_transpose( - subs3: "int[:]", subs_cum3: "int[:]", w3: "float[:,:]", f_int: "float[:,:,:]", mat_f: "float[:,:,:,:]" + subs3: "int[:]", + subs_cum3: "int[:]", + w3: "float[:,:]", + f_int: "float[:,:,:]", + mat_f: "float[:,:,:,:]", ): n1, n2, n3 = shape(f_int) @@ -655,7 +693,11 @@ def kernel_int_3d_eta1_eta2_eta3_transpose( # ========= kernel for integration in eta1-eta2-eta3 cell, reducing to a 3d array ======================= def kernel_int_3d_eta1_eta2_eta3_old( - w1: "float[:,:]", w2: "float[:,:]", w3: "float[:,:]", mat_f: "float[:,:,:,:,:,:]", f_int: "float[:,:,:]" + w1: "float[:,:]", + w2: "float[:,:]", + w3: "float[:,:]", + mat_f: "float[:,:,:,:,:,:]", + f_int: "float[:,:,:]", ): ne1, nq1, ne2, nq2, ne3, nq3 = shape(mat_f) diff --git a/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py b/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py index 9fbc38ce8..a16b44e3d 100644 --- a/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py +++ b/src/struphy/eigenvalue_solvers/legacy/MHD_eigenvalues_cylinder_1D.py @@ -120,17 +120,17 @@ def solve_ev_problem(rho, B_phi, dB_phi, B_z, p, gamma, a, k, m, num_params, bcZ ) W_dXZ = lambda eta: B_phi(r(eta)) * gamma * m * p(r(eta)) / r(eta) ** 2 + B_z(r(eta)) * gamma * k * p(r(eta)) / r( - eta + eta, ) W_ZdX = lambda eta: B_phi(r(eta)) * gamma * m * p(r(eta)) / r(eta) ** 2 + B_z(r(eta)) * gamma * k * p(r(eta)) / r( - eta + eta, ) W_VZ = lambda eta: B_phi(r(eta)) * gamma * m**2 * p(r(eta)) / r(eta) ** 2 + B_z(r(eta)) * gamma * k * m * p( - r(eta) + r(eta), ) / r(eta) W_ZV = lambda eta: B_phi(r(eta)) * gamma * m**2 * p(r(eta)) / r(eta) ** 2 + B_z(r(eta)) * gamma * k * m * p( - r(eta) + r(eta), ) / r(eta) # compute matrices @@ -289,10 +289,16 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, # evaluate basis functions on quadrature points in format (interval, local quad. point, global basis function) basis_his_N = bsp.collocation_matrix(splines.T, splines.p, pts.flatten(), False, normalize=kind_splines[0]).reshape( - pts.shape[0], pts.shape[1], splines.NbaseN + pts.shape[0], + pts.shape[1], + splines.NbaseN, ) basis_his_D = bsp.collocation_matrix( - splines.t, splines.p - 1, pts.flatten(), False, normalize=kind_splines[1] + splines.t, + splines.p - 1, + pts.flatten(), + False, + normalize=kind_splines[1], ).reshape(pts.shape[0], pts.shape[1], splines.NbaseD) # shift first interpolation point away from pole @@ -459,12 +465,14 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, ) rhs0_N_phi = spa.csr_matrix( - (rhs0_N_phi, (pi0_N_i[0], pi0_N_i[1])), shape=(splines.NbaseN, splines.NbaseN) + (rhs0_N_phi, (pi0_N_i[0], pi0_N_i[1])), + shape=(splines.NbaseN, splines.NbaseN), ).toarray() rhs0_N_z = spa.csr_matrix((rhs0_N_z, (pi0_N_i[0], pi0_N_i[1])), shape=(splines.NbaseN, splines.NbaseN)).toarray() rhs1_D_phi = spa.csr_matrix( - (rhs1_D_phi, (pi1_D_i[0], pi1_D_i[1])), shape=(splines.NbaseD, splines.NbaseD) + (rhs1_D_phi, (pi1_D_i[0], pi1_D_i[1])), + shape=(splines.NbaseD, splines.NbaseD), ).toarray() rhs1_D_z = spa.csr_matrix((rhs1_D_z, (pi1_D_i[0], pi1_D_i[1])), shape=(splines.NbaseD, splines.NbaseD)).toarray() @@ -474,10 +482,12 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, rhs1_D_pr = spa.csr_matrix((rhs1_D_pr, (pi1_D_i[0], pi1_D_i[1])), shape=(splines.NbaseD, splines.NbaseD)).toarray() rhs0_N_rho = spa.csr_matrix( - (rhs0_N_rho, (pi0_N_i[0], pi0_N_i[1])), shape=(splines.NbaseN, splines.NbaseN) + (rhs0_N_rho, (pi0_N_i[0], pi0_N_i[1])), + shape=(splines.NbaseN, splines.NbaseN), ).toarray() rhs1_D_rho = spa.csr_matrix( - (rhs1_D_rho, (pi1_D_i[0], pi1_D_i[1])), shape=(splines.NbaseD, splines.NbaseD) + (rhs1_D_rho, (pi1_D_i[0], pi1_D_i[1])), + shape=(splines.NbaseD, splines.NbaseD), ).toarray() pi0_N_phi = xp.linalg.inv(proj.N.toarray()[1:-1, 1:-1]).dot(rhs0_N_phi[1:-1, 1:-1]) @@ -514,7 +524,7 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, [I_11, xp.zeros((len(u2_r) - 2, len(u2_phi))), xp.zeros((len(u2_r) - 2, len(u2_z) - 1))], [I_21, I_22, I_23[:, 1:]], [I_31[1:, :], I_32[1:, :], I_33[1:, 1:]], - ] + ], ) # ======= matrices in strong pressure equation ================ @@ -534,7 +544,7 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, [A_1, xp.zeros((A_1.shape[0], A_2.shape[1])), xp.zeros((A_1.shape[0], A_3.shape[1]))], [xp.zeros((A_2.shape[0], A_1.shape[1])), A_2, xp.zeros((A_2.shape[0], A_3.shape[1]))], [xp.zeros((A_3.shape[0], A_1.shape[1])), xp.zeros((A_3.shape[0], A_2.shape[1])), A_3], - ] + ], ) MB_11 = 2 * xp.pi * n * pi0_N_z.T.dot(M2_r) + 2 * xp.pi * m * pi0_N_phi.T.dot(M2_r) @@ -553,7 +563,7 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, MB_34 = 2 * xp.pi * n * M3 MB_b_all = xp.block( - [[MB_11, MB_12, MB_13[:, 1:]], [MB_21, MB_22, MB_23[:, 1:]], [MB_31[1:, :], MB_32[1:, :], MB_33[1:, 1:]]] + [[MB_11, MB_12, MB_13[:, 1:]], [MB_21, MB_22, MB_23[:, 1:]], [MB_31[1:, :], MB_32[1:, :], MB_33[1:, 1:]]], ) MB_p_all = xp.block([[MB_14[:, 1:]], [MB_24[:, 1:]], [MB_34[1:, 1:]]]) @@ -664,7 +674,7 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, xp.zeros((len(p3) - 1, A_all.shape[1])), xp.identity(len(p3) - 1), ], - ] + ], ) RHS = xp.block( @@ -672,7 +682,7 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, [xp.zeros((MB_b_all.shape[0], I_all.shape[1])), MB_b_all, MB_p_all], [I_all, xp.zeros((I_all.shape[0], MB_b_all.shape[1])), xp.zeros((I_all.shape[0], MB_p_all.shape[1]))], [P_all, xp.zeros((P_all.shape[0], MB_b_all.shape[1])), xp.zeros((P_all.shape[0], MB_p_all.shape[1]))], - ] + ], ) dt = 0.05 @@ -716,13 +726,14 @@ def solve_ev_problem_FEEC(Rho, B_phi, dB_phi, B_z, dB_z, P, gamma, a, R0, n, m, b2_phi_all[n, :], b2_z_all[n, 1:], p3_all[n, 1:], - ) + ), ) new = UPDATE.dot(old) # extract components unew, bnew, pnew = xp.split( - new, [len(u2_r) - 2 + len(u2_phi) + len(u2_z) - 1, 2 * (len(u2_r) - 2 + len(u2_phi) + len(u2_z) - 1)] + new, + [len(u2_r) - 2 + len(u2_phi) + len(u2_z) - 1, 2 * (len(u2_r) - 2 + len(u2_phi) + len(u2_z) - 1)], ) u2_r_all[n + 1, :] = xp.array([0.0] + list(unew[: (splines.NbaseN - 2)]) + [0.0]) diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py index 49be8c3ef..39156f985 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_control_variate.py @@ -204,7 +204,7 @@ def bv_right( ) # ========================= C.T =========================== return tensor_space_FEM.C.T.dot( - xp.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())) + xp.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())), ) diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_kernels_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_kernels_control_variate.py index 9220cdc69..7f15931ec 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_kernels_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fB_massless_kernels_control_variate.py @@ -584,13 +584,49 @@ def vv( bd3[:] = b3[pd3, :pn3] * d3[:] vel[0] = eva.evaluation_kernel( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, NbaseD[0], NbaseN[1], NbaseN[2], bb1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + NbaseD[0], + NbaseN[1], + NbaseN[2], + bb1, ) vel[1] = eva.evaluation_kernel( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, NbaseN[0], NbaseD[1], NbaseN[2], bb2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + NbaseN[0], + NbaseD[1], + NbaseN[2], + bb2, ) vel[2] = eva.evaluation_kernel( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, NbaseN[0], NbaseN[1], NbaseD[2], bb3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + NbaseN[0], + NbaseN[1], + NbaseD[2], + bb3, ) # ======= here we use the linear hat function =========== ie1 = int(eta1 * Nel[0]) diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py index d52a3e90e..5e9c04eb0 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_control_variate.py @@ -248,7 +248,7 @@ def bv_right( ) # ========================= C.T =========================== return tensor_space_FEM.C.T.dot( - xp.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())) + xp.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())), ) @@ -429,7 +429,7 @@ def uv_right( ) # ========================= C.T =========================== temp_final = temp_final_0.flatten() + tensor_space_FEM.G.T.dot( - xp.concatenate((temp_final_1.flatten(), temp_final_2.flatten(), temp_final_3.flatten())) + xp.concatenate((temp_final_1.flatten(), temp_final_2.flatten(), temp_final_3.flatten())), ) return temp_final diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_kernels_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_kernels_control_variate.py index f3b6fca0e..965a7af33 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_kernels_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/fnB_massless_kernels_control_variate.py @@ -212,17 +212,65 @@ def vv( bd3[:] = b3[pd3, :pn3] * d3[:] vel[0] = eva.evaluation_kernel( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, NbaseD[0], NbaseN[1], NbaseN[2], bb1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + NbaseD[0], + NbaseN[1], + NbaseN[2], + bb1, ) vel[1] = eva.evaluation_kernel( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, NbaseN[0], NbaseD[1], NbaseN[2], bb2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + NbaseN[0], + NbaseD[1], + NbaseN[2], + bb2, ) vel[2] = eva.evaluation_kernel( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, NbaseN[0], NbaseN[1], NbaseD[2], bb3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + NbaseN[0], + NbaseN[1], + NbaseD[2], + bb3, ) tt = eva.evaluation_kernel( - pn1, pn2, pn3, bn1, bn2, bn3, span1, span2, span3, NbaseN[0], NbaseN[1], NbaseN[2], n + pn1, + pn2, + pn3, + bn1, + bn2, + bn3, + span1, + span2, + span3, + NbaseN[0], + NbaseN[1], + NbaseN[2], + n, ) if abs(tt) > tol: U_value = 1.0 / tt diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py index 4e97422f0..3459ff7b2 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_control_variate.py @@ -247,7 +247,7 @@ def bv_right( ) # ========================= C.T =========================== return tensor_space_FEM.C.T.dot( - xp.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())) + xp.concatenate((temp_twoform1.flatten(), temp_twoform2.flatten(), temp_twoform3.flatten())), ) @@ -430,7 +430,7 @@ def uv_right( ) # ========================= C.T =========================== temp_final = temp_final_0.flatten() + tensor_space_FEM.G.T.dot( - xp.concatenate((temp_final_1.flatten(), temp_final_2.flatten(), temp_final_3.flatten())) + xp.concatenate((temp_final_1.flatten(), temp_final_2.flatten(), temp_final_3.flatten())), ) return temp_final diff --git a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_kernels_control_variate.py b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_kernels_control_variate.py index c56b67711..b85a74758 100644 --- a/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_kernels_control_variate.py +++ b/src/struphy/eigenvalue_solvers/legacy/control_variates/kinetic_extended/massless_kernels_control_variate.py @@ -195,34 +195,84 @@ def uvright( for q2 in range(nq2): for q3 in range(nq3): dft[0, 0] = DFI_11[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.df_inv(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map, components[0, 0]) dft[0, 1] = DFI_21[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.df_inv(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map, components[0, 1]) dft[0, 2] = DFI_31[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.df_inv(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map, components[0, 2]) dft[1, 0] = DFI_12[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.df_inv(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map, components[1, 0]) dft[1, 1] = DFI_22[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.df_inv(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map, components[1, 1]) dft[1, 2] = DFI_32[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.df_inv(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map, components[1, 2]) dft[2, 0] = DFI_13[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.df_inv(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map, components[2, 0]) dft[2, 1] = DFI_23[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.df_inv(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map, components[2, 1]) dft[2, 2] = DFI_33[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.df_inv(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map, components[2, 2]) detdet = df_det[ - ie1, ie2, ie3, q1, q2, q3 + ie1, + ie2, + ie3, + q1, + q2, + q3, ] # mappings_analytical.det_df(pts1[ie1, q1], pts2[ie2,q2], pts3[ie3,q3], kind_map, params_map) Jeq[0] = Jeqx[ie1, ie2, ie3, q1, q2, q3] Jeq[1] = Jeqy[ie1, ie2, ie3, q1, q2, q3] @@ -705,19 +755,67 @@ def vv( bd3[:] = b3[pd3, :pn3] * d3[:] vel[0] = eva.evaluation_kernel( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, NbaseD[0], NbaseN[1], NbaseN[2], bb1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + NbaseD[0], + NbaseN[1], + NbaseN[2], + bb1, ) vel[1] = eva.evaluation_kernel( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, NbaseN[0], NbaseD[1], NbaseN[2], bb2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + NbaseN[0], + NbaseD[1], + NbaseN[2], + bb2, ) vel[2] = eva.evaluation_kernel( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, NbaseN[0], NbaseN[1], NbaseD[2], bb3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + NbaseN[0], + NbaseN[1], + NbaseD[2], + bb3, ) U_value = exp( -eva.evaluation_kernel( - pn1, pn2, pn3, bn1, bn2, bn3, span1, span2, span3, NbaseN[0], NbaseN[1], NbaseN[2], u - ) + pn1, + pn2, + pn3, + bn1, + bn2, + bn3, + span1, + span2, + span3, + NbaseN[0], + NbaseN[1], + NbaseN[2], + u, + ), ) # ========= mapping evaluation ============= diff --git a/src/struphy/eigenvalue_solvers/legacy/emw_operators.py b/src/struphy/eigenvalue_solvers/legacy/emw_operators.py index e2fd4ed22..9e8c95b3f 100755 --- a/src/struphy/eigenvalue_solvers/legacy/emw_operators.py +++ b/src/struphy/eigenvalue_solvers/legacy/emw_operators.py @@ -198,12 +198,14 @@ def __assemble_M1_cross(self, weight): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M[a][b] = spa.csr_matrix( - (M[a][b].flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (M[a][b].flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M[a][b].eliminate_zeros() M = spa.bmat( - [[M[0][0], M[0][1], M[0][2]], [M[1][0], M[1][1], M[1][2]], [M[2][0], M[2][1], M[2][2]]], format="csr" + [[M[0][0], M[0][1], M[0][2]], [M[1][0], M[1][1], M[1][2]], [M[2][0], M[2][1], M[2][2]]], + format="csr", ) self.R1_mat = -self.SPACES.E1_0.dot(M.dot(self.SPACES.E1_0.T)).tocsr() diff --git a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py index eecb68293..a46097a8f 100644 --- a/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py +++ b/src/struphy/eigenvalue_solvers/legacy/mass_matrices_3d_pre.py @@ -273,10 +273,12 @@ def get_M1_PRE_3(tensor_space_FEM, mats_pol=None): def solve(x): x1 = x[: tensor_space_FEM.E1_pol_0.shape[0] * tensor_space_FEM.NbaseN[2]].reshape( - tensor_space_FEM.E1_pol_0.shape[0], tensor_space_FEM.NbaseN[2] + tensor_space_FEM.E1_pol_0.shape[0], + tensor_space_FEM.NbaseN[2], ) x2 = x[tensor_space_FEM.E1_pol_0.shape[0] * tensor_space_FEM.NbaseN[2] :].reshape( - tensor_space_FEM.E0_pol_0.shape[0], tensor_space_FEM.NbaseD[2] + tensor_space_FEM.E0_pol_0.shape[0], + tensor_space_FEM.NbaseD[2], ) r1 = linkron.kron_fftsolve_2d(M1_pol_0_11_LU, tor_vec0, x1).flatten() @@ -311,10 +313,12 @@ def get_M2_PRE_3(tensor_space_FEM, mats_pol=None): def solve(x): x1 = x[: tensor_space_FEM.E2_pol_0.shape[0] * tensor_space_FEM.NbaseD[2]].reshape( - tensor_space_FEM.E2_pol_0.shape[0], tensor_space_FEM.NbaseD[2] + tensor_space_FEM.E2_pol_0.shape[0], + tensor_space_FEM.NbaseD[2], ) x2 = x[tensor_space_FEM.E2_pol_0.shape[0] * tensor_space_FEM.NbaseD[2] :].reshape( - tensor_space_FEM.E3_pol_0.shape[0], tensor_space_FEM.NbaseN[2] + tensor_space_FEM.E3_pol_0.shape[0], + tensor_space_FEM.NbaseN[2], ) r1 = linkron.kron_fftsolve_2d(M2_pol_0_11_LU, tor_vec1, x1).flatten() @@ -373,10 +377,12 @@ def get_Mv_PRE_3(tensor_space_FEM, mats_pol=None): def solve(x): x1 = x[: tensor_space_FEM.Ev_pol_0.shape[0] * tensor_space_FEM.NbaseN[2]].reshape( - tensor_space_FEM.Ev_pol_0.shape[0], tensor_space_FEM.NbaseN[2] + tensor_space_FEM.Ev_pol_0.shape[0], + tensor_space_FEM.NbaseN[2], ) x2 = x[tensor_space_FEM.Ev_pol_0.shape[0] * tensor_space_FEM.NbaseN[2] :].reshape( - tensor_space_FEM.E0_pol.shape[0], tensor_space_FEM.NbaseN[2] + tensor_space_FEM.E0_pol.shape[0], + tensor_space_FEM.NbaseN[2], ) r1 = linkron.kron_fftsolve_2d(Mv_pol_0_11_LU, tor_vec0, x1).flatten() diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_bv_kernel.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_bv_kernel.py index 5cf3830a4..e477940e6 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_bv_kernel.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_bv_kernel.py @@ -226,7 +226,9 @@ def right_hand2( * bd2[ie2, il2, 0, q2] * bd3[ie3, il3, 0, q3] * temp_vector_1[ - N_index_x[ie1, il1], D_index_y[ie2, il2], D_index_z[ie3, il3] + N_index_x[ie1, il1], + D_index_y[ie2, il2], + D_index_z[ie3, il3], ] ) @@ -252,7 +254,9 @@ def right_hand2( * bn2[ie2, il2, 0, q2] * bd3[ie3, il3, 0, q3] * temp_vector_2[ - D_index_x[ie1, il1], N_index_y[ie2, il2], D_index_z[ie3, il3] + D_index_x[ie1, il1], + N_index_y[ie2, il2], + D_index_z[ie3, il3], ] ) @@ -278,7 +282,9 @@ def right_hand2( * bd2[ie2, il2, 0, q2] * bn3[ie3, il3, 0, q3] * temp_vector_3[ - D_index_x[ie1, il1], D_index_y[ie2, il2], N_index_z[ie3, il3] + D_index_x[ie1, il1], + D_index_y[ie2, il2], + N_index_z[ie3, il3], ] ) @@ -338,7 +344,9 @@ def right_hand1( * bn2[ie2, il2, 0, q2] * bn3[ie3, il3, 0, q3] * temp_vector_1[ - D_index_x[ie1, il1], N_index_y[ie2, il2], N_index_z[ie3, il3] + D_index_x[ie1, il1], + N_index_y[ie2, il2], + N_index_z[ie3, il3], ] ) @@ -364,7 +372,9 @@ def right_hand1( * bd2[ie2, il2, 0, q2] * bn3[ie3, il3, 0, q3] * temp_vector_2[ - N_index_x[ie1, il1], D_index_y[ie2, il2], N_index_z[ie3, il3] + N_index_x[ie1, il1], + D_index_y[ie2, il2], + N_index_z[ie3, il3], ] ) @@ -390,7 +400,9 @@ def right_hand1( * bn2[ie2, il2, 0, q2] * bd3[ie3, il3, 0, q3] * temp_vector_3[ - N_index_x[ie1, il1], N_index_y[ie2, il2], D_index_z[ie3, il3] + N_index_x[ie1, il1], + N_index_y[ie2, il2], + D_index_z[ie3, il3], ] ) diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py index 2ddddc64a..bdf01bf8b 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_massless_linear_operators.py @@ -725,7 +725,8 @@ def linearoperator_step3( # ========================= C =========================== # time1 = time.time() twoform_temp1_long[:], twoform_temp2_long[:], twoform_temp3_long[:] = xp.split( - tensor_space_FEM.C.dot(input_vector), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]] + tensor_space_FEM.C.dot(input_vector), + [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]], ) temp_vector_1[:, :, :] = twoform_temp1_long.reshape(Nbase_2form[0]) temp_vector_2[:, :, :] = twoform_temp2_long.reshape(Nbase_2form[1]) @@ -826,7 +827,7 @@ def linearoperator_step3( # ========================= C.T =========================== # time1 = time.time() temp_final = tensor_space_FEM.M1.dot(input_vector) - dt / 2.0 * tensor_space_FEM.C.T.dot( - xp.concatenate((temp_vector_1.flatten(), temp_vector_2.flatten(), temp_vector_3.flatten())) + xp.concatenate((temp_vector_1.flatten(), temp_vector_2.flatten(), temp_vector_3.flatten())), ) # time2 = time.time() # print('second_curl_time', time2 - time1) @@ -929,7 +930,8 @@ def linearoperator_right_step3( # ================================================================== # ========================= C =========================== twoform_temp1_long[:], twoform_temp2_long[:], twoform_temp3_long[:] = xp.split( - tensor_space_FEM.C.dot(input_vector), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]] + tensor_space_FEM.C.dot(input_vector), + [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]], ) temp_vector_1[:, :, :] = twoform_temp1_long.reshape(Nbase_2form[0]) temp_vector_2[:, :, :] = twoform_temp2_long.reshape(Nbase_2form[1]) @@ -1082,7 +1084,7 @@ def linearoperator_right_step3( # print('final_bb', time2 - time1) # ========================= C.T =========================== temp_final = tensor_space_FEM.M1.dot(input_vector) + dt / 2.0 * tensor_space_FEM.C.T.dot( - xp.concatenate((temp_vector_1.flatten(), temp_vector_2.flatten(), temp_vector_3.flatten())) + xp.concatenate((temp_vector_1.flatten(), temp_vector_2.flatten(), temp_vector_3.flatten())), ) return temp_final @@ -1146,7 +1148,8 @@ def substep4_linear_operator( # ========================================== acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = xp.split( - tensor_space_FEM.C.dot(input), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]] + tensor_space_FEM.C.dot(input), + [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]], ) acc.twoform_temp1[:, :, :] = acc.twoform_temp1_long.reshape(Nbase_2form[0]) acc.twoform_temp2[:, :, :] = acc.twoform_temp2_long.reshape(Nbase_2form[1]) @@ -1359,7 +1362,7 @@ def substep4_linear_operator( ) return M1.dot(input) + dt**2 / 4.0 * tensor_space_FEM.C.T.dot( - xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) + xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())), ) # ========================================================================================================== @@ -1532,7 +1535,7 @@ def substep4_linear_operator_right( xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())), tol=10 ** (-10), M=M1_PRE, - )[0] + )[0], ) acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = xp.split( @@ -1640,7 +1643,7 @@ def substep4_linear_operator_right( ) return M1.dot(xp.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))) - CURL.T.dot( - xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) + xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())), ) # ========================================================================================================== @@ -1961,7 +1964,8 @@ def substep4_localproj_linear_operator( # ========================================== acc.twoform_temp1_long[:], acc.twoform_temp2_long[:], acc.twoform_temp3_long[:] = xp.split( - tensor_space_FEM.C.dot(input), [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]] + tensor_space_FEM.C.dot(input), + [Ntot_2form[0], Ntot_2form[0] + Ntot_2form[1]], ) acc.twoform_temp1[:, :, :] = acc.twoform_temp1_long.reshape(Nbase_2form[0]) acc.twoform_temp2[:, :, :] = acc.twoform_temp2_long.reshape(Nbase_2form[1]) @@ -2065,7 +2069,7 @@ def substep4_localproj_linear_operator( acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = xp.split( mat.dot( - xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())) + xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())), ), [Ntot_1form[0], Ntot_1form[0] + Ntot_1form[1]], ) @@ -2171,7 +2175,7 @@ def substep4_localproj_linear_operator( ) return M1.dot(input) + dt**2 / 4.0 * tensor_space_FEM.C.T.dot( - xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) + xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())), ) # ========================================================================================================== @@ -2339,11 +2343,12 @@ def substep4_localproj_linear_operator_right( tensor_space_FEM.basisD[2], ) acc.oneform_temp_long[:] = mat.dot( - xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())) + xp.concatenate((acc.oneform_temp1.flatten(), acc.oneform_temp2.flatten(), acc.oneform_temp3.flatten())), ) acc.oneform_temp1_long[:], acc.oneform_temp2_long[:], acc.oneform_temp3_long[:] = xp.split( - (dt**2.0 / 4.0 * acc.oneform_temp_long + dt * vec), [Ntot_1form[0], Ntot_1form[0] + Ntot_1form[1]] + (dt**2.0 / 4.0 * acc.oneform_temp_long + dt * vec), + [Ntot_1form[0], Ntot_1form[0] + Ntot_1form[1]], ) acc.oneform_temp1[:, :, :] = acc.oneform_temp1_long.reshape(Nbase_1form[0]) @@ -2447,7 +2452,7 @@ def substep4_localproj_linear_operator_right( ) return M1.dot(xp.concatenate((bb1.flatten(), bb2.flatten(), bb3.flatten()))) - CURL.T.dot( - xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())) + xp.concatenate((acc.twoform_temp1.flatten(), acc.twoform_temp2.flatten(), acc.twoform_temp3.flatten())), ) # ========================================================================================================== diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py index 4a5a2dbe0..1ef3376dc 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_vv_kernel.py @@ -542,7 +542,8 @@ def piecewise_gather( for q2 in range(n_quad[1]): for q3 in range(n_quad[2]): temp1[0] = (cell_left[0] + il1) / Nel[0] + pts1[ - 0, q1 + 0, + q1, ] # quadrature points in the cell x direction temp4[0] = abs(temp1[0] - eta1) - compact[0] / 2.0 # if > 0, result is 0 @@ -741,7 +742,8 @@ def piecewise_scatter( for q2 in range(n_quad[1]): for q3 in range(n_quad[2]): temp1[0] = (cell_left[0] + il1) / Nel[0] + pts1[ - 0, q1 + 0, + q1, ] # quadrature points in the cell x direction temp4[0] = abs(temp1[0] - eta1) - compact[0] / 2 # if > 0, result is 0 diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py index 0c96616de..49464aa58 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py @@ -505,7 +505,9 @@ def __init__(self, tensor_space, n_quad): # quadrature points and weights self.pts[a], self.wts[a] = bsp.quadrature_grid( - xp.unique(self.x_his[a].flatten()), self.pts_loc[a], self.wts_loc[a] + xp.unique(self.x_his[a].flatten()), + self.pts_loc[a], + self.wts_loc[a], ) else: @@ -535,7 +537,9 @@ def __init__(self, tensor_space, n_quad): # quadrature points and weights self.pts[a], self.wts[a] = bsp.quadrature_grid( - xp.append(xp.unique(self.x_his[a].flatten() % 1.0), 1.0), self.pts_loc[a], self.wts_loc[a] + xp.append(xp.unique(self.x_his[a].flatten() % 1.0), 1.0), + self.pts_loc[a], + self.wts_loc[a], ) # evaluate N basis functions at interpolation and quadrature points @@ -556,7 +560,9 @@ def __init__(self, tensor_space, n_quad): self.basisD_his = [ bsp.collocation_matrix(T[1:-1], p - 1, pts.flatten(), bc, normalize=True).reshape( - pts[:, 0].size, pts[0, :].size, NbaseD + pts[:, 0].size, + pts[0, :].size, + NbaseD, ) for T, p, pts, bc, NbaseD in zip(self.T, self.p, self.pts, self.bc, self.NbaseD) ] @@ -802,7 +808,7 @@ def projection_Q_0form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_N[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -827,7 +833,7 @@ def projection_Q_0form(self, domain): self.n_his_nvcof_N[0], self.n_int_nvcof_N[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -852,7 +858,7 @@ def projection_Q_0form(self, domain): self.n_his_nvcof_N[0], self.n_his_nvcof_N[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -1111,7 +1117,7 @@ def projection_Q_2form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_D[1], self.n_his_nvcof_D[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -1136,7 +1142,7 @@ def projection_Q_2form(self, domain): self.n_his_nvcof_D[0], self.n_int_nvcof_N[1], self.n_his_nvcof_D[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -1161,7 +1167,7 @@ def projection_Q_2form(self, domain): self.n_his_nvcof_D[0], self.n_his_nvcof_D[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -1292,7 +1298,7 @@ def projection_W_0form(self, domain): self.n_int_nvcof_N[0], self.n_int_nvcof_N[1], self.n_int_nvcof_N[2], - ) + ), ) # row indices @@ -1735,7 +1741,7 @@ def projection_T_0form(self, domain): self.n_his_nvcof_N[0], self.n_int_nvcof_N[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -1759,7 +1765,7 @@ def projection_T_0form(self, domain): self.n_his_nvcof_N[0], self.n_int_nvcof_N[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -1784,7 +1790,7 @@ def projection_T_0form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_N[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -1808,7 +1814,7 @@ def projection_T_0form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_N[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -1833,7 +1839,7 @@ def projection_T_0form(self, domain): self.n_int_nvcof_N[0], self.n_int_nvcof_N[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -1857,7 +1863,7 @@ def projection_T_0form(self, domain): self.n_int_nvcof_N[0], self.n_int_nvcof_N[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -2268,7 +2274,7 @@ def projection_T_1form(self, domain): self.n_his_nvcof_N[0], self.n_int_nvcof_D[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -2292,7 +2298,7 @@ def projection_T_1form(self, domain): self.n_his_nvcof_N[0], self.n_int_nvcof_N[1], self.n_int_nvcof_D[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -2317,7 +2323,7 @@ def projection_T_1form(self, domain): self.n_int_nvcof_D[0], self.n_his_nvcof_N[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -2341,7 +2347,7 @@ def projection_T_1form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_N[1], self.n_int_nvcof_D[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -2366,7 +2372,7 @@ def projection_T_1form(self, domain): self.n_int_nvcof_D[0], self.n_int_nvcof_N[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -2390,7 +2396,7 @@ def projection_T_1form(self, domain): self.n_int_nvcof_N[0], self.n_int_nvcof_D[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -2801,7 +2807,7 @@ def projection_T_2form(self, domain): self.n_his_nvcof_D[0], self.n_int_nvcof_N[1], self.n_int_nvcof_D[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -2825,7 +2831,7 @@ def projection_T_2form(self, domain): self.n_his_nvcof_D[0], self.n_int_nvcof_D[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -2850,7 +2856,7 @@ def projection_T_2form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_D[1], self.n_int_nvcof_D[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -2874,7 +2880,7 @@ def projection_T_2form(self, domain): self.n_int_nvcof_D[0], self.n_his_nvcof_D[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -2899,7 +2905,7 @@ def projection_T_2form(self, domain): self.n_int_nvcof_N[0], self.n_int_nvcof_D[1], self.n_his_nvcof_D[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -2923,7 +2929,7 @@ def projection_T_2form(self, domain): self.n_int_nvcof_D[0], self.n_int_nvcof_N[1], self.n_his_nvcof_D[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -3182,7 +3188,7 @@ def projection_S_0form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_N[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -3207,7 +3213,7 @@ def projection_S_0form(self, domain): self.n_his_nvcof_N[0], self.n_int_nvcof_N[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -3232,7 +3238,7 @@ def projection_S_0form(self, domain): self.n_his_nvcof_N[0], self.n_his_nvcof_N[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -3491,7 +3497,7 @@ def projection_S_2form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_D[1], self.n_his_nvcof_D[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -3516,7 +3522,7 @@ def projection_S_2form(self, domain): self.n_his_nvcof_D[0], self.n_int_nvcof_N[1], self.n_his_nvcof_D[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -3541,7 +3547,7 @@ def projection_S_2form(self, domain): self.n_his_nvcof_D[0], self.n_his_nvcof_D[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -3664,7 +3670,7 @@ def projection_K_3form(self, domain): self.n_his_nvcof_D[0], self.n_his_nvcof_D[1], self.n_his_nvcof_D[2], - ) + ), ) # row indices @@ -3927,7 +3933,7 @@ def projection_N_0form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_N[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -3952,7 +3958,7 @@ def projection_N_0form(self, domain): self.n_his_nvcof_N[0], self.n_int_nvcof_N[1], self.n_his_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -3977,7 +3983,7 @@ def projection_N_0form(self, domain): self.n_his_nvcof_N[0], self.n_his_nvcof_N[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -4212,7 +4218,7 @@ def projection_N_2form(self, domain): self.n_int_nvcof_N[0], self.n_his_nvcof_D[1], self.n_his_nvcof_D[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -4237,7 +4243,7 @@ def projection_N_2form(self, domain): self.n_his_nvcof_D[0], self.n_int_nvcof_N[1], self.n_his_nvcof_D[2], - ) + ), ) row = self.NbaseN[1] * self.NbaseD[2] * indices[0] + self.NbaseD[2] * indices[1] + indices[2] @@ -4262,7 +4268,7 @@ def projection_N_2form(self, domain): self.n_his_nvcof_D[0], self.n_his_nvcof_D[1], self.n_int_nvcof_N[2], - ) + ), ) row = self.NbaseD[1] * self.NbaseN[2] * indices[0] + self.NbaseN[2] * indices[1] + indices[2] @@ -4323,7 +4329,15 @@ class term_curl_beq: """ def __init__( - self, tensor_space, mapping, kind_map=None, params_map=None, tensor_space_F=None, cx=None, cy=None, cz=None + self, + tensor_space, + mapping, + kind_map=None, + params_map=None, + tensor_space_F=None, + cx=None, + cy=None, + cz=None, ): self.p = tensor_space.p # spline degrees self.Nel = tensor_space.Nel # number of elements @@ -4357,13 +4371,16 @@ def __init__( # ============= evaluation of background magnetic field at quadrature points ========= self.mat_curl_beq_1 = xp.empty( - (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float + (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), + dtype=float, ) self.mat_curl_beq_2 = xp.empty( - (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float + (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), + dtype=float, ) self.mat_curl_beq_3 = xp.empty( - (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float + (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), + dtype=float, ) if mapping == 0: @@ -4455,13 +4472,16 @@ def __init__( # ====================== perturbed magnetic field at quadrature points ========== self.B1 = xp.empty( - (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float + (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), + dtype=float, ) self.B2 = xp.empty( - (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float + (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), + dtype=float, ) self.B3 = xp.empty( - (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), dtype=float + (self.Nel[0], self.Nel[1], self.Nel[2], self.n_quad[0], self.n_quad[1], self.n_quad[2]), + dtype=float, ) # ========================== inner products ===================================== diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py index dacc4f243..3b27b1b5f 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py @@ -137,10 +137,12 @@ def __init__(self, spline_space, n_quad): self.x_int = xp.zeros((n_lambda_int, self.n_int), dtype=float) # interpolation points for each coeff. self.int_global_N = xp.zeros( - (n_lambda_int, self.n_int_locbf_N), dtype=int + (n_lambda_int, self.n_int_locbf_N), + dtype=int, ) # global indices of non-vanishing N bf self.int_global_D = xp.zeros( - (n_lambda_int, self.n_int_locbf_D), dtype=int + (n_lambda_int, self.n_int_locbf_D), + dtype=int, ) # global indices of non-vanishing D bf self.int_loccof_N = xp.zeros((n_lambda_int, self.n_int_locbf_N), dtype=int) # index of non-vanishing coeff. (N) @@ -423,7 +425,9 @@ def __init__(self, spline_space, n_quad): # quadrature points and weights self.pts, self.wts = bsp.quadrature_grid( - xp.append(xp.unique(self.x_his.flatten() % 1.0), 1.0), self.pts_loc, self.wts_loc + xp.append(xp.unique(self.x_his.flatten() % 1.0), 1.0), + self.pts_loc, + self.wts_loc, ) # quasi interpolation @@ -1075,7 +1079,9 @@ def __init__(self, tensor_space, n_quad): # quadrature points and weights self.pts[a], self.wts[a] = bsp.quadrature_grid( - xp.unique(self.x_his[a].flatten()), self.pts_loc[a], self.wts_loc[a] + xp.unique(self.x_his[a].flatten()), + self.pts_loc[a], + self.wts_loc[a], ) else: @@ -1105,7 +1111,9 @@ def __init__(self, tensor_space, n_quad): # quadrature points and weights self.pts[a], self.wts[a] = bsp.quadrature_grid( - xp.append(xp.unique(self.x_his[a].flatten() % 1.0), 1.0), self.pts_loc[a], self.wts_loc[a] + xp.append(xp.unique(self.x_his[a].flatten() % 1.0), 1.0), + self.pts_loc[a], + self.wts_loc[a], ) # projector on space V0 (interpolation) @@ -1427,7 +1435,11 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): self.wts[1], self.wts[2], mat_f.reshape( - x_int1.size, self.pts[1].shape[0], self.pts[1].shape[1], self.pts[2].shape[0], self.pts[2].shape[1] + x_int1.size, + self.pts[1].shape[0], + self.pts[1].shape[1], + self.pts[2].shape[0], + self.pts[2].shape[1], ), lambdas1, ) @@ -1478,7 +1490,11 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): self.wts[0], self.wts[2], mat_f.reshape( - self.pts[0].shape[0], self.pts[0].shape[1], x_int2.size, self.pts[2].shape[0], self.pts[2].shape[1] + self.pts[0].shape[0], + self.pts[0].shape[1], + x_int2.size, + self.pts[2].shape[0], + self.pts[2].shape[1], ), lambdas2, ) @@ -1529,7 +1545,11 @@ def pi_2(self, fun, include_bc=True, eval_kind="meshgrid"): self.wts[0], self.wts[1], mat_f.reshape( - self.pts[0].shape[0], self.pts[0].shape[1], self.pts[1].shape[0], self.pts[1].shape[1], x_int3.size + self.pts[0].shape[0], + self.pts[0].shape[1], + self.pts[1].shape[0], + self.pts[1].shape[1], + x_int3.size, ), lambdas3, ) @@ -1560,7 +1580,8 @@ def pi_3(self, fun, include_bc=True, eval_kind="meshgrid"): # evaluation of function at quadrature points mat_f = xp.empty( - (self.pts[0].flatten().size, self.pts[1].flatten().size, self.pts[2].flatten().size), dtype=float + (self.pts[0].flatten().size, self.pts[1].flatten().size, self.pts[2].flatten().size), + dtype=float, ) # external function call if a callable is passed @@ -1568,7 +1589,10 @@ def pi_3(self, fun, include_bc=True, eval_kind="meshgrid"): # create a meshgrid and evaluate function on point set if eval_kind == "meshgrid": pts1, pts2, pts3 = xp.meshgrid( - self.pts[0].flatten(), self.pts[1].flatten(), self.pts[2].flatten(), indexing="ij" + self.pts[0].flatten(), + self.pts[1].flatten(), + self.pts[2].flatten(), + indexing="ij", ) mat_f[:, :, :] = fun(pts1, pts2, pts3) @@ -1582,7 +1606,9 @@ def pi_3(self, fun, include_bc=True, eval_kind="meshgrid"): for i2 in range(self.pts[1].size): for i3 in range(self.pts[2].size): mat_f[i1, i2, i3] = fun( - self.pts[0].flatten()[i1], self.pts[1].flatten()[i2], self.pts[2].flatten()[i3] + self.pts[0].flatten()[i1], + self.pts[1].flatten()[i2], + self.pts[2].flatten()[i3], ) # internal function call diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_L2_projector_kernel.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_L2_projector_kernel.py index 9aa26b243..0e711dbcf 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_L2_projector_kernel.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_L2_projector_kernel.py @@ -1390,33 +1390,42 @@ def vv_1_form( # evaluation of function at interpolation/quadrature points mat_11 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_21 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_31 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_12 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_22 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_32 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_13 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_23 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_33 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) for i1 in range(cell_number[0]): @@ -1922,33 +1931,42 @@ def vv_push( # evaluation of function at interpolation/quadrature points mat_11 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_21 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_31 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_12 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_22 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_32 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_13 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_23 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_33 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) for i1 in range(cell_number[0]): diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py index 20814c8ac..8978e2464 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py @@ -86,7 +86,7 @@ def __init__(self, tensor_space, p_shape, p_size, NbaseN, NbaseD, mpi_comm): for a in range(3): # self.related[a] = int(xp.floor(NbaseN[a]/2.0)) self.related[a] = int( - xp.floor((3 * int((self.p_size[a] * (self.p_shape[a] + 1)) * self.Nel[a] + 1) + 3 * self.p[a]) / 2.0) + xp.floor((3 * int((self.p_size[a] * (self.p_shape[a] + 1)) * self.Nel[a] + 1) + 3 * self.p[a]) / 2.0), ) if (2 * self.related[a] + 1) > NbaseN[a]: self.related[a] = int(xp.floor(NbaseN[a] / 2.0)) @@ -302,7 +302,7 @@ def assemble_0_form(self, tensor_space_FEM, mpi_comm): # conversion to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -316,7 +316,8 @@ def assemble_0_form(self, tensor_space_FEM, mpi_comm): col = Nj[1] * Ni[2] * col1 + Ni[2] * col2 + col3 M = spa.csr_matrix( - (self.kernel_0.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_0.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M.eliminate_zeros() @@ -359,7 +360,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -373,7 +374,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M11 = spa.csr_matrix( - (self.kernel_1_11.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_11.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M11.eliminate_zeros() @@ -385,7 +387,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -399,7 +401,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M12 = spa.csr_matrix( - (self.kernel_1_12.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_12.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M12.eliminate_zeros() @@ -411,7 +414,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -425,7 +428,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M13 = spa.csr_matrix( - (self.kernel_1_13.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_13.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M13.eliminate_zeros() @@ -437,7 +441,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -451,7 +455,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M22 = spa.csr_matrix( - (self.kernel_1_22.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_22.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M22.eliminate_zeros() @@ -463,7 +468,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -477,7 +482,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M23 = spa.csr_matrix( - (self.kernel_1_23.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_23.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M23.eliminate_zeros() @@ -489,7 +495,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -503,7 +509,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M33 = spa.csr_matrix( - (self.kernel_1_33.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_33.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M33.eliminate_zeros() diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py index 7c9425e47..43c8c8ff9 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py @@ -87,7 +87,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co for a in range(3): # self.related[a] = int(xp.floor(NbaseN[a]/2.0)) self.related[a] = int( - xp.floor((3 * int((self.p_size[a] * (self.p_shape[a] + 1)) * self.Nel[a] + 1) + 3 * self.p[a]) / 2.0) + xp.floor((3 * int((self.p_size[a] * (self.p_shape[a] + 1)) * self.Nel[a] + 1) + 3 * self.p[a]) / 2.0), ) if (2 * self.related[a] + 1) > NbaseN[a]: self.related[a] = int(xp.floor(NbaseN[a] / 2.0)) @@ -294,7 +294,9 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co self.wts = [0, 0, 0] for a in range(3): self.pts[a], self.wts[a] = bsp.quadrature_grid( - [0, 1.0 / 2.0 / self.Nel[a]], self.pts_loc[a], self.wts_loc[a] + [0, 1.0 / 2.0 / self.Nel[a]], + self.pts_loc[a], + self.wts_loc[a], ) # print('check_pts', self.pts[0].shape, self.pts[1].shape, self.pts[2].shape) # print('check_pts', self.wts) @@ -403,7 +405,7 @@ def assemble_0_form(self, tensor_space_FEM, mpi_comm): # conversion to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -417,7 +419,8 @@ def assemble_0_form(self, tensor_space_FEM, mpi_comm): col = Nj[1] * Ni[2] * col1 + Ni[2] * col2 + col3 M = spa.csr_matrix( - (self.kernel_0.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_0.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M.eliminate_zeros() @@ -460,7 +463,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -474,7 +477,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M11 = spa.csr_matrix( - (self.kernel_1_11.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_11.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M11.eliminate_zeros() @@ -486,7 +490,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -500,7 +504,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M12 = spa.csr_matrix( - (self.kernel_1_12.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_12.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M12.eliminate_zeros() @@ -512,7 +517,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -526,7 +531,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M13 = spa.csr_matrix( - (self.kernel_1_13.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_13.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M13.eliminate_zeros() @@ -538,7 +544,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -552,7 +558,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M22 = spa.csr_matrix( - (self.kernel_1_22.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_22.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M22.eliminate_zeros() @@ -564,7 +571,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -578,7 +585,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M23 = spa.csr_matrix( - (self.kernel_1_23.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_23.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M23.eliminate_zeros() @@ -590,7 +598,7 @@ def assemble_1_form(self, tensor_space_FEM): # convert to sparse matrix indices = xp.indices( - (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1) + (Ni[0], Ni[1], Ni[2], 2 * self.related[0] + 1, 2 * self.related[1] + 1, 2 * self.related[2] + 1), ) shift = [xp.arange(Ni) - offset for Ni, offset in zip(Ni, self.related)] @@ -604,7 +612,8 @@ def assemble_1_form(self, tensor_space_FEM): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M33 = spa.csr_matrix( - (self.kernel_1_33.flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (self.kernel_1_33.flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M33.eliminate_zeros() diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_local_projector_kernel.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_local_projector_kernel.py index c0ebc624d..6db315daa 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_local_projector_kernel.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_local_projector_kernel.py @@ -108,7 +108,8 @@ def kernel_0_form( width[il1] = p[il1] + cell_number[il1] - 1 mat_f = empty( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], num_cell[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], num_cell[2]), + dtype=float, ) mat_f[:, :, :, :, :, :] = 0.0 @@ -335,7 +336,8 @@ def potential_kernel_0_form( width[il1] = p[il1] + cell_number[il1] - 1 mat_f = empty( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], num_cell[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], num_cell[2]), + dtype=float, ) mat_f[:, :, :, :, :, :] = 0.0 @@ -539,33 +541,42 @@ def kernel_1_form( # evaluation of function at interpolation/quadrature points mat_11 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_21 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_31 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_12 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_22 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_32 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_13 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_23 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_33 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) for i1 in range(cell_number[0]): @@ -1259,33 +1270,42 @@ def bv_localproj_push( # evaluation of function at interpolation/quadrature points mat_11 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_21 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_31 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_12 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_22 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_32 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_13 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_23 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_33 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) for i1 in range(cell_number[0]): @@ -1805,33 +1825,42 @@ def kernel_1_heavy( # evaluation of function at interpolation/quadrature points mat_11 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_21 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_31 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_12 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_22 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_32 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_13 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_23 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_33 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) for i1 in range(cell_number[0]): @@ -2373,33 +2402,42 @@ def vv_1_form( # evaluation of function at interpolation/quadrature points mat_11 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_21 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_31 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_12 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_22 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_32 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_13 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_23 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_33 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) for i1 in range(cell_number[0]): @@ -2905,33 +2943,42 @@ def vv_push( # evaluation of function at interpolation/quadrature points mat_11 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_21 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_31 = zeros( - (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], 2, num_cell[1], num_cell[2], quad[0]), + dtype=float, ) mat_12 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_22 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_32 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], 2, num_cell[2], quad[1]), + dtype=float, ) mat_13 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_23 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) mat_33 = zeros( - (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), dtype=float + (cell_number[0], cell_number[1], cell_number[2], num_cell[0], num_cell[1], 2, quad[2]), + dtype=float, ) for i1 in range(cell_number[0]): diff --git a/src/struphy/eigenvalue_solvers/mass_matrices_3d.py b/src/struphy/eigenvalue_solvers/mass_matrices_3d.py index 05f019e13..d3dc4cad2 100644 --- a/src/struphy/eigenvalue_solvers/mass_matrices_3d.py +++ b/src/struphy/eigenvalue_solvers/mass_matrices_3d.py @@ -209,7 +209,8 @@ def get_M1(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M[a][b] = spa.csr_matrix( - (M[a][b].flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (M[a][b].flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M[a][b].eliminate_zeros() @@ -328,7 +329,8 @@ def get_M2(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M[a][b] = spa.csr_matrix( - (M[a][b].flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (M[a][b].flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M[a][b].eliminate_zeros() @@ -563,7 +565,8 @@ def get_Mv(tensor_space_FEM, domain, apply_boundary_ops=False, weight=None): col = Nj[1] * Nj[2] * col1 + Nj[2] * col2 + col3 M[a][b] = spa.csr_matrix( - (M[a][b].flatten(), (row, col.flatten())), shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]) + (M[a][b].flatten(), (row, col.flatten())), + shape=(Ni[0] * Ni[1] * Ni[2], Nj[0] * Nj[1] * Nj[2]), ) M[a][b].eliminate_zeros() diff --git a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py index c44c97335..b8d4aaf81 100644 --- a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py +++ b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py @@ -72,7 +72,12 @@ def solve_mhd_ev_problem_2d(num_params, eq_mhd, n_tor, basis_tor="i", path_out=N # set up 2d tensor-product space space_2d = Tensor_spline_space( - [space_1d_1, space_1d_2], polar_ck, eq_mhd.domain.cx[:, :, 0], eq_mhd.domain.cy[:, :, 0], n_tor, basis_tor + [space_1d_1, space_1d_2], + polar_ck, + eq_mhd.domain.cx[:, :, 0], + eq_mhd.domain.cy[:, :, 0], + n_tor, + basis_tor, ) # set up 2d projectors @@ -141,7 +146,7 @@ def solve_mhd_ev_problem_2d(num_params, eq_mhd, n_tor, basis_tor="i", path_out=N .dot( EF.T.dot(space_2d.C0.conjugate().T.dot(M2_0.dot(space_2d.C0.dot(EF)))) + mhd_ops.MJ_mat.dot(space_2d.C0.dot(EF)) - - space_2d.D0.conjugate().T.dot(M3_0.dot(L)) + - space_2d.D0.conjugate().T.dot(M3_0.dot(L)), ) .toarray() ) @@ -162,7 +167,8 @@ def solve_mhd_ev_problem_2d(num_params, eq_mhd, n_tor, basis_tor="i", path_out=N n_tor_str = "+" + str(n_tor) xp.save( - os.path.join(path_out, "spec_n_" + n_tor_str + ".npy"), xp.vstack((omega2.reshape(1, omega2.size), U2_eig)) + os.path.join(path_out, "spec_n_" + n_tor_str + ".npy"), + xp.vstack((omega2.reshape(1, omega2.size), U2_eig)), ) # or return eigenfrequencies, eigenvectors and system matrix @@ -180,7 +186,7 @@ def solve_mhd_ev_problem_2d(num_params, eq_mhd, n_tor, basis_tor="i", path_out=N # parse arguments parser = argparse.ArgumentParser( - description="Computes the complete eigenspectrum for a given axisymmetric MHD equilibrium." + description="Computes the complete eigenspectrum for a given axisymmetric MHD equilibrium.", ) parser.add_argument("n_tor", type=int, help="the toroidal mode number") diff --git a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py index 10e9e274e..1685147e1 100644 --- a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py +++ b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_pproc.py @@ -21,7 +21,10 @@ def main(): ) parser.add_argument( - "--input-abs", type=str, metavar="DIR", help="directory with eigenspectrum (.npy) file, absolute path" + "--input-abs", + type=str, + metavar="DIR", + help="directory with eigenspectrum (.npy) file, absolute path", ) parser.add_argument("lower", type=float, help="lower range of squared eigenfrequency") diff --git a/src/struphy/eigenvalue_solvers/mhd_operators.py b/src/struphy/eigenvalue_solvers/mhd_operators.py index 5f1462c24..6f7325c6b 100644 --- a/src/struphy/eigenvalue_solvers/mhd_operators.py +++ b/src/struphy/eigenvalue_solvers/mhd_operators.py @@ -658,11 +658,13 @@ def __Mn(self, u): if self.Mn_as_tensor: if self.core.basis_u == 0: out = self.core.space.apply_Mv_ten( - u, [[self.Mn_mat[0], self.core.space.M0_tor], [self.Mn_mat[1], self.core.space.M0_tor]] + u, + [[self.Mn_mat[0], self.core.space.M0_tor], [self.Mn_mat[1], self.core.space.M0_tor]], ) elif self.core.basis_u == 2: out = self.core.space.apply_M2_ten( - u, [[self.Mn_mat[0], self.core.space.M1_tor], [self.Mn_mat[1], self.core.space.M0_tor]] + u, + [[self.Mn_mat[0], self.core.space.M1_tor], [self.Mn_mat[1], self.core.space.M0_tor]], ) else: @@ -681,7 +683,8 @@ def __MJ(self, b): out = xp.zeros(self.core.space.Ev_0.shape[0], dtype=float) elif self.core.basis_u == 2: out = self.core.space.apply_M2_ten( - b, [[self.MJ_mat[0], self.core.space.M1_tor], [self.MJ_mat[1], self.core.space.M0_tor]] + b, + [[self.MJ_mat[0], self.core.space.M1_tor], [self.MJ_mat[1], self.core.space.M0_tor]], ) else: @@ -700,7 +703,7 @@ def __L(self, u): if self.core.basis_u == 0: out = -self.core.space.D0.dot(self.__PF(u)) - (self.gamma - 1) * self.__PR( - self.core.space.D0.dot(self.__JF(u)) + self.core.space.D0.dot(self.__JF(u)), ) elif self.core.basis_u == 2: out = -self.core.space.D0.dot(self.__PF(u)) - (self.gamma - 1) * self.__PR(self.core.space.D0.dot(u)) @@ -782,27 +785,32 @@ def set_operators(self, dt_2=1.0, dt_6=1.0): if hasattr(self, "dofs_Mn"): self.Mn = spa.linalg.LinearOperator( - (self.core.space.Ev_0.shape[0], self.core.space.Ev_0.shape[0]), matvec=self.__Mn + (self.core.space.Ev_0.shape[0], self.core.space.Ev_0.shape[0]), + matvec=self.__Mn, ) if hasattr(self, "dofs_MJ"): self.MJ = spa.linalg.LinearOperator( - (self.core.space.Ev_0.shape[0], self.core.space.E2_0.shape[0]), matvec=self.__MJ + (self.core.space.Ev_0.shape[0], self.core.space.E2_0.shape[0]), + matvec=self.__MJ, ) if hasattr(self, "dofs_PF") and hasattr(self, "dofs_PR") and hasattr(self, "dofs_JF"): self.L = spa.linalg.LinearOperator( - (self.core.space.E3_0.shape[0], self.core.space.Ev_0.shape[0]), matvec=self.__L + (self.core.space.E3_0.shape[0], self.core.space.Ev_0.shape[0]), + matvec=self.__L, ) if hasattr(self, "Mn_mat") and hasattr(self, "dofs_EF"): self.S2 = spa.linalg.LinearOperator( - (self.core.space.Ev_0.shape[0], self.core.space.Ev_0.shape[0]), matvec=self.__S2 + (self.core.space.Ev_0.shape[0], self.core.space.Ev_0.shape[0]), + matvec=self.__S2, ) if hasattr(self, "Mn_mat") and hasattr(self, "L"): self.S6 = spa.linalg.LinearOperator( - (self.core.space.Ev_0.shape[0], self.core.space.Ev_0.shape[0]), matvec=self.__S6 + (self.core.space.Ev_0.shape[0], self.core.space.Ev_0.shape[0]), + matvec=self.__S6, ) elif self.core.basis_u == 2: @@ -836,27 +844,32 @@ def set_operators(self, dt_2=1.0, dt_6=1.0): if hasattr(self, "Mn_mat"): self.Mn = spa.linalg.LinearOperator( - (self.core.space.E2_0.shape[0], self.core.space.E2_0.shape[0]), matvec=self.__Mn + (self.core.space.E2_0.shape[0], self.core.space.E2_0.shape[0]), + matvec=self.__Mn, ) if hasattr(self, "MJ_mat"): self.MJ = spa.linalg.LinearOperator( - (self.core.space.E2_0.shape[0], self.core.space.E2_0.shape[0]), matvec=self.__MJ + (self.core.space.E2_0.shape[0], self.core.space.E2_0.shape[0]), + matvec=self.__MJ, ) if hasattr(self, "dofs_PF") and hasattr(self, "dofs_PR"): self.L = spa.linalg.LinearOperator( - (self.core.space.E3_0.shape[0], self.core.space.E2_0.shape[0]), matvec=self.__L + (self.core.space.E3_0.shape[0], self.core.space.E2_0.shape[0]), + matvec=self.__L, ) if hasattr(self, "Mn_mat") and hasattr(self, "dofs_EF"): self.S2 = spa.linalg.LinearOperator( - (self.core.space.E2_0.shape[0], self.core.space.E2_0.shape[0]), matvec=self.__S2 + (self.core.space.E2_0.shape[0], self.core.space.E2_0.shape[0]), + matvec=self.__S2, ) if hasattr(self, "Mn_mat") and hasattr(self, "L"): self.S6 = spa.linalg.LinearOperator( - (self.core.space.E2_0.shape[0], self.core.space.E2_0.shape[0]), matvec=self.__S6 + (self.core.space.E2_0.shape[0], self.core.space.E2_0.shape[0]), + matvec=self.__S6, ) # ====================================== @@ -1024,7 +1037,7 @@ def set_preconditioner_S2(self, which, tol_inv=1e-15, drop_tol=1e-4, fill_fac=10 # assemble approximate S2 matrix S2_approx = Mn + self.dt_2**2 / 4 * EF_approx.T.dot( - self.core.space.C0.T.dot(M2_0.dot(self.core.space.C0.dot(EF_approx))) + self.core.space.C0.T.dot(M2_0.dot(self.core.space.C0.dot(EF_approx))), ) del Mn, EF_approx, M2_0 @@ -1123,7 +1136,7 @@ def set_preconditioner_S6(self, which, tol_inv=1e-15, drop_tol=1e-4, fill_fac=10 # assemble approximate L matrix if self.core.basis_u == 0: L_approx = -self.core.space.D0.dot(PF_approx) - (self.gamma - 1) * PR_approx.dot( - self.core.space.D0.dot(JF_approx) + self.core.space.D0.dot(JF_approx), ) del PF_approx, PR_approx diff --git a/src/struphy/eigenvalue_solvers/mhd_operators_core.py b/src/struphy/eigenvalue_solvers/mhd_operators_core.py index 76752ccc3..61d534148 100644 --- a/src/struphy/eigenvalue_solvers/mhd_operators_core.py +++ b/src/struphy/eigenvalue_solvers/mhd_operators_core.py @@ -139,7 +139,8 @@ def get_blocks_EF(self, pol=True): ) EF_12 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[0] // self.N3, self.space.Ntot_0form // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_1form[0] // self.N3, self.space.Ntot_0form // self.N3), ) EF_12.eliminate_zeros() # ---------------------------------------------------- @@ -173,7 +174,8 @@ def get_blocks_EF(self, pol=True): ) EF_13 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[0] // self.N3, self.space.Ntot_0form // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_1form[0] // self.N3, self.space.Ntot_0form // self.N3), ) EF_13.eliminate_zeros() # ---------------------------------------------------- @@ -207,7 +209,8 @@ def get_blocks_EF(self, pol=True): ) EF_21 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[1] // self.N3, self.space.Ntot_0form // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_1form[1] // self.N3, self.space.Ntot_0form // self.N3), ) EF_21.eliminate_zeros() # ---------------------------------------------------- @@ -241,7 +244,8 @@ def get_blocks_EF(self, pol=True): ) EF_23 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[1] // self.N3, self.space.Ntot_0form // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_1form[1] // self.N3, self.space.Ntot_0form // self.N3), ) EF_23.eliminate_zeros() # ---------------------------------------------------- @@ -269,7 +273,8 @@ def get_blocks_EF(self, pol=True): ) EF_31 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[2] // self.D3, self.space.Ntot_0form // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_1form[2] // self.D3, self.space.Ntot_0form // self.N3), ) EF_31.eliminate_zeros() # ---------------------------------------------------- @@ -297,7 +302,8 @@ def get_blocks_EF(self, pol=True): ) EF_32 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[2] // self.D3, self.space.Ntot_0form // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_1form[2] // self.D3, self.space.Ntot_0form // self.N3), ) EF_32.eliminate_zeros() # ---------------------------------------------------- @@ -310,13 +316,16 @@ def get_blocks_EF(self, pol=True): # assemble sparse matrix val = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float + self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) ker.rhs11( @@ -351,13 +360,16 @@ def get_blocks_EF(self, pol=True): # assemble sparse matrix val = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float + self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) ker.rhs11( @@ -392,13 +404,16 @@ def get_blocks_EF(self, pol=True): # assemble sparse matrix val = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float + self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) ker.rhs12( @@ -433,13 +448,16 @@ def get_blocks_EF(self, pol=True): # assemble sparse matrix val = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float + self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) ker.rhs12( @@ -474,13 +492,16 @@ def get_blocks_EF(self, pol=True): # assemble sparse matrix val = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=float + self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=int, ) ker.rhs13( @@ -515,13 +536,16 @@ def get_blocks_EF(self, pol=True): # assemble sparse matrix val = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=float + self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=int, ) ker.rhs13( @@ -584,7 +608,8 @@ def get_blocks_EF(self, pol=True): ) EF_12 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[0] // self.N3, self.space.Ntot_2form[1] // self.D3) + (val, (row, col)), + shape=(self.space.Ntot_1form[0] // self.N3, self.space.Ntot_2form[1] // self.D3), ) EF_12.eliminate_zeros() # ---------------------------------------------------- @@ -622,7 +647,8 @@ def get_blocks_EF(self, pol=True): ) EF_13 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[0] // self.N3, self.space.Ntot_2form[2] // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_1form[0] // self.N3, self.space.Ntot_2form[2] // self.N3), ) EF_13.eliminate_zeros() # ---------------------------------------------------- @@ -660,7 +686,8 @@ def get_blocks_EF(self, pol=True): ) EF_21 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[1] // self.N3, self.space.Ntot_2form[0] // self.D3) + (val, (row, col)), + shape=(self.space.Ntot_1form[1] // self.N3, self.space.Ntot_2form[0] // self.D3), ) EF_21.eliminate_zeros() # ---------------------------------------------------- @@ -698,7 +725,8 @@ def get_blocks_EF(self, pol=True): ) EF_23 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[1] // self.N3, self.space.Ntot_2form[2] // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_1form[1] // self.N3, self.space.Ntot_2form[2] // self.N3), ) EF_23.eliminate_zeros() # ---------------------------------------------------- @@ -729,7 +757,8 @@ def get_blocks_EF(self, pol=True): ) EF_31 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[2] // self.D3, self.space.Ntot_2form[0] // self.D3) + (val, (row, col)), + shape=(self.space.Ntot_1form[2] // self.D3, self.space.Ntot_2form[0] // self.D3), ) EF_31.eliminate_zeros() # ---------------------------------------------------- @@ -760,7 +789,8 @@ def get_blocks_EF(self, pol=True): ) EF_32 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_1form[2] // self.D3, self.space.Ntot_2form[1] // self.D3) + (val, (row, col)), + shape=(self.space.Ntot_1form[2] // self.D3, self.space.Ntot_2form[1] // self.D3), ) EF_32.eliminate_zeros() # ---------------------------------------------------- @@ -773,19 +803,22 @@ def get_blocks_EF(self, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( - self.equilibrium.domain.jacobian_det(self.eta_his[0].flatten(), self.eta_int[1], self.eta_int[2]) + self.equilibrium.domain.jacobian_det(self.eta_his[0].flatten(), self.eta_int[1], self.eta_int[2]), ) det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nint[2]) # assemble sparse matrix val = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=float + self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_D_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_D_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_0_D_i[2][0].size, + dtype=int, ) ker.rhs11( @@ -820,19 +853,22 @@ def get_blocks_EF(self, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( - self.equilibrium.domain.jacobian_det(self.eta_his[0].flatten(), self.eta_int[1], self.eta_int[2]) + self.equilibrium.domain.jacobian_det(self.eta_his[0].flatten(), self.eta_int[1], self.eta_int[2]), ) det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nint[2]) # assemble sparse matrix val = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float + self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) ker.rhs11( @@ -867,19 +903,22 @@ def get_blocks_EF(self, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( - self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_his[1].flatten(), self.eta_int[2]) + self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_his[1].flatten(), self.eta_int[2]), ) det_dF = det_dF.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix val = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=float + self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_D_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_D_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_D_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_D_i[2][0].size, + dtype=int, ) ker.rhs12( @@ -914,19 +953,22 @@ def get_blocks_EF(self, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( - self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_his[1].flatten(), self.eta_int[2]) + self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_his[1].flatten(), self.eta_int[2]), ) det_dF = det_dF.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix val = xp.empty( - self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float + self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_0_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) ker.rhs12( @@ -961,19 +1003,22 @@ def get_blocks_EF(self, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( - self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_int[1], self.eta_his[2].flatten()) + self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_int[1], self.eta_his[2].flatten()), ) det_dF = det_dF.reshape(self.nint[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix val = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float + self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_0_D_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) ker.rhs13( @@ -1008,19 +1053,22 @@ def get_blocks_EF(self, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( - self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_int[1], self.eta_his[2].flatten()) + self.equilibrium.domain.jacobian_det(self.eta_int[0], self.eta_int[1], self.eta_his[2].flatten()), ) det_dF = det_dF.reshape(self.nint[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix val = xp.empty( - self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float + self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_0_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) ker.rhs13( @@ -1116,7 +1164,8 @@ def get_blocks_FL(self, which, pol=True): ) F_11 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_2form[0] // self.D3, self.space.Ntot_0form // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_2form[0] // self.D3, self.space.Ntot_0form // self.N3), ) F_11.eliminate_zeros() # ------------------------------------------------------------ @@ -1156,7 +1205,8 @@ def get_blocks_FL(self, which, pol=True): ) F_22 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_2form[1] // self.D3, self.space.Ntot_0form // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_2form[1] // self.D3, self.space.Ntot_0form // self.N3), ) F_22.eliminate_zeros() # ------------------------------------------------------------ @@ -1199,7 +1249,8 @@ def get_blocks_FL(self, which, pol=True): ) F_33 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_2form[2] // self.N3, self.space.Ntot_0form // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_2form[2] // self.N3, self.space.Ntot_0form // self.N3), ) F_33.eliminate_zeros() # ------------------------------------------------------------ @@ -1213,20 +1264,25 @@ def get_blocks_FL(self, which, pol=True): EQ = self.equilibrium.p3(self.eta_int[0], self.eta_his[1].flatten(), self.eta_his[2].flatten()) else: EQ = self.equilibrium.domain.jacobian_det( - self.eta_int[0], self.eta_his[1].flatten(), self.eta_his[2].flatten() + self.eta_int[0], + self.eta_his[1].flatten(), + self.eta_his[2].flatten(), ) EQ = EQ.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nhis[2], self.nq[2]) # assemble sparse matrix val = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=float + self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=int, ) ker.rhs21( @@ -1265,20 +1321,25 @@ def get_blocks_FL(self, which, pol=True): EQ = self.equilibrium.p3(self.eta_his[0].flatten(), self.eta_int[1], self.eta_his[2].flatten()) else: EQ = self.equilibrium.domain.jacobian_det( - self.eta_his[0].flatten(), self.eta_int[1], self.eta_his[2].flatten() + self.eta_his[0].flatten(), + self.eta_int[1], + self.eta_his[2].flatten(), ) EQ = EQ.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix val = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=float + self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int + self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, dtype=int + self.dofs_1_N_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_N_i[2][0].size, + dtype=int, ) ker.rhs22( @@ -1317,20 +1378,25 @@ def get_blocks_FL(self, which, pol=True): EQ = self.equilibrium.p3(self.eta_his[0].flatten(), self.eta_his[1].flatten(), self.eta_int[2]) else: EQ = self.equilibrium.domain.jacobian_det( - self.eta_his[0].flatten(), self.eta_his[1].flatten(), self.eta_int[2] + self.eta_his[0].flatten(), + self.eta_his[1].flatten(), + self.eta_int[2], ) EQ = EQ.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix val = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float + self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_N_i[0][0].size * self.dofs_1_N_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) ker.rhs23( @@ -1400,7 +1466,8 @@ def get_blocks_FL(self, which, pol=True): ) F_11 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_2form[0] // self.D3, self.space.Ntot_2form[0] // self.D3) + (val, (row, col)), + shape=(self.space.Ntot_2form[0] // self.D3, self.space.Ntot_2form[0] // self.D3), ) F_11.eliminate_zeros() # ------------------------------------------------------------ @@ -1442,7 +1509,8 @@ def get_blocks_FL(self, which, pol=True): ) F_22 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_2form[1] // self.D3, self.space.Ntot_2form[1] // self.D3) + (val, (row, col)), + shape=(self.space.Ntot_2form[1] // self.D3, self.space.Ntot_2form[1] // self.D3), ) F_22.eliminate_zeros() # ------------------------------------------------------------ @@ -1458,7 +1526,7 @@ def get_blocks_FL(self, which, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( - self.equilibrium.domain.jacobian_det(self.eta_his[0].flatten(), self.eta_his[1].flatten(), 0.0) + self.equilibrium.domain.jacobian_det(self.eta_his[0].flatten(), self.eta_his[1].flatten(), 0.0), ) det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1]) @@ -1489,7 +1557,8 @@ def get_blocks_FL(self, which, pol=True): ) F_33 = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_2form[2] // self.N3, self.space.Ntot_2form[2] // self.N3) + (val, (row, col)), + shape=(self.space.Ntot_2form[2] // self.N3, self.space.Ntot_2form[2] // self.N3), ) F_33.eliminate_zeros() # ------------------------------------------------------------ @@ -1507,20 +1576,25 @@ def get_blocks_FL(self, which, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( self.equilibrium.domain.jacobian_det( - self.eta_int[0], self.eta_his[1].flatten(), self.eta_his[2].flatten() - ) + self.eta_int[0], + self.eta_his[1].flatten(), + self.eta_his[2].flatten(), + ), ) det_dF = det_dF.reshape(self.nint[0], self.nhis[1], self.nq[1], self.nhis[2], self.nq[2]) # assemble sparse matrix val = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float + self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_0_N_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) ker.rhs21( @@ -1563,20 +1637,25 @@ def get_blocks_FL(self, which, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( self.equilibrium.domain.jacobian_det( - self.eta_his[0].flatten(), self.eta_int[1], self.eta_his[2].flatten() - ) + self.eta_his[0].flatten(), + self.eta_int[1], + self.eta_his[2].flatten(), + ), ) det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nint[1], self.nhis[2], self.nq[2]) # assemble sparse matrix val = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float + self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_0_N_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) ker.rhs22( @@ -1619,20 +1698,25 @@ def get_blocks_FL(self, which, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( self.equilibrium.domain.jacobian_det( - self.eta_his[0].flatten(), self.eta_his[1].flatten(), self.eta_int[2] - ) + self.eta_his[0].flatten(), + self.eta_his[1].flatten(), + self.eta_int[2], + ), ) det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1], self.nint[2]) # assemble sparse matrix val = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=float + self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_0_N_i[2][0].size, + dtype=int, ) ker.rhs23( @@ -1691,7 +1775,7 @@ def get_blocks_PR(self, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( - self.equilibrium.domain.jacobian_det(self.eta_his[0].flatten(), self.eta_his[1].flatten(), 0.0) + self.equilibrium.domain.jacobian_det(self.eta_his[0].flatten(), self.eta_his[1].flatten(), 0.0), ) det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1]) @@ -1722,7 +1806,8 @@ def get_blocks_PR(self, pol=True): ) PR = spa.csr_matrix( - (val, (row, col)), shape=(self.space.Ntot_3form // self.D3, self.space.Ntot_3form // self.D3) + (val, (row, col)), + shape=(self.space.Ntot_3form // self.D3, self.space.Ntot_3form // self.D3), ) PR.eliminate_zeros() # ----------------------------------------------------- @@ -1731,7 +1816,9 @@ def get_blocks_PR(self, pol=True): # --------------- ([his, his, his] of DDD) ------------ # evaluate equilibrium pressure at quadrature points P3_pts = self.equilibrium.p3( - self.eta_his[0].flatten(), self.eta_his[1].flatten(), self.eta_his[2].flatten() + self.eta_his[0].flatten(), + self.eta_his[1].flatten(), + self.eta_his[2].flatten(), ) P3_pts = P3_pts.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1], self.nhis[2], self.nq[2]) @@ -1739,20 +1826,25 @@ def get_blocks_PR(self, pol=True): # evaluate Jacobian determinant at at interpolation and quadrature points det_dF = abs( self.equilibrium.domain.jacobian_det( - self.eta_his[0].flatten(), self.eta_his[1].flatten(), self.eta_his[2].flatten() - ) + self.eta_his[0].flatten(), + self.eta_his[1].flatten(), + self.eta_his[2].flatten(), + ), ) det_dF = det_dF.reshape(self.nhis[0], self.nq[0], self.nhis[1], self.nq[1], self.nhis[2], self.nq[2]) # assemble sparse matrix val = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=float + self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=float, ) row = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) col = xp.empty( - self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, dtype=int + self.dofs_1_D_i[0][0].size * self.dofs_1_D_i[1][0].size * self.dofs_1_D_i[2][0].size, + dtype=int, ) ker.rhs3( diff --git a/src/struphy/eigenvalue_solvers/projectors_global.py b/src/struphy/eigenvalue_solvers/projectors_global.py index fa7b69958..ca67c66e6 100644 --- a/src/struphy/eigenvalue_solvers/projectors_global.py +++ b/src/struphy/eigenvalue_solvers/projectors_global.py @@ -271,10 +271,18 @@ def __init__(self, spline_space, n_quad=6): BM_splines = [False, True] self.N_int = bsp.collocation_matrix( - spline_space.T, spline_space.p - 0, self.x_int, spline_space.spl_kind, BM_splines[0] + spline_space.T, + spline_space.p - 0, + self.x_int, + spline_space.spl_kind, + BM_splines[0], ) self.D_int = bsp.collocation_matrix( - spline_space.t, spline_space.p - 1, self.x_int, spline_space.spl_kind, BM_splines[1] + spline_space.t, + spline_space.p - 1, + self.x_int, + spline_space.spl_kind, + BM_splines[1], ) self.N_int[self.N_int < 1e-12] = 0.0 @@ -284,10 +292,18 @@ def __init__(self, spline_space, n_quad=6): self.D_int = spa.csr_matrix(self.D_int) self.N_pts = bsp.collocation_matrix( - spline_space.T, spline_space.p - 0, self.pts.flatten(), spline_space.spl_kind, BM_splines[0] + spline_space.T, + spline_space.p - 0, + self.pts.flatten(), + spline_space.spl_kind, + BM_splines[0], ) self.D_pts = bsp.collocation_matrix( - spline_space.t, spline_space.p - 1, self.pts.flatten(), spline_space.spl_kind, BM_splines[1] + spline_space.t, + spline_space.p - 1, + self.pts.flatten(), + spline_space.spl_kind, + BM_splines[1], ) self.N_pts = spa.csr_matrix(self.N_pts) @@ -542,29 +558,29 @@ def D_jD_k(eta): dofs_1_DD_i_red[i] = xp.nonzero(un == nv[i])[0] dofs_0_NN_indices = xp.vstack( - (dofs_0_NN_indices[0], dofs_0_NN_indices[1], dofs_0_NN_indices[2], dofs_0_NN_i_red) + (dofs_0_NN_indices[0], dofs_0_NN_indices[1], dofs_0_NN_indices[2], dofs_0_NN_i_red), ) dofs_0_DN_indices = xp.vstack( - (dofs_0_DN_indices[0], dofs_0_DN_indices[1], dofs_0_DN_indices[2], dofs_0_DN_i_red) + (dofs_0_DN_indices[0], dofs_0_DN_indices[1], dofs_0_DN_indices[2], dofs_0_DN_i_red), ) dofs_0_ND_indices = xp.vstack( - (dofs_0_ND_indices[0], dofs_0_ND_indices[1], dofs_0_ND_indices[2], dofs_0_ND_i_red) + (dofs_0_ND_indices[0], dofs_0_ND_indices[1], dofs_0_ND_indices[2], dofs_0_ND_i_red), ) dofs_0_DD_indices = xp.vstack( - (dofs_0_DD_indices[0], dofs_0_DD_indices[1], dofs_0_DD_indices[2], dofs_0_DD_i_red) + (dofs_0_DD_indices[0], dofs_0_DD_indices[1], dofs_0_DD_indices[2], dofs_0_DD_i_red), ) dofs_1_NN_indices = xp.vstack( - (dofs_1_NN_indices[0], dofs_1_NN_indices[1], dofs_1_NN_indices[2], dofs_1_NN_i_red) + (dofs_1_NN_indices[0], dofs_1_NN_indices[1], dofs_1_NN_indices[2], dofs_1_NN_i_red), ) dofs_1_DN_indices = xp.vstack( - (dofs_1_DN_indices[0], dofs_1_DN_indices[1], dofs_1_DN_indices[2], dofs_1_DN_i_red) + (dofs_1_DN_indices[0], dofs_1_DN_indices[1], dofs_1_DN_indices[2], dofs_1_DN_i_red), ) dofs_1_ND_indices = xp.vstack( - (dofs_1_ND_indices[0], dofs_1_ND_indices[1], dofs_1_ND_indices[2], dofs_1_ND_i_red) + (dofs_1_ND_indices[0], dofs_1_ND_indices[1], dofs_1_ND_indices[2], dofs_1_ND_i_red), ) dofs_1_DD_indices = xp.vstack( - (dofs_1_DD_indices[0], dofs_1_DD_indices[1], dofs_1_DD_indices[2], dofs_1_DD_i_red) + (dofs_1_DD_indices[0], dofs_1_DD_indices[1], dofs_1_DD_indices[2], dofs_1_DD_i_red), ) return ( @@ -1119,7 +1135,8 @@ def dofs(self, comp, mat_f): if comp == "0": dofs = kron_matvec_3d( - [spa.identity(mat_f.shape[0]), spa.identity(mat_f.shape[1]), spa.identity(mat_f.shape[2])], mat_f + [spa.identity(mat_f.shape[0]), spa.identity(mat_f.shape[1]), spa.identity(mat_f.shape[2])], + mat_f, ) elif comp == "11": @@ -1172,15 +1189,18 @@ def dofs_T(self, comp, mat_dofs): elif comp == "11": rhs = kron_matvec_3d( - [self.Q1.T, spa.identity(mat_dofs.shape[1]), spa.identity(mat_dofs.shape[2])], mat_dofs + [self.Q1.T, spa.identity(mat_dofs.shape[1]), spa.identity(mat_dofs.shape[2])], + mat_dofs, ) elif comp == "12": rhs = kron_matvec_3d( - [spa.identity(mat_dofs.shape[0]), self.Q2.T, spa.identity(mat_dofs.shape[2])], mat_dofs + [spa.identity(mat_dofs.shape[0]), self.Q2.T, spa.identity(mat_dofs.shape[2])], + mat_dofs, ) elif comp == "13": rhs = kron_matvec_3d( - [spa.identity(mat_dofs.shape[0]), spa.identity(mat_dofs.shape[1]), self.Q3.T], mat_dofs + [spa.identity(mat_dofs.shape[0]), spa.identity(mat_dofs.shape[1]), self.Q3.T], + mat_dofs, ) elif comp == "21": @@ -1836,10 +1856,12 @@ def solve_V1(self, dofs_1, include_bc): # with boundary splines if include_bc: dofs_11 = dofs_1[: self.P1_pol.shape[0] * self.I_tor.shape[0]].reshape( - self.P1_pol.shape[0], self.I_tor.shape[0] + self.P1_pol.shape[0], + self.I_tor.shape[0], ) dofs_12 = dofs_1[self.P1_pol.shape[0] * self.I_tor.shape[0] :].reshape( - self.P0_pol.shape[0], self.H_tor.shape[0] + self.P0_pol.shape[0], + self.H_tor.shape[0], ) coeffs1 = self.I_tor_LU.solve(self.I1_pol_LU.solve(dofs_11).T).T @@ -1848,10 +1870,12 @@ def solve_V1(self, dofs_1, include_bc): # without boundary splines else: dofs_11 = dofs_1[: self.P1_pol_0.shape[0] * self.I0_tor.shape[0]].reshape( - self.P1_pol_0.shape[0], self.I0_tor.shape[0] + self.P1_pol_0.shape[0], + self.I0_tor.shape[0], ) dofs_12 = dofs_1[self.P1_pol_0.shape[0] * self.I0_tor.shape[0] :].reshape( - self.P0_pol_0.shape[0], self.H0_tor.shape[0] + self.P0_pol_0.shape[0], + self.H0_tor.shape[0], ) coeffs1 = self.I0_tor_LU.solve(self.I1_pol_0_LU.solve(dofs_11).T).T @@ -1864,10 +1888,12 @@ def solve_V2(self, dofs_2, include_bc): # with boundary splines if include_bc: dofs_21 = dofs_2[: self.P2_pol.shape[0] * self.H_tor.shape[0]].reshape( - self.P2_pol.shape[0], self.H_tor.shape[0] + self.P2_pol.shape[0], + self.H_tor.shape[0], ) dofs_22 = dofs_2[self.P2_pol.shape[0] * self.H_tor.shape[0] :].reshape( - self.P3_pol.shape[0], self.I_tor.shape[0] + self.P3_pol.shape[0], + self.I_tor.shape[0], ) coeffs1 = self.H_tor_LU.solve(self.I2_pol_LU.solve(dofs_21).T).T @@ -1876,10 +1902,12 @@ def solve_V2(self, dofs_2, include_bc): # without boundary splines else: dofs_21 = dofs_2[: self.P2_pol_0.shape[0] * self.H0_tor.shape[0]].reshape( - self.P2_pol_0.shape[0], self.H0_tor.shape[0] + self.P2_pol_0.shape[0], + self.H0_tor.shape[0], ) dofs_22 = dofs_2[self.P2_pol_0.shape[0] * self.H0_tor.shape[0] :].reshape( - self.P3_pol_0.shape[0], self.I0_tor.shape[0] + self.P3_pol_0.shape[0], + self.I0_tor.shape[0], ) coeffs1 = self.H0_tor_LU.solve(self.I2_pol_0_LU.solve(dofs_21).T).T @@ -1938,10 +1966,12 @@ def apply_IinvT_V1(self, rhs, include_bc=False): # without boundary splines else: rhs1 = rhs[: self.P1_pol_0.shape[0] * self.I0_tor.shape[0]].reshape( - self.P1_pol_0.shape[0], self.I0_tor.shape[0] + self.P1_pol_0.shape[0], + self.I0_tor.shape[0], ) rhs2 = rhs[self.P1_pol_0.shape[0] * self.I0_tor.shape[0] :].reshape( - self.P0_pol_0.shape[0], self.H0_tor.shape[0] + self.P0_pol_0.shape[0], + self.H0_tor.shape[0], ) rhs1 = self.I1_pol_0_T_LU.solve(self.I0_tor_T_LU.solve(rhs1.T).T) @@ -1968,10 +1998,12 @@ def apply_IinvT_V2(self, rhs, include_bc=False): # without boundary splines else: rhs1 = rhs[: self.P2_pol_0.shape[0] * self.H0_tor.shape[0]].reshape( - self.P2_pol_0.shape[0], self.H0_tor.shape[0] + self.P2_pol_0.shape[0], + self.H0_tor.shape[0], ) rhs2 = rhs[self.P2_pol_0.shape[0] * self.H0_tor.shape[0] :].reshape( - self.P3_pol_0.shape[0], self.I0_tor.shape[0] + self.P3_pol_0.shape[0], + self.I0_tor.shape[0], ) rhs1 = self.I2_pol_0_T_LU.solve(self.H0_tor_T_LU.solve(rhs1.T).T) @@ -2004,7 +2036,8 @@ def dofs_0(self, fun, include_bc=True, eval_kind="meshgrid"): # get dofs on tensor-product grid dofs = kron_matvec_3d( - [spa.identity(dofs.shape[0]), spa.identity(dofs.shape[1]), spa.identity(dofs.shape[2])], dofs + [spa.identity(dofs.shape[0]), spa.identity(dofs.shape[1]), spa.identity(dofs.shape[2])], + dofs, ) # apply extraction operator for dofs diff --git a/src/struphy/eigenvalue_solvers/spline_space.py b/src/struphy/eigenvalue_solvers/spline_space.py index b36ad73db..c10124e57 100644 --- a/src/struphy/eigenvalue_solvers/spline_space.py +++ b/src/struphy/eigenvalue_solvers/spline_space.py @@ -712,27 +712,33 @@ def __init__(self, spline_spaces, ck=-1, cx=None, cy=None, n_tor=0, basis_tor="r # extraction operators for 3D diagram: without boundary conditions self.E0 = spa.kron(self.E0_pol, self.E0_tor, format="csr") self.E1 = spa.bmat( - [[spa.kron(self.E1_pol, self.E0_tor), None], [None, spa.kron(self.E0_pol, self.E1_tor)]], format="csr" + [[spa.kron(self.E1_pol, self.E0_tor), None], [None, spa.kron(self.E0_pol, self.E1_tor)]], + format="csr", ) self.E2 = spa.bmat( - [[spa.kron(self.E2_pol, self.E1_tor), None], [None, spa.kron(self.E3_pol, self.E0_tor)]], format="csr" + [[spa.kron(self.E2_pol, self.E1_tor), None], [None, spa.kron(self.E3_pol, self.E0_tor)]], + format="csr", ) self.E3 = spa.kron(self.E3_pol, self.E1_tor, format="csr") self.Ev = spa.bmat( - [[spa.kron(self.Ev_pol, self.E0_tor), None], [None, spa.kron(self.E0_pol, self.E0_tor)]], format="csr" + [[spa.kron(self.Ev_pol, self.E0_tor), None], [None, spa.kron(self.E0_pol, self.E0_tor)]], + format="csr", ) # boundary operators for 3D diagram self.B0 = spa.kron(self.B0_pol, self.B0_tor, format="csr") self.B1 = spa.bmat( - [[spa.kron(self.B1_pol, self.B0_tor), None], [None, spa.kron(self.B0_pol, self.B1_tor)]], format="csr" + [[spa.kron(self.B1_pol, self.B0_tor), None], [None, spa.kron(self.B0_pol, self.B1_tor)]], + format="csr", ) self.B2 = spa.bmat( - [[spa.kron(self.B2_pol, self.B1_tor), None], [None, spa.kron(self.B3_pol, self.B0_tor)]], format="csr" + [[spa.kron(self.B2_pol, self.B1_tor), None], [None, spa.kron(self.B3_pol, self.B0_tor)]], + format="csr", ) self.B3 = spa.kron(self.B3_pol, self.B1_tor, format="csr") self.Bv = spa.bmat( - [[spa.kron(self.Bv_pol, self.E0_tor), None], [None, spa.kron(Bv3, self.B0_tor)]], format="csr" + [[spa.kron(self.Bv_pol, self.E0_tor), None], [None, spa.kron(Bv3, self.B0_tor)]], + format="csr", ) # extraction operators for 3D diagram: with boundary conditions @@ -841,10 +847,10 @@ def apply_M1_0_ten(self, x, mats): x1, x2 = self.reshape_pol_1(x) out1 = self.B0_tor.dot( - mats[0][1].dot(self.B0_tor.T.dot(self.B1_pol.dot(mats[0][0].dot(self.B1_pol.T.dot(x1))).T)) + mats[0][1].dot(self.B0_tor.T.dot(self.B1_pol.dot(mats[0][0].dot(self.B1_pol.T.dot(x1))).T)), ).T out2 = self.B1_tor.dot( - mats[1][1].dot(self.B1_tor.T.dot(self.B0_pol.dot(mats[1][0].dot(self.B0_pol.T.dot(x2))).T)) + mats[1][1].dot(self.B1_tor.T.dot(self.B0_pol.dot(mats[1][0].dot(self.B0_pol.T.dot(x2))).T)), ).T return xp.concatenate((out1.flatten(), out2.flatten())) @@ -857,10 +863,10 @@ def apply_M2_0_ten(self, x, mats): x1, x2 = self.reshape_pol_2(x) out1 = self.B1_tor.dot( - mats[0][1].dot(self.B1_tor.T.dot(self.B2_pol.dot(mats[0][0].dot(self.B2_pol.T.dot(x1))).T)) + mats[0][1].dot(self.B1_tor.T.dot(self.B2_pol.dot(mats[0][0].dot(self.B2_pol.T.dot(x1))).T)), ).T out2 = self.B0_tor.dot( - mats[1][1].dot(self.B0_tor.T.dot(self.B3_pol.dot(mats[1][0].dot(self.B3_pol.T.dot(x2))).T)) + mats[1][1].dot(self.B0_tor.T.dot(self.B3_pol.dot(mats[1][0].dot(self.B3_pol.T.dot(x2))).T)), ).T return xp.concatenate((out1.flatten(), out2.flatten())) @@ -928,10 +934,12 @@ def __assemble_M1(self, domain, as_tensor=False): self.M1_pol_mat = mass_2d.get_M1(self, domain) matvec = lambda x: self.apply_M1_ten( - x, [[self.M1_pol_mat[0], self.M0_tor], [self.M1_pol_mat[1], self.M1_tor]] + x, + [[self.M1_pol_mat[0], self.M0_tor], [self.M1_pol_mat[1], self.M1_tor]], ) matvec_0 = lambda x: self.apply_M1_0_ten( - x, [[self.M1_pol_mat[0], self.M0_tor], [self.M1_pol_mat[1], self.M1_tor]] + x, + [[self.M1_pol_mat[0], self.M0_tor], [self.M1_pol_mat[1], self.M1_tor]], ) # 3D @@ -939,7 +947,8 @@ def __assemble_M1(self, domain, as_tensor=False): if self.dim == 2: M11, M22 = mass_2d.get_M1(self, domain) self.M1_mat = spa.bmat( - [[spa.kron(M11, self.M0_tor), None], [None, spa.kron(M22, self.M1_tor)]], format="csr" + [[spa.kron(M11, self.M0_tor), None], [None, spa.kron(M22, self.M1_tor)]], + format="csr", ) else: self.M1_mat = mass_3d.get_M1(self, domain) @@ -963,10 +972,12 @@ def __assemble_M2(self, domain, as_tensor=False): self.M2_pol_mat = mass_2d.get_M2(self, domain) matvec = lambda x: self.apply_M2_ten( - x, [[self.M2_pol_mat[0], self.M1_tor], [self.M2_pol_mat[1], self.M0_tor]] + x, + [[self.M2_pol_mat[0], self.M1_tor], [self.M2_pol_mat[1], self.M0_tor]], ) matvec_0 = lambda x: self.apply_M2_0_ten( - x, [[self.M2_pol_mat[0], self.M1_tor], [self.M2_pol_mat[1], self.M0_tor]] + x, + [[self.M2_pol_mat[0], self.M1_tor], [self.M2_pol_mat[1], self.M0_tor]], ) # 3D @@ -974,7 +985,8 @@ def __assemble_M2(self, domain, as_tensor=False): if self.dim == 2: M11, M22 = mass_2d.get_M2(self, domain) self.M2_mat = spa.bmat( - [[spa.kron(M11, self.M1_tor), None], [None, spa.kron(M22, self.M0_tor)]], format="csr" + [[spa.kron(M11, self.M1_tor), None], [None, spa.kron(M22, self.M0_tor)]], + format="csr", ) else: self.M2_mat = mass_3d.get_M2(self, domain) @@ -1026,10 +1038,12 @@ def __assemble_Mv(self, domain, as_tensor=False): self.Mv_pol_mat = mass_2d.get_Mv(self, domain) matvec = lambda x: self.apply_Mv_ten( - x, [[self.Mv_pol_mat[0], self.M0_tor], [self.Mv_pol_mat[1], self.M0_tor]] + x, + [[self.Mv_pol_mat[0], self.M0_tor], [self.Mv_pol_mat[1], self.M0_tor]], ) matvec_0 = lambda x: self.apply_Mv_0_ten( - x, [[self.Mv_pol_mat[0], self.M0_tor], [self.Mv_pol_mat[1], self.M0_tor]] + x, + [[self.Mv_pol_mat[0], self.M0_tor], [self.Mv_pol_mat[1], self.M0_tor]], ) # 3D @@ -1037,7 +1051,8 @@ def __assemble_Mv(self, domain, as_tensor=False): if self.dim == 2: M11, M22 = mass_2d.get_Mv(self, domain) self.Mv_mat = spa.bmat( - [[spa.kron(M11, self.M0_tor), None], [None, spa.kron(M22, self.M0_tor)]], format="csr" + [[spa.kron(M11, self.M0_tor), None], [None, spa.kron(M22, self.M0_tor)]], + format="csr", ) else: self.Mv_mat = mass_3d.get_Mv(self, domain) @@ -1093,17 +1108,21 @@ def reshape_pol_1(self, coeff): if c_size == self.E1.shape[0]: coeff1_pol_1 = coeff[: self.E1_pol.shape[0] * self.E0_tor.shape[0]].reshape( - self.E1_pol.shape[0], self.E0_tor.shape[0] + self.E1_pol.shape[0], + self.E0_tor.shape[0], ) coeff1_pol_3 = coeff[self.E1_pol.shape[0] * self.E0_tor.shape[0] :].reshape( - self.E0_pol.shape[0], self.E1_tor.shape[0] + self.E0_pol.shape[0], + self.E1_tor.shape[0], ) else: coeff1_pol_1 = coeff[: self.E1_pol_0.shape[0] * self.E0_tor_0.shape[0]].reshape( - self.E1_pol_0.shape[0], self.E0_tor_0.shape[0] + self.E1_pol_0.shape[0], + self.E0_tor_0.shape[0], ) coeff1_pol_3 = coeff[self.E1_pol_0.shape[0] * self.E0_tor_0.shape[0] :].reshape( - self.E0_pol_0.shape[0], self.E1_tor_0.shape[0] + self.E0_pol_0.shape[0], + self.E1_tor_0.shape[0], ) return coeff1_pol_1, coeff1_pol_3 @@ -1119,17 +1138,21 @@ def reshape_pol_2(self, coeff): if c_size == self.E2.shape[0]: coeff2_pol_1 = coeff[: self.E2_pol.shape[0] * self.E1_tor.shape[0]].reshape( - self.E2_pol.shape[0], self.E1_tor.shape[0] + self.E2_pol.shape[0], + self.E1_tor.shape[0], ) coeff2_pol_3 = coeff[self.E2_pol.shape[0] * self.E1_tor.shape[0] :].reshape( - self.E3_pol.shape[0], self.E0_tor.shape[0] + self.E3_pol.shape[0], + self.E0_tor.shape[0], ) else: coeff2_pol_1 = coeff[: self.E2_pol_0.shape[0] * self.E1_tor_0.shape[0]].reshape( - self.E2_pol_0.shape[0], self.E1_tor_0.shape[0] + self.E2_pol_0.shape[0], + self.E1_tor_0.shape[0], ) coeff2_pol_3 = coeff[self.E2_pol_0.shape[0] * self.E1_tor_0.shape[0] :].reshape( - self.E3_pol_0.shape[0], self.E0_tor_0.shape[0] + self.E3_pol_0.shape[0], + self.E0_tor_0.shape[0], ) return coeff2_pol_1, coeff2_pol_3 @@ -1161,18 +1184,22 @@ def reshape_pol_v(self, coeff): if c_size == self.Ev.shape[0]: coeffv_pol_1 = coeff[: self.Ev_pol.shape[0] * self.E0_tor.shape[0]].reshape( - self.Ev_pol.shape[0], self.E0_tor.shape[0] + self.Ev_pol.shape[0], + self.E0_tor.shape[0], ) coeffv_pol_3 = coeff[self.Ev_pol.shape[0] * self.E0_tor.shape[0] :].reshape( - self.E0_pol.shape[0], self.E0_tor.shape[0] + self.E0_pol.shape[0], + self.E0_tor.shape[0], ) else: coeffv_pol_1 = coeff[: self.Ev_pol_0.shape[0] * self.E0_tor.shape[0]].reshape( - self.Ev_pol_0.shape[0], self.E0_tor.shape[0] + self.Ev_pol_0.shape[0], + self.E0_tor.shape[0], ) coeffv_pol_3 = coeff[self.Ev_pol_0.shape[0] * self.E0_tor.shape[0] :].reshape( - self.E0_pol.shape[0], self.E0_tor_0.shape[0] + self.E0_pol.shape[0], + self.E0_tor_0.shape[0], ) return coeffv_pol_1, coeffv_pol_3 @@ -1527,10 +1554,26 @@ def evaluate_NN(self, eta1, eta2, eta3, coeff, which="V0", part="r"): if self.n_tor != 0 and self.basis_tor == "r": real_2 = eva_2d.evaluate_n_n( - self.T[0], self.T[1], self.p[0], self.p[1], self.indN[0], self.indN[1], coeff_r[:, :, 1], eta1, eta2 + self.T[0], + self.T[1], + self.p[0], + self.p[1], + self.indN[0], + self.indN[1], + coeff_r[:, :, 1], + eta1, + eta2, ) imag_2 = eva_2d.evaluate_n_n( - self.T[0], self.T[1], self.p[0], self.p[1], self.indN[0], self.indN[1], coeff_i[:, :, 1], eta1, eta2 + self.T[0], + self.T[1], + self.p[0], + self.p[1], + self.indN[0], + self.indN[1], + coeff_i[:, :, 1], + eta1, + eta2, ) # multiply with Fourier basis in third direction if |n_tor| > 0 diff --git a/src/struphy/examples/_draw_parallel.py b/src/struphy/examples/_draw_parallel.py index e95598b3c..b31ba19c4 100644 --- a/src/struphy/examples/_draw_parallel.py +++ b/src/struphy/examples/_draw_parallel.py @@ -81,7 +81,7 @@ def main(): print( f"rank {rank} | markers not on correct process: {xp.nonzero(xp.logical_and(~stay, ~holes))} \ - \n corresponding positions:\n {error_mks[:, :3]}" + \n corresponding positions:\n {error_mks[:, :3]}", ) assert error_mks.size == 0 diff --git a/src/struphy/feec/basis_projection_ops.py b/src/struphy/feec/basis_projection_ops.py index 3bce4701f..d76dc7ca9 100644 --- a/src/struphy/feec/basis_projection_ops.py +++ b/src/struphy/feec/basis_projection_ops.py @@ -147,7 +147,7 @@ def K3(self): e3, ) / self.sqrt_g(e1, e2, e3), - ] + ], ] self._K3 = self.create_basis_op( fun, @@ -263,7 +263,7 @@ def Q3(self): e3, ) / self.sqrt_g(e1, e2, e3), - ] + ], ] self._Q3 = self.create_basis_op( fun, @@ -1064,7 +1064,7 @@ def __init__( self._V1ds[1][2].nbasis, ], [self._V1ds[2][0].nbasis, self._V1ds[2][1].nbasis, self._V1ds[2][2].nbasis], - ] + ], ) # output space: 3d StencilVectorSpaces and 1d SplineSpaces of each component @@ -1363,7 +1363,7 @@ def assemble(self, verbose=False): col0, col1, col2, - ] + ], ), self._VNbasis[hh], Aux._data, @@ -1443,7 +1443,7 @@ def assemble(self, verbose=False): col0, col1, col2, - ] + ], ), self._VNbasis, Aux[h]._data, @@ -1544,7 +1544,7 @@ def assemble(self, verbose=False): col0, col1, col2, - ] + ], ), self._VNbasis[hh], Aux[h]._data, @@ -2043,7 +2043,7 @@ def assemble(self, weights=None, verbose=False): getattr( basis_projection_kernels, "assemble_dofs_for_weighted_basisfuns_" + str(V.ldim) + "d", - ) + ), ) if rank == 0 and verbose: diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 3a29ea821..7469d68e9 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -136,7 +136,7 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): itterables = [] for i in range(ndim[h]): itterables.append( - range(allstarts[currentrank][i + npredim], allends[currentrank][i + npredim] + 1) + range(allstarts[currentrank][i + npredim], allends[currentrank][i + npredim] + 1), ) # We iterate over all the entries that belong to rank number currentrank for i in itertools.product(*itterables): @@ -293,7 +293,7 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): return sparse.csr_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols)).todia() else: raise Exception( - "The selected sparse matrix format must be one of the following : csr, csc, bsr, lil, dok, coo or dia." + "The selected sparse matrix format must be one of the following : csr, csc, bsr, lil, dok, coo or dia.", ) diff --git a/src/struphy/feec/local_projectors_kernels.py b/src/struphy/feec/local_projectors_kernels.py index 09b4c182f..f1eb285c9 100644 --- a/src/struphy/feec/local_projectors_kernels.py +++ b/src/struphy/feec/local_projectors_kernels.py @@ -168,7 +168,10 @@ def get_dofs_local_1_form_ec_component_weighted( @stack_array("shp") def get_dofs_local_1_form_ec_component( - args_solve: LocalProjectorsArguments, f3: "float[:,:,:]", f_eval_aux: "float[:,:,:]", c: int + args_solve: LocalProjectorsArguments, + f3: "float[:,:,:]", + f_eval_aux: "float[:,:,:]", + c: int, ): """Kernel for evaluating the degrees of freedom for the c-th component of 1-forms. This function is for local commuting projetors. @@ -220,7 +223,10 @@ def get_dofs_local_1_form_ec_component( @stack_array("shp") def get_dofs_local_2_form_ec_component( - args_solve: LocalProjectorsArguments, fc: "float[:,:,:]", f_eval_aux: "float[:,:,:]", c: int + args_solve: LocalProjectorsArguments, + fc: "float[:,:,:]", + f_eval_aux: "float[:,:,:]", + c: int, ): """Kernel for evaluating the degrees of freedom for the c-th component of 2-forms. This function is for local commuting projetors. @@ -1037,7 +1043,15 @@ def get_rows_periodic(starts: int, ends: int, modl: int, modr: int, Nbasis: int, def get_rows( - col: int, starts: int, ends: int, p: int, Nbasis: int, periodic: bool, IoH: bool, BoD: bool, aux: "int[:]" + col: int, + starts: int, + ends: int, + p: int, + Nbasis: int, + periodic: bool, + IoH: bool, + BoD: bool, + aux: "int[:]", ): """Kernel for getting the list of rows that are non-zero for the current BasisProjectionLocal column, within the start and end indices of the current MPI rank. diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 76a26bd4d..16f0109d9 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -449,7 +449,7 @@ def M2B_div0(self): self.weights[self.selected_weight].a1_1, self.weights[self.selected_weight].a1_2, self.weights[self.selected_weight].a1_3, - ] + ], ) tmp_b2 = self.derham.curl.dot(a_eq) @@ -555,7 +555,7 @@ def M2Bn(self): self.weights[self.selected_weight].a1_1, self.weights[self.selected_weight].a1_2, self.weights[self.selected_weight].a1_3, - ] + ], ) tmp_b2 = self.derham.curl.dot(a_eq) @@ -918,7 +918,11 @@ def DFinvT(e1, e2, e3): for n in range(3): fun[-1] += [ lambda e1, e2, e3, m=m, n=n: self._matrix_operate(e1, e2, e3, *weights_rank2)[ - :, :, :, m, n + :, + :, + :, + m, + n, ], ] # Scalar operations second @@ -945,14 +949,14 @@ def DFinvT(e1, e2, e3): fun = [ [ lambda e1, e2, e3: 1.0 / weights_rank0[0](e1, e2, e3), - ] + ], ] for f2, op in zip(weights_rank0[1:], operations[1:]): fun = [ [ lambda e1, e2, e3, f=fun[0][0], op=op, f2=f2: self._operate(f, f2, op, e1, e2, e3), - ] + ], ] V_id = self.derham.space_to_form[V_id] @@ -1626,7 +1630,7 @@ def M2B_div0(self): self.weights[self.selected_weight].a1_1, self.weights[self.selected_weight].a1_2, self.weights[self.selected_weight].a1_3, - ] + ], ) tmp_b2 = self.derham.curl.dot(a_eq) @@ -1740,7 +1744,7 @@ def M2Bn(self): self.weights[self.selected_weight].a1_1, self.weights[self.selected_weight].a1_2, self.weights[self.selected_weight].a1_3, - ] + ], ) tmp_b2 = self.derham.curl.dot(a_eq) @@ -1861,7 +1865,11 @@ def M1perp(self): for n in range(3): fun[-1] += [ lambda e1, e2, e3, m=m, n=n: (self.DFinv(e1, e2, e3) @ self.D @ self.DFinv(e1, e2, e3))[ - :, :, :, m, n + :, + :, + :, + m, + n, ] * self.sqrt_g( e1, @@ -1898,7 +1906,7 @@ def M0ad(self): e3, ) * self.sqrt_g(e1, e2, e3), - ] + ], ] self._M0ad = self._assemble_weighted_mass( @@ -2359,7 +2367,8 @@ def __init__( pts = [ quad_grid[nquad].points.flatten() for quad_grid, nquad in zip( - self.derham.get_quad_grids(wspace, nquads=self.nquads), self.nquads + self.derham.get_quad_grids(wspace, nquads=self.nquads), + self.nquads, ) ] @@ -2473,7 +2482,7 @@ def __init__( getattr( mass_kernels, "kernel_" + str(self._V.ldim) + "d_mat", - ) + ), ) @property @@ -2779,7 +2788,8 @@ def assemble(self, weights=None, clear=True, verbose=True): codomain_spans = [ quad_grid[nquad].spans for quad_grid, nquad in zip( - self.derham.get_quad_grids(codomain_space, nquads=self.nquads), self.nquads + self.derham.get_quad_grids(codomain_space, nquads=self.nquads), + self.nquads, ) ] @@ -2792,7 +2802,8 @@ def assemble(self, weights=None, clear=True, verbose=True): pts = [ quad_grid[nquad].points.flatten() for quad_grid, nquad in zip( - self.derham.get_quad_grids(codomain_space, nquads=self.nquads), self.nquads + self.derham.get_quad_grids(codomain_space, nquads=self.nquads), + self.nquads, ) ] wts = [ @@ -2807,7 +2818,8 @@ def assemble(self, weights=None, clear=True, verbose=True): codomain_basis = [ quad_grid[nquad].basis for quad_grid, nquad in zip( - self.derham.get_quad_grids(codomain_space, nquads=self.nquads), self.nquads + self.derham.get_quad_grids(codomain_space, nquads=self.nquads), + self.nquads, ) ] @@ -2850,7 +2862,8 @@ def assemble(self, weights=None, clear=True, verbose=True): domain_basis = [ quad_grid[nquad].basis for quad_grid, nquad in zip( - self.derham.get_quad_grids(domain_space, nquads=self.nquads), self.nquads + self.derham.get_quad_grids(domain_space, nquads=self.nquads), + self.nquads, ) ] @@ -3036,7 +3049,8 @@ def eval_quad(W, coeffs, out=None): [ q_grid[nquad].points.size for q_grid, nquad in zip( - self.derham.get_quad_grids(space, nquads=self.nquads), self.nquads + self.derham.get_quad_grids(space, nquads=self.nquads), + self.nquads, ) ], dtype=float, @@ -3150,14 +3164,14 @@ def __init__(self, derham, V, W, weights=None, nquads=None): getattr( mass_kernels, "kernel_" + str(self._V.ldim) + "d_matrixfree", - ) + ), ) self._diag_kernel = Pyccelkernel( getattr( mass_kernels, "kernel_" + str(self._V.ldim) + "d_diag", - ) + ), ) shape = tuple(e - s + 1 for s, e in zip(V.coeff_space.starts, V.coeff_space.ends)) @@ -3245,7 +3259,11 @@ def toarray(self): def transpose(self, conjugate=False): return StencilMatrixFreeMassOperator( - self._derham, self._codomain, self._domain, self._weights, nquads=self._nquads + self._derham, + self._codomain, + self._domain, + self._weights, + nquads=self._nquads, ) @property diff --git a/src/struphy/feec/mass_kernels.py b/src/struphy/feec/mass_kernels.py index 7e62b3248..7b4f09720 100644 --- a/src/struphy/feec/mass_kernels.py +++ b/src/struphy/feec/mass_kernels.py @@ -341,7 +341,9 @@ def kernel_3d_mat( for iel2 in range(ne2): for iel3 in range(ne3): tmp_mat_fun[:, :, :] = mat_fun[ - iel1 * nq1 : (iel1 + 1) * nq1, iel2 * nq2 : (iel2 + 1) * nq2, iel3 * nq3 : (iel3 + 1) * nq3 + iel1 * nq1 : (iel1 + 1) * nq1, + iel2 * nq2 : (iel2 + 1) * nq2, + iel3 * nq3 : (iel3 + 1) * nq3, ] tmp_w1[:] = w1[iel1, :] @@ -600,7 +602,9 @@ def kernel_3d_matrixfree( for iel2 in range(ne2): for iel3 in range(ne3): tmp_mat_fun[:, :, :] = mat_fun[ - iel1 * nq1 : (iel1 + 1) * nq1, iel2 * nq2 : (iel2 + 1) * nq2, iel3 * nq3 : (iel3 + 1) * nq3 + iel1 * nq1 : (iel1 + 1) * nq1, + iel2 * nq2 : (iel2 + 1) * nq2, + iel3 * nq3 : (iel3 + 1) * nq3, ] tmp_w1[:] = w1[iel1, :] @@ -713,7 +717,9 @@ def kernel_3d_diag( for iel2 in range(ne2): for iel3 in range(ne3): tmp_mat_fun[:, :, :] = mat_fun[ - iel1 * nq1 : (iel1 + 1) * nq1, iel2 * nq2 : (iel2 + 1) * nq2, iel3 * nq3 : (iel3 + 1) * nq3 + iel1 * nq1 : (iel1 + 1) * nq1, + iel2 * nq2 : (iel2 + 1) * nq2, + iel3 * nq3 : (iel3 + 1) * nq3, ] tmp_w1[:] = w1[iel1, :] diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index f3016b532..87b7e89fb 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -915,7 +915,7 @@ def solve(self, rhs, out=None, transposed=False): out[:] = solve_circulant(self._column, rhs.T).T except xp.linalg.LinAlgError: eps = 1e-4 - print(f"Stabilizing singular preconditioning FFTSolver with {eps = }:") + print(f"Stabilizing singular preconditioning FFTSolver with {eps =}:") self._column[0] *= 1.0 + eps out[:] = solve_circulant(self._column, rhs.T).T diff --git a/src/struphy/feec/projectors.py b/src/struphy/feec/projectors.py index a291d18bd..115ed0aa1 100644 --- a/src/struphy/feec/projectors.py +++ b/src/struphy/feec/projectors.py @@ -80,7 +80,11 @@ class CommutingProjector: """ def __init__( - self, projector_tensor: GlobalProjector, dofs_extraction_op=None, base_extraction_op=None, boundary_op=None + self, + projector_tensor: GlobalProjector, + dofs_extraction_op=None, + base_extraction_op=None, + boundary_op=None, ): self._projector_tensor = projector_tensor @@ -746,7 +750,7 @@ def __init__( # List that will contain the LocalProjectorsArguments for each value of h = 0,1,2. self._solve_args = [] else: - raise TypeError(f"{fem_space = } is not of type FemSpace.") + raise TypeError(f"{fem_space =} is not of type FemSpace.") if isinstance(fem_space, TensorFemSpace): if space_id == "H1": @@ -1425,7 +1429,7 @@ def get_dofs(self, fun, dofs=None): fh = fun(*self._meshgrid[h])[h] # Case in which fun is a list of three functions, each one with one output. else: - assert len(fun) == 3, f"List input only for vector-valued spaces of size 3, but {len(fun) = }." + assert len(fun) == 3, f"List input only for vector-valued spaces of size 3, but {len(fun) =}." # Evaluation of the function to compute the h component fh = fun[h](*self._meshgrid[h]) @@ -1461,7 +1465,7 @@ def get_dofs(self, fun, dofs=None): fun, ) == 3 - ), f"List input only for vector-valued spaces of size 3, but {len(fun) = }." + ), f"List input only for vector-valued spaces of size 3, but {len(fun) =}." for h in range(3): f_eval.append(fun[h](*self._meshgrid[h])) @@ -1481,7 +1485,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non pre_computed_dofs = [fun(*self._meshgrid)] elif self._space_key == "1" or self._space_key == "2": - assert len(fun) == 3, f"List input only for vector-valued spaces of size 3, but {len(fun) = }." + assert len(fun) == 3, f"List input only for vector-valued spaces of size 3, but {len(fun) =}." self._do_nothing = xp.zeros(3, dtype=int) f_eval = [] @@ -1561,7 +1565,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non ) elif self._space_key == "v": - assert len(fun) == 3, f"List input only for vector-valued spaces of size 3, but {len(fun) = }." + assert len(fun) == 3, f"List input only for vector-valued spaces of size 3, but {len(fun) =}." self._do_nothing = xp.zeros(3, dtype=int) for h in range(3): @@ -1668,7 +1672,8 @@ def __call__( return self.solve_weighted(rhs, out=out), rhs_weights else: return self.solve_weighted( - self.get_dofs_weighted(fun, dofs=dofs, first_go=False, pre_computed_dofs=pre_computed_dofs), out=out + self.get_dofs_weighted(fun, dofs=dofs, first_go=False, pre_computed_dofs=pre_computed_dofs), + out=out, ) def get_translation_b(self, i, h): @@ -2027,7 +2032,7 @@ def get_dofs(self, fun, dofs=None, apply_bc=False, clear=True): fun_weights = fun(*self._quad_grid_mesh) elif isinstance(fun, xp.ndarray): assert fun.shape == self._quad_grid_mesh[0].shape, ( - f"Expected shape {self._quad_grid_mesh[0].shape}, got {fun.shape = } instead." + f"Expected shape {self._quad_grid_mesh[0].shape}, got {fun.shape =} instead." ) fun_weights = fun else: @@ -2036,7 +2041,7 @@ def get_dofs(self, fun, dofs=None, apply_bc=False, clear=True): fun, ) == 3 - ), f"List input only for vector-valued spaces of size 3, but {len(fun) = }." + ), f"List input only for vector-valued spaces of size 3, but {len(fun) =}." fun_weights = [] # loop over rows (different meshes) for mesh in self._quad_grid_mesh: @@ -2046,11 +2051,11 @@ def get_dofs(self, fun, dofs=None, apply_bc=False, clear=True): if callable(f): fun_weights[-1] += [f(*mesh)] elif isinstance(f, xp.ndarray): - assert f.shape == mesh[0].shape, f"Expected shape {mesh[0].shape}, got {f.shape = } instead." + assert f.shape == mesh[0].shape, f"Expected shape {mesh[0].shape}, got {f.shape =} instead." fun_weights[-1] += [f] else: raise ValueError( - f"Expected callable or numpy array, got {type(f) = } instead.", + f"Expected callable or numpy array, got {type(f) =} instead.", ) # check output vector diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index f296d95a9..523a5bb97 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -356,7 +356,7 @@ def __init__( self._spline_types_pyccel[sp_form], ) else: - raise TypeError(f"{fem_space = } is not a valid type.") + raise TypeError(f"{fem_space =} is not a valid type.") # break points self._breaks = [space.breaks for space in _derham.spaces[0].spaces] @@ -1093,7 +1093,7 @@ def _discretize_space( elif V == "L2": Wh = Vh.reduce_degree(axes=[0, 1, 2], multiplicity=Vh.multiplicity, basis=basis) else: - raise ValueError(f"V must be one of H1, Hcurl, Hdiv or L2, but is {V = }.") + raise ValueError(f"V must be one of H1, Hcurl, Hdiv or L2, but is {V =}.") Wh.symbolic_space = V for key in Wh._refined_space: @@ -1573,7 +1573,9 @@ def vector(self, value): e1, e2, e3 = self.ends self._vector.tp[s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1] = value[1][ - s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1 + s1 : e1 + 1, + s2 : e2 + 1, + s3 : e3 + 1, ] else: for n in range(3): @@ -1589,7 +1591,9 @@ def vector(self, value): e1, e2, e3 = self.ends[n] self._vector.tp[n][s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1] = value[n][1][ - s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1 + s1 : e1 + 1, + s2 : e2 + 1, + s3 : e3 + 1, ] self._vector.update_ghost_regions() @@ -1732,13 +1736,13 @@ def f_tmp(e1, e2, e3): else: assert equil is not None var = fb.variable - assert var in dir(MHDequilibrium), f"{var = } is not an attribute of any fields background." + assert var in dir(MHDequilibrium), f"{var =} is not an attribute of any fields background." if self.space_id in {"H1", "L2"}: fun = getattr(equil, var) else: assert (var + "_1") in dir(MHDequilibrium), ( - f"{(var + '_1') = } is not an attribute of any fields background." + f"{(var + '_1') =} is not an attribute of any fields background." ) fun = [ getattr(equil, var + "_1"), diff --git a/src/struphy/feec/tests/test_basis_ops.py b/src/struphy/feec/tests/test_basis_ops.py index 7b06e09e1..7ba56aefa 100644 --- a/src/struphy/feec/tests/test_basis_ops.py +++ b/src/struphy/feec/tests/test_basis_ops.py @@ -503,7 +503,7 @@ def test_basis_ops_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=Fal "n2": 4.0, "na": 0.0, "beta": 0.1, - } + }, ) if show_plots: @@ -800,8 +800,8 @@ def assert_ops(mpi_rank, res_PSY, res_STR, verbose=False, MPI_COMM=None): res_PSY.starts[0] : res_PSY.ends[0] + 1, res_PSY.starts[1] : res_PSY.ends[1] + 1, res_PSY.starts[2] : res_PSY.ends[2] + 1, - ] - ) + ], + ), ), ) @@ -834,5 +834,10 @@ def assert_ops(mpi_rank, res_PSY, res_STR, verbose=False, MPI_COMM=None): # mapping=["Cuboid", {"l1": 0.0, "r1": 1.0, "l2": 0.0, "r2": 1.0, "l3": 0.0, "r3": 1.0}], # ) test_basis_ops_polar( - [6, 9, 7], [2, 2, 3], [False, True, True], None, ["IGAPolarCylinder", {"a": 1.0, "Lz": 3.0}], False + [6, 9, 7], + [2, 2, 3], + [False, True, True], + None, + ["IGAPolarCylinder", {"a": 1.0, "Lz": 3.0}], + False, ) diff --git a/src/struphy/feec/tests/test_derham.py b/src/struphy/feec/tests/test_derham.py index c5c0e57ea..1e857b5a2 100644 --- a/src/struphy/feec/tests/test_derham.py +++ b/src/struphy/feec/tests/test_derham.py @@ -70,7 +70,9 @@ def test_psydac_derham(Nel, p, spl_kind): # Assign from start to end index + 1 x0_PSY[s0[0] : e0[0] + 1, s0[1] : e0[1] + 1, s0[2] : e0[2] + 1] = DR_STR.extract_0(x0)[ - s0[0] : e0[0] + 1, s0[1] : e0[1] + 1, s0[2] : e0[2] + 1 + s0[0] : e0[0] + 1, + s0[1] : e0[1] + 1, + s0[2] : e0[2] + 1, ] # Block of StencilVecttors @@ -87,13 +89,19 @@ def test_psydac_derham(Nel, p, spl_kind): x11, x12, x13 = DR_STR.extract_1(x1) x1_PSY[0][s11[0] : e11[0] + 1, s11[1] : e11[1] + 1, s11[2] : e11[2] + 1] = x11[ - s11[0] : e11[0] + 1, s11[1] : e11[1] + 1, s11[2] : e11[2] + 1 + s11[0] : e11[0] + 1, + s11[1] : e11[1] + 1, + s11[2] : e11[2] + 1, ] x1_PSY[1][s12[0] : e12[0] + 1, s12[1] : e12[1] + 1, s12[2] : e12[2] + 1] = x12[ - s12[0] : e12[0] + 1, s12[1] : e12[1] + 1, s12[2] : e12[2] + 1 + s12[0] : e12[0] + 1, + s12[1] : e12[1] + 1, + s12[2] : e12[2] + 1, ] x1_PSY[2][s13[0] : e13[0] + 1, s13[1] : e13[1] + 1, s13[2] : e13[2] + 1] = x13[ - s13[0] : e13[0] + 1, s13[1] : e13[1] + 1, s13[2] : e13[2] + 1 + s13[0] : e13[0] + 1, + s13[1] : e13[1] + 1, + s13[2] : e13[2] + 1, ] x2_PSY = BlockVector(derham.Vh["2"]) @@ -109,13 +117,19 @@ def test_psydac_derham(Nel, p, spl_kind): x21, x22, x23 = DR_STR.extract_2(x2) x2_PSY[0][s21[0] : e21[0] + 1, s21[1] : e21[1] + 1, s21[2] : e21[2] + 1] = x21[ - s21[0] : e21[0] + 1, s21[1] : e21[1] + 1, s21[2] : e21[2] + 1 + s21[0] : e21[0] + 1, + s21[1] : e21[1] + 1, + s21[2] : e21[2] + 1, ] x2_PSY[1][s22[0] : e22[0] + 1, s22[1] : e22[1] + 1, s22[2] : e22[2] + 1] = x22[ - s22[0] : e22[0] + 1, s22[1] : e22[1] + 1, s22[2] : e22[2] + 1 + s22[0] : e22[0] + 1, + s22[1] : e22[1] + 1, + s22[2] : e22[2] + 1, ] x2_PSY[2][s23[0] : e23[0] + 1, s23[1] : e23[1] + 1, s23[2] : e23[2] + 1] = x23[ - s23[0] : e23[0] + 1, s23[1] : e23[1] + 1, s23[2] : e23[2] + 1 + s23[0] : e23[0] + 1, + s23[1] : e23[1] + 1, + s23[2] : e23[2] + 1, ] x3_PSY = StencilVector(derham.Vh["3"]) @@ -130,7 +144,9 @@ def test_psydac_derham(Nel, p, spl_kind): e3 = x3_PSY.ends x3_PSY[s3[0] : e3[0] + 1, s3[1] : e3[1] + 1, s3[2] : e3[2] + 1] = DR_STR.extract_3(x3)[ - s3[0] : e3[0] + 1, s3[1] : e3[1] + 1, s3[2] : e3[2] + 1 + s3[0] : e3[0] + 1, + s3[1] : e3[1] + 1, + s3[2] : e3[2] + 1, ] ######################## diff --git a/src/struphy/feec/tests/test_eval_field.py b/src/struphy/feec/tests/test_eval_field.py index 6148fa41e..f9a00c18d 100644 --- a/src/struphy/feec/tests/test_eval_field.py +++ b/src/struphy/feec/tests/test_eval_field.py @@ -234,13 +234,13 @@ def test_eval_field(Nel, p, spl_kind): m_vals_ref_3 = E1(0.0, 0.0, eta3, squeeze_out=True) assert xp.all( - [xp.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)] + [xp.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)], ) assert xp.all( - [xp.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)] + [xp.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)], ) assert xp.all( - [xp.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] + [xp.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)], ) ###### @@ -354,13 +354,13 @@ def test_eval_field(Nel, p, spl_kind): m_vals_ref_3 = B2(0.0, 0.0, eta3, squeeze_out=True) assert xp.all( - [xp.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)] + [xp.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)], ) assert xp.all( - [xp.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)] + [xp.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)], ) assert xp.all( - [xp.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] + [xp.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)], ) ###### @@ -526,13 +526,13 @@ def test_eval_field(Nel, p, spl_kind): m_vals_ref_3 = uv(0.0, 0.0, eta3, squeeze_out=True) assert xp.all( - [xp.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)] + [xp.allclose(m_vals_1_i, m_vals_ref_1_i) for m_vals_1_i, m_vals_ref_1_i in zip(m_vals_1, m_vals_ref_1)], ) assert xp.all( - [xp.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)] + [xp.allclose(m_vals_2_i, m_vals_ref_2_i) for m_vals_2_i, m_vals_ref_2_i in zip(m_vals_2, m_vals_ref_2)], ) assert xp.all( - [xp.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)] + [xp.allclose(m_vals_3_i, m_vals_ref_3_i) for m_vals_3_i, m_vals_ref_3_i in zip(m_vals_3, m_vals_ref_3)], ) print("\nAll assertions passed.") diff --git a/src/struphy/feec/tests/test_field_init.py b/src/struphy/feec/tests/test_field_init.py index 292edf76a..2f0da1611 100644 --- a/src/struphy/feec/tests/test_field_init.py +++ b/src/struphy/feec/tests/test_field_init.py @@ -40,7 +40,7 @@ def test_bckgr_init_const(Nel, p, spl_kind, spaces, vec_comps): background = FieldsBackground(type="LogicalConst", values=(val,)) field.initialize_coeffs(backgrounds=background) print( - f"\n{rank = }, {space = }, after init:\n {xp.max(xp.abs(field(*meshgrids) - val)) = }", + f"\n{rank =}, {space =}, after init:\n {xp.max(xp.abs(field(*meshgrids) - val)) =}", ) # print(f'{field(*meshgrids) = }') assert xp.allclose(field(*meshgrids), val) @@ -50,7 +50,7 @@ def test_bckgr_init_const(Nel, p, spl_kind, spaces, vec_comps): for j, val in enumerate(background.values): if val is not None: print( - f"\n{rank = }, {space = }, after init:\n {j = }, {xp.max(xp.abs(field(*meshgrids)[j] - val)) = }", + f"\n{rank =}, {space =}, after init:\n {j =}, {xp.max(xp.abs(field(*meshgrids)[j] - val)) =}", ) # print(f'{field(*meshgrids)[i] = }') assert xp.allclose(field(*meshgrids)[j], val) @@ -96,20 +96,20 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show # test for key, val in inspect.getmembers(equils): if inspect.isclass(val) and val.__module__ == equils.__name__: - print(f"{key = }") + print(f"{key =}") if "DESC" in key and not with_desc: - print(f"Attention: {with_desc = }, DESC not tested here !!") + print(f"Attention: {with_desc =}, DESC not tested here !!") continue if "GVEC" in key and not with_gvec: - print(f"Attention: {with_gvec = }, GVEC not tested here !!") + print(f"Attention: {with_gvec =}, GVEC not tested here !!") continue mhd_equil = val() if not isinstance(mhd_equil, FluidEquilibriumWithB): continue - print(f"{mhd_equil.params = }") + print(f"{mhd_equil.params =}") if "AdhocTorus" in key: mhd_equil.domain = domains.HollowTorus( @@ -186,7 +186,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show # scalar spaces print( - f"{xp.max(xp.abs(field_3(*meshgrids) - mhd_equil.p3(*meshgrids))) / xp.max(xp.abs(mhd_equil.p3(*meshgrids)))}" + f"{xp.max(xp.abs(field_3(*meshgrids) - mhd_equil.p3(*meshgrids))) / xp.max(xp.abs(mhd_equil.p3(*meshgrids)))}", ) assert ( xp.max( @@ -198,7 +198,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show if isinstance(mhd_equil, FluidEquilibriumWithB): print( - f"{xp.max(xp.abs(field_0(*meshgrids) - mhd_equil.absB0(*meshgrids))) / xp.max(xp.abs(mhd_equil.absB0(*meshgrids)))}" + f"{xp.max(xp.abs(field_0(*meshgrids) - mhd_equil.absB0(*meshgrids))) / xp.max(xp.abs(mhd_equil.absB0(*meshgrids)))}", ) assert ( xp.max( @@ -216,7 +216,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show else: denom = xp.max(xp.abs(ref[0])) print( - f"{xp.max(xp.abs(field_1(*meshgrids)[0] - ref[0])) / denom = }", + f"{xp.max(xp.abs(field_1(*meshgrids)[0] - ref[0])) / denom =}", ) assert xp.max(xp.abs(field_1(*meshgrids)[0] - ref[0])) / denom < 0.28 if xp.max(xp.abs(ref[1])) < 1e-11: @@ -224,7 +224,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show else: denom = xp.max(xp.abs(ref[1])) print( - f"{xp.max(xp.abs(field_1(*meshgrids)[1] - ref[1])) / denom = }", + f"{xp.max(xp.abs(field_1(*meshgrids)[1] - ref[1])) / denom =}", ) assert xp.max(xp.abs(field_1(*meshgrids)[1] - ref[1])) / denom < 0.33 if xp.max(xp.abs(ref[2])) < 1e-11: @@ -232,7 +232,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show else: denom = xp.max(xp.abs(ref[2])) print( - f"{xp.max(xp.abs(field_1(*meshgrids)[2] - ref[2])) / denom = }", + f"{xp.max(xp.abs(field_1(*meshgrids)[2] - ref[2])) / denom =}", ) assert ( xp.max( @@ -251,7 +251,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show else: denom = xp.max(xp.abs(ref[0])) print( - f"{xp.max(xp.abs(field_2(*meshgrids)[0] - ref[0])) / denom = }", + f"{xp.max(xp.abs(field_2(*meshgrids)[0] - ref[0])) / denom =}", ) assert xp.max(xp.abs(field_2(*meshgrids)[0] - ref[0])) / denom < 0.86 if xp.max(xp.abs(ref[1])) < 1e-11: @@ -259,7 +259,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show else: denom = xp.max(xp.abs(ref[1])) print( - f"{xp.max(xp.abs(field_2(*meshgrids)[1] - ref[1])) / denom = }", + f"{xp.max(xp.abs(field_2(*meshgrids)[1] - ref[1])) / denom =}", ) assert ( xp.max( @@ -275,7 +275,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show else: denom = xp.max(xp.abs(ref[2])) print( - f"{xp.max(xp.abs(field_2(*meshgrids)[2] - ref[2])) / denom = }", + f"{xp.max(xp.abs(field_2(*meshgrids)[2] - ref[2])) / denom =}", ) assert xp.max(xp.abs(field_2(*meshgrids)[2] - ref[2])) / denom < 0.21 print("u2 asserts passed.") @@ -286,7 +286,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show else: denom = xp.max(xp.abs(ref[0])) print( - f"{xp.max(xp.abs(field_4(*meshgrids)[0] - ref[0])) / denom = }", + f"{xp.max(xp.abs(field_4(*meshgrids)[0] - ref[0])) / denom =}", ) assert xp.max(xp.abs(field_4(*meshgrids)[0] - ref[0])) / denom < 0.6 if xp.max(xp.abs(ref[1])) < 1e-11: @@ -294,7 +294,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show else: denom = xp.max(xp.abs(ref[1])) print( - f"{xp.max(xp.abs(field_4(*meshgrids)[1] - ref[1])) / denom = }", + f"{xp.max(xp.abs(field_4(*meshgrids)[1] - ref[1])) / denom =}", ) assert ( xp.max( @@ -310,7 +310,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show else: denom = xp.max(xp.abs(ref[2])) print( - f"{xp.max(xp.abs(field_4(*meshgrids)[2] - ref[2])) / denom = }", + f"{xp.max(xp.abs(field_4(*meshgrids)[2] - ref[2])) / denom =}", ) assert ( xp.max( @@ -325,27 +325,27 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show # plotting fields with equilibrium if show_plot and rank == 0: - plt.figure(f"0/3-forms top, {mhd_equil = }", figsize=(24, 16)) + plt.figure(f"0/3-forms top, {mhd_equil =}", figsize=(24, 16)) plt.figure( - f"0/3-forms poloidal, {mhd_equil = }", + f"0/3-forms poloidal, {mhd_equil =}", figsize=(24, 16), ) - plt.figure(f"1-forms top, {mhd_equil = }", figsize=(24, 16)) + plt.figure(f"1-forms top, {mhd_equil =}", figsize=(24, 16)) plt.figure( - f"1-forms poloidal, {mhd_equil = }", + f"1-forms poloidal, {mhd_equil =}", figsize=(24, 16), ) - plt.figure(f"2-forms top, {mhd_equil = }", figsize=(24, 16)) + plt.figure(f"2-forms top, {mhd_equil =}", figsize=(24, 16)) plt.figure( - f"2-forms poloidal, {mhd_equil = }", + f"2-forms poloidal, {mhd_equil =}", figsize=(24, 16), ) plt.figure( - f"vector-fields top, {mhd_equil = }", + f"vector-fields top, {mhd_equil =}", figsize=(24, 16), ) plt.figure( - f"vector-fields poloidal, {mhd_equil = }", + f"vector-fields poloidal, {mhd_equil =}", figsize=(24, 16), ) x, y, z = mhd_equil.domain(*meshgrids) @@ -357,7 +357,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show levels = xp.linspace(xp.min(absB0) - 1e-10, xp.max(absB0), 20) - plt.figure(f"0/3-forms top, {mhd_equil = }") + plt.figure(f"0/3-forms top, {mhd_equil =}") plt.subplot(2, 3, 1) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -443,7 +443,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show plt.colorbar() plt.title("reference, top view (e1-e3)") - plt.figure(f"0/3-forms poloidal, {mhd_equil = }") + plt.figure(f"0/3-forms poloidal, {mhd_equil =}") plt.subplot(2, 3, 1) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -495,7 +495,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show levels = xp.linspace(xp.min(p3) - 1e-10, xp.max(p3), 20) - plt.figure(f"0/3-forms top, {mhd_equil = }") + plt.figure(f"0/3-forms top, {mhd_equil =}") plt.subplot(2, 3, 2) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -581,7 +581,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show plt.colorbar() plt.title("reference, top view (e1-e3)") - plt.figure(f"0/3-forms poloidal, {mhd_equil = }") + plt.figure(f"0/3-forms poloidal, {mhd_equil =}") plt.subplot(2, 3, 2) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -642,7 +642,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show for i, (bh, b) in enumerate(zip(b1h, b1)): levels = xp.linspace(xp.min(b) - 1e-10, xp.max(b), 20) - plt.figure(f"1-forms top, {mhd_equil = }") + plt.figure(f"1-forms top, {mhd_equil =}") plt.subplot(2, 3, 1 + i) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -728,7 +728,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show plt.colorbar() plt.title("reference, top view (e1-e3)") - plt.figure(f"1-forms poloidal, {mhd_equil = }") + plt.figure(f"1-forms poloidal, {mhd_equil =}") plt.subplot(2, 3, 1 + i) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -791,7 +791,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show for i, (bh, b) in enumerate(zip(b2h, b2)): levels = xp.linspace(xp.min(b) - 1e-10, xp.max(b), 20) - plt.figure(f"2-forms top, {mhd_equil = }") + plt.figure(f"2-forms top, {mhd_equil =}") plt.subplot(2, 3, 1 + i) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -877,7 +877,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show plt.colorbar() plt.title("reference, top view (e1-e3)") - plt.figure(f"2-forms poloidal, {mhd_equil = }") + plt.figure(f"2-forms poloidal, {mhd_equil =}") plt.subplot(2, 3, 1 + i) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -940,7 +940,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show for i, (bh, b) in enumerate(zip(bvh, bv)): levels = xp.linspace(xp.min(b) - 1e-10, xp.max(b), 20) - plt.figure(f"vector-fields top, {mhd_equil = }") + plt.figure(f"vector-fields top, {mhd_equil =}") plt.subplot(2, 3, 1 + i) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -1026,7 +1026,7 @@ def test_bckgr_init_mhd(Nel, p, spl_kind, with_desc=False, with_gvec=False, show plt.colorbar() plt.title("reference, top view (e1-e3)") - plt.figure(f"vector-fields poloidal, {mhd_equil = }") + plt.figure(f"vector-fields poloidal, {mhd_equil =}") plt.subplot(2, 3, 1 + i) if "Slab" in key or "Pinch" in key: plt.contourf( @@ -1162,7 +1162,10 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): field_0 = derham.create_spline_function("name_0", "H1", backgrounds=bckgr_0, perturbations=[f_sin_0, f_cos_0]) field_1 = derham.create_spline_function( - "name_1", "Hcurl", backgrounds=bckgr_1, perturbations=[f_sin_11, f_sin_13, f_cos_11, f_cos_12] + "name_1", + "Hcurl", + backgrounds=bckgr_1, + perturbations=[f_sin_11, f_sin_13, f_cos_11, f_cos_12], ) field_2 = derham.create_spline_function("name_2", "Hdiv", backgrounds=bckgr_2, perturbations=[f_cos_22]) @@ -1189,13 +1192,13 @@ def test_sincos_init_const(Nel, p, spl_kind, show_plot=False): f1_h = field_1(*meshgrids) f2_h = field_2(*meshgrids) - print(f"{xp.max(xp.abs(fun_0 - f0_h)) = }") - print(f"{xp.max(xp.abs(fun_1[0] - f1_h[0])) = }") - print(f"{xp.max(xp.abs(fun_1[1] - f1_h[1])) = }") - print(f"{xp.max(xp.abs(fun_1[2] - f1_h[2])) = }") - print(f"{xp.max(xp.abs(fun_2[0] - f2_h[0])) = }") - print(f"{xp.max(xp.abs(fun_2[1] - f2_h[1])) = }") - print(f"{xp.max(xp.abs(fun_2[2] - f2_h[2])) = }") + print(f"{xp.max(xp.abs(fun_0 - f0_h)) =}") + print(f"{xp.max(xp.abs(fun_1[0] - f1_h[0])) =}") + print(f"{xp.max(xp.abs(fun_1[1] - f1_h[1])) =}") + print(f"{xp.max(xp.abs(fun_1[2] - f1_h[2])) =}") + print(f"{xp.max(xp.abs(fun_2[0] - f2_h[0])) =}") + print(f"{xp.max(xp.abs(fun_2[1] - f2_h[1])) =}") + print(f"{xp.max(xp.abs(fun_2[2] - f2_h[2])) =}") assert xp.max(xp.abs(fun_0 - f0_h)) < 3e-5 assert xp.max(xp.abs(fun_1[0] - f1_h[0])) < 3e-5 diff --git a/src/struphy/feec/tests/test_l2_projectors.py b/src/struphy/feec/tests/test_l2_projectors.py index 6d3695caa..7da42eff4 100644 --- a/src/struphy/feec/tests/test_l2_projectors.py +++ b/src/struphy/feec/tests/test_l2_projectors.py @@ -47,11 +47,11 @@ def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_desc, do_plo for dom_type, dom_class in zip(dom_types, dom_classes): print("#" * 80) - print(f"Testing {dom_class = }") + print(f"Testing {dom_class =}") print("#" * 80) if "DESC" in dom_type and not with_desc: - print(f"Attention: {with_desc = }, DESC not tested here !!") + print(f"Attention: {with_desc =}, DESC not tested here !!") continue domain = dom_class() @@ -103,7 +103,7 @@ def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_desc, do_plo err = [xp.max(xp.abs(exact(ee1, ee2, ee3) - field_v)) for exact, field_v in zip(f_analytic, field_vals)] f_plot = field_vals[0] - print(f"{sp_id = }, {xp.max(err) = }") + print(f"{sp_id =}, {xp.max(err) =}") if sp_id in ("H1", "H1vec"): assert xp.max(err) < 0.004 else: @@ -233,7 +233,7 @@ def f(x, y, z): line_for_rate_p0 = [Ne ** (-rate_p0) * errors[sp_id][0] / Nels[0] ** (-rate_p0) for Ne in Nels] m, _ = xp.polyfit(xp.log(Nels), xp.log(errors[sp_id]), deg=1) - print(f"{sp_id = }, fitted convergence rate = {-m}, degree = {pi}") + print(f"{sp_id =}, fitted convergence rate = {-m}, degree = {pi}") if sp_id in ("H1", "H1vec"): assert -m > (pi + 1 - 0.05) else: @@ -247,7 +247,7 @@ def f(x, y, z): plt.loglog(Nels, line_for_rate_p0, "k--") plt.text(Nels[-2], line_for_rate_p1[-2], f"1/Nel^{rate_p1}") plt.text(Nels[-2], line_for_rate_p0[-2], f"1/Nel^{rate_p0}") - plt.title(f"{sp_id = }, degree = {pi}") + plt.title(f"{sp_id =}, degree = {pi}") plt.xlabel("Nel") if do_plot and rank == 0: diff --git a/src/struphy/feec/tests/test_local_projectors.py b/src/struphy/feec/tests/test_local_projectors.py index ce2abda01..4cf2d401c 100644 --- a/src/struphy/feec/tests/test_local_projectors.py +++ b/src/struphy/feec/tests/test_local_projectors.py @@ -142,7 +142,7 @@ def f(e1, e2, e3): errg[1] = xp.max(xp.abs(fieldg_vals[1] - field_vals[1])) errg[2] = xp.max(xp.abs(fieldg_vals[2] - field_vals[2])) - print(f"{sp_id = }, {xp.max(err) = }, {xp.max(errg) = },{exectime = }") + print(f"{sp_id =}, {xp.max(err) =}, {xp.max(errg) =},{exectime =}") if sp_id in ("H1", "H1vec"): assert xp.max(err) < 0.011 assert xp.max(errg) < 0.011 @@ -264,14 +264,14 @@ def f(x, y, z): # for those cases is better to compute the convergance rate using only the information of Nel with smaller number if -m <= (pi + 1 - 0.1): m = -xp.log2(errors[sp_id][1] / errors[sp_id][2]) - print(f"{sp_id = }, fitted convergence rate = {-m}, degree = {pi}") + print(f"{sp_id =}, fitted convergence rate = {-m}, degree = {pi}") assert -m > (pi + 1 - 0.1) else: # Sometimes for very large number of elements the convergance rate falls of a bit since the error is already so small floating point impressions become relevant # for those cases is better to compute the convergance rate using only the information of Nel with smaller number if -m <= (pi - 0.1): m = -xp.log2(errors[sp_id][1] / errors[sp_id][2]) - print(f"{sp_id = }, fitted convergence rate = {-m}, degree = {pi}") + print(f"{sp_id =}, fitted convergence rate = {-m}, degree = {pi}") assert -m > (pi - 0.1) if do_plot: @@ -282,7 +282,7 @@ def f(x, y, z): plt.loglog(Nels, line_for_rate_p0, "k--") plt.text(Nels[-2], line_for_rate_p1[-2], f"1/Nel^{rate_p1}") plt.text(Nels[-2], line_for_rate_p0[-2], f"1/Nel^{rate_p0}") - plt.title(f"{sp_id = }, degree = {pi}") + plt.title(f"{sp_id =}, degree = {pi}") plt.xlabel("Nel") if do_plot and rank == 0: @@ -500,7 +500,7 @@ def fun(eta1, eta2, eta3): VFEM1ds[1][2].nbasis, ], [VFEM1ds[2][0].nbasis, VFEM1ds[2][1].nbasis, VFEM1ds[2][2].nbasis], - ] + ], ) if in_sp_key == "0" or in_sp_key == "3": @@ -578,7 +578,7 @@ def basis3(i3, h=None): npts_in[0][0] * npts_in[0][1] * npts_in[0][2] + npts_in[1][0] * npts_in[1][1] * npts_in[1][2] + npts_in[2][0] * npts_in[2][1] * npts_in[2][2], - ) + ), ) else: @@ -591,57 +591,57 @@ def basis3(i3, h=None): ( npts_out[0][0] * npts_out[0][1] * npts_out[0][2], npts_in[0][0] * npts_in[0][1] * npts_in[0][2], - ) + ), ) matrix10 = xp.zeros( ( npts_out[1][0] * npts_out[1][1] * npts_out[1][2], npts_in[0][0] * npts_in[0][1] * npts_in[0][2], - ) + ), ) matrix20 = xp.zeros( ( npts_out[2][0] * npts_out[2][1] * npts_out[2][2], npts_in[0][0] * npts_in[0][1] * npts_in[0][2], - ) + ), ) matrix01 = xp.zeros( ( npts_out[0][0] * npts_out[0][1] * npts_out[0][2], npts_in[1][0] * npts_in[1][1] * npts_in[1][2], - ) + ), ) matrix11 = xp.zeros( ( npts_out[1][0] * npts_out[1][1] * npts_out[1][2], npts_in[1][0] * npts_in[1][1] * npts_in[1][2], - ) + ), ) matrix21 = xp.zeros( ( npts_out[2][0] * npts_out[2][1] * npts_out[2][2], npts_in[1][0] * npts_in[1][1] * npts_in[1][2], - ) + ), ) matrix02 = xp.zeros( ( npts_out[0][0] * npts_out[0][1] * npts_out[0][2], npts_in[2][0] * npts_in[2][1] * npts_in[2][2], - ) + ), ) matrix12 = xp.zeros( ( npts_out[1][0] * npts_out[1][1] * npts_out[1][2], npts_in[2][0] * npts_in[2][1] * npts_in[2][2], - ) + ), ) matrix22 = xp.zeros( ( npts_out[2][0] * npts_out[2][1] * npts_out[2][2], npts_in[2][0] * npts_in[2][1] * npts_in[2][2], - ) + ), ) # We build the BasisProjectionOperator by hand @@ -1233,7 +1233,7 @@ def f_analytic3(e1, e2, e3): * basis1(random_i0)(*meshgrid) * basis2(random_i1)(*meshgrid) * basis3(random_i2)(*meshgrid), - ] + ], ) else: @@ -1307,7 +1307,7 @@ def f_analytic22(e1, e2, e3): * basis2(random_i1, random_h)(*meshgrid) * basis3(random_i2, random_h)(*meshgrid) for dim in range(3) - ] + ], ) FE_loc = matrix_new.dot(input) @@ -1367,10 +1367,10 @@ def f_analytic22(e1, e2, e3): if rank == 0: assert reducemeanlocal < 10.0 * reducemeanglobal or reducemeanlocal < 10.0**-5 - print(f"{reducemeanlocal = }") - print(f"{reducemaxlocal = }") - print(f"{reducemeanglobal = }") - print(f"{reducemaxglobal = }") + print(f"{reducemeanlocal =}") + print(f"{reducemaxlocal =}") + print(f"{reducemeanglobal =}") + print(f"{reducemaxglobal =}") if do_plot: if out_sp_key == "0" or out_sp_key == "3": @@ -1502,7 +1502,7 @@ def error(e1, e2, e3): maxerrorB = auxerror inputB[col0, col1, col2] = 0.0 - print(f"{maxerrorB = }") + print(f"{maxerrorB =}") assert maxerrorB < 10.0**-13 maxerrorD = 0.0 @@ -1530,7 +1530,7 @@ def error(e1, e2, e3): maxerrorD = auxerror inputD[col0, col1, col2] = 0.0 - print(f"{maxerrorD = }") + print(f"{maxerrorD =}") assert maxerrorD < 10.0**-13 print("Test spline evaluation passed.") diff --git a/src/struphy/feec/tests/test_lowdim_nel_is_1.py b/src/struphy/feec/tests/test_lowdim_nel_is_1.py index fdbe6be3d..cefcddf61 100644 --- a/src/struphy/feec/tests/test_lowdim_nel_is_1.py +++ b/src/struphy/feec/tests/test_lowdim_nel_is_1.py @@ -161,7 +161,7 @@ def div_f(x, y, z): # a) projection error err_f0 = xp.max(xp.abs(f(e1, e2, e3) - field_f0_vals)) - print(f"\n{err_f0 = }") + print(f"\n{err_f0 =}") assert err_f0 < 1e-2 # b) commuting property @@ -174,7 +174,7 @@ def div_f(x, y, z): field_df0_vals = field_df0(e1, e2, e3, squeeze_out=True) err_df0 = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(grad_f, field_df0_vals)] - print(f"{err_df0 = }") + print(f"{err_df0 =}") assert xp.max(err_df0) < 0.64 # d) plotting @@ -203,7 +203,7 @@ def div_f(x, y, z): # a) projection error err_f1 = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip([f, f, f], field_f1_vals)] - print(f"{err_f1 = }") + print(f"{err_f1 =}") assert xp.max(err_f1) < 0.09 # b) commuting property @@ -216,7 +216,7 @@ def div_f(x, y, z): field_df1_vals = field_df1(e1, e2, e3, squeeze_out=True) err_df1 = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip(curl_f, field_df1_vals)] - print(f"{err_df1 = }") + print(f"{err_df1 =}") assert xp.max(err_df1) < 0.64 # d) plotting @@ -250,7 +250,7 @@ def div_f(x, y, z): # a) projection error err_f2 = [xp.max(xp.abs(exact(e1, e2, e3) - field_v)) for exact, field_v in zip([f, f, f], field_f2_vals)] - print(f"{err_f2 = }") + print(f"{err_f2 =}") assert xp.max(err_f2) < 0.09 # b) commuting property @@ -263,7 +263,7 @@ def div_f(x, y, z): field_df2_vals = field_df2(e1, e2, e3, squeeze_out=True) err_df2 = xp.max(xp.abs(div_f(e1, e2, e3) - field_df2_vals)) - print(f"{err_df2 = }") + print(f"{err_df2 =}") assert xp.max(err_df2) < 0.64 # d) plotting @@ -292,7 +292,7 @@ def div_f(x, y, z): # a) projection error err_f3 = xp.max(xp.abs(f(e1, e2, e3) - field_f3_vals)) - print(f"{err_f3 = }") + print(f"{err_f3 =}") assert err_f3 < 0.09 # d) plotting diff --git a/src/struphy/feec/tests/test_mass_matrices.py b/src/struphy/feec/tests/test_mass_matrices.py index a57d67cdc..e1d629c2e 100644 --- a/src/struphy/feec/tests/test_mass_matrices.py +++ b/src/struphy/feec/tests/test_mass_matrices.py @@ -56,7 +56,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): "n2": 4.0, "na": 0.0, "beta": 0.1, - } + }, ) elif mapping[0] == "Colella": @@ -71,7 +71,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): "n2": 4.0, "na": 0.0, "beta": 0.1, - } + }, ) if show_plots: @@ -89,7 +89,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): "n2": 4.0, "na": 0.0, "beta": 0.1, - } + }, ) if show_plots: @@ -106,7 +106,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): dirichlet_bc = [(False, False)] * 3 dirichlet_bc = tuple(dirichlet_bc) - print(f"{dirichlet_bc = }") + print(f"{dirichlet_bc =}") # derham object derham = Derham(Nel, p, spl_kind, comm=mpi_comm, dirichlet_bc=dirichlet_bc) @@ -124,7 +124,7 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): # test calling the diagonal method aaa = mass_mats.M0.matrix.diagonal() bbb = mass_mats.M1.matrix.diagonal() - print(f"{aaa = }, {bbb[0, 0] = }, {bbb[0, 1] = }") + print(f"{aaa =}, {bbb[0, 0] =}, {bbb[0, 1] =}") # compare to old STRUPHY bc_old = [[None, None], [None, None], [None, None]] @@ -220,7 +220,11 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): # Change order of input in callable rM1ninvswitch_psy = mass_mats.create_weighted_mass( - "Hcurl", "Hcurl", weights=["sqrt_g", "1/eq_n0", "Ginv"], name="M1ninv", assemble=True + "Hcurl", + "Hcurl", + weights=["sqrt_g", "1/eq_n0", "Ginv"], + name="M1ninv", + assemble=True, ).dot(x1_psy, apply_bc=True) rot_B = RotationMatrix( @@ -229,7 +233,11 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): mass_mats.weights[mass_mats.selected_weight].b2_3, ) rM1Bninvswitch_psy = mass_mats.create_weighted_mass( - "Hcurl", "Hcurl", weights=["1/eq_n0", "sqrt_g", "Ginv", rot_B, "Ginv"], name="M1Bninv", assemble=True + "Hcurl", + "Hcurl", + weights=["1/eq_n0", "sqrt_g", "Ginv", rot_B, "Ginv"], + name="M1Bninv", + assemble=True, ).dot(x1_psy, apply_bc=True) # Test matrix free operators @@ -255,7 +263,11 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): # Change order of input in callable rM1ninvswitch_fre = mass_mats_free.create_weighted_mass( - "Hcurl", "Hcurl", weights=["sqrt_g", "1/eq_n0", "Ginv"], name="M1ninvswitch", assemble=True + "Hcurl", + "Hcurl", + weights=["sqrt_g", "1/eq_n0", "Ginv"], + name="M1ninvswitch", + assemble=True, ).dot(x1_psy, apply_bc=True) rot_B = RotationMatrix( mass_mats_free.weights[mass_mats_free.selected_weight].b2_1, @@ -264,7 +276,11 @@ def test_mass(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): ) rM1Bninvswitch_fre = mass_mats_free.create_weighted_mass( - "Hcurl", "Hcurl", weights=["1/eq_n0", "sqrt_g", "Ginv", rot_B, "Ginv"], name="M1Bninvswitch", assemble=True + "Hcurl", + "Hcurl", + weights=["1/eq_n0", "sqrt_g", "Ginv", rot_B, "Ginv"], + name="M1Bninvswitch", + assemble=True, ).dot(x1_psy, apply_bc=True) # compare output arrays @@ -420,7 +436,7 @@ def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): "n2": 4.0, "na": 0.0, "beta": 0.1, - } + }, ) if show_plots: @@ -440,7 +456,14 @@ def test_mass_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots=False): # derham object derham = Derham( - Nel, p, spl_kind, comm=mpi_comm, dirichlet_bc=dirichlet_bc, with_projectors=False, polar_ck=1, domain=domain + Nel, + p, + spl_kind, + comm=mpi_comm, + dirichlet_bc=dirichlet_bc, + with_projectors=False, + polar_ck=1, + domain=domain, ) print(f"Rank {mpi_rank} | Local domain : " + str(derham.domain_array[mpi_rank])) @@ -619,7 +642,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots "n2": 4.0, "na": 0.0, "beta": 0.1, - } + }, ) elif mapping[0] == "Colella": @@ -634,7 +657,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots "n2": 4.0, "na": 0.0, "beta": 0.1, - } + }, ) if show_plots: @@ -652,7 +675,7 @@ def test_mass_preconditioner(Nel, p, spl_kind, dirichlet_bc, mapping, show_plots "n2": 4.0, "na": 0.0, "beta": 0.1, - } + }, ) if show_plots: @@ -926,7 +949,7 @@ def test_mass_preconditioner_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show "n2": 4.0, "na": 0.0, "beta": 0.1, - } + }, ) if show_plots: @@ -946,7 +969,14 @@ def test_mass_preconditioner_polar(Nel, p, spl_kind, dirichlet_bc, mapping, show # derham object derham = Derham( - Nel, p, spl_kind, comm=mpi_comm, dirichlet_bc=dirichlet_bc, with_projectors=False, polar_ck=1, domain=domain + Nel, + p, + spl_kind, + comm=mpi_comm, + dirichlet_bc=dirichlet_bc, + with_projectors=False, + polar_ck=1, + domain=domain, ) print(f"Rank {mpi_rank} | Local domain : " + str(derham.domain_array[mpi_rank])) diff --git a/src/struphy/feec/tests/test_toarray_struphy.py b/src/struphy/feec/tests/test_toarray_struphy.py index 0279293ea..90427d8e4 100644 --- a/src/struphy/feec/tests/test_toarray_struphy.py +++ b/src/struphy/feec/tests/test_toarray_struphy.py @@ -5,7 +5,8 @@ @pytest.mark.parametrize("p", [[3, 2, 1]]) @pytest.mark.parametrize("spl_kind", [[False, True, True], [True, False, False]]) @pytest.mark.parametrize( - "mapping", [["Cuboid", {"l1": 1.0, "r1": 2.0, "l2": 10.0, "r2": 20.0, "l3": 100.0, "r3": 200.0}]] + "mapping", + [["Cuboid", {"l1": 1.0, "r1": 2.0, "l2": 10.0, "r2": 20.0, "l3": 100.0, "r3": 200.0}]], ) def test_toarray_struphy(Nel, p, spl_kind, mapping): """ diff --git a/src/struphy/feec/tests/test_tosparse_struphy.py b/src/struphy/feec/tests/test_tosparse_struphy.py index 020ab9e29..48cbfd7a2 100644 --- a/src/struphy/feec/tests/test_tosparse_struphy.py +++ b/src/struphy/feec/tests/test_tosparse_struphy.py @@ -7,7 +7,8 @@ @pytest.mark.parametrize("p", [[3, 2, 1]]) @pytest.mark.parametrize("spl_kind", [[False, True, True], [True, False, False]]) @pytest.mark.parametrize( - "mapping", [["Cuboid", {"l1": 1.0, "r1": 2.0, "l2": 10.0, "r2": 20.0, "l3": 100.0, "r3": 200.0}]] + "mapping", + [["Cuboid", {"l1": 1.0, "r1": 2.0, "l2": 10.0, "r2": 20.0, "l3": 100.0, "r3": 200.0}]], ) def test_tosparse_struphy(Nel, p, spl_kind, mapping): """ @@ -115,14 +116,26 @@ def test_tosparse_struphy(Nel, p, spl_kind, mapping): if __name__ == "__main__": test_tosparse_struphy( - [32, 2, 2], [2, 1, 1], [True, True, True], ["Colella", {"Lx": 1.0, "Ly": 2.0, "alpha": 0.5, "Lz": 3.0}] + [32, 2, 2], + [2, 1, 1], + [True, True, True], + ["Colella", {"Lx": 1.0, "Ly": 2.0, "alpha": 0.5, "Lz": 3.0}], ) test_tosparse_struphy( - [2, 32, 2], [1, 2, 1], [True, True, True], ["Colella", {"Lx": 1.0, "Ly": 2.0, "alpha": 0.5, "Lz": 3.0}] + [2, 32, 2], + [1, 2, 1], + [True, True, True], + ["Colella", {"Lx": 1.0, "Ly": 2.0, "alpha": 0.5, "Lz": 3.0}], ) test_tosparse_struphy( - [2, 2, 32], [1, 1, 2], [True, True, True], ["Colella", {"Lx": 1.0, "Ly": 2.0, "alpha": 0.5, "Lz": 3.0}] + [2, 2, 32], + [1, 1, 2], + [True, True, True], + ["Colella", {"Lx": 1.0, "Ly": 2.0, "alpha": 0.5, "Lz": 3.0}], ) test_tosparse_struphy( - [2, 2, 32], [1, 1, 2], [False, False, False], ["Colella", {"Lx": 1.0, "Ly": 2.0, "alpha": 0.5, "Lz": 3.0}] + [2, 2, 32], + [1, 1, 2], + [False, False, False], + ["Colella", {"Lx": 1.0, "Ly": 2.0, "alpha": 0.5, "Lz": 3.0}], ) diff --git a/src/struphy/feec/utilities.py b/src/struphy/feec/utilities.py index aa3912857..2fffe38e3 100644 --- a/src/struphy/feec/utilities.py +++ b/src/struphy/feec/utilities.py @@ -45,7 +45,7 @@ def __call__(self, e1, e2, e3): [ [self._cross_mask[m][n] * fun(e1, e2, e3) for n, fun in enumerate(row)] for m, row in enumerate(self._funs) - ] + ], ) # numpy operates on the last two indices with @ @@ -99,7 +99,9 @@ def create_equal_random_arrays(V, seed=123, flattened=False): e = arr_psy.ends arr_psy[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = arr[-1][ - s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1 + s[0] : e[0] + 1, + s[1] : e[1] + 1, + s[2] : e[2] + 1, ] if flattened: @@ -117,7 +119,9 @@ def create_equal_random_arrays(V, seed=123, flattened=False): e = block.ends arr_psy[d][s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = arr[-1][ - s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1 + s[0] : e[0] + 1, + s[1] : e[1] + 1, + s[2] : e[2] + 1, ] if flattened: @@ -126,7 +130,7 @@ def create_equal_random_arrays(V, seed=123, flattened=False): arr[0].flatten(), arr[1].flatten(), arr[2].flatten(), - ) + ), ) arr_psy.update_ghost_regions() diff --git a/src/struphy/feec/utilities_local_projectors.py b/src/struphy/feec/utilities_local_projectors.py index 9ff91a120..3ce6590f2 100644 --- a/src/struphy/feec/utilities_local_projectors.py +++ b/src/struphy/feec/utilities_local_projectors.py @@ -592,10 +592,10 @@ def build_translation_list_for_non_zero_spline_indices( for h in range(3): translation_indices_B_or_D_splines[h]["B"][Basis_functions_indices_B[h]] = xp.arange( - len(Basis_functions_indices_B[h]) + len(Basis_functions_indices_B[h]), ) translation_indices_B_or_D_splines[h]["D"][Basis_functions_indices_D[h]] = xp.arange( - len(Basis_functions_indices_D[h]) + len(Basis_functions_indices_D[h]), ) if sp_id in {"Hcurl", "Hdiv", "H1vec"}: @@ -610,7 +610,11 @@ def build_translation_list_for_non_zero_spline_indices( def evaluate_relevant_splines_at_relevant_points( - localpts, Bspaces_1d, Dspaces_1d, Basis_functions_indices_B, Basis_functions_indices_D + localpts, + Bspaces_1d, + Dspaces_1d, + Basis_functions_indices_B, + Basis_functions_indices_D, ): """This function evaluates all the B and D-splines that produce non-zeros in the BasisProjectionOperatorLocal's rows that belong to the current MPI rank over all the local evaluation points. They are store as float arrays in a dictionary of lists. @@ -693,7 +697,15 @@ def evaluate_relevant_splines_at_relevant_points( def determine_non_zero_rows_for_each_spline( - Basis_functions_indices_B, Basis_functions_indices_D, starts, ends, p, B_nbasis, D_nbasis, periodic, IoH + Basis_functions_indices_B, + Basis_functions_indices_D, + starts, + ends, + p, + B_nbasis, + D_nbasis, + periodic, + IoH, ): """This function determines for which rows (amongst those belonging to the current MPI rank) of the BasisProjectionOperatorLocal each B and D spline, of relevance for the current MPI rank, produces non-zero entries and annotates this regions of non-zeros by saving the rows at which each region starts and ends. @@ -792,7 +804,8 @@ def process_splines(indices, nbasis, is_D, h): def get_splines_that_are_relevant_for_at_least_one_block( - Basis_function_indices_agreggated_B, Basis_function_indices_agreggated_D + Basis_function_indices_agreggated_B, + Basis_function_indices_agreggated_D, ): """This function builds one list with all the B-spline indices (and another one for the D-splines) that are required for at least one block of the FE coefficients the current MPI rank needs to build its share of the BasisProjectionOperatorLocal. diff --git a/src/struphy/feec/variational_utilities.py b/src/struphy/feec/variational_utilities.py index 12695bfa2..d03a75e3d 100644 --- a/src/struphy/feec/variational_utilities.py +++ b/src/struphy/feec/variational_utilities.py @@ -264,19 +264,27 @@ def dot(self, v, out=None): self.gv3f.vector = grad_3_v vf_values = self.vf.eval_tp_fixed_loc( - self.interpolation_grid_spans, [self.interpolation_grid_bn] * 3, out=self._vf_values + self.interpolation_grid_spans, + [self.interpolation_grid_bn] * 3, + out=self._vf_values, ) gvf1_values = self.gv1f.eval_tp_fixed_loc( - self.interpolation_grid_spans, self.interpolation_grid_gradient, out=self._gvf1_values + self.interpolation_grid_spans, + self.interpolation_grid_gradient, + out=self._gvf1_values, ) gvf2_values = self.gv2f.eval_tp_fixed_loc( - self.interpolation_grid_spans, self.interpolation_grid_gradient, out=self._gvf2_values + self.interpolation_grid_spans, + self.interpolation_grid_gradient, + out=self._gvf2_values, ) gvf3_values = self.gv3f.eval_tp_fixed_loc( - self.interpolation_grid_spans, self.interpolation_grid_gradient, out=self._gvf3_values + self.interpolation_grid_spans, + self.interpolation_grid_gradient, + out=self._gvf3_values, ) self.PiuT.update_weights([[vf_values[0], vf_values[1], vf_values[2]]]) @@ -1383,7 +1391,12 @@ def update_weight(self, coeffs): self._pc.update_mass_operator(self._massop) def _create_inv( - self, type="pcg", pc_type="MassMatrixDiagonalPreconditioner", tol=1e-16, maxiter=500, verbose=False + self, + type="pcg", + pc_type="MassMatrixDiagonalPreconditioner", + tol=1e-16, + maxiter=500, + verbose=False, ): """Inverse the weighted mass matrix""" if pc_type is None: diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index 4696acf54..7ad6e3887 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -783,19 +783,28 @@ def u_cart(self, *etas, squeeze_out=False): def curl_unit_b1(self, *etas, squeeze_out=False): """1-form components of curl of unit magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z).""" return self.domain.pull( - self.curl_unit_b_cart(*etas, squeeze_out=False)[0], *etas, kind="1", squeeze_out=squeeze_out + self.curl_unit_b_cart(*etas, squeeze_out=False)[0], + *etas, + kind="1", + squeeze_out=squeeze_out, ) def curl_unit_b2(self, *etas, squeeze_out=False): """2-form components of curl of unit magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z).""" return self.domain.pull( - self.curl_unit_b_cart(*etas, squeeze_out=False)[0], *etas, kind="2", squeeze_out=squeeze_out + self.curl_unit_b_cart(*etas, squeeze_out=False)[0], + *etas, + kind="2", + squeeze_out=squeeze_out, ) def curl_unit_bv(self, *etas, squeeze_out=False): """Contra-variant components of curl of unit magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z).""" return self.domain.pull( - self.curl_unit_b_cart(*etas, squeeze_out=False)[0], *etas, kind="v", squeeze_out=squeeze_out + self.curl_unit_b_cart(*etas, squeeze_out=False)[0], + *etas, + kind="v", + squeeze_out=squeeze_out, ) def curl_unit_b_cart(self, *etas, squeeze_out=False): diff --git a/src/struphy/fields_background/coil_fields/base.py b/src/struphy/fields_background/coil_fields/base.py index 3b068c62d..331e89e7d 100644 --- a/src/struphy/fields_background/coil_fields/base.py +++ b/src/struphy/fields_background/coil_fields/base.py @@ -71,7 +71,11 @@ def bv(self, *etas, squeeze_out=False): def b_cart(self, *etas, squeeze_out=False): """Cartesian components of equilibrium magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z).""" b_out = self.domain.push( - self.bv(*etas, squeeze_out=False), *etas, kind="v", a_kwargs={"squeeze_out": False}, squeeze_out=squeeze_out + self.bv(*etas, squeeze_out=False), + *etas, + kind="v", + a_kwargs={"squeeze_out": False}, + squeeze_out=squeeze_out, ) return b_out, self.domain(*etas, squeeze_out=squeeze_out) diff --git a/src/struphy/fields_background/coil_fields/coil_fields.py b/src/struphy/fields_background/coil_fields/coil_fields.py index e1f66c71d..1b5c66a15 100644 --- a/src/struphy/fields_background/coil_fields/coil_fields.py +++ b/src/struphy/fields_background/coil_fields/coil_fields.py @@ -15,7 +15,9 @@ def __init__(self, csv_path=None, Nel=[16, 16, 16], p=[3, 3, 3], domain=None, ** self._ratgui_csv_data = load_csv_data(csv_path) derham = Derham( - Nel=Nel, p=p, spl_kind=[False, False, True] + Nel=Nel, + p=p, + spl_kind=[False, False, True], ) # Assuming (R=eta1, Z=eta2, phi=eta3) coordinates for csv data (periodic in eta3 only). self._interpolate = derham.P[ "v" @@ -34,14 +36,14 @@ def __init__(self, csv_path=None, Nel=[16, 16, 16], p=[3, 3, 3], domain=None, ** self.rhs[1][:] = B_Z self.rhs[2][:] = B_phi - print(f"{self.rhs = }") - print(f"{derham.nbasis['v'] = }") - print(f"{self.rhs[0] = }") - print(f"{self.rhs[1] = }") - print(f"{self.rhs[2] = }") - print(f"{self.rhs[0][:].shape = }") - print(f"{self.rhs[1][:].shape = }") - print(f"{self.rhs[2][:].shape = }") + print(f"{self.rhs =}") + print(f"{derham.nbasis['v'] =}") + print(f"{self.rhs[0] =}") + print(f"{self.rhs[1] =}") + print(f"{self.rhs[2] =}") + print(f"{self.rhs[0][:].shape =}") + print(f"{self.rhs[1][:].shape =}") + print(f"{self.rhs[2][:].shape =}") # We need to choose Nel and p such that the csv_data fits into this vector. # For a periodic direction, the size of the vector is Nel, for non-periodic (spl_kind=False) the size is Nel + p. # See the Tutorial on FEEC data structures https://struphy.pages.mpcdf.de/struphy/tutorials/tutorial_06_data_structures.html#FEEC-data-structures on how to address such a vector diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 93688b55f..f1afacd35 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -1759,7 +1759,7 @@ def __init__( units["p"] = 1.0 units["n"] = 1e20 warnings.warn( - f"{units = }, no rescaling performed in EQDSK output.", + f"{units =}, no rescaling performed in EQDSK output.", ) self._units = units @@ -2133,7 +2133,7 @@ def __init__( with pytest.raises(SystemExit) as exc: print("Simulation aborted, gvec must be installed (pip install gvec)!") sys.exit(1) - print(f"{exc.value.code = }") + print(f"{exc.value.code =}") import gvec @@ -2148,7 +2148,7 @@ def __init__( units["p"] = 1.0 units["n"] = 1e20 warnings.warn( - f"{units = }, no rescaling performed in GVEC output.", + f"{units =}, no rescaling performed in GVEC output.", ) self._units = units @@ -2428,7 +2428,7 @@ def __init__( units["p"] = 1.0 units["n"] = 1e20 warnings.warn( - f"{units = }, no rescaling performed in DESC output.", + f"{units =}, no rescaling performed in DESC output.", ) self._units = units @@ -2880,23 +2880,23 @@ def desc_eval( if verbose: # import sys - print(f"\n{nfp = }") - print(f"{self.eq.axis = }") - print(f"{rho.size = }") - print(f"{theta.size = }") - print(f"{zeta.size = }") - print(f"{grid_3d.num_rho = }") - print(f"{grid_3d.num_theta = }") - print(f"{grid_3d.num_zeta = }") + print(f"\n{nfp =}") + print(f"{self.eq.axis =}") + print(f"{rho.size =}") + print(f"{theta.size =}") + print(f"{zeta.size =}") + print(f"{grid_3d.num_rho =}") + print(f"{grid_3d.num_theta =}") + print(f"{grid_3d.num_zeta =}") # print(f'\n{grid_3d.nodes[:, 0] = }') # print(f'\n{grid_3d.nodes[:, 1] = }') # print(f'\n{grid_3d.nodes[:, 2] = }') - print(f"\n{rho = }") - print(f"{rho1 = }") - print(f"\n{theta = }") - print(f"{theta1 = }") - print(f"\n{zeta = }") - print(f"{zeta1 = }") + print(f"\n{rho =}") + print(f"{rho1 =}") + print(f"\n{theta =}") + print(f"{theta1 =}") + print(f"\n{zeta =}") + print(f"{zeta1 =}") # make c-contiguous out = xp.ascontiguousarray(out) diff --git a/src/struphy/fields_background/mhd_equil/eqdsk/readeqdsk.py b/src/struphy/fields_background/mhd_equil/eqdsk/readeqdsk.py index 67578ce17..812381e87 100644 --- a/src/struphy/fields_background/mhd_equil/eqdsk/readeqdsk.py +++ b/src/struphy/fields_background/mhd_equil/eqdsk/readeqdsk.py @@ -173,7 +173,11 @@ def main(): action="store_true", ) parser.add_option( - "-v", "--vars", dest="vars", help="comma separated list of variables (use '-v \"*\"' for all)", default="*" + "-v", + "--vars", + dest="vars", + help="comma separated list of variables (use '-v \"*\"' for all)", + default="*", ) parser.add_option( "-p", diff --git a/src/struphy/fields_background/tests/test_desc_equil.py b/src/struphy/fields_background/tests/test_desc_equil.py index 5aca31b8d..c7130f0a3 100644 --- a/src/struphy/fields_background/tests/test_desc_equil.py +++ b/src/struphy/fields_background/tests/test_desc_equil.py @@ -167,7 +167,7 @@ def test_desc_equil(do_plot=False): err_lim = 0.09 for nfp in nfps: - print(f"\n{nfp = }") + print(f"\n{nfp =}") for var in vars: if var in ("B_R", "B_phi", "B_Z", "J_R", "J_phi", "J_Z"): continue @@ -179,7 +179,7 @@ def test_desc_equil(do_plot=False): assert err < err_lim print( - f"compare {var}: {err = }", + f"compare {var}: {err =}", ) if do_plot: @@ -193,7 +193,7 @@ def test_desc_equil(do_plot=False): plt.subplot(2, 2, 1) map1 = plt.contourf(R, Z, outs[nfp][var][:, :, 0], levels=levels) - plt.title(f"DESC, {var = }, {nfp = }") + plt.title(f"DESC, {var =}, {nfp =}") plt.xlabel("$R$") plt.ylabel("$Z$") plt.axis("equal") @@ -201,7 +201,7 @@ def test_desc_equil(do_plot=False): plt.subplot(2, 2, 2) map2 = plt.contourf(R, Z, outs_struphy[nfp][var][:, :, 0], levels=levels) - plt.title(f"Struphy, {err = }") + plt.title(f"Struphy, {err =}") plt.xlabel("$R$") plt.ylabel("$Z$") plt.axis("equal") @@ -217,7 +217,7 @@ def test_desc_equil(do_plot=False): plt.subplot(2, 2, 3) map3 = plt.contourf(x1, y1, outs[nfp][var][:, 0, :], levels=levels) map3b = plt.contourf(x2, y2, outs[nfp][var][:, n2 // 2, :], levels=levels) - plt.title(f"DESC, {var = }, {nfp = }") + plt.title(f"DESC, {var =}, {nfp =}") plt.xlabel("$x$") plt.ylabel("$y$") plt.axis("equal") @@ -226,7 +226,7 @@ def test_desc_equil(do_plot=False): plt.subplot(2, 2, 4) map4 = plt.contourf(x1, y1, outs_struphy[nfp][var][:, 0, :], levels=levels) map4b = plt.contourf(x2, y2, outs_struphy[nfp][var][:, n2 // 2, :], levels=levels) - plt.title(f"Struphy, {err = }") + plt.title(f"Struphy, {err =}") plt.xlabel("$x$") plt.ylabel("$y$") plt.axis("equal") diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index a789913e7..d2b21688e 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -1287,9 +1287,9 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): elif isinstance(component, xp.ndarray): if flat_eval: - assert component.ndim == 1, print(f"{component.ndim = }") + assert component.ndim == 1, print(f"{component.ndim =}") else: - assert component.ndim == 3, print(f"{component.ndim = }") + assert component.ndim == 3, print(f"{component.ndim =}") a_out += [component] @@ -1312,7 +1312,7 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): else: raise ValueError( "Input array a_in must be either 1d (scalar) or \ - 2d (vector-valued, shape (3,:)) for flat evaluation!" + 2d (vector-valued, shape (3,:)) for flat evaluation!", ) else: @@ -1331,13 +1331,13 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): else: raise ValueError( "Input array a_in must be either 3d (scalar) or \ - 4d (vector-valued, shape (3,:,:,:)) for non-flat evaluation!" + 4d (vector-valued, shape (3,:,:,:)) for non-flat evaluation!", ) else: raise TypeError( "Argument a must be either a float OR a list/tuple of 1 or 3 callable(s)/numpy array(s)/float(s) \ - OR a single numpy array OR a single callable!" + OR a single numpy array OR a single callable!", ) # make sure that output array is 2d and of shape (:, 1) or (:, 3) for flat evaluation diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index 024b980c9..7b2c25064 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -743,7 +743,7 @@ def __init__( self.params = copy.deepcopy(locals()) self.params_numpy = self.get_params_numpy() - assert a2 <= R0, f"The minor radius must be smaller or equal than the major radius! {a2 = }, {R0 = }" + assert a2 <= R0, f"The minor radius must be smaller or equal than the major radius! {a2 =}, {R0 =}" if sfl: assert pol_period == 1, ( diff --git a/src/struphy/geometry/mappings_kernels.py b/src/struphy/geometry/mappings_kernels.py index 8b643855d..83e6275fd 100644 --- a/src/struphy/geometry/mappings_kernels.py +++ b/src/struphy/geometry/mappings_kernels.py @@ -49,13 +49,40 @@ def spline_3d( tmp3 = ind3[span3 - int(p[2]), :] f_out[0] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), b1, b2, b3, tmp1, tmp2, tmp3, args.cx + int(p[0]), + int(p[1]), + int(p[2]), + b1, + b2, + b3, + tmp1, + tmp2, + tmp3, + args.cx, ) f_out[1] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), b1, b2, b3, tmp1, tmp2, tmp3, args.cy + int(p[0]), + int(p[1]), + int(p[2]), + b1, + b2, + b3, + tmp1, + tmp2, + tmp3, + args.cy, ) f_out[2] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), b1, b2, b3, tmp1, tmp2, tmp3, args.cz + int(p[0]), + int(p[1]), + int(p[2]), + b1, + b2, + b3, + tmp1, + tmp2, + tmp3, + args.cz, ) @@ -97,31 +124,112 @@ def spline_3d_df( tmp3 = ind3[span3 - int(p[2]), :] df_out[0, 0] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), der1, b2, b3, tmp1, tmp2, tmp3, args.cx + int(p[0]), + int(p[1]), + int(p[2]), + der1, + b2, + b3, + tmp1, + tmp2, + tmp3, + args.cx, ) df_out[0, 1] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), b1, der2, b3, tmp1, tmp2, tmp3, args.cx + int(p[0]), + int(p[1]), + int(p[2]), + b1, + der2, + b3, + tmp1, + tmp2, + tmp3, + args.cx, ) df_out[0, 2] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), b1, b2, der3, tmp1, tmp2, tmp3, args.cx + int(p[0]), + int(p[1]), + int(p[2]), + b1, + b2, + der3, + tmp1, + tmp2, + tmp3, + args.cx, ) df_out[1, 0] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), der1, b2, b3, tmp1, tmp2, tmp3, args.cy + int(p[0]), + int(p[1]), + int(p[2]), + der1, + b2, + b3, + tmp1, + tmp2, + tmp3, + args.cy, ) df_out[1, 1] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), b1, der2, b3, tmp1, tmp2, tmp3, args.cy + int(p[0]), + int(p[1]), + int(p[2]), + b1, + der2, + b3, + tmp1, + tmp2, + tmp3, + args.cy, ) df_out[1, 2] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), b1, b2, der3, tmp1, tmp2, tmp3, args.cy + int(p[0]), + int(p[1]), + int(p[2]), + b1, + b2, + der3, + tmp1, + tmp2, + tmp3, + args.cy, ) df_out[2, 0] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), der1, b2, b3, tmp1, tmp2, tmp3, args.cz + int(p[0]), + int(p[1]), + int(p[2]), + der1, + b2, + b3, + tmp1, + tmp2, + tmp3, + args.cz, ) df_out[2, 1] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), b1, der2, b3, tmp1, tmp2, tmp3, args.cz + int(p[0]), + int(p[1]), + int(p[2]), + b1, + der2, + b3, + tmp1, + tmp2, + tmp3, + args.cz, ) df_out[2, 2] = evaluation_kernels_3d.evaluation_kernel_3d( - int(p[0]), int(p[1]), int(p[2]), b1, b2, der3, tmp1, tmp2, tmp3, args.cz + int(p[0]), + int(p[1]), + int(p[2]), + b1, + b2, + der3, + tmp1, + tmp2, + tmp3, + args.cz, ) @@ -276,7 +384,7 @@ def spline_2d_torus( tmp2 = ind2[span2 - int(p[1]), :] f_out[0] = evaluation_kernels_2d.evaluation_kernel_2d(int(p[0]), int(p[1]), b1, b2, tmp1, tmp2, cx) * cos( - 2 * pi * eta3 / tor_period + 2 * pi * eta3 / tor_period, ) f_out[1] = ( evaluation_kernels_2d.evaluation_kernel_2d(int(p[0]), int(p[1]), b1, b2, tmp1, tmp2, cx) @@ -329,10 +437,10 @@ def spline_2d_torus_df( tmp2 = ind2[span2 - int(p[1]), :] df_out[0, 0] = evaluation_kernels_2d.evaluation_kernel_2d(int(p[0]), int(p[1]), der1, b2, tmp1, tmp2, cx) * cos( - 2 * pi * eta3 / tor_period + 2 * pi * eta3 / tor_period, ) df_out[0, 1] = evaluation_kernels_2d.evaluation_kernel_2d(int(p[0]), int(p[1]), b1, der2, tmp1, tmp2, cx) * cos( - 2 * pi * eta3 / tor_period + 2 * pi * eta3 / tor_period, ) df_out[0, 2] = ( evaluation_kernels_2d.evaluation_kernel_2d(int(p[0]), int(p[1]), b1, b2, tmp1, tmp2, cx) @@ -621,7 +729,14 @@ def hollow_cyl_df(eta1: float, eta2: float, a1: float, a2: float, lz: float, poc @pure def powered_ellipse( - eta1: float, eta2: float, eta3: float, rx: float, ry: float, lz: float, s: float, f_out: "float[:]" + eta1: float, + eta2: float, + eta3: float, + rx: float, + ry: float, + lz: float, + s: float, + f_out: "float[:]", ): r""" Point-wise evaluation of @@ -664,7 +779,14 @@ def powered_ellipse( @pure def powered_ellipse_df( - eta1: float, eta2: float, eta3: float, rx: float, ry: float, lz: float, s: float, df_out: "float[:,:]" + eta1: float, + eta2: float, + eta3: float, + rx: float, + ry: float, + lz: float, + s: float, + df_out: "float[:,:]", ): """Jacobian matrix for :meth:`struphy.geometry.mappings_kernels.powered_ellipse`.""" @@ -843,7 +965,14 @@ def hollow_torus_df( @pure def shafranov_shift( - eta1: float, eta2: float, eta3: float, rx: float, ry: float, lz: float, de: float, f_out: "float[:]" + eta1: float, + eta2: float, + eta3: float, + rx: float, + ry: float, + lz: float, + de: float, + f_out: "float[:]", ): r""" Point-wise evaluation of @@ -887,7 +1016,14 @@ def shafranov_shift( @pure def shafranov_shift_df( - eta1: float, eta2: float, eta3: float, rx: float, ry: float, lz: float, de: float, df_out: "float[:,:]" + eta1: float, + eta2: float, + eta3: float, + rx: float, + ry: float, + lz: float, + de: float, + df_out: "float[:,:]", ): """Jacobian matrix for :meth:`struphy.geometry.mappings_kernels.shafranov_shift`.""" @@ -904,7 +1040,14 @@ def shafranov_shift_df( @pure def shafranov_sqrt( - eta1: float, eta2: float, eta3: float, rx: float, ry: float, lz: float, de: float, f_out: "float[:]" + eta1: float, + eta2: float, + eta3: float, + rx: float, + ry: float, + lz: float, + de: float, + f_out: "float[:]", ): r""" Point-wise evaluation of @@ -946,7 +1089,14 @@ def shafranov_sqrt( @pure def shafranov_sqrt_df( - eta1: float, eta2: float, eta3: float, rx: float, ry: float, lz: float, de: float, df_out: "float[:,:]" + eta1: float, + eta2: float, + eta3: float, + rx: float, + ry: float, + lz: float, + de: float, + df_out: "float[:,:]", ): """Jacobian matrix for :meth:`struphy.geometry.mappings_kernels.shafranov_sqrt`.""" diff --git a/src/struphy/geometry/transform_kernels.py b/src/struphy/geometry/transform_kernels.py index c95156b96..f9e6d8077 100644 --- a/src/struphy/geometry/transform_kernels.py +++ b/src/struphy/geometry/transform_kernels.py @@ -54,7 +54,13 @@ @stack_array("dfmat1", "dfmat2") def pull( - a: "float[:]", eta1: float, eta2: float, eta3: float, kind_fun: int, args_domain: "DomainArguments", out: "float[:]" + a: "float[:]", + eta1: float, + eta2: float, + eta3: float, + kind_fun: int, + args_domain: "DomainArguments", + out: "float[:]", ): """ Pull-back of a Cartesian scalar/vector field to a differential p-form. @@ -114,7 +120,13 @@ def pull( @stack_array("dfmat1", "dfmat2", "dfmat3") def push( - a: "float[:]", eta1: float, eta2: float, eta3: float, kind_fun: int, args_domain: "DomainArguments", out: "float[:]" + a: "float[:]", + eta1: float, + eta2: float, + eta3: float, + kind_fun: int, + args_domain: "DomainArguments", + out: "float[:]", ): """ Pushforward of a differential p-forms to a Cartesian scalar/vector field. @@ -172,7 +184,13 @@ def push( @stack_array("dfmat1", "dfmat2", "dfmat3", "vec1", "vec2") def tran( - a: "float[:]", eta1: float, eta2: float, eta3: float, kind_fun: int, args_domain: "DomainArguments", out: "float[:]" + a: "float[:]", + eta1: float, + eta2: float, + eta3: float, + kind_fun: int, + args_domain: "DomainArguments", + out: "float[:]", ): """ Transformations between differential p-forms and/or vector fields. diff --git a/src/struphy/initial/eigenfunctions.py b/src/struphy/initial/eigenfunctions.py index 2f66fcacb..239ae24a9 100644 --- a/src/struphy/initial/eigenfunctions.py +++ b/src/struphy/initial/eigenfunctions.py @@ -67,13 +67,16 @@ def __init__(self, derham, **params): nnz_tor = derham.boundary_ops["2"].dim_nz_tor eig_vec_1 = U2_eig[ - 0 * nnz_pol[0] + 0 * nnz_pol[1] + 0 * nnz_pol[2] : 1 * nnz_pol[0] + 0 * nnz_pol[1] + 0 * nnz_pol[2], mode + 0 * nnz_pol[0] + 0 * nnz_pol[1] + 0 * nnz_pol[2] : 1 * nnz_pol[0] + 0 * nnz_pol[1] + 0 * nnz_pol[2], + mode, ] eig_vec_2 = U2_eig[ - 1 * nnz_pol[0] + 0 * nnz_pol[1] + 0 * nnz_pol[2] : 1 * nnz_pol[0] + 1 * nnz_pol[1] + 0 * nnz_pol[2], mode + 1 * nnz_pol[0] + 0 * nnz_pol[1] + 0 * nnz_pol[2] : 1 * nnz_pol[0] + 1 * nnz_pol[1] + 0 * nnz_pol[2], + mode, ] eig_vec_3 = U2_eig[ - 1 * nnz_pol[0] + 1 * nnz_pol[1] + 0 * nnz_pol[2] : 1 * nnz_pol[0] + 1 * nnz_pol[1] + 1 * nnz_pol[2], mode + 1 * nnz_pol[0] + 1 * nnz_pol[1] + 0 * nnz_pol[2] : 1 * nnz_pol[0] + 1 * nnz_pol[1] + 1 * nnz_pol[2], + mode, ] del omega2, U2_eig @@ -182,7 +185,7 @@ def __init__(self, derham, **params): ] eigvec_1_ten[derham.Vh_pol["2"].n_rings[0] : derham.nbasis["2"][0][0] - bc1_2, :, :] = eig_vec_1[1].reshape( - n_v2_0[0] + n_v2_0[0], ) eigvec_2_ten[derham.Vh_pol["2"].n_rings[1] :, :, :] = eig_vec_2[1].reshape(n_v2_0[1]) eigvec_3_ten[derham.Vh_pol["2"].n_rings[2] :, :, :] = eig_vec_3[1].reshape(n_v2_0[2]) diff --git a/src/struphy/initial/perturbations.py b/src/struphy/initial/perturbations.py index 767f899c9..4c113d02d 100644 --- a/src/struphy/initial/perturbations.py +++ b/src/struphy/initial/perturbations.py @@ -418,7 +418,7 @@ def __call__(self, eta1, eta2, eta3): z = eta3 val += (self._a * scipy.special.jv(self._m, r) + self._b * scipy.special.yn(self._m, r)) * xp.cos( - self._m * theta + self._m * theta, ) return val diff --git a/src/struphy/initial/tests/test_init_perturbations.py b/src/struphy/initial/tests/test_init_perturbations.py index 2f5fa3176..dd391cf56 100644 --- a/src/struphy/initial/tests/test_init_perturbations.py +++ b/src/struphy/initial/tests/test_init_perturbations.py @@ -222,7 +222,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False params = { key: { "given_in_basis": [fun_form] * 3, - } + }, } if "Modes" in key: @@ -267,12 +267,20 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False fun3_xyz = perturbation(eee1, eee2, eee3) elif fun_form == "norm": tmp1, tmp2, tmp3 = domain.transform( - [perturbation, perturbation, perturbation], eee1, eee2, eee3, kind=fun_form + "_to_v" + [perturbation, perturbation, perturbation], + eee1, + eee2, + eee3, + kind=fun_form + "_to_v", ) fun1_xyz, fun2_xyz, fun3_xyz = domain.push([tmp1, tmp2, tmp3], eee1, eee2, eee3, kind="v") else: fun1_xyz, fun2_xyz, fun3_xyz = domain.push( - [perturbation, perturbation, perturbation], eee1, eee2, eee3, kind=fun_form + [perturbation, perturbation, perturbation], + eee1, + eee2, + eee3, + kind=fun_form, ) fun_xyz_vec = [fun1_xyz, fun2_xyz, fun3_xyz] @@ -299,7 +307,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False plt.ylabel("y") plt.colorbar() plt.title( - f"component {c + 1}, init was {fun_form}, (m,n)=({kwargs['ms'][0]},{kwargs['ns'][0]})" + f"component {c + 1}, init was {fun_form}, (m,n)=({kwargs['ms'][0]},{kwargs['ns'][0]})", ) ax = plt.gca() ax.set_aspect("equal", adjustable="box") @@ -316,7 +324,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False plt.ylabel("z") plt.colorbar() plt.title( - f"component {c + 1}, init was {fun_form}, (m,n)=({kwargs['ms'][0]},{kwargs['ns'][0]})" + f"component {c + 1}, init was {fun_form}, (m,n)=({kwargs['ms'][0]},{kwargs['ns'][0]})", ) ax = plt.gca() ax.set_aspect("equal", adjustable="box") diff --git a/src/struphy/io/output_handling.py b/src/struphy/io/output_handling.py index 734a2d1a3..808c2bcf3 100644 --- a/src/struphy/io/output_handling.py +++ b/src/struphy/io/output_handling.py @@ -54,7 +54,7 @@ def __init__(self, path_out, file_name=None, comm=None): dataset_keys = [] self._file.visit( - lambda key: dataset_keys.append(key) if isinstance(self._file[key], h5py.Dataset) else None + lambda key: dataset_keys.append(key) if isinstance(self._file[key], h5py.Dataset) else None, ) for key in dataset_keys: @@ -110,7 +110,11 @@ def add_data(self, data_dict): self._file[key][0] = val[0] else: self._file.create_dataset( - key, (1,) + val.shape, maxshape=(None,) + val.shape, dtype=val.dtype, chunks=True + key, + (1,) + val.shape, + maxshape=(None,) + val.shape, + dtype=val.dtype, + chunks=True, ) self._file[key][0] = val diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index 78a295b35..f38654160 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -233,35 +233,35 @@ def descend_options_dict( out = copy.deepcopy(d) if verbose: - print(f"{d = }") - print(f"{out = }") - print(f"{d_default = }") - print(f"{d_opts = }") - print(f"{keys = }") - print(f"{depth = }") - print(f"{pop_again = }") + print(f"{d =}") + print(f"{out =}") + print(f"{d_default =}") + print(f"{d_opts =}") + print(f"{keys =}") + print(f"{depth =}") + print(f"{pop_again =}") if verbose: - print(f"{d = }") - print(f"{out = }") - print(f"{d_default = }") - print(f"{d_opts = }") - print(f"{keys = }") - print(f"{depth = }") - print(f"{pop_again = }") + print(f"{d =}") + print(f"{out =}") + print(f"{d_default =}") + print(f"{d_opts =}") + print(f"{keys =}") + print(f"{depth =}") + print(f"{pop_again =}") count = 0 for key, val in d.items(): count += 1 if verbose: - print(f"\n{keys = } | {key = }, {type(val) = }, {count = }\n") + print(f"\n{keys =} | {key =}, {type(val) =}, {count =}\n") if isinstance(val, list): # create default parameter dict "out" if verbose: - print(f"{val = }") + print(f"{val =}") if d_default is None: if len(keys) == 0: @@ -299,10 +299,10 @@ def descend_options_dict( out += [out_sublist] if verbose: - print(f"{out = }") + print(f"{out =}") if verbose: - print(f"{out = }") + print(f"{out =}") # recurse if necessary elif isinstance(val, dict): diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 41c7b6575..765ee1508 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -395,7 +395,7 @@ def gaussian(self, v, u=0.0, vth=1.0, polar=False, volume_form=False): """ if isinstance(v, xp.ndarray) and isinstance(u, xp.ndarray): - assert v.shape == u.shape, f"{v.shape = } but {u.shape = }" + assert v.shape == u.shape, f"{v.shape =} but {u.shape =}" if not polar: out = 1.0 / vth * 1.0 / xp.sqrt(2.0 * xp.pi) * xp.exp(-((v - u) ** 2) / (2.0 * vth**2)) @@ -434,9 +434,9 @@ def __call__(self, *args): # Check that all args have the same shape shape0 = xp.shape(args[0]) for i, arg in enumerate(args): - assert xp.shape(arg) == shape0, f"Argument {i} has {xp.shape(arg) = }, but must be {shape0 = }." + assert xp.shape(arg) == shape0, f"Argument {i} has {xp.shape(arg) =}, but must be {shape0 =}." assert xp.ndim(arg) == 1 or xp.ndim(arg) == 3 + self.vdim, ( - f"{xp.ndim(arg) = } not allowed for Maxwellian evaluation." + f"{xp.ndim(arg) =} not allowed for Maxwellian evaluation." ) # flat or meshgrid evaluation # Get result evaluated at eta's diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index c7dea067a..d9e34dcab 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -431,7 +431,7 @@ def gaussian(self, e, vth=1.0): """ if isinstance(vth, xp.ndarray): - assert e.shape == vth.shape, f"{e.shape = } but {vth.shape = }" + assert e.shape == vth.shape, f"{e.shape =} but {vth.shape =}" return 2.0 * xp.sqrt(e / xp.pi) / vth**3 * xp.exp(-e / vth**2) @@ -462,9 +462,9 @@ def __call__(self, *args): # Check that all args have the same shape shape0 = xp.shape(args[0]) for i, arg in enumerate(args): - assert xp.shape(arg) == shape0, f"Argument {i} has {xp.shape(arg) = }, but must be {shape0 = }." + assert xp.shape(arg) == shape0, f"Argument {i} has {xp.shape(arg) =}, but must be {shape0 =}." assert xp.ndim(arg) == 1 or xp.ndim(arg) == 3, ( - f"{xp.ndim(arg) = } not allowed for canonical Maxwellian evaluation." + f"{xp.ndim(arg) =} not allowed for canonical Maxwellian evaluation." ) # flat or meshgrid evaluation # Get result evaluated with each particles' psic diff --git a/src/struphy/kinetic_background/tests/test_maxwellians.py b/src/struphy/kinetic_background/tests/test_maxwellians.py index edf24af4c..710a88262 100644 --- a/src/struphy/kinetic_background/tests/test_maxwellians.py +++ b/src/struphy/kinetic_background/tests/test_maxwellians.py @@ -287,10 +287,10 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): for key, val in inspect.getmembers(equils): if inspect.isclass(val) and val.__module__ == equils.__name__: - print(f"{key = }") + print(f"{key =}") if "DESCequilibrium" in key and not with_desc: - print(f"Attention: {with_desc = }, DESC not tested here !!") + print(f"Attention: {with_desc =}, DESC not tested here !!") continue if "GVECequilibrium" in key: @@ -298,16 +298,22 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): mhd_equil = val() assert isinstance(mhd_equil, FluidEquilibrium) - print(f"{mhd_equil.params = }") + print(f"{mhd_equil.params =}") if "AdhocTorus" in key: mhd_equil.domain = domains.HollowTorus( - a1=1e-3, a2=mhd_equil.params["a"], R0=mhd_equil.params["R0"], tor_period=1 + a1=1e-3, + a2=mhd_equil.params["a"], + R0=mhd_equil.params["R0"], + tor_period=1, ) elif "EQDSKequilibrium" in key: mhd_equil.domain = domains.Tokamak(equilibrium=mhd_equil) elif "CircularTokamak" in key: mhd_equil.domain = domains.HollowTorus( - a1=1e-3, a2=mhd_equil.params["a"], R0=mhd_equil.params["R0"], tor_period=1 + a1=1e-3, + a2=mhd_equil.params["a"], + R0=mhd_equil.params["R0"], + tor_period=1, ) elif "HomogenSlab" in key: mhd_equil.domain = domains.Cuboid() @@ -319,11 +325,15 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): ) elif "ShearFluid" in key: mhd_equil.domain = domains.Cuboid( - r1=mhd_equil.params["a"], r2=mhd_equil.params["b"], r3=mhd_equil.params["c"] + r1=mhd_equil.params["a"], + r2=mhd_equil.params["b"], + r3=mhd_equil.params["c"], ) elif "ScrewPinch" in key: mhd_equil.domain = domains.HollowCylinder( - a1=1e-3, a2=mhd_equil.params["a"], Lz=mhd_equil.params["R0"] * 2 * xp.pi + a1=1e-3, + a2=mhd_equil.params["a"], + Lz=mhd_equil.params["R0"] * 2 * xp.pi, ) else: try: @@ -354,11 +364,13 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): # test meshgrid evaluation n0 = mhd_equil.n0(*e_meshgrids) assert xp.allclose( - maxwellian(*meshgrids)[:, :, :, 0, 0, 0], n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 0, 0] + maxwellian(*meshgrids)[:, :, :, 0, 0, 0], + n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 0, 0], ) assert xp.allclose( - maxwellian(*meshgrids)[:, :, :, 0, 1, 2], n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 1, 2] + maxwellian(*meshgrids)[:, :, :, 0, 1, 2], + n0 * maxwellian_1(*meshgrids)[:, :, :, 0, 1, 2], ) # test flat evaluation @@ -378,7 +390,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): # plotting moments if show_plot: - plt.figure(f"{mhd_equil = }", figsize=(24, 16)) + plt.figure(f"{mhd_equil =}", figsize=(24, 16)) x, y, z = mhd_equil.domain(*e_meshgrids) # density plots @@ -390,14 +402,20 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): if "Slab" in key or "Pinch" in key: plt.contourf(x[:, 0, :], z[:, 0, :], n_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], z[:, Nel[1] // 2 - 1, :], n_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + z[:, Nel[1] // 2 - 1, :], + n_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("z") else: plt.contourf(x[:, 0, :], y[:, 0, :], n_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], y[:, Nel[1] // 2 - 1, :], n_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + y[:, Nel[1] // 2 - 1, :], + n_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("y") @@ -459,14 +477,20 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): if "Slab" in key or "Pinch" in key: plt.contourf(x[:, 0, :], z[:, 0, :], vth_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], z[:, Nel[1] // 2 - 1, :], vth_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + z[:, Nel[1] // 2 - 1, :], + vth_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("z") else: plt.contourf(x[:, 0, :], y[:, 0, :], vth_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], y[:, Nel[1] // 2 - 1, :], vth_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + y[:, Nel[1] // 2 - 1, :], + vth_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("y") @@ -496,7 +520,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): if inspect.isclass(val_2) and val_2.__module__ == perturbations.__name__: pert = val_2() assert isinstance(pert, Perturbation) - print(f"{pert = }") + print(f"{pert =}") if isinstance(pert, perturbations.Noise): continue @@ -550,14 +574,20 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): if "Slab" in key or "Pinch" in key: plt.contourf(x[:, 0, :], z[:, 0, :], n_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], z[:, Nel[1] // 2, :], n_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + z[:, Nel[1] // 2, :], + n_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("z") else: plt.contourf(x[:, 0, :], y[:, 0, :], n_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], y[:, Nel[1] // 2, :], n_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + y[:, Nel[1] // 2, :], + n_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("y") @@ -586,14 +616,20 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): if "Slab" in key or "Pinch" in key: plt.contourf(x[:, 0, :], z[:, 0, :], u[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], z[:, Nel[1] // 2, :], u[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + z[:, Nel[1] // 2, :], + u[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("z") else: plt.contourf(x[:, 0, :], y[:, 0, :], u[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], y[:, Nel[1] // 2, :], u[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + y[:, Nel[1] // 2, :], + u[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("y") @@ -1047,10 +1083,10 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): for key, val in inspect.getmembers(equils): if inspect.isclass(val) and val.__module__ == equils.__name__: - print(f"{key = }") + print(f"{key =}") if "DESCequilibrium" in key and not with_desc: - print(f"Attention: {with_desc = }, DESC not tested here !!") + print(f"Attention: {with_desc =}, DESC not tested here !!") continue if "GVECequilibrium" in key: @@ -1060,16 +1096,22 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): if not isinstance(mhd_equil, FluidEquilibriumWithB): continue - print(f"{mhd_equil.params = }") + print(f"{mhd_equil.params =}") if "AdhocTorus" in key: mhd_equil.domain = domains.HollowTorus( - a1=1e-3, a2=mhd_equil.params["a"], R0=mhd_equil.params["R0"], tor_period=1 + a1=1e-3, + a2=mhd_equil.params["a"], + R0=mhd_equil.params["R0"], + tor_period=1, ) elif "EQDSKequilibrium" in key: mhd_equil.domain = domains.Tokamak(equilibrium=mhd_equil) elif "CircularTokamak" in key: mhd_equil.domain = domains.HollowTorus( - a1=1e-3, a2=mhd_equil.params["a"], R0=mhd_equil.params["R0"], tor_period=1 + a1=1e-3, + a2=mhd_equil.params["a"], + R0=mhd_equil.params["R0"], + tor_period=1, ) elif "HomogenSlab" in key: mhd_equil.domain = domains.Cuboid() @@ -1081,11 +1123,15 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): ) elif "ShearFluid" in key: mhd_equil.domain = domains.Cuboid( - r1=mhd_equil.params["a"], r2=mhd_equil.params["b"], r3=mhd_equil.params["c"] + r1=mhd_equil.params["a"], + r2=mhd_equil.params["b"], + r3=mhd_equil.params["c"], ) elif "ScrewPinch" in key: mhd_equil.domain = domains.HollowCylinder( - a1=1e-3, a2=mhd_equil.params["a"], Lz=mhd_equil.params["R0"] * 2 * xp.pi + a1=1e-3, + a2=mhd_equil.params["a"], + Lz=mhd_equil.params["R0"] * 2 * xp.pi, ) else: try: @@ -1135,7 +1181,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): # plotting moments if show_plot: - plt.figure(f"{mhd_equil = }", figsize=(24, 16)) + plt.figure(f"{mhd_equil =}", figsize=(24, 16)) x, y, z = mhd_equil.domain(*e_meshgrids) # density plots @@ -1147,14 +1193,20 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): if "Slab" in key or "Pinch" in key: plt.contourf(x[:, 0, :], z[:, 0, :], n_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], z[:, Nel[1] // 2 - 1, :], n_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + z[:, Nel[1] // 2 - 1, :], + n_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("z") else: plt.contourf(x[:, 0, :], y[:, 0, :], n_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], y[:, Nel[1] // 2 - 1, :], n_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + y[:, Nel[1] // 2 - 1, :], + n_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("y") @@ -1216,14 +1268,20 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): if "Slab" in key or "Pinch" in key: plt.contourf(x[:, 0, :], z[:, 0, :], vth_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], z[:, Nel[1] // 2 - 1, :], vth_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + z[:, Nel[1] // 2 - 1, :], + vth_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("z") else: plt.contourf(x[:, 0, :], y[:, 0, :], vth_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], y[:, Nel[1] // 2 - 1, :], vth_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + y[:, Nel[1] // 2 - 1, :], + vth_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("y") @@ -1250,7 +1308,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): for key_2, val_2 in inspect.getmembers(perturbations): if inspect.isclass(val_2) and val_2.__module__ == perturbations.__name__: pert = val_2() - print(f"{pert = }") + print(f"{pert =}") assert isinstance(pert, Perturbation) if isinstance(pert, perturbations.Noise): @@ -1301,14 +1359,20 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): if "Slab" in key or "Pinch" in key: plt.contourf(x[:, 0, :], z[:, 0, :], n_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], z[:, Nel[1] // 2, :], n_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + z[:, Nel[1] // 2, :], + n_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("z") else: plt.contourf(x[:, 0, :], y[:, 0, :], n_cart[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], y[:, Nel[1] // 2, :], n_cart[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + y[:, Nel[1] // 2, :], + n_cart[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("y") @@ -1337,14 +1401,20 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): if "Slab" in key or "Pinch" in key: plt.contourf(x[:, 0, :], z[:, 0, :], u[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], z[:, Nel[1] // 2, :], u[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + z[:, Nel[1] // 2, :], + u[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("z") else: plt.contourf(x[:, 0, :], y[:, 0, :], u[:, 0, :], levels=levels) plt.contourf( - x[:, Nel[1] // 2, :], y[:, Nel[1] // 2, :], u[:, Nel[1] // 2, :], levels=levels + x[:, Nel[1] // 2, :], + y[:, Nel[1] // 2, :], + u[:, Nel[1] // 2, :], + levels=levels, ) plt.xlabel("x") plt.ylabel("y") diff --git a/src/struphy/linear_algebra/linalg_kron.py b/src/struphy/linear_algebra/linalg_kron.py index 2e0dd57dc..fedd05979 100644 --- a/src/struphy/linear_algebra/linalg_kron.py +++ b/src/struphy/linear_algebra/linalg_kron.py @@ -82,8 +82,9 @@ def kron_matvec_3d(kmat, vec3d): ( kmat[2].dot( ((kmat[1].dot(((kmat[0].dot(vec3d.reshape(v0, v1 * v2))).T).reshape(v1, v2 * k0))).T).reshape( - v2, k0 * k1 - ) + v2, + k0 * k1, + ), ) ).T ).reshape(k0, k1, k2) @@ -278,8 +279,9 @@ def kron_lusolve_3d(kmatlu, rhs): ( kmatlu[2].solve( ((kmatlu[1].solve(((kmatlu[0].solve(rhs.reshape(r0, r1 * r2))).T).reshape(r1, r2 * r0))).T).reshape( - r2, r0 * r1 - ) + r2, + r0 * r1, + ), ) ).T ).reshape(r0, r1, r2) @@ -320,7 +322,7 @@ def kron_solve_3d(kmat, rhs): splu(kmat[2]).solve( ( (splu(kmat[1]).solve(((splu(kmat[0]).solve(rhs.reshape(r0, r1 * r2))).T).reshape(r1, r2 * r0))).T - ).reshape(r2, r0 * r1) + ).reshape(r2, r0 * r1), ) ).T ).reshape(r0, r1, r2) @@ -361,7 +363,8 @@ def kron_fftsolve_3d(cvec, rhs): ( ( solve_circulant( - cvec[1], ((solve_circulant(cvec[0], rhs.reshape(r0, r1 * r2))).T).reshape(r1, r2 * r0) + cvec[1], + ((solve_circulant(cvec[0], rhs.reshape(r0, r1 * r2))).T).reshape(r1, r2 * r0), ) ).T ).reshape(r2, r0 * r1), diff --git a/src/struphy/linear_algebra/saddle_point.py b/src/struphy/linear_algebra/saddle_point.py index 47dd36190..3a191cde4 100644 --- a/src/struphy/linear_algebra/saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point.py @@ -413,7 +413,10 @@ def _setup_inverses(self): # === Inverse for A[1] if hasattr(self, "_Aenpinv") and self._is_inverse_still_valid( - self._Aenpinv, A1, "A[1]", pre=self._A22npinv + self._Aenpinv, + A1, + "A[1]", + pre=self._A22npinv, ): pass else: @@ -497,7 +500,7 @@ def _spectral_analysis(self): # print(f'{minbeforeA11_abs = }') # print(f'{minbeforeA11 = }') # print(f'{specA11_bef = }') - print(f"{specA11_bef_abs = }") + print(f"{specA11_bef_abs =}") # A22 before if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): @@ -517,8 +520,8 @@ def _spectral_analysis(self): # print(f'{minbeforeA22_abs = }') # print(f'{minbeforeA22 = }') # print(f'{specA22_bef = }') - print(f"{specA22_bef_abs = }") - print(f"{condA22_before = }") + print(f"{specA22_bef_abs =}") + print(f"{condA22_before =}") if self._preconditioner == True: # A11 after preconditioning with its inverse @@ -537,7 +540,7 @@ def _spectral_analysis(self): # print(f'{minafterA11_abs_prec = }') # print(f'{minafterA11_prec = }') # print(f'{specA11_aft_prec = }') - print(f"{specA11_aft_abs_prec = }") + print(f"{specA11_aft_abs_prec =}") # A22 after preconditioning with its inverse if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): @@ -557,7 +560,7 @@ def _spectral_analysis(self): # print(f'{minafterA22_abs_prec = }') # print(f'{minafterA22_prec = }') # print(f'{specA22_aft_prec = }') - print(f"{specA22_aft_abs_prec = }") + print(f"{specA22_aft_abs_prec =}") return condA22_before, specA22_bef_abs, condA11_before, condA22_after, specA22_aft_abs_prec diff --git a/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py b/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py index 06482f958..3aa3f4ab0 100644 --- a/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py +++ b/src/struphy/linear_algebra/tests/test_saddle_point_propagator.py @@ -306,7 +306,7 @@ def test_propagator2D(Nel, p, spl_kind, dirichlet_bc, mapping, epsilon, dt): "ManufacturedSolutionPotential": { "given_in_basis": "physical", "dimension": "2D", - } + }, } uvec.initialize_coeffs(domain=domain, pert_params=pp_u) diff --git a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py index ed5ff4d8a..42d3ae8d3 100644 --- a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py +++ b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py @@ -323,7 +323,7 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m compare_arrays(y1_rdm, y_uzawa, mpi_rank, atol=1e-5) compare_arrays(x1, x_uzawa[0], mpi_rank, atol=1e-5) compare_arrays(x2, x_uzawa[1], mpi_rank, atol=1e-5) - print(f"{info = }") + print(f"{info =}") elif isinstance(x_uzawa[0], BlockVector): # Output as Blockvector Rx1 = x1 - x_uzawa[0] diff --git a/src/struphy/main.py b/src/struphy/main.py index 524e19eb7..4b7b65645 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -338,7 +338,8 @@ def run( t1 = time.time() if rank == 0 and verbose: message = "Particles sorted | wall clock [s]: {0:8.4f} | sorting duration [s]: {1:8.4f}".format( - run_time_now * 60, t1 - t0 + run_time_now * 60, + t1 - t0, ) print(message, end="\n") print() @@ -383,10 +384,12 @@ def run( message = "time step: " + step + "/" + str(total_steps) message += " | " + "time: {0:10.5f}/{1:10.5f}".format(time_state["value"][0], Tend) message += " | " + "phys. time [s]: {0:12.10f}/{1:12.10f}".format( - time_state["value_sec"][0], Tend * model.units.t + time_state["value_sec"][0], + Tend * model.units.t, ) message += " | " + "wall clock [s]: {0:8.4f} | last step duration [s]: {1:8.4f}".format( - run_time_now * 60, t1 - t0 + run_time_now * 60, + t1 - t0, ) print(message, end="\n") @@ -513,7 +516,10 @@ def pproc( if physical: point_data_phy, grids_log, grids_phy = eval_femfields( - params_in, fields, celldivide=[celldivide] * 3, physical=True + params_in, + fields, + celldivide=[celldivide] * 3, + physical=True, ) # directory for field data @@ -699,7 +705,7 @@ def load_data(path: str) -> SimData: path_pproc = os.path.join(path, "post_processing") assert os.path.exists(path_pproc), f"Path {path_pproc} does not exist, run 'pproc' first?" print("\n*** Loading post-processed simulation data:") - print(f"{path = }") + print(f"{path =}") simdata = SimData(path) @@ -738,7 +744,7 @@ def load_data(path: str) -> SimData: if os.path.exists(path_kinetic): # species folders species = next(os.walk(path_kinetic))[1] - print(f"{species = }") + print(f"{species =}") for spec in species: path_spec = os.path.join(path_kinetic, spec) wlk = os.walk(path_spec) @@ -792,20 +798,20 @@ def load_data(path: str) -> SimData: simdata._n_sph[spec][sli][name] = tmp else: - print(f"{folder = }") + print(f"{folder =}") raise NotImplementedError print("\nThe following data has been loaded:") print(f"\ngrids:") - print(f"{simdata.t_grid.shape = }") + print(f"{simdata.t_grid.shape =}") if simdata.grids_log is not None: - print(f"{simdata.grids_log[0].shape = }") - print(f"{simdata.grids_log[1].shape = }") - print(f"{simdata.grids_log[2].shape = }") + print(f"{simdata.grids_log[0].shape =}") + print(f"{simdata.grids_log[1].shape =}") + print(f"{simdata.grids_log[2].shape =}") if simdata.grids_phy is not None: - print(f"{simdata.grids_phy[0].shape = }") - print(f"{simdata.grids_phy[1].shape = }") - print(f"{simdata.grids_phy[2].shape = }") + print(f"{simdata.grids_phy[0].shape =}") + print(f"{simdata.grids_phy[1].shape =}") + print(f"{simdata.grids_phy[2].shape =}") print(f"\nsimdata.spline_values:") for k, v in simdata.spline_values.items(): print(f" {k}") diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index d18942154..e6eb5b665 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -174,7 +174,7 @@ def allocate_feec(self, grid: TensorProductGrid, derham_opts: DerhamOptions): if grid is None or derham_opts is None: if MPI.COMM_WORLD.Get_rank() == 0: - print(f"\n{grid = }, {derham_opts = }: no Derham object set up.") + print(f"\n{grid =}, {derham_opts =}: no Derham object set up.") self._derham = None else: self._derham = setup_derham( @@ -472,7 +472,7 @@ def add_scalar(self, name: str, variable: PICVariable | SPHVariable = None, comp assert isinstance(name, str), "name must be a string" if compute == "from_particles": - assert isinstance(variable, (PICVariable, SPHVariable)), f"Variable is needed when {compute = }" + assert isinstance(variable, (PICVariable, SPHVariable)), f"Variable is needed when {compute =}" if not hasattr(self, "_scalar_quantities"): self._scalar_quantities = {} @@ -1172,7 +1172,7 @@ def show_options(cls): print( 'Options are given under the keyword "options" for each species dict. \ -Available options stand in lists as dict values.\nThe first entry of a list denotes the default value.' +Available options stand in lists as dict values.\nThe first entry of a list denotes the default value.', ) tab = " " @@ -1408,7 +1408,7 @@ def generate_default_parameter_file( BoundaryParameters,\n\ BinningPlot,\n\ KernelDensityPlot,\n\ - )\n" + )\n", ) file.write("from struphy import main\n") @@ -1484,14 +1484,14 @@ def generate_default_parameter_file( grid=grid,\n\ derham_opts=derham_opts,\n\ verbose=verbose,\n\ - )" + )", ) file.close() print( f"\nDefault parameter file for '{self.__class__.__name__}' has been created in the cwd ({path}).\n\ -You can now launch a simulation with 'python params_{self.__class__.__name__}.py'" +You can now launch a simulation with 'python params_{self.__class__.__name__}.py'", ) return path diff --git a/src/struphy/models/fluid.py b/src/struphy/models/fluid.py index 2b832ba00..a4916e39d 100644 --- a/src/struphy/models/fluid.py +++ b/src/struphy/models/fluid.py @@ -157,7 +157,7 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "mag_sonic.Options" in line: new_file += [ - "model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n" + "model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n", ] else: new_file += [line] @@ -332,7 +332,7 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "hall.Options" in line: new_file += [ - "model.propagators.hall.options = model.propagators.hall.Options(epsilon_from=model.mhd)\n" + "model.propagators.hall.options = model.propagators.hall.Options(epsilon_from=model.mhd)\n", ] else: new_file += [line] @@ -652,25 +652,25 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full',\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full',\n", ] new_file += [ - " s=model.mhd.entropy)\n" + " s=model.mhd.entropy)\n", ] elif "variat_ent.Options" in line: new_file += [ - "model.propagators.variat_ent.options = model.propagators.variat_ent.Options(model='full',\n" + "model.propagators.variat_ent.options = model.propagators.variat_ent.Options(model='full',\n", ] new_file += [ - " rho=model.mhd.density)\n" + " rho=model.mhd.density)\n", ] elif "variat_viscous.Options" in line: new_file += [ - "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(rho=model.mhd.density)\n" + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(rho=model.mhd.density)\n", ] elif "variat_resist.Options" in line: new_file += [ - "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(rho=model.mhd.density)\n" + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(rho=model.mhd.density)\n", ] elif "entropy.add_background" in line: new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] @@ -842,21 +842,21 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full',\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full',\n", ] new_file += [ - " s=model.fluid.entropy)\n" + " s=model.fluid.entropy)\n", ] elif "variat_ent.Options" in line: new_file += [ - "model.propagators.variat_ent.options = model.propagators.variat_ent.Options(model='full',\n" + "model.propagators.variat_ent.options = model.propagators.variat_ent.Options(model='full',\n", ] new_file += [ - " rho=model.fluid.density)\n" + " rho=model.fluid.density)\n", ] elif "variat_viscous.Options" in line: new_file += [ - "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(rho=model.fluid.density)\n" + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(rho=model.fluid.density)\n", ] elif "entropy.add_background" in line: new_file += ["model.fluid.density.add_background(FieldsBackground())\n"] @@ -1042,18 +1042,18 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_pb.Options" in line: new_file += [ - "model.propagators.variat_pb.options = model.propagators.variat_pb.Options(div_u=model.diagnostics.div_u,\n" + "model.propagators.variat_pb.options = model.propagators.variat_pb.Options(div_u=model.diagnostics.div_u,\n", ] new_file += [ - " u2=model.diagnostics.u2)\n" + " u2=model.diagnostics.u2)\n", ] elif "variat_viscous.Options" in line: new_file += [ - "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(rho=model.mhd.density)\n" + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(rho=model.mhd.density)\n", ] elif "variat_resist.Options" in line: new_file += [ - "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(rho=model.mhd.density)\n" + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(rho=model.mhd.density)\n", ] elif "pressure.add_background" in line: new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] @@ -1255,40 +1255,40 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='linear')\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='linear')\n", ] elif "variat_pb.Options" in line: new_file += [ - "model.propagators.variat_pb.options = model.propagators.variat_pb.Options(model='linear',\n" + "model.propagators.variat_pb.options = model.propagators.variat_pb.Options(model='linear',\n", ] new_file += [ - " div_u=model.diagnostics.div_u,\n" + " div_u=model.diagnostics.div_u,\n", ] new_file += [ - " u2=model.diagnostics.u2,\n" + " u2=model.diagnostics.u2,\n", ] new_file += [ - " pt3=model.diagnostics.pt3,\n" + " pt3=model.diagnostics.pt3,\n", ] new_file += [ - " bt2=model.diagnostics.bt2)\n" + " bt2=model.diagnostics.bt2)\n", ] elif "variat_viscous.Options" in line: new_file += [ - "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='linear_p',\n" + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='linear_p',\n", ] new_file += [ - " rho=model.mhd.density)\n" + " rho=model.mhd.density)\n", ] elif "variat_resist.Options" in line: new_file += [ - "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='linear_p',\n" + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='linear_p',\n", ] new_file += [ - " rho=model.mhd.density,\n" + " rho=model.mhd.density,\n", ] new_file += [ - " pt3=model.diagnostics.pt3)\n" + " pt3=model.diagnostics.pt3)\n", ] elif "pressure.add_background" in line: new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] @@ -1493,31 +1493,31 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='deltaf')\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='deltaf')\n", ] elif "variat_pb.Options" in line: new_file += [ - "model.propagators.variat_pb.options = model.propagators.variat_pb.Options(model='deltaf',\n" + "model.propagators.variat_pb.options = model.propagators.variat_pb.Options(model='deltaf',\n", ] new_file += [ - " pt3=model.diagnostics.pt3,\n" + " pt3=model.diagnostics.pt3,\n", ] new_file += [ - " bt2=model.diagnostics.bt2)\n" + " bt2=model.diagnostics.bt2)\n", ] elif "variat_viscous.Options" in line: new_file += [ - "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='full_p',\n" + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='full_p',\n", ] new_file += [ - " rho=model.mhd.density)\n" + " rho=model.mhd.density)\n", ] elif "variat_resist.Options" in line: new_file += [ - "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='full_p',\n" + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='full_p',\n", ] new_file += [ - " rho=model.mhd.density)\n" + " rho=model.mhd.density)\n", ] elif "pressure.add_background" in line: new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] @@ -1705,25 +1705,25 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full_q')\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full_q')\n", ] elif "variat_qb.Options" in line: new_file += [ - "model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='full_q')\n" + "model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='full_q')\n", ] elif "variat_viscous.Options" in line: new_file += [ - "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='full_q',\n" + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='full_q',\n", ] new_file += [ - " rho=model.mhd.density)\n" + " rho=model.mhd.density)\n", ] elif "variat_resist.Options" in line: new_file += [ - "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='full_q',\n" + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='full_q',\n", ] new_file += [ - " rho=model.mhd.density)\n" + " rho=model.mhd.density)\n", ] elif "sqrt_p.add_background" in line: new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] @@ -1909,43 +1909,43 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='linear_q')\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='linear_q')\n", ] elif "variat_qb.Options" in line: new_file += [ - "model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='linear_q',\n" + "model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='linear_q',\n", ] new_file += [ - " div_u=model.diagnostics.div_u,\n" + " div_u=model.diagnostics.div_u,\n", ] new_file += [ - " u2=model.diagnostics.u2,\n" + " u2=model.diagnostics.u2,\n", ] new_file += [ - " qt3=model.diagnostics.qt3,\n" + " qt3=model.diagnostics.qt3,\n", ] new_file += [ - " bt2=model.diagnostics.bt2)\n" + " bt2=model.diagnostics.bt2)\n", ] elif "variat_viscous.Options" in line: new_file += [ - "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='linear_q',\n" + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='linear_q',\n", ] new_file += [ - " rho=model.mhd.density,\n" + " rho=model.mhd.density,\n", ] new_file += [ - " pt3=model.diagnostics.qt3)\n" + " pt3=model.diagnostics.qt3)\n", ] elif "variat_resist.Options" in line: new_file += [ - "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='linear_q',\n" + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='linear_q',\n", ] new_file += [ - " rho=model.mhd.density,\n" + " rho=model.mhd.density,\n", ] new_file += [ - " pt3=model.diagnostics.qt3)\n" + " pt3=model.diagnostics.qt3)\n", ] elif "sqrt_p.add_background" in line: new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] @@ -2134,43 +2134,43 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='deltaf_q')\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='deltaf_q')\n", ] elif "variat_qb.Options" in line: new_file += [ - "model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='deltaf_q',\n" + "model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='deltaf_q',\n", ] new_file += [ - " div_u=model.diagnostics.div_u,\n" + " div_u=model.diagnostics.div_u,\n", ] new_file += [ - " u2=model.diagnostics.u2,\n" + " u2=model.diagnostics.u2,\n", ] new_file += [ - " qt3=model.diagnostics.qt3,\n" + " qt3=model.diagnostics.qt3,\n", ] new_file += [ - " bt2=model.diagnostics.bt2)\n" + " bt2=model.diagnostics.bt2)\n", ] elif "variat_viscous.Options" in line: new_file += [ - "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='deltaf_q',\n" + "model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='deltaf_q',\n", ] new_file += [ - " rho=model.mhd.density,\n" + " rho=model.mhd.density,\n", ] new_file += [ - " pt3=model.diagnostics.qt3)\n" + " pt3=model.diagnostics.qt3)\n", ] elif "variat_resist.Options" in line: new_file += [ - "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='deltaf_q',\n" + "model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='deltaf_q',\n", ] new_file += [ - " rho=model.mhd.density,\n" + " rho=model.mhd.density,\n", ] new_file += [ - " pt3=model.diagnostics.qt3)\n" + " pt3=model.diagnostics.qt3)\n", ] elif "sqrt_p.add_background" in line: new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] @@ -2280,7 +2280,7 @@ def update_scalar_quantities(self): particles = self.euler_fluid.var.particles valid_markers = particles.markers_wo_holes_and_ghost en_kin = valid_markers[:, 6].dot( - valid_markers[:, 3] ** 2 + valid_markers[:, 4] ** 2 + valid_markers[:, 5] ** 2 + valid_markers[:, 3] ** 2 + valid_markers[:, 4] ** 2 + valid_markers[:, 5] ** 2, ) / (2.0 * particles.Np) self.update_scalar("en_kin", en_kin) @@ -2425,7 +2425,7 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "hw.Options" in line: new_file += [ - "model.propagators.hw.options = model.propagators.hw.Options(phi=model.em_fields.phi)\n" + "model.propagators.hw.options = model.propagators.hw.Options(phi=model.em_fields.phi)\n", ] elif "vorticity.add_background" in line: new_file += ["model.plasma.density.add_background(FieldsBackground())\n"] diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index b5e7c2547..7e158bfb2 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -211,18 +211,18 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "mag_sonic.Options" in line: new_file += [ - "model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n" + "model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n", ] elif "couple_dens.Options" in line: new_file += [ - "model.propagators.couple_dens.options = model.propagators.couple_dens.Options(energetic_ions=model.energetic_ions.var,\n" + "model.propagators.couple_dens.options = model.propagators.couple_dens.Options(energetic_ions=model.energetic_ions.var,\n", ] new_file += [ - " b_tilde=model.em_fields.b_field)\n" + " b_tilde=model.em_fields.b_field)\n", ] elif "couple_curr.Options" in line: new_file += [ - "model.propagators.couple_curr.options = model.propagators.couple_curr.Options(b_tilde=model.em_fields.b_field)\n" + "model.propagators.couple_curr.options = model.propagators.couple_curr.Options(b_tilde=model.em_fields.b_field)\n", ] elif "set_save_data" in line: new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] @@ -390,7 +390,7 @@ def __init__(self, params, comm, clone_config=None): self.equil.b2_1, self.equil.b2_2, self.equil.b2_3, - ] + ], ) self._p_eq = self.derham.P["3"](self.equil.p3) self._ones = self._p_eq.space.zeros() @@ -760,44 +760,44 @@ def generate_default_parameter_file(self, path=None, prompt=True): if "shearalfen_cc5d.Options" in line: new_file += [ """model.propagators.shearalfen_cc5d.options = model.propagators.shearalfen_cc5d.Options( - energetic_ions = model.energetic_ions.var,)\n""" + energetic_ions = model.energetic_ions.var,)\n""", ] elif "magnetosonic.Options" in line: new_file += [ """model.propagators.magnetosonic.options = model.propagators.magnetosonic.Options( - b_field=model.em_fields.b_field,)\n""" + b_field=model.em_fields.b_field,)\n""", ] elif "cc5d_density.Options" in line: new_file += [ """model.propagators.cc5d_density.options = model.propagators.cc5d_density.Options( energetic_ions = model.energetic_ions.var, - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] elif "cc5d_curlb.Options" in line: new_file += [ """model.propagators.cc5d_curlb.options = model.propagators.cc5d_curlb.Options( - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] elif "cc5d_gradb.Options" in line: new_file += [ """model.propagators.cc5d_gradb.options = model.propagators.cc5d_gradb.Options( - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] elif "push_bxe.Options" in line: new_file += [ """model.propagators.push_bxe.options = model.propagators.push_bxe.Options( - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] elif "push_parallel.Options" in line: new_file += [ """model.propagators.push_parallel.options = model.propagators.push_parallel.Options( - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] else: diff --git a/src/struphy/models/kinetic.py b/src/struphy/models/kinetic.py index 9a01e38fd..3f327e9b2 100644 --- a/src/struphy/models/kinetic.py +++ b/src/struphy/models/kinetic.py @@ -477,7 +477,7 @@ def generate_default_parameter_file(self, path=None, prompt=True): new_file += ["model.initial_poisson.options = model.initial_poisson.Options()\n"] elif "push_vxb.Options" in line: new_file += [ - "model.propagators.push_vxb.options = model.propagators.push_vxb.Options(b2_var=model.em_fields.b_field)\n" + "model.propagators.push_vxb.options = model.propagators.push_vxb.Options(b2_var=model.em_fields.b_field)\n", ] elif "set_save_data" in line: new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] @@ -1079,11 +1079,11 @@ def generate_default_parameter_file(self, path=None, prompt=True): new_file += ["base_units = BaseUnits(kBT=1.0)\n"] elif "push_gc_bxe.Options" in line: new_file += [ - "model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options(phi=model.em_fields.phi)\n" + "model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options(phi=model.em_fields.phi)\n", ] elif "push_gc_para.Options" in line: new_file += [ - "model.propagators.push_gc_para.options = model.propagators.push_gc_para.Options(phi=model.em_fields.phi)\n" + "model.propagators.push_gc_para.options = model.propagators.push_gc_para.Options(phi=model.em_fields.phi)\n", ] elif "set_save_data" in line: new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 56a338ca9..8a9f11008 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -91,21 +91,21 @@ def __init__( else: self.alpha = alpha if MPI.COMM_WORLD.Get_rank() == 0: - warnings.warn(f"Override equation parameter {self.alpha = }") + warnings.warn(f"Override equation parameter {self.alpha =}") if epsilon is None: self.epsilon = 1.0 / (om_c * units.t) else: self.epsilon = epsilon if MPI.COMM_WORLD.Get_rank() == 0: - warnings.warn(f"Override equation parameter {self.epsilon = }") + warnings.warn(f"Override equation parameter {self.epsilon =}") if kappa is None: self.kappa = om_p * units.t else: self.kappa = kappa if MPI.COMM_WORLD.Get_rank() == 0: - warnings.warn(f"Override equation parameter {self.kappa = }") + warnings.warn(f"Override equation parameter {self.kappa =}") if verbose and MPI.COMM_WORLD.Get_rank() == 0: print(f"\nSet normalization parameters for species {species.__class__.__name__}:") diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index aa2af4013..e473fd9de 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -19,21 +19,21 @@ if inspect.isclass(obj) and "models.toy" in obj.__module__: toy_models += [name] if rank == 0: - print(f"\n{toy_models = }") + print(f"\n{toy_models =}") fluid_models = [] for name, obj in inspect.getmembers(fluid): if inspect.isclass(obj) and "models.fluid" in obj.__module__: fluid_models += [name] if rank == 0: - print(f"\n{fluid_models = }") + print(f"\n{fluid_models =}") kinetic_models = [] for name, obj in inspect.getmembers(kinetic): if inspect.isclass(obj) and "models.kinetic" in obj.__module__: kinetic_models += [name] if rank == 0: - print(f"\n{kinetic_models = }") + print(f"\n{kinetic_models =}") hybrid_models = [ "LinearMHDDriftkineticCC", @@ -44,7 +44,7 @@ # if inspect.isclass(obj) and "models.hybrid" in obj.__module__: # hybrid_models += [name] if rank == 0: - print(f"\n{hybrid_models = }") + print(f"\n{hybrid_models =}") # folder for test simulations @@ -58,7 +58,7 @@ def call_test(model_name: str, module: ModuleType = None, verbose=True): # exceptions if model_name == "TwoFluidQuasiNeutralToy" and MPI.COMM_WORLD.Get_size() > 1: - print(f"WARNING: Model {model_name} cannot be tested for {MPI.COMM_WORLD.Get_size() = }") + print(f"WARNING: Model {model_name} cannot be tested for {MPI.COMM_WORLD.Get_size() =}") return if module is None: diff --git a/src/struphy/models/tests/test_verif_EulerSPH.py b/src/struphy/models/tests/test_verif_EulerSPH.py index 79a248ac9..48eb8a7a8 100644 --- a/src/struphy/models/tests/test_verif_EulerSPH.py +++ b/src/struphy/models/tests/test_verif_EulerSPH.py @@ -133,7 +133,7 @@ def test_soundwave_1d(nx: int, plot_pts: int, do_plot: bool = False): plot_ct = 0 for i in range(0, Nt + 1): if i % interval == 0: - print(f"{i = }") + print(f"{i =}") plot_ct += 1 ax = plt.gca() @@ -150,14 +150,14 @@ def test_soundwave_1d(nx: int, plot_pts: int, do_plot: bool = False): plt.xlabel("x") plt.ylabel(r"$\rho$") - plt.title(f"standing sound wave ($c_s = 1$) for {nx = } and {ppb = }") + plt.title(f"standing sound wave ($c_s = 1$) for {nx =} and {ppb =}") if plot_ct == 11: break plt.show() error = xp.max(xp.abs(n_sph[0] - n_sph[-1])) - print(f"SPH sound wave {error = }.") + print(f"SPH sound wave {error =}.") assert error < 6e-4 print("Assertion passed.") diff --git a/src/struphy/models/tests/test_verif_LinearMHD.py b/src/struphy/models/tests/test_verif_LinearMHD.py index 5cbbbc9fd..475b11aef 100644 --- a/src/struphy/models/tests/test_verif_LinearMHD.py +++ b/src/struphy/models/tests/test_verif_LinearMHD.py @@ -115,7 +115,7 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): # assert vA = xp.sqrt(Bsquare / n0) v_alfven = vA * B0z / xp.sqrt(Bsquare) - print(f"{v_alfven = }") + print(f"{v_alfven =}") assert xp.abs(coeffs[0][0] - v_alfven) < 0.07 # second fft @@ -144,8 +144,8 @@ def test_slab_waves_1d(algo: str, do_plot: bool = False): delta = (4 * B0z**2 * cS**2 * vA**2) / ((cS**2 + vA**2) ** 2 * Bsquare) v_slow = xp.sqrt(1 / 2 * (cS**2 + vA**2) * (1 - xp.sqrt(1 - delta))) v_fast = xp.sqrt(1 / 2 * (cS**2 + vA**2) * (1 + xp.sqrt(1 - delta))) - print(f"{v_slow = }") - print(f"{v_fast = }") + print(f"{v_slow =}") + print(f"{v_fast =}") assert xp.abs(coeffs[0][0] - v_slow) < 0.05 assert xp.abs(coeffs[1][0] - v_fast) < 0.19 diff --git a/src/struphy/models/tests/test_verif_Maxwell.py b/src/struphy/models/tests/test_verif_Maxwell.py index e97675e7b..ccea67c18 100644 --- a/src/struphy/models/tests/test_verif_Maxwell.py +++ b/src/struphy/models/tests/test_verif_Maxwell.py @@ -203,7 +203,7 @@ def E_theta(X, Y, Z, m, t): r = (X**2 + Y**2) ** 0.5 theta = xp.arctan2(Y, X) return ((m / r * jv(m, r) - jv(m + 1, r)) - 0.28 * (m / r * yn(m, r) - yn(m + 1, r))) * xp.sin( - m * theta - t + m * theta - t, ) def to_E_r(X, Y, E_x, E_y): @@ -222,7 +222,13 @@ def to_E_theta(X, Y, E_x, E_y): vmax = E_theta(X, Y, grids_phy[0], modes, 0).max() fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4)) plot_exac = ax1.contourf( - X, Y, E_theta(X, Y, grids_phy[0], modes, t_grid[-1]), cmap="plasma", levels=100, vmin=vmin, vmax=vmax + X, + Y, + E_theta(X, Y, grids_phy[0], modes, t_grid[-1]), + cmap="plasma", + levels=100, + vmin=vmin, + vmax=vmax, ) ax2.contourf( X, @@ -256,12 +262,12 @@ def to_E_theta(X, Y, E_x, E_y): rel_err_Bz = error_Bz / xp.max(xp.abs(Bz_exact)) print("") - assert rel_err_Bz < 0.0021, f"Assertion for magnetic field Maxwell failed: {rel_err_Bz = }" - print(f"Assertion for magnetic field Maxwell passed ({rel_err_Bz = }).") - assert rel_err_Etheta < 0.0021, f"Assertion for electric (E_theta) field Maxwell failed: {rel_err_Etheta = }" - print(f"Assertion for electric field Maxwell passed ({rel_err_Etheta = }).") - assert rel_err_Er < 0.0021, f"Assertion for electric (E_r) field Maxwell failed: {rel_err_Er = }" - print(f"Assertion for electric field Maxwell passed ({rel_err_Er = }).") + assert rel_err_Bz < 0.0021, f"Assertion for magnetic field Maxwell failed: {rel_err_Bz =}" + print(f"Assertion for magnetic field Maxwell passed ({rel_err_Bz =}).") + assert rel_err_Etheta < 0.0021, f"Assertion for electric (E_theta) field Maxwell failed: {rel_err_Etheta =}" + print(f"Assertion for electric field Maxwell passed ({rel_err_Etheta =}).") + assert rel_err_Er < 0.0021, f"Assertion for electric (E_r) field Maxwell failed: {rel_err_Er =}" + print(f"Assertion for electric field Maxwell passed ({rel_err_Er =}).") if __name__ == "__main__": diff --git a/src/struphy/models/tests/test_verif_Poisson.py b/src/struphy/models/tests/test_verif_Poisson.py index 5b62d61ab..e82ea22c7 100644 --- a/src/struphy/models/tests/test_verif_Poisson.py +++ b/src/struphy/models/tests/test_verif_Poisson.py @@ -124,14 +124,14 @@ def test_poisson_1d(do_plot=False): plt.subplot(5, 2, 2 * c + 1) plt.plot(x, phi_h, label="phi") plt.plot(x, phi_e, "r--", label="exact") - plt.title(f"phi at {t = }") + plt.title(f"phi at {t =}") plt.ylim(-amp / (l * 2 * xp.pi / Lx) ** 2, amp / (l * 2 * xp.pi / Lx) ** 2) plt.legend() plt.subplot(5, 2, 2 * c + 2) plt.plot(x, source[t][0][:, 0, 0], label="rhs") plt.plot(x, rhs_exact(x, 0, 0, t), "r--", label="exact") - plt.title(f"source at {t = }") + plt.title(f"source at {t =}") plt.ylim(-amp, amp) plt.legend() @@ -140,7 +140,7 @@ def test_poisson_1d(do_plot=False): break plt.show() - print(f"{err = }") + print(f"{err =}") assert err < 0.0057 diff --git a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py index 585a9776a..a2625ba17 100644 --- a/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/test_verif_VlasovAmpereOneSpecies.py @@ -159,8 +159,8 @@ def E_exact(t): # assert rel_error = xp.abs(gamma_num - gamma) / xp.abs(gamma) - assert rel_error < 0.22, f"Assertion for weak Landau damping failed: {gamma_num = } vs. {gamma = }." - print(f"Assertion for weak Landau damping passed ({rel_error = }).") + assert rel_error < 0.22, f"Assertion for weak Landau damping failed: {gamma_num =} vs. {gamma =}." + print(f"Assertion for weak Landau damping passed ({rel_error =}).") if __name__ == "__main__": diff --git a/src/struphy/models/tests/test_xxpproc.py b/src/struphy/models/tests/test_xxpproc.py index acb8a6590..3d4fef2f0 100644 --- a/src/struphy/models/tests/test_xxpproc.py +++ b/src/struphy/models/tests/test_xxpproc.py @@ -49,7 +49,7 @@ def test_pproc_codes(model: str = None, group: str = None): elif group == "toy": list_models = list_toy else: - raise ValueError(f"{group = } is not a valid group specification.") + raise ValueError(f"{group =} is not a valid group specification.") if comm.Get_rank() == 0: if model is None: diff --git a/src/struphy/models/tests/verification.py b/src/struphy/models/tests/verification.py index 75b0a1047..d28fc3d6e 100644 --- a/src/struphy/models/tests/verification.py +++ b/src/struphy/models/tests/verification.py @@ -85,8 +85,8 @@ def E_exact(t): # assert rel_error = xp.abs(gamma_num - gamma) / xp.abs(gamma) - assert rel_error < 0.25, f"{rank = }: Assertion for weak Landau damping failed: {gamma_num = } vs. {gamma = }." - print(f"{rank = }: Assertion for weak Landau damping passed ({rel_error = }).") + assert rel_error < 0.25, f"{rank =}: Assertion for weak Landau damping failed: {gamma_num =} vs. {gamma =}." + print(f"{rank =}: Assertion for weak Landau damping passed ({rel_error =}).") def LinearVlasovAmpereOneSpecies_weakLandau( @@ -161,8 +161,8 @@ def E_exact(t): # assert rel_error = xp.abs(gamma_num - gamma) / xp.abs(gamma) - assert rel_error < 0.25, f"{rank = }: Assertion for weak Landau damping failed: {gamma_num = } vs. {gamma = }." - print(f"{rank = }: Assertion for weak Landau damping passed ({rel_error = }).") + assert rel_error < 0.25, f"{rank =}: Assertion for weak Landau damping failed: {gamma_num =} vs. {gamma =}." + print(f"{rank =}: Assertion for weak Landau damping passed ({rel_error =}).") def IsothermalEulerSPH_soundwave( @@ -207,7 +207,7 @@ def IsothermalEulerSPH_soundwave( plot_ct = 0 for i in range(0, Nt + 1): if i % interval == 0: - print(f"{i = }") + print(f"{i =}") plot_ct += 1 ax = plt.gca() @@ -224,7 +224,7 @@ def IsothermalEulerSPH_soundwave( plt.xlabel("x") plt.ylabel(r"$\rho$") - plt.title(f"standing sound wave ($c_s = 1$) for {nx = } and {ppb = }") + plt.title(f"standing sound wave ($c_s = 1$) for {nx =} and {ppb =}") if plot_ct == 11: break @@ -232,7 +232,7 @@ def IsothermalEulerSPH_soundwave( # assert error = xp.max(xp.abs(n_sph[0] - n_sph[-1])) - print(f"{rank = }: Assertion for SPH sound wave passed ({error = }).") + print(f"{rank =}: Assertion for SPH sound wave passed ({error =}).") assert error < 1.3e-3 @@ -311,7 +311,13 @@ def to_E_theta(X, Y, E_x, E_y): vmax = E_theta(X, Y, grids_phy[0], modes, 0).max() fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4)) plot_exac = ax1.contourf( - X, Y, E_theta(X, Y, grids_phy[0], modes, t_grid[-1]), cmap="plasma", levels=100, vmin=vmin, vmax=vmax + X, + Y, + E_theta(X, Y, grids_phy[0], modes, t_grid[-1]), + cmap="plasma", + levels=100, + vmin=vmin, + vmax=vmax, ) ax2.contourf( X, @@ -344,18 +350,18 @@ def to_E_theta(X, Y, E_x, E_y): rel_err_Etheta = error_Etheta / xp.max(xp.abs(Etheta_exact)) rel_err_Bz = error_Bz / xp.max(xp.abs(Bz_exact)) - print(f"{rel_err_Er = }") - print(f"{rel_err_Etheta = }") - print(f"{rel_err_Bz = }") + print(f"{rel_err_Er =}") + print(f"{rel_err_Etheta =}") + print(f"{rel_err_Bz =}") - assert rel_err_Bz < 0.0021, f"{rank = }: Assertion for magnetic field Maxwell failed: {rel_err_Bz = }" - print(f"{rank = }: Assertion for magnetic field Maxwell passed ({rel_err_Bz = }).") + assert rel_err_Bz < 0.0021, f"{rank =}: Assertion for magnetic field Maxwell failed: {rel_err_Bz =}" + print(f"{rank =}: Assertion for magnetic field Maxwell passed ({rel_err_Bz =}).") assert rel_err_Etheta < 0.0021, ( - f"{rank = }: Assertion for electric (E_theta) field Maxwell failed: {rel_err_Etheta = }" + f"{rank =}: Assertion for electric (E_theta) field Maxwell failed: {rel_err_Etheta =}" ) - print(f"{rank = }: Assertion for electric field Maxwell passed ({rel_err_Etheta = }).") - assert rel_err_Er < 0.0021, f"{rank = }: Assertion for electric (E_r) field Maxwell failed: {rel_err_Er = }" - print(f"{rank = }: Assertion for electric field Maxwell passed ({rel_err_Er = }).") + print(f"{rank =}: Assertion for electric field Maxwell passed ({rel_err_Etheta =}).") + assert rel_err_Er < 0.0021, f"{rank =}: Assertion for electric (E_r) field Maxwell failed: {rel_err_Er =}" + print(f"{rank =}: Assertion for electric field Maxwell passed ({rel_err_Er =}).") if __name__ == "__main__": diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index bad2b7916..fd36b5d5f 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -81,10 +81,12 @@ def allocate_helpers(self): def update_scalar_quantities(self): en_E = 0.5 * self.mass_ops.M1.dot_inner( - self.em_fields.e_field.spline.vector, self.em_fields.e_field.spline.vector + self.em_fields.e_field.spline.vector, + self.em_fields.e_field.spline.vector, ) en_B = 0.5 * self.mass_ops.M2.dot_inner( - self.em_fields.b_field.spline.vector, self.em_fields.b_field.spline.vector + self.em_fields.b_field.spline.vector, + self.em_fields.b_field.spline.vector, ) self.update_scalar("electric energy", en_E) @@ -336,7 +338,7 @@ def allocate_helpers(self): self.equil.b2_1, self.equil.b2_2, self.equil.b2_3, - ] + ], ) # temporary vectors for scalar quantities @@ -371,7 +373,8 @@ def update_scalar_quantities(self): # perturbed fields en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.mhd.velocity.spline.vector, self.mhd.velocity.spline.vector) en_B = 0.5 * self.mass_ops.M2.dot_inner( - self.em_fields.b_field.spline.vector, self.em_fields.b_field.spline.vector + self.em_fields.b_field.spline.vector, + self.em_fields.b_field.spline.vector, ) self.update_scalar("en_U", en_U) @@ -477,7 +480,7 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='pressureless')\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='pressureless')\n", ] else: new_file += [line] @@ -583,7 +586,7 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='barotropic')\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='barotropic')\n", ] else: new_file += [line] @@ -703,17 +706,17 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "variat_dens.Options" in line: new_file += [ - "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full',\n" + "model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='full',\n", ] new_file += [ - " s=model.fluid.entropy)\n" + " s=model.fluid.entropy)\n", ] elif "variat_ent.Options" in line: new_file += [ - "model.propagators.variat_ent.options = model.propagators.variat_ent.Options(model='full',\n" + "model.propagators.variat_ent.options = model.propagators.variat_ent.Options(model='full',\n", ] new_file += [ - " rho=model.fluid.density)\n" + " rho=model.fluid.density)\n", ] elif "entropy.add_background" in line: new_file += ["model.fluid.density.add_background(FieldsBackground())\n"] @@ -859,7 +862,7 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "poisson.Options" in line: new_file += [ - "model.propagators.poisson.options = model.propagators.poisson.Options(rho=model.em_fields.source)\n" + "model.propagators.poisson.options = model.propagators.poisson.Options(rho=model.em_fields.source)\n", ] else: new_file += [line] diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index d71f4644d..e47ab1cd0 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -82,7 +82,7 @@ def add_background(self, background, verbose=True): if verbose and MPI.COMM_WORLD.Get_rank() == 0: print( - f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added background '{background.__class__.__name__}' with:" + f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added background '{background.__class__.__name__}' with:", ) for k, v in background.__dict__.items(): print(f" {k}: {v}") @@ -120,7 +120,7 @@ def add_perturbation(self, perturbation: Perturbation, verbose=True): if verbose and MPI.COMM_WORLD.Get_rank() == 0: print( - f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added perturbation '{perturbation.__class__.__name__}' with:" + f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added perturbation '{perturbation.__class__.__name__}' with:", ) for k, v in perturbation.__dict__.items(): print(f" {k}: {v}") @@ -175,7 +175,7 @@ def add_initial_condition(self, init: KineticBackground, verbose=True): self._initial_condition = init if verbose and MPI.COMM_WORLD.Get_rank() == 0: print( - f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added initial condition '{init.__class__.__name__}' with:" + f"\nVariable '{self.__name__}' of species '{self.species.__class__.__name__}' - added initial condition '{init.__class__.__name__}' with:", ) for k, v in init.__dict__.items(): print(f" {k}: {v}") diff --git a/src/struphy/ode/tests/test_ode_feec.py b/src/struphy/ode/tests/test_ode_feec.py index dadd3aad3..c0ef51b08 100644 --- a/src/struphy/ode/tests/test_ode_feec.py +++ b/src/struphy/ode/tests/test_ode_feec.py @@ -100,9 +100,9 @@ def f(t, y1, y2, y3, out=out): vector_field[var] = f - print(f"{vector_field = }") + print(f"{vector_field =}") butcher = ButcherTableau(algo=algo) - print(f"{butcher = }") + print(f"{butcher =}") solver = ODEsolverFEEC(vector_field, butcher=butcher) @@ -118,7 +118,7 @@ def f(t, y1, y2, y3, out=out): for i, h in enumerate(hs): errors[h] = {} time = xp.linspace(0, Tend, int(Tend / h) + 1) - print(f"{h = }, {time.size = }") + print(f"{h =}, {time.size =}") yvec = y_exact(time) ymax = {} for var in vector_field: @@ -139,7 +139,7 @@ def f(t, y1, y2, y3, out=out): # checks for var in vector_field: errors[h][var] = h * xp.sum(xp.abs(yvec - ymax[var])) / (h * xp.sum(xp.abs(yvec))) - print(f"{errors[h][var] = }") + print(f"{errors[h][var] =}") assert errors[h][var] < 0.31 if rank == 0: @@ -162,9 +162,9 @@ def f(t, y1, y2, y3, out=out): err_vec += [dct[var]] m, _ = xp.polyfit(xp.log(h_vec), xp.log(err_vec), deg=1) - print(f"{spaces[j]}-space, fitted convergence rate = {m} for {algo = } with {solver.butcher.conv_rate = }") + print(f"{spaces[j]}-space, fitted convergence rate = {m} for {algo =} with {solver.butcher.conv_rate =}") assert xp.abs(m - solver.butcher.conv_rate) < 0.1 - print(f"Convergence check passed on {rank = }.") + print(f"Convergence check passed on {rank =}.") if rank == 0: plt.loglog(h_vec, h_vec, "--", label=f"h") diff --git a/src/struphy/pic/accumulation/accum_kernels_gc.py b/src/struphy/pic/accumulation/accum_kernels_gc.py index c5e836f81..fecf6a255 100644 --- a/src/struphy/pic/accumulation/accum_kernels_gc.py +++ b/src/struphy/pic/accumulation/accum_kernels_gc.py @@ -240,7 +240,16 @@ def cc_lin_mhd_5d_D( # call the appropriate matvec filler particle_to_mat_kernels.mat_fill_v0vec_asym( - args_derham, span1, span2, span3, mat12, mat13, mat23, filling_m12, filling_m13, filling_m23 + args_derham, + span1, + span2, + span3, + mat12, + mat13, + mat23, + filling_m12, + filling_m13, + filling_m23, ) elif basis_u == 1: @@ -257,7 +266,16 @@ def cc_lin_mhd_5d_D( # call the appropriate matvec filler particle_to_mat_kernels.mat_fill_v1_asym( - args_derham, span1, span2, span3, mat12, mat13, mat23, filling_m12, filling_m13, filling_m23 + args_derham, + span1, + span2, + span3, + mat12, + mat13, + mat23, + filling_m12, + filling_m13, + filling_m23, ) elif basis_u == 2: @@ -268,7 +286,16 @@ def cc_lin_mhd_5d_D( # call the appropriate matvec filler particle_to_mat_kernels.mat_fill_v2_asym( - args_derham, span1, span2, span3, mat12, mat13, mat23, filling_m12, filling_m13, filling_m23 + args_derham, + span1, + span2, + span3, + mat12, + mat13, + mat23, + filling_m12, + filling_m13, + filling_m23, ) # -- removed omp: #$ omp end parallel @@ -579,7 +606,16 @@ def cc_lin_mhd_5d_M( filling_v[:] = weight * mu / det_df * scale_vec * norm_b1 particle_to_mat_kernels.vec_fill_v2( - args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + args_derham, + span1, + span2, + span3, + vec1, + vec2, + vec3, + filling_v[0], + filling_v[1], + filling_v[2], ) vec1 /= Np @@ -759,7 +795,16 @@ def cc_lin_mhd_5d_gradB( # call the appropriate matvec filler particle_to_mat_kernels.vec_fill_v0vec( - args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + args_derham, + span1, + span2, + span3, + vec1, + vec2, + vec3, + filling_v[0], + filling_v[1], + filling_v[2], ) elif basis_u == 2: @@ -770,7 +815,16 @@ def cc_lin_mhd_5d_gradB( # call the appropriate matvec filler particle_to_mat_kernels.vec_fill_v2( - args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + args_derham, + span1, + span2, + span3, + vec1, + vec2, + vec3, + filling_v[0], + filling_v[1], + filling_v[2], ) vec1 /= Np vec2 /= Np @@ -955,7 +1009,16 @@ def cc_lin_mhd_5d_gradB_dg_init( # call the appropriate matvec filler particle_to_mat_kernels.vec_fill_v0vec( - args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + args_derham, + span1, + span2, + span3, + vec1, + vec2, + vec3, + filling_v[0], + filling_v[1], + filling_v[2], ) elif basis_u == 2: @@ -981,7 +1044,16 @@ def cc_lin_mhd_5d_gradB_dg_init( # call the appropriate matvec filler particle_to_mat_kernels.vec_fill_v2( - args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + args_derham, + span1, + span2, + span3, + vec1, + vec2, + vec3, + filling_v[0], + filling_v[1], + filling_v[2], ) vec1 /= Np @@ -1180,7 +1252,16 @@ def cc_lin_mhd_5d_gradB_dg( # call the appropriate matvec filler particle_to_mat_kernels.vec_fill_v0vec( - args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + args_derham, + span1, + span2, + span3, + vec1, + vec2, + vec3, + filling_v[0], + filling_v[1], + filling_v[2], ) elif basis_u == 2: @@ -1218,7 +1299,16 @@ def cc_lin_mhd_5d_gradB_dg( # call the appropriate matvec filler particle_to_mat_kernels.vec_fill_v2( - args_derham, span1, span2, span3, vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2] + args_derham, + span1, + span2, + span3, + vec1, + vec2, + vec3, + filling_v[0], + filling_v[1], + filling_v[2], ) vec1 /= Np diff --git a/src/struphy/pic/accumulation/particle_to_mat_kernels.py b/src/struphy/pic/accumulation/particle_to_mat_kernels.py index 576d4571c..bc9364f6a 100644 --- a/src/struphy/pic/accumulation/particle_to_mat_kernels.py +++ b/src/struphy/pic/accumulation/particle_to_mat_kernels.py @@ -5834,7 +5834,12 @@ def m_v_fill_v2_full( def mat_fill_b_v0( - args_derham: "DerhamArguments", eta1: float, eta2: float, eta3: float, mat: "float[:,:,:,:,:,:]", fill: float + args_derham: "DerhamArguments", + eta1: float, + eta2: float, + eta3: float, + mat: "float[:,:,:,:,:,:]", + fill: float, ): """ Adds the contribution of one particle to the elements of an accumulation matrix V0 -> V0. The result is returned in mat. @@ -5969,7 +5974,12 @@ def m_v_fill_b_v0( def mat_fill_b_v3( - args_derham: "DerhamArguments", eta1: float, eta2: float, eta3: float, mat: "float[:,:,:,:,:,:]", fill: float + args_derham: "DerhamArguments", + eta1: float, + eta2: float, + eta3: float, + mat: "float[:,:,:,:,:,:]", + fill: float, ): """ Adds the contribution of one particle to the elements of an accumulation matrix V3 -> V3. The result is returned in mat. @@ -6112,7 +6122,12 @@ def m_v_fill_b_v3( def mat_fill_v0( - args_derham: "DerhamArguments", span1: int, span2: int, span3: int, mat: "float[:,:,:,:,:,:]", fill: float + args_derham: "DerhamArguments", + span1: int, + span2: int, + span3: int, + mat: "float[:,:,:,:,:,:]", + fill: float, ): """ Adds the contribution of one particle to the elements of an accumulation matrix V0 -> V0. The result is returned in mat. @@ -6239,7 +6254,12 @@ def m_v_fill_v0( def mat_fill_v3( - args_derham: "DerhamArguments", span1: int, span2: int, span3: int, mat: "float[:,:,:,:,:,:]", fill: float + args_derham: "DerhamArguments", + span1: int, + span2: int, + span3: int, + mat: "float[:,:,:,:,:,:]", + fill: float, ): """ Adds the contribution of one particle to the elements of an accumulation block matrix V3 -> V3. The result is returned in mat. @@ -12949,7 +12969,12 @@ def vec_fill_v0vec( def vec_fill_b_v0( - args_derham: "DerhamArguments", eta1: float, eta2: float, eta3: float, vec: "float[:,:,:]", fill: float + args_derham: "DerhamArguments", + eta1: float, + eta2: float, + eta3: float, + vec: "float[:,:,:]", + fill: float, ): """TODO""" @@ -13127,7 +13152,12 @@ def vec_fill_b_v2( def vec_fill_b_v3( - args_derham: "DerhamArguments", eta1: float, eta2: float, eta3: float, vec: "float[:,:,:]", fill: float + args_derham: "DerhamArguments", + eta1: float, + eta2: float, + eta3: float, + vec: "float[:,:,:]", + fill: float, ): """TODO""" diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 0f92323ae..84900418d 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -230,10 +230,10 @@ def __init__( n_boxes = self.mpi_size * self.num_clones else: assert all([nboxes >= nproc for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)]), ( - f"There must be at least one box {self.boxes_per_dim = } on each process {self.nprocs = } in each direction." + f"There must be at least one box {self.boxes_per_dim =} on each process {self.nprocs =} in each direction." ) assert all([nboxes % nproc == 0 for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)]), ( - f"Number of boxes {self.boxes_per_dim = } must be divisible by number of processes {self.nprocs = } in each direction." + f"Number of boxes {self.boxes_per_dim =} must be divisible by number of processes {self.nprocs =} in each direction." ) n_boxes = xp.prod(self.boxes_per_dim, dtype=int) * self.num_clones @@ -776,7 +776,7 @@ def velocities(self): @velocities.setter def velocities(self, new): assert isinstance(new, xp.ndarray) - assert new.shape == (self.n_mks_loc, self.vdim), f"{self.n_mks_loc = } and {self.vdim = } but {new.shape = }" + assert new.shape == (self.n_mks_loc, self.vdim), f"{self.n_mks_loc =} and {self.vdim =} but {new.shape =}" self._markers[self.valid_mks, self.index["vel"]] = new @property @@ -1126,7 +1126,7 @@ def _allocate_marker_array(self): # Have at least 3 spare places in markers array assert self.args_markers.first_free_idx + 2 < self.n_cols - 1, ( - f"{self.args_markers.first_free_idx + 2} is not smaller than {self.n_cols - 1 = }; not enough columns in marker array !!" + f"{self.args_markers.first_free_idx + 2} is not smaller than {self.n_cols - 1 =}; not enough columns in marker array !!" ) def _initialize_sorting_boxes(self): @@ -1565,7 +1565,7 @@ def draw_markers( num_loaded_particles_loc += num_valid # make sure all particles are loaded - assert self.Np == int(num_loaded_particles_glob), f"{self.Np = }, {int(num_loaded_particles_glob) = }" + assert self.Np == int(num_loaded_particles_glob), f"{self.Np =}, {int(num_loaded_particles_glob) =}" # set new n_mks_load self._gather_scalar_in_subcomm_array(num_loaded_particles_loc, out=self.n_mks_load) @@ -1843,7 +1843,7 @@ def initialize_weights( self.update_holes() self.reset_marker_ids() print( - f"\nWeights < {self.threshold} have been rejected, number of valid markers on process {self.mpi_rank} is {self.n_mks_loc}." + f"\nWeights < {self.threshold} have been rejected, number of valid markers on process {self.mpi_rank} is {self.n_mks_loc}.", ) # compute (time-dependent) weights at vdim + 3 @@ -2464,7 +2464,7 @@ def _set_boundary_boxes(self): self._bnd_boxes_x_p.append(flatten_index(self.nx, j, k, self.nx, self.ny, self.nz)) if self._verbose: - print(f"eta1 boundary on {self._rank = }:\n{self._bnd_boxes_x_m = }\n{self._bnd_boxes_x_p = }") + print(f"eta1 boundary on {self._rank =}:\n{self._bnd_boxes_x_m =}\n{self._bnd_boxes_x_p =}") # y boundary # negative direction @@ -2479,7 +2479,7 @@ def _set_boundary_boxes(self): self._bnd_boxes_y_p.append(flatten_index(i, self.ny, k, self.nx, self.ny, self.nz)) if self._verbose: - print(f"eta2 boundary on {self._rank = }:\n{self._bnd_boxes_y_m = }\n{self._bnd_boxes_y_p = }") + print(f"eta2 boundary on {self._rank =}:\n{self._bnd_boxes_y_m =}\n{self._bnd_boxes_y_p =}") # z boundary # negative direction @@ -2494,7 +2494,7 @@ def _set_boundary_boxes(self): self._bnd_boxes_z_p.append(flatten_index(i, j, self.nz, self.nx, self.ny, self.nz)) if self._verbose: - print(f"eta3 boundary on {self._rank = }:\n{self._bnd_boxes_z_m = }\n{self._bnd_boxes_z_p = }") + print(f"eta3 boundary on {self._rank =}:\n{self._bnd_boxes_z_m =}\n{self._bnd_boxes_z_p =}") # x-y edges self._bnd_boxes_x_m_y_m = [] @@ -2512,11 +2512,11 @@ def _set_boundary_boxes(self): if self._verbose: print( ( - f"eta1-eta2 edge on {self._rank = }:\n{self._bnd_boxes_x_m_y_m = }" - f"\n{self._bnd_boxes_x_m_y_p = }" - f"\n{self._bnd_boxes_x_p_y_m = }" - f"\n{self._bnd_boxes_x_p_y_p = }" - ) + f"eta1-eta2 edge on {self._rank =}:\n{self._bnd_boxes_x_m_y_m =}" + f"\n{self._bnd_boxes_x_m_y_p =}" + f"\n{self._bnd_boxes_x_p_y_m =}" + f"\n{self._bnd_boxes_x_p_y_p =}" + ), ) # x-z edges @@ -2535,11 +2535,11 @@ def _set_boundary_boxes(self): if self._verbose: print( ( - f"eta1-eta3 edge on {self._rank = }:\n{self._bnd_boxes_x_m_z_m = }" - f"\n{self._bnd_boxes_x_m_z_p = }" - f"\n{self._bnd_boxes_x_p_z_m = }" - f"\n{self._bnd_boxes_x_p_z_p = }" - ) + f"eta1-eta3 edge on {self._rank =}:\n{self._bnd_boxes_x_m_z_m =}" + f"\n{self._bnd_boxes_x_m_z_p =}" + f"\n{self._bnd_boxes_x_p_z_m =}" + f"\n{self._bnd_boxes_x_p_z_p =}" + ), ) # y-z edges @@ -2558,11 +2558,11 @@ def _set_boundary_boxes(self): if self._verbose: print( ( - f"eta2-eta3 edge on {self._rank = }:\n{self._bnd_boxes_y_m_z_m = }" - f"\n{self._bnd_boxes_y_m_z_p = }" - f"\n{self._bnd_boxes_y_p_z_m = }" - f"\n{self._bnd_boxes_y_p_z_p = }" - ) + f"eta2-eta3 edge on {self._rank =}:\n{self._bnd_boxes_y_m_z_m =}" + f"\n{self._bnd_boxes_y_m_z_p =}" + f"\n{self._bnd_boxes_y_p_z_m =}" + f"\n{self._bnd_boxes_y_p_z_p =}" + ), ) # corners @@ -2588,15 +2588,15 @@ def _set_boundary_boxes(self): if self._verbose: print( ( - f"corners on {self._rank = }:\n{self._bnd_boxes_x_m_y_m_z_m = }" - f"\n{self._bnd_boxes_x_m_y_m_z_p = }" - f"\n{self._bnd_boxes_x_m_y_p_z_m = }" - f"\n{self._bnd_boxes_x_p_y_m_z_m = }" - f"\n{self._bnd_boxes_x_m_y_p_z_p = }" - f"\n{self._bnd_boxes_x_p_y_m_z_p = }" - f"\n{self._bnd_boxes_x_p_y_p_z_m = }" - f"\n{self._bnd_boxes_x_p_y_p_z_p = }" - ) + f"corners on {self._rank =}:\n{self._bnd_boxes_x_m_y_m_z_m =}" + f"\n{self._bnd_boxes_x_m_y_m_z_p =}" + f"\n{self._bnd_boxes_x_m_y_p_z_m =}" + f"\n{self._bnd_boxes_x_p_y_m_z_m =}" + f"\n{self._bnd_boxes_x_m_y_p_z_p =}" + f"\n{self._bnd_boxes_x_p_y_m_z_p =}" + f"\n{self._bnd_boxes_x_p_y_p_z_m =}" + f"\n{self._bnd_boxes_x_p_y_p_z_p =}" + ), ) def _sort_boxed_particles_numpy(self): @@ -2650,7 +2650,7 @@ def check_and_assign_particles_to_boxes(self): f'Strong load imbalance detected in sorting boxes: \ max number of markers in a box ({max_in_box}) on rank {self.mpi_rank} \ exceeds the column-size of the box array ({self._sorting_boxes.boxes.shape[1]}). \ -Increasing the value of "box_bufsize" in the markers parameters for the next run.' +Increasing the value of "box_bufsize" in the markers parameters for the next run.', ) self.mpi_comm.Abort() @@ -2734,17 +2734,23 @@ def prepare_ghost_particles(self): # Mirror position for boundary condition if self.bc_sph[0] in ("mirror", "fixed"): self._mirror_particles( - "_markers_x_m", "_markers_x_p", is_domain_boundary=self.sorting_boxes.is_domain_boundary + "_markers_x_m", + "_markers_x_p", + is_domain_boundary=self.sorting_boxes.is_domain_boundary, ) if self.bc_sph[1] in ("mirror", "fixed"): self._mirror_particles( - "_markers_y_m", "_markers_y_p", is_domain_boundary=self.sorting_boxes.is_domain_boundary + "_markers_y_m", + "_markers_y_p", + is_domain_boundary=self.sorting_boxes.is_domain_boundary, ) if self.bc_sph[2] in ("mirror", "fixed"): self._mirror_particles( - "_markers_z_m", "_markers_z_p", is_domain_boundary=self.sorting_boxes.is_domain_boundary + "_markers_z_m", + "_markers_z_p", + is_domain_boundary=self.sorting_boxes.is_domain_boundary, ) ## Edges x-y @@ -2899,7 +2905,8 @@ def _mirror_particles(self, *marker_array_names, is_domain_boundary=None): arr[:, 0] *= -1.0 if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, flat_eval=True + *arr[:, :3].T, + flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = -boundary_values / self.s0( *arr[:, :3].T, @@ -2911,7 +2918,8 @@ def _mirror_particles(self, *marker_array_names, is_domain_boundary=None): arr[:, 0] = 2.0 - arr[:, 0] if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, flat_eval=True + *arr[:, :3].T, + flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = -boundary_values / self.s0( *arr[:, :3].T, @@ -2926,7 +2934,8 @@ def _mirror_particles(self, *marker_array_names, is_domain_boundary=None): arr[:, 1] *= -1.0 if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, flat_eval=True + *arr[:, :3].T, + flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = -boundary_values / self.s0( *arr[:, :3].T, @@ -2938,7 +2947,8 @@ def _mirror_particles(self, *marker_array_names, is_domain_boundary=None): arr[:, 1] = 2.0 - arr[:, 1] if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, flat_eval=True + *arr[:, :3].T, + flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = -boundary_values / self.s0( *arr[:, :3].T, @@ -2953,7 +2963,8 @@ def _mirror_particles(self, *marker_array_names, is_domain_boundary=None): arr[:, 2] *= -1.0 if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, flat_eval=True + *arr[:, :3].T, + flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = -boundary_values / self.s0( *arr[:, :3].T, @@ -2965,7 +2976,8 @@ def _mirror_particles(self, *marker_array_names, is_domain_boundary=None): arr[:, 2] = 2.0 - arr[:, 2] if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, flat_eval=True + *arr[:, :3].T, + flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! arr[:, self.index["weights"]] = -boundary_values / self.s0( *arr[:, :3].T, @@ -3018,124 +3030,124 @@ def get_destinations_box(self): # if self._x_m_y_m_proc is not None: self._send_info_box[self._x_m_y_m_proc] += len(self._markers_x_m_y_m) self._send_list_box[self._x_m_y_m_proc] = xp.concatenate( - (self._send_list_box[self._x_m_y_m_proc], self._markers_x_m_y_m) + (self._send_list_box[self._x_m_y_m_proc], self._markers_x_m_y_m), ) # if self._x_m_y_p_proc is not None: self._send_info_box[self._x_m_y_p_proc] += len(self._markers_x_m_y_p) self._send_list_box[self._x_m_y_p_proc] = xp.concatenate( - (self._send_list_box[self._x_m_y_p_proc], self._markers_x_m_y_p) + (self._send_list_box[self._x_m_y_p_proc], self._markers_x_m_y_p), ) # if self._x_p_y_m_proc is not None: self._send_info_box[self._x_p_y_m_proc] += len(self._markers_x_p_y_m) self._send_list_box[self._x_p_y_m_proc] = xp.concatenate( - (self._send_list_box[self._x_p_y_m_proc], self._markers_x_p_y_m) + (self._send_list_box[self._x_p_y_m_proc], self._markers_x_p_y_m), ) # if self._x_p_y_p_proc is not None: self._send_info_box[self._x_p_y_p_proc] += len(self._markers_x_p_y_p) self._send_list_box[self._x_p_y_p_proc] = xp.concatenate( - (self._send_list_box[self._x_p_y_p_proc], self._markers_x_p_y_p) + (self._send_list_box[self._x_p_y_p_proc], self._markers_x_p_y_p), ) # x-z edges # if self._x_m_z_m_proc is not None: self._send_info_box[self._x_m_z_m_proc] += len(self._markers_x_m_z_m) self._send_list_box[self._x_m_z_m_proc] = xp.concatenate( - (self._send_list_box[self._x_m_z_m_proc], self._markers_x_m_z_m) + (self._send_list_box[self._x_m_z_m_proc], self._markers_x_m_z_m), ) # if self._x_m_z_p_proc is not None: self._send_info_box[self._x_m_z_p_proc] += len(self._markers_x_m_z_p) self._send_list_box[self._x_m_z_p_proc] = xp.concatenate( - (self._send_list_box[self._x_m_z_p_proc], self._markers_x_m_z_p) + (self._send_list_box[self._x_m_z_p_proc], self._markers_x_m_z_p), ) # if self._x_p_z_m_proc is not None: self._send_info_box[self._x_p_z_m_proc] += len(self._markers_x_p_z_m) self._send_list_box[self._x_p_z_m_proc] = xp.concatenate( - (self._send_list_box[self._x_p_z_m_proc], self._markers_x_p_z_m) + (self._send_list_box[self._x_p_z_m_proc], self._markers_x_p_z_m), ) # if self._x_p_z_p_proc is not None: self._send_info_box[self._x_p_z_p_proc] += len(self._markers_x_p_z_p) self._send_list_box[self._x_p_z_p_proc] = xp.concatenate( - (self._send_list_box[self._x_p_z_p_proc], self._markers_x_p_z_p) + (self._send_list_box[self._x_p_z_p_proc], self._markers_x_p_z_p), ) # y-z edges # if self._y_m_z_m_proc is not None: self._send_info_box[self._y_m_z_m_proc] += len(self._markers_y_m_z_m) self._send_list_box[self._y_m_z_m_proc] = xp.concatenate( - (self._send_list_box[self._y_m_z_m_proc], self._markers_y_m_z_m) + (self._send_list_box[self._y_m_z_m_proc], self._markers_y_m_z_m), ) # if self._y_m_z_p_proc is not None: self._send_info_box[self._y_m_z_p_proc] += len(self._markers_y_m_z_p) self._send_list_box[self._y_m_z_p_proc] = xp.concatenate( - (self._send_list_box[self._y_m_z_p_proc], self._markers_y_m_z_p) + (self._send_list_box[self._y_m_z_p_proc], self._markers_y_m_z_p), ) # if self._y_p_z_m_proc is not None: self._send_info_box[self._y_p_z_m_proc] += len(self._markers_y_p_z_m) self._send_list_box[self._y_p_z_m_proc] = xp.concatenate( - (self._send_list_box[self._y_p_z_m_proc], self._markers_y_p_z_m) + (self._send_list_box[self._y_p_z_m_proc], self._markers_y_p_z_m), ) # if self._y_p_z_p_proc is not None: self._send_info_box[self._y_p_z_p_proc] += len(self._markers_y_p_z_p) self._send_list_box[self._y_p_z_p_proc] = xp.concatenate( - (self._send_list_box[self._y_p_z_p_proc], self._markers_y_p_z_p) + (self._send_list_box[self._y_p_z_p_proc], self._markers_y_p_z_p), ) # corners # if self._x_m_y_m_z_m_proc is not None: self._send_info_box[self._x_m_y_m_z_m_proc] += len(self._markers_x_m_y_m_z_m) self._send_list_box[self._x_m_y_m_z_m_proc] = xp.concatenate( - (self._send_list_box[self._x_m_y_m_z_m_proc], self._markers_x_m_y_m_z_m) + (self._send_list_box[self._x_m_y_m_z_m_proc], self._markers_x_m_y_m_z_m), ) # if self._x_m_y_m_z_p_proc is not None: self._send_info_box[self._x_m_y_m_z_p_proc] += len(self._markers_x_m_y_m_z_p) self._send_list_box[self._x_m_y_m_z_p_proc] = xp.concatenate( - (self._send_list_box[self._x_m_y_m_z_p_proc], self._markers_x_m_y_m_z_p) + (self._send_list_box[self._x_m_y_m_z_p_proc], self._markers_x_m_y_m_z_p), ) # if self._x_m_y_p_z_m_proc is not None: self._send_info_box[self._x_m_y_p_z_m_proc] += len(self._markers_x_m_y_p_z_m) self._send_list_box[self._x_m_y_p_z_m_proc] = xp.concatenate( - (self._send_list_box[self._x_m_y_p_z_m_proc], self._markers_x_m_y_p_z_m) + (self._send_list_box[self._x_m_y_p_z_m_proc], self._markers_x_m_y_p_z_m), ) # if self._x_m_y_p_z_p_proc is not None: self._send_info_box[self._x_m_y_p_z_p_proc] += len(self._markers_x_m_y_p_z_p) self._send_list_box[self._x_m_y_p_z_p_proc] = xp.concatenate( - (self._send_list_box[self._x_m_y_p_z_p_proc], self._markers_x_m_y_p_z_p) + (self._send_list_box[self._x_m_y_p_z_p_proc], self._markers_x_m_y_p_z_p), ) # if self._x_p_y_m_z_m_proc is not None: self._send_info_box[self._x_p_y_m_z_m_proc] += len(self._markers_x_p_y_m_z_m) self._send_list_box[self._x_p_y_m_z_m_proc] = xp.concatenate( - (self._send_list_box[self._x_p_y_m_z_m_proc], self._markers_x_p_y_m_z_m) + (self._send_list_box[self._x_p_y_m_z_m_proc], self._markers_x_p_y_m_z_m), ) # if self._x_p_y_m_z_p_proc is not None: self._send_info_box[self._x_p_y_m_z_p_proc] += len(self._markers_x_p_y_m_z_p) self._send_list_box[self._x_p_y_m_z_p_proc] = xp.concatenate( - (self._send_list_box[self._x_p_y_m_z_p_proc], self._markers_x_p_y_m_z_p) + (self._send_list_box[self._x_p_y_m_z_p_proc], self._markers_x_p_y_m_z_p), ) # if self._x_p_y_p_z_m_proc is not None: self._send_info_box[self._x_p_y_p_z_m_proc] += len(self._markers_x_p_y_p_z_m) self._send_list_box[self._x_p_y_p_z_m_proc] = xp.concatenate( - (self._send_list_box[self._x_p_y_p_z_m_proc], self._markers_x_p_y_p_z_m) + (self._send_list_box[self._x_p_y_p_z_m_proc], self._markers_x_p_y_p_z_m), ) # if self._x_p_y_p_z_p_proc is not None: self._send_info_box[self._x_p_y_p_z_p_proc] += len(self._markers_x_p_y_p_z_p) self._send_list_box[self._x_p_y_p_z_p_proc] = xp.concatenate( - (self._send_list_box[self._x_p_y_p_z_p_proc], self._markers_x_p_y_p_z_p) + (self._send_list_box[self._x_p_y_p_z_p_proc], self._markers_x_p_y_p_z_p), ) def self_communication_boxes(self): @@ -3151,7 +3163,7 @@ def self_communication_boxes(self): f'Strong load imbalance detected: \ number of holes ({holes_inds.size}) on rank {self.mpi_rank} \ is smaller than number of incoming particles ({self._send_info_box[self.mpi_rank]}). \ -Increasing the value of "bufsize" in the markers parameters for the next run.' +Increasing the value of "bufsize" in the markers parameters for the next run.', ) self.mpi_comm.Abort() @@ -3242,7 +3254,7 @@ def sendrecv_markers_boxes(self): f'Strong load imbalance detected: \ number of holes ({hole_inds.size}) on rank {self.mpi_rank} \ is smaller than number of incoming particles ({first_hole[i] + self._recv_info_box[i]}). \ -Increasing the value of "bufsize" in the markers parameters for the next run.' +Increasing the value of "bufsize" in the markers parameters for the next run.', ) self.mpi_comm.Abort() # exit() @@ -4007,7 +4019,7 @@ def sendrecv_markers(self, recv_info, hole_inds_after_send): f'Strong load imbalance detected: \ number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ is smaller than number of incoming particles ({first_hole[i] + recv_info[i]}). \ -Increasing the value of "bufsize" in the markers parameters for the next run.' +Increasing the value of "bufsize" in the markers parameters for the next run.', ) self.mpi_comm.Abort() diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 21e50ffda..79634d13a 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -219,7 +219,8 @@ def save_constants_of_motion(self): # send particles to the guiding center positions self.markers[~self.holes, self.first_pusher_idx : self.first_pusher_idx + 3] = self.markers[ - ~self.holes, slice_gc + ~self.holes, + slice_gc, ] if self.mpi_comm is not None: self.mpi_sort_markers(alpha=1) diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 190525de9..14c756b31 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -136,7 +136,7 @@ def __init__( # check marker array column number assert isinstance(comps, xp.ndarray) assert column_nr + comps.size < particles.n_cols, ( - f"{column_nr + comps.size} not smaller than {particles.n_cols = }; not enough columns in marker array !!" + f"{column_nr + comps.size} not smaller than {particles.n_cols =}; not enough columns in marker array !!" ) # prepare and check eval_kernels @@ -148,7 +148,7 @@ def __init__( # check marker array column number assert isinstance(comps, xp.ndarray) assert column_nr + comps.size < particles.n_cols, ( - f"{column_nr + comps.size} not smaller than {particles.n_cols = }; not enough columns in marker array !!" + f"{column_nr + comps.size} not smaller than {particles.n_cols =}; not enough columns in marker array !!" ) self._init_kernels = init_kernels @@ -178,10 +178,10 @@ def __call__(self, dt: float): residual_idx = self.particles.residual_idx if self.verbose: - print(f"{first_pusher_idx = }") - print(f"{first_shift_idx = }") - print(f"{residual_idx = }") - print(f"{self.particles.n_cols = }") + print(f"{first_pusher_idx =}") + print(f"{first_shift_idx =}") + print(f"{residual_idx =}") + print(f"{self.particles.n_cols =}") init_slice = slice(first_pusher_idx, first_shift_idx) shift_slice = slice(first_shift_idx, residual_idx) @@ -231,7 +231,7 @@ def __call__(self, dt: float): if self.verbose and self.maxiter > 1: max_res = 1.0 print( - f"rank {rank}: {k = }, tol: {self._tol}, {n_not_converged[0] = }, {max_res = }", + f"rank {rank}: {k =}, tol: {self._tol}, {n_not_converged[0] =}, {max_res =}", ) if self.particles.mpi_comm is not None: self.particles.mpi_comm.Barrier() @@ -309,7 +309,7 @@ def __call__(self, dt: float): if self.verbose: print( - f"rank {rank}: {k = }, tol: {self._tol}, {n_not_converged[0] = }, {max_res = }", + f"rank {rank}: {k =}, tol: {self._tol}, {n_not_converged[0] =}, {max_res =}", ) if self.particles.mpi_comm is not None: self.particles.mpi_comm.Barrier() @@ -329,7 +329,7 @@ def __call__(self, dt: float): if self.maxiter > 1: rank = self.particles.mpi_rank print( - f"rank {rank}: {k = }, maxiter={self.maxiter} reached! tol: {self._tol}, {n_not_converged[0] = }, {max_res = }", + f"rank {rank}: {k =}, maxiter={self.maxiter} reached! tol: {self._tol}, {n_not_converged[0] =}, {max_res =}", ) # sort markers according to domain decomposition if self.mpi_sort == "each": diff --git a/src/struphy/pic/sobol_seq.py b/src/struphy/pic/sobol_seq.py index ce965cc8f..f4c01347a 100644 --- a/src/struphy/pic/sobol_seq.py +++ b/src/struphy/pic/sobol_seq.py @@ -264,7 +264,7 @@ def i4_sobol(dim_num, seed): 1, 1, 1, - ] + ], ) v[2:40, 1] = xp.transpose( @@ -307,7 +307,7 @@ def i4_sobol(dim_num, seed): 3, 1, 3, - ] + ], ) v[3:40, 2] = xp.transpose( @@ -349,7 +349,7 @@ def i4_sobol(dim_num, seed): 1, 3, 3, - ] + ], ) v[5:40, 3] = xp.transpose( @@ -389,7 +389,7 @@ def i4_sobol(dim_num, seed): 1, 7, 9, - ] + ], ) v[7:40, 4] = xp.transpose( @@ -427,15 +427,15 @@ def i4_sobol(dim_num, seed): 9, 31, 9, - ] + ], ) v[13:40, 5] = xp.transpose( - [37, 33, 7, 5, 11, 39, 63, 27, 17, 15, 23, 29, 3, 21, 13, 31, 25, 9, 49, 33, 19, 29, 11, 19, 27, 15, 25] + [37, 33, 7, 5, 11, 39, 63, 27, 17, 15, 23, 29, 3, 21, 13, 31, 25, 9, 49, 33, 19, 29, 11, 19, 27, 15, 25], ) v[19:40, 6] = xp.transpose( - [13, 33, 115, 41, 79, 17, 29, 119, 75, 73, 105, 7, 59, 65, 21, 3, 113, 61, 89, 45, 107] + [13, 33, 115, 41, 79, 17, 29, 119, 75, 73, 105, 7, 59, 65, 21, 3, 113, 61, 89, 45, 107], ) v[37:40, 7] = xp.transpose([7, 23, 39]) diff --git a/src/struphy/pic/sph_eval_kernels.py b/src/struphy/pic/sph_eval_kernels.py index 37414f447..4c63e0156 100644 --- a/src/struphy/pic/sph_eval_kernels.py +++ b/src/struphy/pic/sph_eval_kernels.py @@ -297,7 +297,19 @@ def naive_evaluation_meshgrid( e2 = eta2[i, j, k] e3 = eta3[i, j, k] out[i, j, k] = naive_evaluation_kernel( - args_markers, e1, e2, e3, holes, periodic1, periodic2, periodic3, index, kernel_type, h1, h2, h3 + args_markers, + e1, + e2, + e3, + holes, + periodic1, + periodic2, + periodic3, + index, + kernel_type, + h1, + h2, + h3, ) diff --git a/src/struphy/pic/tests/test_accum_vec_H1.py b/src/struphy/pic/tests/test_accum_vec_H1.py index 7ed52b153..cb5cbb17e 100644 --- a/src/struphy/pic/tests/test_accum_vec_H1.py +++ b/src/struphy/pic/tests/test_accum_vec_H1.py @@ -6,7 +6,8 @@ @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @pytest.mark.parametrize("p", [[2, 3, 4]]) @pytest.mark.parametrize( - "spl_kind", [[False, False, True], [False, True, True], [True, False, True], [True, True, True]] + "spl_kind", + [[False, False, True], [False, True, True], [True, False, True], [True, True, True]], ) @pytest.mark.parametrize( "mapping", @@ -156,7 +157,7 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): if clone_config is not None: clone_config.sub_comm.Allreduce(MPI.IN_PLACE, _sum_within_clone, op=MPI.SUM) - print(f"rank {mpi_rank}: {_sum_within_clone = }, {_sqrtg = }") + print(f"rank {mpi_rank}: {_sum_within_clone =}, {_sqrtg =}") # Check within clone assert xp.isclose(_sum_within_clone, _sqrtg) @@ -169,7 +170,7 @@ def test_accum_poisson(Nel, p, spl_kind, mapping, num_clones, Np=1000): mpi_comm.Allreduce(MPI.IN_PLACE, _sum_between_clones, op=MPI.SUM) clone_config.inter_comm.Allreduce(MPI.IN_PLACE, _sqrtg, op=MPI.SUM) - print(f"rank {mpi_rank}: {_sum_between_clones = }, {_sqrtg = }") + print(f"rank {mpi_rank}: {_sum_between_clones =}, {_sqrtg =}") # Check within clone assert xp.isclose(_sum_between_clones, _sqrtg) diff --git a/src/struphy/pic/tests/test_accumulation.py b/src/struphy/pic/tests/test_accumulation.py index 012c73e33..14cadf638 100644 --- a/src/struphy/pic/tests/test_accumulation.py +++ b/src/struphy/pic/tests/test_accumulation.py @@ -6,7 +6,8 @@ @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @pytest.mark.parametrize("p", [[2, 3, 4]]) @pytest.mark.parametrize( - "spl_kind", [[False, False, True], [False, True, False], [True, False, True], [True, True, False]] + "spl_kind", + [[False, False, True], [False, True, False], [True, False, True], [True, True, False]], ) @pytest.mark.parametrize( "mapping", diff --git a/src/struphy/pic/tests/test_binning.py b/src/struphy/pic/tests/test_binning.py index a6a1dde6e..cda2524e7 100644 --- a/src/struphy/pic/tests/test_binning.py +++ b/src/struphy/pic/tests/test_binning.py @@ -631,8 +631,8 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): "given_in_basis": "0", "ls": [l_n1], "amps": [amp_n1], - } - } + }, + }, }, "Maxwellian3D_2": { "n": { @@ -640,8 +640,8 @@ def test_binning_6D_full_f_mpi(mapping, show_plot=False): "given_in_basis": "0", "ls": [l_n2], "amps": [amp_n2], - } - } + }, + }, }, } pert_1 = perturbations.ModesCos(ls=(l_n1,), amps=(amp_n1,)) @@ -814,8 +814,8 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): "given_in_basis": "0", "ls": [l_n], "amps": [amp_n], - } - } + }, + }, } pert = perturbations.ModesCos(ls=(l_n,), amps=(amp_n,)) background = Maxwellian3D(n=(1.0, pert)) @@ -891,7 +891,7 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): "given_in_basis": "0", "ls": [l_n1], "amps": [amp_n1], - } + }, }, }, "Maxwellian3D_2": { @@ -901,7 +901,7 @@ def test_binning_6D_delta_f_mpi(mapping, show_plot=False): "given_in_basis": "0", "ls": [l_n2], "amps": [amp_n2], - } + }, }, }, } diff --git a/src/struphy/pic/tests/test_mat_vec_filler.py b/src/struphy/pic/tests/test_mat_vec_filler.py index 491e1f20e..c6bee1faa 100644 --- a/src/struphy/pic/tests/test_mat_vec_filler.py +++ b/src/struphy/pic/tests/test_mat_vec_filler.py @@ -70,8 +70,11 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): for j in range(3): mat["v1"][-1] += [ StencilMatrix( - DR.Vh["1"].spaces[i], DR.Vh["1"].spaces[j], backend=PSYDAC_BACKEND_GPYCCEL, precompiled=True - )._data + DR.Vh["1"].spaces[i], + DR.Vh["1"].spaces[j], + backend=PSYDAC_BACKEND_GPYCCEL, + precompiled=True, + )._data, ] vec["v1"] = [] @@ -84,8 +87,11 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): for j in range(3): mat["v2"][-1] += [ StencilMatrix( - DR.Vh["2"].spaces[i], DR.Vh["2"].spaces[j], backend=PSYDAC_BACKEND_GPYCCEL, precompiled=True - )._data + DR.Vh["2"].spaces[i], + DR.Vh["2"].spaces[j], + backend=PSYDAC_BACKEND_GPYCCEL, + precompiled=True, + )._data, ] vec["v2"] = [] @@ -213,7 +219,13 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): for n, ij in enumerate(ind_pairs): assert_mat( - args[n], rows, cols, basis[space][ij[0]], basis[space][ij[1]], rank, verbose=False + args[n], + rows, + cols, + basis[space][ij[0]], + basis[space][ij[1]], + rank, + verbose=False, ) # assertion test of mat if mv == "m_v": for i in range(3): @@ -230,7 +242,13 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): for n, ij in enumerate(ind_pairs): assert_mat( - args[n], rows, cols, basis[space][ij[0]], basis[space][ij[1]], rank, verbose=False + args[n], + rows, + cols, + basis[space][ij[0]], + basis[space][ij[1]], + rank, + verbose=False, ) # assertion test of mat if mv == "m_v": for i in range(3): diff --git a/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py b/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py index 2ee5ba7b2..6bb225571 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/accumulation.py @@ -159,7 +159,8 @@ def to_sparse_step1(self): # final block matrix M = spa.bmat( - [[None, M[0][1], M[0][2]], [-M[0][1].T, None, M[1][2]], [-M[0][2].T, -M[1][2].T, None]], format="csr" + [[None, M[0][1], M[0][2]], [-M[0][1].T, None, M[1][2]], [-M[0][2].T, -M[1][2].T, None]], + format="csr", ) # apply extraction operator @@ -226,7 +227,8 @@ def to_sparse_step3(self): # final block matrix M = spa.bmat( - [[M[0][0], M[0][1], M[0][2]], [M[0][1].T, M[1][1], M[1][2]], [M[0][2].T, M[1][2].T, M[2][2]]], format="csr" + [[M[0][0], M[0][1], M[0][2]], [M[0][1].T, M[1][1], M[1][2]], [M[0][2].T, M[1][2].T, M[2][2]]], + format="csr", ) # apply extraction operator @@ -528,15 +530,15 @@ def assemble_step3(self, b2_eq, b2): # build global sparse matrix and global vector if self.basis_u == 0: return self.to_sparse_step3(), self.space.Ev_0.dot( - xp.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())) + xp.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())), ) elif self.basis_u == 1: return self.to_sparse_step3(), self.space.E1_0.dot( - xp.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())) + xp.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())), ) elif self.basis_u == 2: return self.to_sparse_step3(), self.space.E2_0.dot( - xp.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())) + xp.concatenate((self.vecs[0].flatten(), self.vecs[1].flatten(), self.vecs[2].flatten())), ) diff --git a/src/struphy/pic/tests/test_pic_legacy_files/accumulation_kernels_3d.py b/src/struphy/pic/tests/test_pic_legacy_files/accumulation_kernels_3d.py index c70261023..349cca379 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/accumulation_kernels_3d.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/accumulation_kernels_3d.py @@ -185,13 +185,49 @@ def kernel_step1( bsp.b_d_splines_slim(t3, int(pn3), eta3, int(span3), bn3, bd3) b[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], b2_1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + b2_1, ) b[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], b2_2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + b2_2, ) b[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], b2_3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + b2_3, ) b_prod[0, 1] = -b[2] @@ -554,13 +590,49 @@ def kernel_step3( bsp.b_d_splines_slim(t3, int(pn3), eta3, int(span3), bn3, bd3) b[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], b2_1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + b2_1, ) b[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], b2_2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + b2_2, ) b[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], b2_3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + b2_3, ) b_prod[0, 1] = -b[2] @@ -1130,7 +1202,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat11[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] for jl1 in range(pn1 + 1): @@ -1142,7 +1221,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat12[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] for jl1 in range(pn1 + 1): @@ -1154,7 +1240,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat13[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] # add contribution to 22 component (NDN NDN) and 23 component (NDN NND) @@ -1179,7 +1272,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat22[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] for jl1 in range(pn1 + 1): @@ -1191,7 +1291,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat23[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] # add contribution to 33 component (NND NND) @@ -1216,7 +1323,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat33[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] elif basis_u == 2: @@ -1242,7 +1356,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat11[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] for jl1 in range(pd1 + 1): @@ -1254,7 +1375,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat12[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] for jl1 in range(pd1 + 1): @@ -1266,7 +1394,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat13[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] # add contribution to 22 component (DND DND) and 23 component (DND DDN) @@ -1291,7 +1426,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat22[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] for jl1 in range(pd1 + 1): @@ -1303,7 +1445,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat23[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] # add contribution to 33 component (DDN DDN) @@ -1328,7 +1477,14 @@ def kernel_step_ph_full( for vp in range(3): for vq in range(3): mat33[ - i1, i2, i3, pn1 + jl1 - il1, pn2 + jl2 - il2, pn3 + jl3 - il3, vp, vq + i1, + i2, + i3, + pn1 + jl1 - il1, + pn2 + jl2 - il2, + pn3 + jl3 - il3, + vp, + vq, ] += bj3 * v[vp] * v[vq] # -- removed omp: #$ omp end parallel diff --git a/src/struphy/pic/tests/test_pic_legacy_files/mappings_3d.py b/src/struphy/pic/tests/test_pic_legacy_files/mappings_3d.py index 587b8b15f..2e54d34dd 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/mappings_3d.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/mappings_3d.py @@ -74,17 +74,53 @@ def f( if kind_map == 0: if component == 1: value = eva_3d.evaluate_n_n_n( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cx, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cx, + eta1, + eta2, + eta3, ) elif component == 2: value = eva_3d.evaluate_n_n_n( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cy, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cy, + eta1, + eta2, + eta3, ) elif component == 3: value = eva_3d.evaluate_n_n_n( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cz, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cz, + eta1, + eta2, + eta3, ) # ==== 2d spline (straight in 3rd direction) === @@ -110,7 +146,7 @@ def f( elif kind_map == 2: if component == 1: value = eva_2d.evaluate_n_n(tn1, tn2, pn[0], pn[1], nbase_n[0], nbase_n[1], cx[:, :, 0], eta1, eta2) * cos( - 2 * pi * eta3 + 2 * pi * eta3, ) if eta1 == 0.0 and cx[0, 0, 0] == cx[0, 1, 0]: @@ -124,7 +160,7 @@ def f( elif component == 3: value = eva_2d.evaluate_n_n(tn1, tn2, pn[0], pn[1], nbase_n[0], nbase_n[1], cx[:, :, 0], eta1, eta2) * sin( - 2 * pi * eta3 + 2 * pi * eta3, ) if eta1 == 0.0 and cx[0, 0, 0] == cx[0, 1, 0]: @@ -299,39 +335,147 @@ def df( if kind_map == 0: if component == 11: value = eva_3d.evaluate_diffn_n_n( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cx, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cx, + eta1, + eta2, + eta3, ) elif component == 12: value = eva_3d.evaluate_n_diffn_n( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cx, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cx, + eta1, + eta2, + eta3, ) elif component == 13: value = eva_3d.evaluate_n_n_diffn( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cx, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cx, + eta1, + eta2, + eta3, ) elif component == 21: value = eva_3d.evaluate_diffn_n_n( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cy, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cy, + eta1, + eta2, + eta3, ) elif component == 22: value = eva_3d.evaluate_n_diffn_n( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cy, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cy, + eta1, + eta2, + eta3, ) elif component == 23: value = eva_3d.evaluate_n_n_diffn( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cy, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cy, + eta1, + eta2, + eta3, ) elif component == 31: value = eva_3d.evaluate_diffn_n_n( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cz, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cz, + eta1, + eta2, + eta3, ) elif component == 32: value = eva_3d.evaluate_n_diffn_n( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cz, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cz, + eta1, + eta2, + eta3, ) elif component == 33: value = eva_3d.evaluate_n_n_diffn( - tn1, tn2, tn3, pn[0], pn[1], pn[2], nbase_n[0], nbase_n[1], nbase_n[2], cz, eta1, eta2, eta3 + tn1, + tn2, + tn3, + pn[0], + pn[1], + pn[2], + nbase_n[0], + nbase_n[1], + nbase_n[2], + cz, + eta1, + eta2, + eta3, ) # ==== 2d spline (straight in 3rd direction) === @@ -369,11 +513,27 @@ def df( elif kind_map == 2: if component == 11: value = eva_2d.evaluate_diffn_n( - tn1, tn2, pn[0], pn[1], nbase_n[0], nbase_n[1], cx[:, :, 0], eta1, eta2 + tn1, + tn2, + pn[0], + pn[1], + nbase_n[0], + nbase_n[1], + cx[:, :, 0], + eta1, + eta2, ) * cos(2 * pi * eta3) elif component == 12: value = eva_2d.evaluate_n_diffn( - tn1, tn2, pn[0], pn[1], nbase_n[0], nbase_n[1], cx[:, :, 0], eta1, eta2 + tn1, + tn2, + pn[0], + pn[1], + nbase_n[0], + nbase_n[1], + cx[:, :, 0], + eta1, + eta2, ) * cos(2 * pi * eta3) if eta1 == 0.0 and cx[0, 0, 0] == cx[0, 1, 0]: @@ -397,11 +557,27 @@ def df( value = 0.0 elif component == 31: value = eva_2d.evaluate_diffn_n( - tn1, tn2, pn[0], pn[1], nbase_n[0], nbase_n[1], cx[:, :, 0], eta1, eta2 + tn1, + tn2, + pn[0], + pn[1], + nbase_n[0], + nbase_n[1], + cx[:, :, 0], + eta1, + eta2, ) * sin(2 * pi * eta3) elif component == 32: value = eva_2d.evaluate_n_diffn( - tn1, tn2, pn[0], pn[1], nbase_n[0], nbase_n[1], cx[:, :, 0], eta1, eta2 + tn1, + tn2, + pn[0], + pn[1], + nbase_n[0], + nbase_n[1], + cx[:, :, 0], + eta1, + eta2, ) * sin(2 * pi * eta3) if eta1 == 0.0 and cx[0, 0, 0] == cx[0, 1, 0]: diff --git a/src/struphy/pic/tests/test_pic_legacy_files/mappings_3d_fast.py b/src/struphy/pic/tests/test_pic_legacy_files/mappings_3d_fast.py index fbd912b39..f87380685 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/mappings_3d_fast.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/mappings_3d_fast.py @@ -264,19 +264,51 @@ def df_all( if mat_or_vec == 0 or mat_or_vec == 2: # sum-up non-vanishing contributions (line 1: df_11, df_12 and df_13) mat_out[0, 0] = evaluation_kernel_2d( - pn[0], pn[1], der1, b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + der1, + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) mat_out[0, 1] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], der2, span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + der2, + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) mat_out[0, 2] = 0.0 # sum-up non-vanishing contributions (line 2: df_21, df_22 and df_23) mat_out[1, 0] = evaluation_kernel_2d( - pn[0], pn[1], der1, b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cy[:, :, 0] + pn[0], + pn[1], + der1, + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cy[:, :, 0], ) mat_out[1, 1] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], der2, span_n1, span_n2, nbase_n[0], nbase_n[1], cy[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + der2, + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cy[:, :, 0], ) mat_out[1, 2] = 0.0 @@ -288,10 +320,26 @@ def df_all( # evaluate mapping if mat_or_vec == 1 or mat_or_vec == 2: vec_out[0] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) vec_out[1] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cy[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cy[:, :, 0], ) vec_out[2] = lz * eta3 @@ -305,14 +353,38 @@ def df_all( if mat_or_vec == 0 or mat_or_vec == 2: # sum-up non-vanishing contributions (line 1: df_11, df_12 and df_13) mat_out[0, 0] = evaluation_kernel_2d( - pn[0], pn[1], der1, b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + der1, + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) * cos(2 * pi * eta3) mat_out[0, 1] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], der2, span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + der2, + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) * cos(2 * pi * eta3) mat_out[0, 2] = ( evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) * sin(2 * pi * eta3) * (-2 * pi) @@ -320,23 +392,63 @@ def df_all( # sum-up non-vanishing contributions (line 2: df_21, df_22 and df_23) mat_out[1, 0] = evaluation_kernel_2d( - pn[0], pn[1], der1, b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cy[:, :, 0] + pn[0], + pn[1], + der1, + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cy[:, :, 0], ) mat_out[1, 1] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], der2, span_n1, span_n2, nbase_n[0], nbase_n[1], cy[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + der2, + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cy[:, :, 0], ) mat_out[1, 2] = 0.0 # sum-up non-vanishing contributions (line 3: df_31, df_32 and df_33) mat_out[2, 0] = evaluation_kernel_2d( - pn[0], pn[1], der1, b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + der1, + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) * sin(2 * pi * eta3) mat_out[2, 1] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], der2, span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + der2, + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) * sin(2 * pi * eta3) mat_out[2, 2] = ( evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) * cos(2 * pi * eta3) * 2 @@ -346,13 +458,37 @@ def df_all( # evaluate mapping if mat_or_vec == 1 or mat_or_vec == 2: vec_out[0] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) * cos(2 * pi * eta3) vec_out[1] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cy[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cy[:, :, 0], ) vec_out[2] = evaluation_kernel_2d( - pn[0], pn[1], b1[pn[0]], b2[pn[1]], span_n1, span_n2, nbase_n[0], nbase_n[1], cx[:, :, 0] + pn[0], + pn[1], + b1[pn[0]], + b2[pn[1]], + span_n1, + span_n2, + nbase_n[0], + nbase_n[1], + cx[:, :, 0], ) * sin(2 * pi * eta3) # analytical mapping @@ -360,33 +496,150 @@ def df_all( # evaluate Jacobian matrix if mat_or_vec == 0 or mat_or_vec == 2: mat_out[0, 0] = mapping.df( - eta1, eta2, eta3, 11, kind_map, params_map, tn1, tn2, tn3, pn, nbase_n, cx, cy, cz + eta1, + eta2, + eta3, + 11, + kind_map, + params_map, + tn1, + tn2, + tn3, + pn, + nbase_n, + cx, + cy, + cz, ) mat_out[0, 1] = mapping.df( - eta1, eta2, eta3, 12, kind_map, params_map, tn1, tn2, tn3, pn, nbase_n, cx, cy, cz + eta1, + eta2, + eta3, + 12, + kind_map, + params_map, + tn1, + tn2, + tn3, + pn, + nbase_n, + cx, + cy, + cz, ) mat_out[0, 2] = mapping.df( - eta1, eta2, eta3, 13, kind_map, params_map, tn1, tn2, tn3, pn, nbase_n, cx, cy, cz + eta1, + eta2, + eta3, + 13, + kind_map, + params_map, + tn1, + tn2, + tn3, + pn, + nbase_n, + cx, + cy, + cz, ) mat_out[1, 0] = mapping.df( - eta1, eta2, eta3, 21, kind_map, params_map, tn1, tn2, tn3, pn, nbase_n, cx, cy, cz + eta1, + eta2, + eta3, + 21, + kind_map, + params_map, + tn1, + tn2, + tn3, + pn, + nbase_n, + cx, + cy, + cz, ) mat_out[1, 1] = mapping.df( - eta1, eta2, eta3, 22, kind_map, params_map, tn1, tn2, tn3, pn, nbase_n, cx, cy, cz + eta1, + eta2, + eta3, + 22, + kind_map, + params_map, + tn1, + tn2, + tn3, + pn, + nbase_n, + cx, + cy, + cz, ) mat_out[1, 2] = mapping.df( - eta1, eta2, eta3, 23, kind_map, params_map, tn1, tn2, tn3, pn, nbase_n, cx, cy, cz + eta1, + eta2, + eta3, + 23, + kind_map, + params_map, + tn1, + tn2, + tn3, + pn, + nbase_n, + cx, + cy, + cz, ) mat_out[2, 0] = mapping.df( - eta1, eta2, eta3, 31, kind_map, params_map, tn1, tn2, tn3, pn, nbase_n, cx, cy, cz + eta1, + eta2, + eta3, + 31, + kind_map, + params_map, + tn1, + tn2, + tn3, + pn, + nbase_n, + cx, + cy, + cz, ) mat_out[2, 1] = mapping.df( - eta1, eta2, eta3, 32, kind_map, params_map, tn1, tn2, tn3, pn, nbase_n, cx, cy, cz + eta1, + eta2, + eta3, + 32, + kind_map, + params_map, + tn1, + tn2, + tn3, + pn, + nbase_n, + cx, + cy, + cz, ) mat_out[2, 2] = mapping.df( - eta1, eta2, eta3, 33, kind_map, params_map, tn1, tn2, tn3, pn, nbase_n, cx, cy, cz + eta1, + eta2, + eta3, + 33, + kind_map, + params_map, + tn1, + tn2, + tn3, + pn, + nbase_n, + cx, + cy, + cz, ) # evaluate mapping diff --git a/src/struphy/pic/tests/test_pic_legacy_files/pusher_pos.py b/src/struphy/pic/tests/test_pic_legacy_files/pusher_pos.py index 78631440e..81b5e1e53 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/pusher_pos.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/pusher_pos.py @@ -532,13 +532,52 @@ def pusher_step4_pcart( # compute old pseudo-cartesian coordinates fx_pseudo[0] = mapping.f( - eta[0], eta[1], eta[2], 1, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 1, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) fx_pseudo[1] = mapping.f( - eta[0], eta[1], eta[2], 2, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 2, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) fx_pseudo[2] = mapping.f( - eta[0], eta[1], eta[2], 3, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 3, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) # evaluate old Jacobian matrix of mapping F @@ -588,33 +627,150 @@ def pusher_step4_pcart( # evaluate old Jacobian matrix of mapping F_pseudo df_pseudo_old[0, 0] = mapping.df( - eta[0], eta[1], eta[2], 11, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 11, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo_old[0, 1] = mapping.df( - eta[0], eta[1], eta[2], 12, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 12, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo_old[0, 2] = mapping.df( - eta[0], eta[1], eta[2], 13, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 13, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo_old[1, 0] = mapping.df( - eta[0], eta[1], eta[2], 21, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 21, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo_old[1, 1] = mapping.df( - eta[0], eta[1], eta[2], 22, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 22, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo_old[1, 2] = mapping.df( - eta[0], eta[1], eta[2], 23, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 23, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo_old[2, 0] = mapping.df( - eta[0], eta[1], eta[2], 31, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 31, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo_old[2, 1] = mapping.df( - eta[0], eta[1], eta[2], 32, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 32, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo_old[2, 2] = mapping.df( - eta[0], eta[1], eta[2], 33, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 33, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) while True: @@ -687,33 +843,150 @@ def pusher_step4_pcart( # evaluate Jacobian matrix of mapping F_pseudo df_pseudo[0, 0] = mapping.df( - eta[0], eta[1], eta[2], 11, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 11, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[0, 1] = mapping.df( - eta[0], eta[1], eta[2], 12, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 12, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[0, 2] = mapping.df( - eta[0], eta[1], eta[2], 13, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 13, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[1, 0] = mapping.df( - eta[0], eta[1], eta[2], 21, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 21, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[1, 1] = mapping.df( - eta[0], eta[1], eta[2], 22, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 22, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[1, 2] = mapping.df( - eta[0], eta[1], eta[2], 23, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 23, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[2, 0] = mapping.df( - eta[0], eta[1], eta[2], 31, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 31, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[2, 1] = mapping.df( - eta[0], eta[1], eta[2], 32, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 32, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[2, 2] = mapping.df( - eta[0], eta[1], eta[2], 33, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 33, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) # compute df_pseudo*df_inv*v @@ -784,33 +1057,150 @@ def pusher_step4_pcart( # evaluate Jacobian matrix of mapping F_pseudo df_pseudo[0, 0] = mapping.df( - eta[0], eta[1], eta[2], 11, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 11, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[0, 1] = mapping.df( - eta[0], eta[1], eta[2], 12, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 12, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[0, 2] = mapping.df( - eta[0], eta[1], eta[2], 13, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 13, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[1, 0] = mapping.df( - eta[0], eta[1], eta[2], 21, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 21, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[1, 1] = mapping.df( - eta[0], eta[1], eta[2], 22, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 22, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[1, 2] = mapping.df( - eta[0], eta[1], eta[2], 23, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 23, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[2, 0] = mapping.df( - eta[0], eta[1], eta[2], 31, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 31, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[2, 1] = mapping.df( - eta[0], eta[1], eta[2], 32, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 32, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[2, 2] = mapping.df( - eta[0], eta[1], eta[2], 33, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 33, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) # compute df_pseudo*df_inv*v @@ -881,33 +1271,150 @@ def pusher_step4_pcart( # evaluate Jacobian matrix of mapping F_pseudo df_pseudo[0, 0] = mapping.df( - eta[0], eta[1], eta[2], 11, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 11, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[0, 1] = mapping.df( - eta[0], eta[1], eta[2], 12, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 12, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[0, 2] = mapping.df( - eta[0], eta[1], eta[2], 13, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 13, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[1, 0] = mapping.df( - eta[0], eta[1], eta[2], 21, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 21, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[1, 1] = mapping.df( - eta[0], eta[1], eta[2], 22, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 22, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[1, 2] = mapping.df( - eta[0], eta[1], eta[2], 23, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 23, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[2, 0] = mapping.df( - eta[0], eta[1], eta[2], 31, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 31, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[2, 1] = mapping.df( - eta[0], eta[1], eta[2], 32, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 32, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) df_pseudo[2, 2] = mapping.df( - eta[0], eta[1], eta[2], 33, map_pseudo, params_pseudo, tf1, tf2, tf3, pf, nbasef, cx, cy, cz + eta[0], + eta[1], + eta[2], + 33, + map_pseudo, + params_pseudo, + tf1, + tf2, + tf3, + pf, + nbasef, + cx, + cy, + cz, ) # compute df_pseudo*df_inv*v @@ -1374,26 +1881,98 @@ def pusher_rk4_pc_full( # velocity field if basis_u == 1: u[0] = eva3.evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, nbase_d[0], nbase_n[1], nbase_n[2], u1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + nbase_d[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, nbase_n[0], nbase_d[1], nbase_n[2], u2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + nbase_n[0], + nbase_d[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, nbase_n[0], nbase_n[1], nbase_d[2], u3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + nbase_n[0], + nbase_n[1], + nbase_d[2], + u3, ) linalg.matrix_vector(Ginv, u, k1_u) elif basis_u == 2: u[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], u1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], u2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], u3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + u3, ) k1_u[:] = u / det_df @@ -1491,26 +2070,98 @@ def pusher_rk4_pc_full( # velocity field if basis_u == 1: u[0] = eva3.evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, nbase_d[0], nbase_n[1], nbase_n[2], u1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + nbase_d[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, nbase_n[0], nbase_d[1], nbase_n[2], u2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + nbase_n[0], + nbase_d[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, nbase_n[0], nbase_n[1], nbase_d[2], u3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + nbase_n[0], + nbase_n[1], + nbase_d[2], + u3, ) linalg.matrix_vector(Ginv, u, k2_u) elif basis_u == 2: u[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], u1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], u2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], u3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + u3, ) k2_u[:] = u / det_df @@ -1608,26 +2259,98 @@ def pusher_rk4_pc_full( # velocity field if basis_u == 1: u[0] = eva3.evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, nbase_d[0], nbase_n[1], nbase_n[2], u1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + nbase_d[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, nbase_n[0], nbase_d[1], nbase_n[2], u2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + nbase_n[0], + nbase_d[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, nbase_n[0], nbase_n[1], nbase_d[2], u3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + nbase_n[0], + nbase_n[1], + nbase_d[2], + u3, ) linalg.matrix_vector(Ginv, u, k3_u) elif basis_u == 2: u[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], u1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], u2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], u3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + u3, ) k3_u[:] = u / det_df @@ -1722,26 +2445,98 @@ def pusher_rk4_pc_full( # velocity field if basis_u == 1: u[0] = eva3.evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, nbase_d[0], nbase_n[1], nbase_n[2], u1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + nbase_d[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, nbase_n[0], nbase_d[1], nbase_n[2], u2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + nbase_n[0], + nbase_d[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, nbase_n[0], nbase_n[1], nbase_d[2], u3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + nbase_n[0], + nbase_n[1], + nbase_d[2], + u3, ) linalg.matrix_vector(Ginv, u, k4_u) elif basis_u == 2: u[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], u1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], u2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], u3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + u3, ) k4_u[:] = u / det_df @@ -1992,26 +2787,98 @@ def pusher_rk4_pc_perp( # velocity field if basis_u == 1: u[0] = eva3.evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, nbase_d[0], nbase_n[1], nbase_n[2], u1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + nbase_d[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, nbase_n[0], nbase_d[1], nbase_n[2], u2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + nbase_n[0], + nbase_d[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, nbase_n[0], nbase_n[1], nbase_d[2], u3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + nbase_n[0], + nbase_n[1], + nbase_d[2], + u3, ) linalg.matrix_vector(Ginv, u, k1_u) elif basis_u == 2: u[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], u1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], u2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], u3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + u3, ) k1_u[:] = u / det_df @@ -2108,26 +2975,98 @@ def pusher_rk4_pc_perp( # velocity field if basis_u == 1: u[0] = eva3.evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, nbase_d[0], nbase_n[1], nbase_n[2], u1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + nbase_d[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, nbase_n[0], nbase_d[1], nbase_n[2], u2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + nbase_n[0], + nbase_d[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, nbase_n[0], nbase_n[1], nbase_d[2], u3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + nbase_n[0], + nbase_n[1], + nbase_d[2], + u3, ) linalg.matrix_vector(Ginv, u, k2_u) elif basis_u == 2: u[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], u1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], u2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], u3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + u3, ) k2_u[:] = u / det_df @@ -2223,26 +3162,98 @@ def pusher_rk4_pc_perp( # velocity field if basis_u == 1: u[0] = eva3.evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, nbase_d[0], nbase_n[1], nbase_n[2], u1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + nbase_d[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, nbase_n[0], nbase_d[1], nbase_n[2], u2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + nbase_n[0], + nbase_d[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, nbase_n[0], nbase_n[1], nbase_d[2], u3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + nbase_n[0], + nbase_n[1], + nbase_d[2], + u3, ) linalg.matrix_vector(Ginv, u, k3_u) elif basis_u == 2: u[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], u1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], u2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], u3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + u3, ) k3_u[:] = u / det_df @@ -2338,26 +3349,98 @@ def pusher_rk4_pc_perp( # velocity field if basis_u == 1: u[0] = eva3.evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, nbase_d[0], nbase_n[1], nbase_n[2], u1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + nbase_d[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, nbase_n[0], nbase_d[1], nbase_n[2], u2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + nbase_n[0], + nbase_d[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, nbase_n[0], nbase_n[1], nbase_d[2], u3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + nbase_n[0], + nbase_n[1], + nbase_d[2], + u3, ) linalg.matrix_vector(Ginv, u, k4_u) elif basis_u == 2: u[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], u1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], u2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], u3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + u3, ) k4_u[:] = u / det_df diff --git a/src/struphy/pic/tests/test_pic_legacy_files/pusher_vel_2d.py b/src/struphy/pic/tests/test_pic_legacy_files/pusher_vel_2d.py index 43e320311..0fcc29751 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/pusher_vel_2d.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/pusher_vel_2d.py @@ -248,19 +248,43 @@ def pusher_step3( for i in range(nbase_n[2]): u[0] += ( eva2.evaluation_kernel_2d( - pd1, pn2, bd1, bn2, span1 - 1, span2 - 0, nbase_d[0], nbase_n[1], u1[:, :, i] + pd1, + pn2, + bd1, + bn2, + span1 - 1, + span2 - 0, + nbase_d[0], + nbase_n[1], + u1[:, :, i], ) * cs[i] ) u[1] += ( eva2.evaluation_kernel_2d( - pn1, pd2, bn1, bd2, span1 - 0, span2 - 1, nbase_n[0], nbase_d[1], u2[:, :, i] + pn1, + pd2, + bn1, + bd2, + span1 - 0, + span2 - 1, + nbase_n[0], + nbase_d[1], + u2[:, :, i], ) * cs[i] ) u[2] += ( eva2.evaluation_kernel_2d( - pn1, pn2, bn1, bn2, span1 - 0, span2 - 0, nbase_n[0], nbase_n[1], u3[:, :, i] + pn1, + pn2, + bn1, + bn2, + span1 - 0, + span2 - 0, + nbase_n[0], + nbase_n[1], + u3[:, :, i], ) * cs[i] ) @@ -274,19 +298,43 @@ def pusher_step3( for i in range(nbase_n[2]): u[0] += ( eva2.evaluation_kernel_2d( - pn1, pd2, bn1, bd2, span1 - 0, span2 - 1, nbase_n[0], nbase_d[1], u1[:, :, i] + pn1, + pd2, + bn1, + bd2, + span1 - 0, + span2 - 1, + nbase_n[0], + nbase_d[1], + u1[:, :, i], ) * cs[i] ) u[1] += ( eva2.evaluation_kernel_2d( - pd1, pn2, bd1, bn2, span1 - 1, span2 - 0, nbase_d[0], nbase_n[1], u2[:, :, i] + pd1, + pn2, + bd1, + bn2, + span1 - 1, + span2 - 0, + nbase_d[0], + nbase_n[1], + u2[:, :, i], ) * cs[i] ) u[2] += ( eva2.evaluation_kernel_2d( - pd1, pd2, bd1, bd2, span1 - 1, span2 - 1, nbase_d[0], nbase_d[1], u3[:, :, i] + pd1, + pd2, + bd1, + bd2, + span1 - 1, + span2 - 1, + nbase_d[0], + nbase_d[1], + u3[:, :, i], ) * cs[i] ) @@ -299,32 +347,80 @@ def pusher_step3( # equilibrium magnetic field (2-form) b[0] = eva2.evaluation_kernel_2d( - pn1, pd2, bn1, bd2, span1 - 0, span2 - 1, nbase_n[0], nbase_d[1], b_eq_1[:, :, 0] + pn1, + pd2, + bn1, + bd2, + span1 - 0, + span2 - 1, + nbase_n[0], + nbase_d[1], + b_eq_1[:, :, 0], ) b[1] = eva2.evaluation_kernel_2d( - pd1, pn2, bd1, bn2, span1 - 1, span2 - 0, nbase_d[0], nbase_n[1], b_eq_2[:, :, 0] + pd1, + pn2, + bd1, + bn2, + span1 - 1, + span2 - 0, + nbase_d[0], + nbase_n[1], + b_eq_2[:, :, 0], ) b[2] = eva2.evaluation_kernel_2d( - pd1, pd2, bd1, bd2, span1 - 1, span2 - 1, nbase_d[0], nbase_d[1], b_eq_3[:, :, 0] + pd1, + pd2, + bd1, + bd2, + span1 - 1, + span2 - 1, + nbase_d[0], + nbase_d[1], + b_eq_3[:, :, 0], ) # perturbed magnetic field (2-form) for i in range(nbase_n[2]): b[0] += ( eva2.evaluation_kernel_2d( - pn1, pd2, bn1, bd2, span1 - 0, span2 - 1, nbase_n[0], nbase_d[1], b_p_1[:, :, i] + pn1, + pd2, + bn1, + bd2, + span1 - 0, + span2 - 1, + nbase_n[0], + nbase_d[1], + b_p_1[:, :, i], ) * cs[i] ) b[1] += ( eva2.evaluation_kernel_2d( - pd1, pn2, bd1, bn2, span1 - 1, span2 - 0, nbase_d[0], nbase_n[1], b_p_2[:, :, i] + pd1, + pn2, + bd1, + bn2, + span1 - 1, + span2 - 0, + nbase_d[0], + nbase_n[1], + b_p_2[:, :, i], ) * cs[i] ) b[2] += ( eva2.evaluation_kernel_2d( - pd1, pd2, bd1, bd2, span1 - 1, span2 - 1, nbase_d[0], nbase_d[1], b_p_3[:, :, i] + pd1, + pd2, + bd1, + bd2, + span1 - 1, + span2 - 1, + nbase_d[0], + nbase_d[1], + b_p_3[:, :, i], ) * cs[i] ) @@ -338,10 +434,26 @@ def pusher_step3( # gradient of absolute value of magnetic field (1-form) b_grad[0] = eva2.evaluation_kernel_2d( - pn1, pn2, der1, bn2, span1, span2, nbase_n[0], nbase_n[1], b_norm[:, :, 0] + pn1, + pn2, + der1, + bn2, + span1, + span2, + nbase_n[0], + nbase_n[1], + b_norm[:, :, 0], ) b_grad[1] = eva2.evaluation_kernel_2d( - pn1, pn2, bn1, der2, span1, span2, nbase_n[0], nbase_n[1], b_norm[:, :, 0] + pn1, + pn2, + bn1, + der2, + span1, + span2, + nbase_n[0], + nbase_n[1], + b_norm[:, :, 0], ) b_grad[2] = 0.0 @@ -562,32 +674,80 @@ def pusher_step5( # equilibrium magnetic field (2-form) b[0] = eva2.evaluation_kernel_2d( - pn1, pd2, bn1, bd2, span1 - 0, span2 - 1, nbase_n[0], nbase_d[1], b_eq_1[:, :, 0] + pn1, + pd2, + bn1, + bd2, + span1 - 0, + span2 - 1, + nbase_n[0], + nbase_d[1], + b_eq_1[:, :, 0], ) b[1] = eva2.evaluation_kernel_2d( - pd1, pn2, bd1, bn2, span1 - 1, span2 - 0, nbase_d[0], nbase_n[1], b_eq_2[:, :, 0] + pd1, + pn2, + bd1, + bn2, + span1 - 1, + span2 - 0, + nbase_d[0], + nbase_n[1], + b_eq_2[:, :, 0], ) b[2] = eva2.evaluation_kernel_2d( - pd1, pd2, bd1, bd2, span1 - 1, span2 - 1, nbase_d[0], nbase_d[1], b_eq_3[:, :, 0] + pd1, + pd2, + bd1, + bd2, + span1 - 1, + span2 - 1, + nbase_d[0], + nbase_d[1], + b_eq_3[:, :, 0], ) # perturbed magnetic field (2-form) for i in range(nbase_n[2]): b[0] += ( eva2.evaluation_kernel_2d( - pn1, pd2, bn1, bd2, span1 - 0, span2 - 1, nbase_n[0], nbase_d[1], b_p_1[:, :, i] + pn1, + pd2, + bn1, + bd2, + span1 - 0, + span2 - 1, + nbase_n[0], + nbase_d[1], + b_p_1[:, :, i], ) * cs[i] ) b[1] += ( eva2.evaluation_kernel_2d( - pd1, pn2, bd1, bn2, span1 - 1, span2 - 0, nbase_d[0], nbase_n[1], b_p_2[:, :, i] + pd1, + pn2, + bd1, + bn2, + span1 - 1, + span2 - 0, + nbase_d[0], + nbase_n[1], + b_p_2[:, :, i], ) * cs[i] ) b[2] += ( eva2.evaluation_kernel_2d( - pd1, pd2, bd1, bd2, span1 - 1, span2 - 1, nbase_d[0], nbase_d[1], b_p_3[:, :, i] + pd1, + pd2, + bd1, + bd2, + span1 - 1, + span2 - 1, + nbase_d[0], + nbase_d[1], + b_p_3[:, :, i], ) * cs[i] ) diff --git a/src/struphy/pic/tests/test_pic_legacy_files/pusher_vel_3d.py b/src/struphy/pic/tests/test_pic_legacy_files/pusher_vel_3d.py index 4eabb26dd..cd3884209 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/pusher_vel_3d.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/pusher_vel_3d.py @@ -227,13 +227,49 @@ def pusher_step3( # velocity field (0-form, push-forward with df) if basis_u == 0: u[0] = eva3.evaluation_kernel_3d( - pn1, pn2, pn3, bn1, bn2, bn3, span1, span2, span3, nbase_n[0], nbase_n[1], nbase_n[2], u1 + pn1, + pn2, + pn3, + bn1, + bn2, + bn3, + span1, + span2, + span3, + nbase_n[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pn2, pn3, bn1, bn2, bn3, span1, span2, span3, nbase_n[0], nbase_n[1], nbase_n[2], u2 + pn1, + pn2, + pn3, + bn1, + bn2, + bn3, + span1, + span2, + span3, + nbase_n[0], + nbase_n[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pn3, bn1, bn2, bn3, span1, span2, span3, nbase_n[0], nbase_n[1], nbase_n[2], u3 + pn1, + pn2, + pn3, + bn1, + bn2, + bn3, + span1, + span2, + span3, + nbase_n[0], + nbase_n[1], + nbase_n[2], + u3, ) linalg.matrix_vector(df, u, u_cart) @@ -241,13 +277,49 @@ def pusher_step3( # velocity field (1-form, push forward with df^(-T)) elif basis_u == 1: u[0] = eva3.evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span1 - 1, span2, span3, nbase_d[0], nbase_n[1], nbase_n[2], u1 + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span1 - 1, + span2, + span3, + nbase_d[0], + nbase_n[1], + nbase_n[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span1, span2 - 1, span3, nbase_n[0], nbase_d[1], nbase_n[2], u2 + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span1, + span2 - 1, + span3, + nbase_n[0], + nbase_d[1], + nbase_n[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span1, span2, span3 - 1, nbase_n[0], nbase_n[1], nbase_d[2], u3 + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span1, + span2, + span3 - 1, + nbase_n[0], + nbase_n[1], + nbase_d[2], + u3, ) linalg.matrix_vector(dfinv_t, u, u_cart) @@ -255,13 +327,49 @@ def pusher_step3( # velocity field (2-form, push forward with df/|det df|) elif basis_u == 2: u[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], u1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + u1, ) u[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], u2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + u2, ) u[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], u3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + u3, ) linalg.matrix_vector(df, u, u_cart) @@ -272,13 +380,49 @@ def pusher_step3( # magnetic field (2-form) b[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], b2_1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + b2_1, ) b[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], b2_2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + b2_2, ) b[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], b2_3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + b2_3, ) # push-forward to physical domain @@ -290,13 +434,49 @@ def pusher_step3( # gradient of absolute value of magnetic field (1-form) b_grad[0] = eva3.evaluation_kernel_3d( - pn1, pn2, pn3, der1, bn2, bn3, span1, span2, span3, nbase_n[0], nbase_n[1], nbase_n[2], b0 + pn1, + pn2, + pn3, + der1, + bn2, + bn3, + span1, + span2, + span3, + nbase_n[0], + nbase_n[1], + nbase_n[2], + b0, ) b_grad[1] = eva3.evaluation_kernel_3d( - pn1, pn2, pn3, bn1, der2, bn3, span1, span2, span3, nbase_n[0], nbase_n[1], nbase_n[2], b0 + pn1, + pn2, + pn3, + bn1, + der2, + bn3, + span1, + span2, + span3, + nbase_n[0], + nbase_n[1], + nbase_n[2], + b0, ) b_grad[2] = eva3.evaluation_kernel_3d( - pn1, pn2, pn3, bn1, bn2, der3, span1, span2, span3, nbase_n[0], nbase_n[1], nbase_n[2], b0 + pn1, + pn2, + pn3, + bn1, + bn2, + der3, + span1, + span2, + span3, + nbase_n[0], + nbase_n[1], + nbase_n[2], + b0, ) # push-forward to physical domain @@ -537,13 +717,49 @@ def pusher_step5_old( # magnetic field (2-form) b[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], b2_1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + b2_1, ) b[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], b2_2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + b2_2, ) b[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], b2_3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + b2_3, ) b_prod[0, 1] = -b[2] @@ -794,13 +1010,49 @@ def pusher_step5( # magnetic field (2-form) b[0] = eva3.evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span1, span2 - 1, span3 - 1, nbase_n[0], nbase_d[1], nbase_d[2], b2_1 + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span1, + span2 - 1, + span3 - 1, + nbase_n[0], + nbase_d[1], + nbase_d[2], + b2_1, ) b[1] = eva3.evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span1 - 1, span2, span3 - 1, nbase_d[0], nbase_n[1], nbase_d[2], b2_2 + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span1 - 1, + span2, + span3 - 1, + nbase_d[0], + nbase_n[1], + nbase_d[2], + b2_2, ) b[2] = eva3.evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span1 - 1, span2 - 1, span3, nbase_d[0], nbase_d[1], nbase_n[2], b2_3 + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span1 - 1, + span2 - 1, + span3, + nbase_d[0], + nbase_d[1], + nbase_n[2], + b2_3, ) # push-forward to physical domain diff --git a/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_3d.py b/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_3d.py index a40d12fc9..7923b3966 100644 --- a/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_3d.py +++ b/src/struphy/pic/tests/test_pic_legacy_files/spline_evaluation_3d.py @@ -127,7 +127,19 @@ def evaluate_n_n_n( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pn1, pn2, pn3, bn1, bn2, bn3, span_n1, span_n2, span_n3, nbase_n1, nbase_n2, nbase_n3, coeff + pn1, + pn2, + pn3, + bn1, + bn2, + bn3, + span_n1, + span_n2, + span_n3, + nbase_n1, + nbase_n2, + nbase_n3, + coeff, ) return value @@ -189,7 +201,19 @@ def evaluate_diffn_n_n( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pn1, pn2, pn3, bn1, bn2, bn3, span_n1, span_n2, span_n3, nbase_n1, nbase_n2, nbase_n3, coeff + pn1, + pn2, + pn3, + bn1, + bn2, + bn3, + span_n1, + span_n2, + span_n3, + nbase_n1, + nbase_n2, + nbase_n3, + coeff, ) return value @@ -251,7 +275,19 @@ def evaluate_n_diffn_n( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pn1, pn2, pn3, bn1, bn2, bn3, span_n1, span_n2, span_n3, nbase_n1, nbase_n2, nbase_n3, coeff + pn1, + pn2, + pn3, + bn1, + bn2, + bn3, + span_n1, + span_n2, + span_n3, + nbase_n1, + nbase_n2, + nbase_n3, + coeff, ) return value @@ -313,7 +349,19 @@ def evaluate_n_n_diffn( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pn1, pn2, pn3, bn1, bn2, bn3, span_n1, span_n2, span_n3, nbase_n1, nbase_n2, nbase_n3, coeff + pn1, + pn2, + pn3, + bn1, + bn2, + bn3, + span_n1, + span_n2, + span_n3, + nbase_n1, + nbase_n2, + nbase_n3, + coeff, ) return value @@ -377,7 +425,19 @@ def evaluate_d_n_n( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pd1, pn2, pn3, bd1, bn2, bn3, span_d1, span_n2, span_n3, nbase_d1, nbase_n2, nbase_n3, coeff + pd1, + pn2, + pn3, + bd1, + bn2, + bn3, + span_d1, + span_n2, + span_n3, + nbase_d1, + nbase_n2, + nbase_n3, + coeff, ) return value @@ -441,7 +501,19 @@ def evaluate_n_d_n( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pn1, pd2, pn3, bn1, bd2, bn3, span_n1, span_d2, span_n3, nbase_n1, nbase_d2, nbase_n3, coeff + pn1, + pd2, + pn3, + bn1, + bd2, + bn3, + span_n1, + span_d2, + span_n3, + nbase_n1, + nbase_d2, + nbase_n3, + coeff, ) return value @@ -505,7 +577,19 @@ def evaluate_n_n_d( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pn1, pn2, pd3, bn1, bn2, bd3, span_n1, span_n2, span_d3, nbase_n1, nbase_n2, nbase_d3, coeff + pn1, + pn2, + pd3, + bn1, + bn2, + bd3, + span_n1, + span_n2, + span_d3, + nbase_n1, + nbase_n2, + nbase_d3, + coeff, ) return value @@ -570,7 +654,19 @@ def evaluate_n_d_d( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pn1, pd2, pd3, bn1, bd2, bd3, span_n1, span_d2, span_d3, nbase_n1, nbase_d2, nbase_d3, coeff + pn1, + pd2, + pd3, + bn1, + bd2, + bd3, + span_n1, + span_d2, + span_d3, + nbase_n1, + nbase_d2, + nbase_d3, + coeff, ) return value @@ -635,7 +731,19 @@ def evaluate_d_n_d( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pd1, pn2, pd3, bd1, bn2, bd3, span_d1, span_n2, span_d3, nbase_d1, nbase_n2, nbase_d3, coeff + pd1, + pn2, + pd3, + bd1, + bn2, + bd3, + span_d1, + span_n2, + span_d3, + nbase_d1, + nbase_n2, + nbase_d3, + coeff, ) return value @@ -700,7 +808,19 @@ def evaluate_d_d_n( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pd1, pd2, pn3, bd1, bd2, bn3, span_d1, span_d2, span_n3, nbase_d1, nbase_d2, nbase_n3, coeff + pd1, + pd2, + pn3, + bd1, + bd2, + bn3, + span_d1, + span_d2, + span_n3, + nbase_d1, + nbase_d2, + nbase_n3, + coeff, ) return value @@ -766,7 +886,19 @@ def evaluate_d_d_d( # sum up non-vanishing contributions value = evaluation_kernel_3d( - pd1, pd2, pd3, bd1, bd2, bd3, span_d1, span_d2, span_d3, nbase_d1, nbase_d2, nbase_d3, coeff + pd1, + pd2, + pd3, + bd1, + bd2, + bd3, + span_d1, + span_d2, + span_d3, + nbase_d1, + nbase_d2, + nbase_d3, + coeff, ) return value @@ -815,41 +947,137 @@ def evaluate_tensor_product( # V0 - space if kind == 0: values[i1, i2, i3] = evaluate_n_n_n( - t1, t2, t3, p1, p2, p3, nbase_1, nbase_2, nbase_3, coeff, eta1[i1], eta2[i2], eta3[i3] + t1, + t2, + t3, + p1, + p2, + p3, + nbase_1, + nbase_2, + nbase_3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) # V1 - space elif kind == 11: values[i1, i2, i3] = evaluate_d_n_n( - t1, t2, t3, p1, p2, p3, nbase_1, nbase_2, nbase_3, coeff, eta1[i1], eta2[i2], eta3[i3] + t1, + t2, + t3, + p1, + p2, + p3, + nbase_1, + nbase_2, + nbase_3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 12: values[i1, i2, i3] = evaluate_n_d_n( - t1, t2, t3, p1, p2, p3, nbase_1, nbase_2, nbase_3, coeff, eta1[i1], eta2[i2], eta3[i3] + t1, + t2, + t3, + p1, + p2, + p3, + nbase_1, + nbase_2, + nbase_3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 13: values[i1, i2, i3] = evaluate_n_n_d( - t1, t2, t3, p1, p2, p3, nbase_1, nbase_2, nbase_3, coeff, eta1[i1], eta2[i2], eta3[i3] + t1, + t2, + t3, + p1, + p2, + p3, + nbase_1, + nbase_2, + nbase_3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) # V2 - space elif kind == 21: values[i1, i2, i3] = evaluate_n_d_d( - t1, t2, t3, p1, p2, p3, nbase_1, nbase_2, nbase_3, coeff, eta1[i1], eta2[i2], eta3[i3] + t1, + t2, + t3, + p1, + p2, + p3, + nbase_1, + nbase_2, + nbase_3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 22: values[i1, i2, i3] = evaluate_d_n_d( - t1, t2, t3, p1, p2, p3, nbase_1, nbase_2, nbase_3, coeff, eta1[i1], eta2[i2], eta3[i3] + t1, + t2, + t3, + p1, + p2, + p3, + nbase_1, + nbase_2, + nbase_3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) elif kind == 23: values[i1, i2, i3] = evaluate_d_d_n( - t1, t2, t3, p1, p2, p3, nbase_1, nbase_2, nbase_3, coeff, eta1[i1], eta2[i2], eta3[i3] + t1, + t2, + t3, + p1, + p2, + p3, + nbase_1, + nbase_2, + nbase_3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) # V3 - space elif kind == 3: values[i1, i2, i3] = evaluate_d_d_d( - t1, t2, t3, p1, p2, p3, nbase_1, nbase_2, nbase_3, coeff, eta1[i1], eta2[i2], eta3[i3] + t1, + t2, + t3, + p1, + p2, + p3, + nbase_1, + nbase_2, + nbase_3, + coeff, + eta1[i1], + eta2[i2], + eta3[i3], ) diff --git a/src/struphy/pic/tests/test_pushers.py b/src/struphy/pic/tests/test_pushers.py index d64076cd1..321ab9aba 100644 --- a/src/struphy/pic/tests/test_pushers.py +++ b/src/struphy/pic/tests/test_pushers.py @@ -6,7 +6,8 @@ @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @pytest.mark.parametrize("p", [[2, 3, 1], [1, 2, 3]]) @pytest.mark.parametrize( - "spl_kind", [[False, True, True], [True, False, True], [False, False, True], [True, True, True]] + "spl_kind", + [[False, True, True], [True, False, True], [False, False, True], [True, True, True]], ) @pytest.mark.parametrize( "mapping", @@ -141,7 +142,8 @@ def test_push_vxb_analytic(Nel, p, spl_kind, mapping, show_plots=False): @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @pytest.mark.parametrize("p", [[2, 3, 1], [1, 2, 3]]) @pytest.mark.parametrize( - "spl_kind", [[False, True, True], [True, False, True], [False, False, True], [True, True, True]] + "spl_kind", + [[False, True, True], [True, False, True], [False, False, True], [True, True, True]], ) @pytest.mark.parametrize( "mapping", @@ -287,7 +289,8 @@ def test_push_bxu_Hdiv(Nel, p, spl_kind, mapping, show_plots=False): @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @pytest.mark.parametrize("p", [[2, 3, 1], [1, 2, 3]]) @pytest.mark.parametrize( - "spl_kind", [[False, True, True], [True, False, True], [False, False, True], [True, True, True]] + "spl_kind", + [[False, True, True], [True, False, True], [False, False, True], [True, True, True]], ) @pytest.mark.parametrize( "mapping", @@ -433,7 +436,8 @@ def test_push_bxu_Hcurl(Nel, p, spl_kind, mapping, show_plots=False): @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @pytest.mark.parametrize("p", [[2, 3, 1], [1, 2, 3]]) @pytest.mark.parametrize( - "spl_kind", [[False, True, True], [True, False, True], [False, False, True], [True, True, True]] + "spl_kind", + [[False, True, True], [True, False, True], [False, False, True], [True, True, True]], ) @pytest.mark.parametrize( "mapping", @@ -579,7 +583,8 @@ def test_push_bxu_H1vec(Nel, p, spl_kind, mapping, show_plots=False): @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @pytest.mark.parametrize("p", [[2, 3, 1], [1, 2, 3]]) @pytest.mark.parametrize( - "spl_kind", [[False, True, True], [True, False, True], [False, False, True], [True, True, True]] + "spl_kind", + [[False, True, True], [True, False, True], [False, False, True], [True, True, True]], ) @pytest.mark.parametrize( "mapping", @@ -727,7 +732,8 @@ def test_push_bxu_Hdiv_pauli(Nel, p, spl_kind, mapping, show_plots=False): @pytest.mark.parametrize("Nel", [[8, 9, 5], [7, 8, 9]]) @pytest.mark.parametrize("p", [[2, 3, 1], [1, 2, 3]]) @pytest.mark.parametrize( - "spl_kind", [[False, True, True], [True, False, True], [False, False, True], [True, True, True]] + "spl_kind", + [[False, True, True], [True, False, True], [False, False, True], [True, True, True]], ) @pytest.mark.parametrize( "mapping", @@ -880,7 +886,11 @@ def test_push_eta_rk4(Nel, p, spl_kind, mapping, show_plots=False): if __name__ == "__main__": test_push_vxb_analytic( - [8, 9, 5], [4, 2, 3], [False, True, True], ["Colella", {"Lx": 2.0, "Ly": 2.0, "alpha": 0.1, "Lz": 4.0}], False + [8, 9, 5], + [4, 2, 3], + [False, True, True], + ["Colella", {"Lx": 2.0, "Ly": 2.0, "alpha": 0.1, "Lz": 4.0}], + False, ) # test_push_bxu_Hdiv([8, 9, 5], [4, 2, 3], [False, True, True], ['Colella', { # 'Lx': 2., 'Ly': 2., 'alpha': 0.1, 'Lz': 4.}], False) diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index 337dfaede..a11c9600e 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -73,7 +73,8 @@ def test_flattening(nx, ny, nz, algo): @pytest.mark.parametrize("Nel", [[8, 9, 10]]) @pytest.mark.parametrize("p", [[2, 3, 4]]) @pytest.mark.parametrize( - "spl_kind", [[False, False, True], [False, True, False], [True, False, True], [True, True, False]] + "spl_kind", + [[False, False, True], [False, True, False], [True, False, True], [True, True, False]], ) @pytest.mark.parametrize( "mapping", diff --git a/src/struphy/pic/tests/test_sph.py b/src/struphy/pic/tests/test_sph.py index 64b7b0cc3..294e7f9dc 100644 --- a/src/struphy/pic/tests/test_sph.py +++ b/src/struphy/pic/tests/test_sph.py @@ -111,9 +111,9 @@ def test_sph_evaluation_1d( err_max_norm = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) if rank == 0: - print(f"\n{boxes_per_dim = }") - print(f"{kernel = }, {derivative =}") - print(f"{bc_x = }, {eval_pts = }, {tesselation = }, {err_max_norm = }") + print(f"\n{boxes_per_dim =}") + print(f"{kernel =}, {derivative =}") + print(f"{bc_x =}, {eval_pts =}, {tesselation =}, {err_max_norm =}") if show_plot: plt.figure(figsize=(12, 8)) plt.plot(ee1.squeeze(), fun_exact(ee1, ee2, ee3).squeeze(), label="exact") @@ -235,9 +235,9 @@ def test_sph_evaluation_2d( err_max_norm = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) if rank == 0: - print(f"\n{boxes_per_dim = }") - print(f"{kernel = }, {derivative =}") - print(f"{bc_x = }, {bc_y = }, {eval_pts = }, {tesselation = }, {err_max_norm = }") + print(f"\n{boxes_per_dim =}") + print(f"{kernel =}, {derivative =}") + print(f"{bc_x =}, {bc_y =}, {eval_pts =}, {tesselation =}, {err_max_norm =}") if show_plot: plt.figure(figsize=(12, 24)) plt.subplot(2, 1, 1) @@ -354,28 +354,28 @@ def test_sph_evaluation_3d( err_max_norm = xp.max(xp.abs(all_eval - exact_eval)) if rank == 0: - print(f"\n{boxes_per_dim = }") - print(f"{kernel = }, {derivative =}") - print(f"{bc_x = }, {bc_y = }, {bc_z = }, {eval_pts = }, {tesselation = }, {err_max_norm = }") + print(f"\n{boxes_per_dim =}") + print(f"{kernel =}, {derivative =}") + print(f"{bc_x =}, {bc_y =}, {bc_z =}, {eval_pts =}, {tesselation =}, {err_max_norm =}") if show_plot: - print(f"\n{fun_exact(ee1, ee2, ee3)[5, 5, 5] = }") - print(f"{ee1[5, 5, 5] = }, {ee2[5, 5, 5] = }, {ee3[5, 5, 5] = }") - print(f"{all_eval[5, 5, 5] = }") + print(f"\n{fun_exact(ee1, ee2, ee3)[5, 5, 5] =}") + print(f"{ee1[5, 5, 5] =}, {ee2[5, 5, 5] =}, {ee3[5, 5, 5] =}") + print(f"{all_eval[5, 5, 5] =}") - print(f"\n{ee1[4, 4, 4] = }, {ee2[4, 4, 4] = }, {ee3[4, 4, 4] = }") - print(f"{all_eval[4, 4, 4] = }") + print(f"\n{ee1[4, 4, 4] =}, {ee2[4, 4, 4] =}, {ee3[4, 4, 4] =}") + print(f"{all_eval[4, 4, 4] =}") - print(f"\n{ee1[3, 3, 3] = }, {ee2[3, 3, 3] = }, {ee3[3, 3, 3] = }") - print(f"{all_eval[3, 3, 3] = }") + print(f"\n{ee1[3, 3, 3] =}, {ee2[3, 3, 3] =}, {ee3[3, 3, 3] =}") + print(f"{all_eval[3, 3, 3] =}") - print(f"\n{ee1[2, 2, 2] = }, {ee2[2, 2, 2] = }, {ee3[2, 2, 2] = }") - print(f"{all_eval[2, 2, 2] = }") + print(f"\n{ee1[2, 2, 2] =}, {ee2[2, 2, 2] =}, {ee3[2, 2, 2] =}") + print(f"{all_eval[2, 2, 2] =}") - print(f"\n{ee1[1, 1, 1] = }, {ee2[1, 1, 1] = }, {ee3[1, 1, 1] = }") - print(f"{all_eval[1, 1, 1] = }") + print(f"\n{ee1[1, 1, 1] =}, {ee2[1, 1, 1] =}, {ee3[1, 1, 1] =}") + print(f"{all_eval[1, 1, 1] =}") - print(f"\n{ee1[0, 0, 0] = }, {ee2[0, 0, 0] = }, {ee3[0, 0, 0] = }") - print(f"{all_eval[0, 0, 0] = }") + print(f"\n{ee1[0, 0, 0] =}, {ee2[0, 0, 0] =}, {ee3[0, 0, 0] =}") + print(f"{all_eval[0, 0, 0] =}") # plt.figure(figsize=(12, 24)) # plt.subplot(2, 1, 1) # plt.pcolor(ee1[0, :, :], ee2[0, :, :], fun_exact(ee1, ee2, ee3)[0, :, :]) @@ -478,12 +478,12 @@ def test_evaluation_SPH_Np_convergence_1d(boxes_per_dim, bc_x, eval_pts, tessela plt.figure() plt.plot(ee1.squeeze(), exact_eval.squeeze(), label="exact") plt.plot(ee1.squeeze(), all_eval.squeeze(), "--.", label="eval_sph") - plt.title(f"{Np = }, {ppb = }") + plt.title(f"{Np =}, {ppb =}") # plt.savefig(f"fun_{Np}_{ppb}.png") diff = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) err_vec += [diff] - print(f"{Np = }, {ppb = }, {diff = }") + print(f"{Np =}, {ppb =}, {diff =}") if tesselation: fit = xp.polyfit(xp.log(ppbs), xp.log(err_vec), 1) @@ -501,7 +501,7 @@ def test_evaluation_SPH_Np_convergence_1d(boxes_per_dim, bc_x, eval_pts, tessela # plt.savefig(f"Convergence_SPH_{tesselation=}") if rank == 0: - print(f"\n{bc_x = }, {eval_pts = }, {tesselation = }, {fit[0] = }") + print(f"\n{bc_x =}, {eval_pts =}, {tesselation =}, {fit[0] =}") if tesselation: assert fit[0] < 2e-3 @@ -594,13 +594,13 @@ def test_evaluation_SPH_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselat plt.figure() plt.plot(ee1.squeeze(), exact_eval.squeeze(), label="exact") plt.plot(ee1.squeeze(), all_eval.squeeze(), "--.", label="eval_sph") - plt.title(f"{h1 = }") + plt.title(f"{h1 =}") # plt.savefig(f"fun_{h1}.png") # error in max-norm diff = xp.max(xp.abs(all_eval - exact_eval)) / xp.max(xp.abs(exact_eval)) - print(f"{h1 = }, {diff = }") + print(f"{h1 =}, {diff =}") if tesselation and h1 < 0.256: assert diff < 0.036 @@ -621,7 +621,7 @@ def test_evaluation_SPH_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, tesselat # plt.savefig("Convergence_SPH") if rank == 0: - print(f"\n{bc_x = }, {eval_pts = }, {tesselation = }, {fit[0] = }") + print(f"\n{bc_x =}, {eval_pts =}, {tesselation =}, {fit[0] =}") if not tesselation: assert xp.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate @@ -718,7 +718,7 @@ def test_evaluation_mc_Np_and_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, te err_vec[-1] += [diff] if rank == 0: - print(f"{Np = }, {ppb = }, {diff = }") + print(f"{Np =}, {ppb =}, {diff =}") # if show_plot: # plt.figure() # plt.plot(ee1.squeeze(), fun_exact(ee1, ee2, ee3).squeeze(), label="exact") @@ -756,7 +756,7 @@ def test_evaluation_mc_Np_and_h_convergence_1d(boxes_per_dim, bc_x, eval_pts, te plt.show() if rank == 0: - print(f"\n{tesselation = }, {bc_x = }, {err_min = }") + print(f"\n{tesselation =}, {bc_x =}, {err_min =}") if tesselation: if bc_x == "periodic": @@ -879,14 +879,14 @@ def test_evaluation_SPH_Np_convergence_2d(boxes_per_dim, bc_x, bc_y, tesselation assert diff < 0.06 if rank == 0: - print(f"{Np = }, {ppb = }, {diff = }") + print(f"{Np =}, {ppb =}, {diff =}") if show_plot: fig, ax = plt.subplots() d = ax.pcolor(ee1.squeeze(), ee2.squeeze(), all_eval.squeeze(), label="eval_sph", vmin=1.0, vmax=2.0) fig.colorbar(d, ax=ax, label="2d_SPH") ax.set_xlabel("ee1") ax.set_ylabel("ee2") - ax.set_title(f"{Np}_{ppb = }") + ax.set_title(f"{Np}_{ppb =}") # fig.savefig(f"2d_sph_{Np}_{ppb}.png") if tesselation: @@ -905,7 +905,7 @@ def test_evaluation_SPH_Np_convergence_2d(boxes_per_dim, bc_x, bc_y, tesselation # plt.savefig(f"Convergence_SPH_{tesselation=}") if rank == 0: - print(f"\n{bc_x = }, {tesselation = }, {fit[0] = }") + print(f"\n{bc_x =}, {tesselation =}, {fit[0] =}") if not tesselation: assert xp.abs(fit[0] + 0.5) < 0.1 # Monte Carlo rate diff --git a/src/struphy/pic/tests/test_tesselation.py b/src/struphy/pic/tests/test_tesselation.py index cf6ed922e..b138af50a 100644 --- a/src/struphy/pic/tests/test_tesselation.py +++ b/src/struphy/pic/tests/test_tesselation.py @@ -176,7 +176,7 @@ def test_cell_average(ppb, nx, ny, nz, n_quad, show_plot=False): plt.show() # test - print(f"\n{rank = }, {xp.max(xp.abs(particles.weights - particles.f_init(particles.positions))) = }") + print(f"\n{rank =}, {xp.max(xp.abs(particles.weights - particles.f_init(particles.positions))) =}") assert xp.max(xp.abs(particles.weights - particles.f_init(particles.positions))) < 0.012 diff --git a/src/struphy/polar/basic.py b/src/struphy/polar/basic.py index 6a1c2c2f2..f737c671e 100644 --- a/src/struphy/polar/basic.py +++ b/src/struphy/polar/basic.py @@ -347,7 +347,7 @@ def toarray(self, allreduce=False): out2, self.pol[2].flatten(), out3, - ) + ), ) return out diff --git a/src/struphy/polar/extraction_operators.py b/src/struphy/polar/extraction_operators.py index f58c88f59..1c2a461ce 100644 --- a/src/struphy/polar/extraction_operators.py +++ b/src/struphy/polar/extraction_operators.py @@ -72,7 +72,7 @@ def __init__(self, domain, derham): ((self.cx[1] - self.pole[0]) * (-2)).max(), ((self.cx[1] - self.pole[0]) - xp.sqrt(3) * (self.cy[1] - self.pole[1])).max(), ((self.cx[1] - self.pole[0]) + xp.sqrt(3) * (self.cy[1] - self.pole[1])).max(), - ] + ], ) # barycentric coordinates @@ -692,7 +692,7 @@ def __init__(self, cx, cy): (-2 * (cx[1] - self.x0)).max(), ((cx[1] - self.x0) - xp.sqrt(3) * (cy[1] - self.y0)).max(), ((cx[1] - self.x0) + xp.sqrt(3) * (cy[1] - self.y0)).max(), - ] + ], ).max() self.Xi_0 = xp.zeros((3, n1), dtype=float) @@ -966,7 +966,7 @@ def __init__(self, tensor_space, cx, cy): # size of control triangle self.tau = xp.array( - [(-2 * cx[1]).max(), (cx[1] - xp.sqrt(3) * cy[1]).max(), (cx[1] + xp.sqrt(3) * cy[1]).max()] + [(-2 * cx[1]).max(), (cx[1] - xp.sqrt(3) * cy[1]).max(), (cx[1] + xp.sqrt(3) * cy[1]).max()], ).max() self.Xi_0 = xp.zeros((3, n1), dtype=float) @@ -982,7 +982,8 @@ def __init__(self, tensor_space, cx, cy): # =========== extraction operators for discrete 0-forms ================== # extraction operator for basis functions self.E0_pol = spa.bmat( - [[xp.hstack((self.Xi_0, self.Xi_1)), None], [None, spa.identity((n0 - 2) * n1)]], format="csr" + [[xp.hstack((self.Xi_0, self.Xi_1)), None], [None, spa.identity((n0 - 2) * n1)]], + format="csr", ) self.E0 = spa.kron(self.E0_pol, spa.identity(n2), format="csr") @@ -1215,7 +1216,7 @@ def __init__(self, tensor_space, cx, cy): [ [None, -spa.kron(spa.identity((n0 - 2) * d1 + 2), grad_1d_3)], [spa.kron(spa.identity((d0 - 1) * n1), grad_1d_3), None], - ] + ], ) # total polar curl diff --git a/src/struphy/polar/linear_operators.py b/src/struphy/polar/linear_operators.py index 03a3e504a..0e37c7e76 100644 --- a/src/struphy/polar/linear_operators.py +++ b/src/struphy/polar/linear_operators.py @@ -334,7 +334,14 @@ class PolarLinearOperator(LinOpWithTransp): """ def __init__( - self, V, W, tp_operator=None, blocks_pol_to_ten=None, blocks_pol_to_pol=None, blocks_e3=None, transposed=False + self, + V, + W, + tp_operator=None, + blocks_pol_to_ten=None, + blocks_pol_to_pol=None, + blocks_e3=None, + transposed=False, ): assert isinstance(V, PolarDerhamSpace) assert isinstance(W, PolarDerhamSpace) diff --git a/src/struphy/post_processing/likwid/plot_likwidproject.py b/src/struphy/post_processing/likwid/plot_likwidproject.py index feda5d3b6..cde2a2b76 100644 --- a/src/struphy/post_processing/likwid/plot_likwidproject.py +++ b/src/struphy/post_processing/likwid/plot_likwidproject.py @@ -818,7 +818,7 @@ def load_projects(data_paths, procs_per_clone="any"): ) if (procs_per_clone != "any") and (procs_per_clone != project.procs_per_clone): print( - f"Incorrect number of procs_per_clone: {project.procs_per_clone = } {procs_per_clone = }", + f"Incorrect number of procs_per_clone: {project.procs_per_clone =} {procs_per_clone =}", ) continue project.read_project() diff --git a/src/struphy/post_processing/likwid/plot_time_traces.py b/src/struphy/post_processing/likwid/plot_time_traces.py index ed0a34010..f97681ffa 100644 --- a/src/struphy/post_processing/likwid/plot_time_traces.py +++ b/src/struphy/post_processing/likwid/plot_time_traces.py @@ -54,7 +54,7 @@ def plot_time_vs_duration( plt.figure(figsize=(10, 6)) for path in paths: - print(f"{path = }") + print(f"{path =}") with open(path, "rb") as file: profiling_data = pickle.load(file) @@ -204,7 +204,7 @@ def plot_gantt_chart_plotly( Start=start_times[i], Finish=end_times[i], Duration=durations[i], - ) + ), ) if len(bars) == 0: @@ -226,7 +226,7 @@ def plot_gantt_chart_plotly( name=bar["Rank"], marker_color=rank_color_map[bar["Rank"]], hovertemplate=f"Rank: {bar['Rank']}
Start: {bar['Start']:.3f}s
Duration: {bar['Duration']:.3f}s", - ) + ), ) fig.update_layout( diff --git a/src/struphy/post_processing/pproc_struphy.py b/src/struphy/post_processing/pproc_struphy.py index 044ec0de8..940c94ab3 100644 --- a/src/struphy/post_processing/pproc_struphy.py +++ b/src/struphy/post_processing/pproc_struphy.py @@ -96,12 +96,17 @@ def main( fields, t_grid = pproc.create_femfields(path, params_in, step=step) point_data, grids_log, grids_phy = pproc.eval_femfields( - params_in, fields, celldivide=[celldivide, celldivide, celldivide] + params_in, + fields, + celldivide=[celldivide, celldivide, celldivide], ) if physical: point_data_phy, grids_log, grids_phy = pproc.eval_femfields( - params_in, fields, celldivide=[celldivide, celldivide, celldivide], physical=True + params_in, + fields, + celldivide=[celldivide, celldivide, celldivide], + physical=True, ) # directory for field data @@ -196,14 +201,19 @@ def main( libpath = struphy.__path__[0] parser = argparse.ArgumentParser( - description="Post-process data of finished Struphy runs to prepare for diagnostics." + description="Post-process data of finished Struphy runs to prepare for diagnostics.", ) # paths of simulation folders parser.add_argument("dir", type=str, metavar="DIR", help="absolute path of simulation ouput folder to post-process") parser.add_argument( - "-s", "--step", type=int, metavar="N", help="do post-processing every N-th time step (default=1)", default=1 + "-s", + "--step", + type=int, + metavar="N", + help="do post-processing every N-th time step (default=1)", + default=1, ) parser.add_argument( @@ -221,11 +231,15 @@ def main( ) parser.add_argument( - "--guiding-center", help="compute guiding-center coordinates (only from Particles6D)", action="store_true" + "--guiding-center", + help="compute guiding-center coordinates (only from Particles6D)", + action="store_true", ) parser.add_argument( - "--classify", help="classify guiding-center trajectories (passing, trapped or lost)", action="store_true" + "--classify", + help="classify guiding-center trajectories (passing, trapped or lost)", + action="store_true", ) parser.add_argument("--no-vtk", help="whether vtk files creation should be skipped", action="store_true") diff --git a/src/struphy/post_processing/profile_struphy.py b/src/struphy/post_processing/profile_struphy.py index da4632555..43d4be47d 100644 --- a/src/struphy/post_processing/profile_struphy.py +++ b/src/struphy/post_processing/profile_struphy.py @@ -93,7 +93,7 @@ def main(): + "ncalls".ljust(15) + "totime".ljust(15) + "percall".ljust(15) - + "cumtime".ljust(15) + + "cumtime".ljust(15), ) print("-" * 154) for position, key in enumerate(dicts[0].keys()): @@ -157,7 +157,7 @@ def main(): plt.xlabel("mpi_size") plt.ylabel("time [s]") plt.title( - "Weak scaling for cells/mpi_size=" + str(xp.prod(val["Nel"][0]) / val["mpi_size"][0]) + "=const." + "Weak scaling for cells/mpi_size=" + str(xp.prod(val["Nel"][0]) / val["mpi_size"][0]) + "=const.", ) plt.legend(loc="upper left") # plt.loglog(val['mpi_size'], val['time'][0]*xp.ones_like(val['time']), 'k--', alpha=0.3) diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index f69d4c7fe..945107bbd 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -259,7 +259,7 @@ def add_init_kernel( column_nr, comps, args_init, - ) + ), ] def add_eval_kernel( @@ -314,5 +314,5 @@ def add_eval_kernel( column_nr, comps, args_eval, - ) + ), ] diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index 76c2396ef..132cd8dbf 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -1029,7 +1029,7 @@ def allocate(self): kernel = Pyccelkernel(pusher_kernels.push_bxu_H1vec) else: raise ValueError( - f'{self.options.u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', + f'{self.options.u_space =} not valid, choose from "Hcurl", "Hdiv" or "H1vec.', ) # instantiate Pusher @@ -1333,7 +1333,7 @@ def allocate(self): pusher_kernel = Pyccelkernel(pusher_kernels_gc.push_gc_cc_J1_H1vec) else: raise ValueError( - f'{self.options.u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', + f'{self.options.u_space =} not valid, choose from "Hcurl", "Hdiv" or "H1vec.', ) args_pusher_kernel = ( @@ -1623,7 +1623,7 @@ def allocate(self): self._pusher_kernel = pusher_kernels_gc.push_gc_cc_J2_stage_H1vec else: raise ValueError( - f'{self.options.u_space = } not valid, choose from "Hdiv" or "H1vec.', + f'{self.options.u_space =} not valid, choose from "Hdiv" or "H1vec.', ) # temp fix due to refactoring of ButcherTableau: @@ -1970,7 +1970,7 @@ def __call__(self, dt): sum_u_diff_loc = xp.sum((u_diff.toarray() ** 2)) sum_H_diff_loc = xp.sum( - (markers[~holes, :3] - markers[~holes, first_init_idx : first_init_idx + 3]) ** 2 + (markers[~holes, :3] - markers[~holes, first_init_idx : first_init_idx + 3]) ** 2, ) buffer_array = xp.array([sum_u_diff_loc]) @@ -2064,7 +2064,7 @@ def __call__(self, dt): ) sum_H_diff_loc = xp.sum( - xp.abs(markers[~holes, 0:3] - markers[~holes, first_free_idx : first_free_idx + 3]) + xp.abs(markers[~holes, 0:3] - markers[~holes, first_free_idx : first_free_idx + 3]), ) if particles.mpi_comm is not None: @@ -2146,7 +2146,7 @@ def __call__(self, dt): if iter_num == self.options.dg_solver_params.maxiter: if self.options.dg_solver_params.info and MPI.COMM_WORLD.Get_rank() == 0: print( - f"{iter_num = }, maxiter={self.options.dg_solver_params.maxiter} reached! diff: {diff}, e_diff: {e_diff}", + f"{iter_num =}, maxiter={self.options.dg_solver_params.maxiter} reached! diff: {diff}, e_diff: {e_diff}", ) if particles.mpi_comm is not None: particles.mpi_comm.Barrier() diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index f7ac3a3c3..c3f3e1381 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -2594,7 +2594,7 @@ def allocate(self): if xp.abs(self.options.sigma_1) < 1e-14: self.options.sigma_1 = 1e-14 if MPI.COMM_WORLD.Get_rank() == 0: - print(f"Stabilizing Poisson solve with {self.options.sigma_1 = }") + print(f"Stabilizing Poisson solve with {self.options.sigma_1 =}") # model parameters self._sigma_1 = self.options.sigma_1 @@ -2617,7 +2617,7 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: elif isinstance(rho, Callable): rhs = L2Projector("H1", self.mass_ops).get_dofs(rho, apply_bc=True) else: - raise TypeError(f"{type(rho) = } is not accepted.") + raise TypeError(f"{type(rho) =} is not accepted.") return rhs @@ -3029,7 +3029,7 @@ def __call_newton(self, dt): if it == self.options.nonlin_solver.maxiter - 1 or xp.isnan(err): print( - f"!!!WARNING: Maximum iteration in VariationalMomentumAdvection reached - not converged \n {err = } \n {tol**2 = }", + f"!!!WARNING: Maximum iteration in VariationalMomentumAdvection reached - not converged \n {err =} \n {tol**2 =}", ) self.update_feec_variables(u=un1) @@ -3075,7 +3075,7 @@ def __call_picard(self, dt): if it == self.options.nonlin_solver.maxiter - 1 or xp.isnan(err): print( - f"!!!WARNING: Maximum iteration in VariationalMomentumAdvection reached - not converged \n {err = } \n {tol**2 = }", + f"!!!WARNING: Maximum iteration in VariationalMomentumAdvection reached - not converged \n {err =} \n {tol**2 =}", ) self.update_feec_variables(u=un1) @@ -3427,7 +3427,7 @@ def __call_newton(self, dt): if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( - f"!!!Warning: Maximum iteration in VariationalDensityEvolve reached - not converged:\n {err = } \n {tol**2 = }", + f"!!!Warning: Maximum iteration in VariationalDensityEvolve reached - not converged:\n {err =} \n {tol**2 =}", ) self._tmp_un_diff = un1 - un @@ -3605,11 +3605,15 @@ def _update_linear_form_dl_drho(self, rhon, rhon1, un, un1, sn): def _compute_init_linear_form(self): if abs(self._gamma - 5 / 3) < 1e-3: self._energy_evaluator.evaluate_exact_de_drho_grid( - self.projected_equil.n3, self.projected_equil.s3_monoatomic, out=self._init_dener_drho + self.projected_equil.n3, + self.projected_equil.s3_monoatomic, + out=self._init_dener_drho, ) elif abs(self._gamma - 7 / 5) < 1e-3: self._energy_evaluator.evaluate_exact_de_drho_grid( - self.projected_equil.n3, self.projected_equil.s3_diatomic, out=self._init_dener_drho + self.projected_equil.n3, + self.projected_equil.s3_diatomic, + out=self._init_dener_drho, ) else: raise ValueError("Gamma should be 7/5 or 5/3 for if you want to linearize") @@ -3890,7 +3894,7 @@ def __call_newton(self, dt): if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( - f"!!!Warning: Maximum iteration in VariationalEntropyEvolve reached - not converged:\n {err = } \n {tol**2 = }", + f"!!!Warning: Maximum iteration in VariationalEntropyEvolve reached - not converged:\n {err =} \n {tol**2 =}", ) self._tmp_sn_diff = sn1 - sn self._tmp_un_diff = un1 - un @@ -4024,11 +4028,15 @@ def _update_linear_form_dl_ds(self, rhon, sn, sn1): def _compute_init_linear_form(self): if abs(self._gamma - 5 / 3) < 1e-3: self._energy_evaluator.evaluate_exact_de_ds_grid( - self.projected_equil.n3, self.projected_equil.s3_monoatomic, out=self._init_dener_ds + self.projected_equil.n3, + self.projected_equil.s3_monoatomic, + out=self._init_dener_ds, ) elif abs(self._gamma - 7 / 5) < 1e-3: self._energy_evaluator.evaluate_exact_de_ds_grid( - self.projected_equil.n3, self.projected_equil.s3_diatomic, out=self._init_dener_ds + self.projected_equil.n3, + self.projected_equil.s3_diatomic, + out=self._init_dener_ds, ) else: raise ValueError("Gamma should be 7/5 or 5/3 for if you want to linearize") @@ -4310,7 +4318,7 @@ def __call_newton(self, dt): if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( - f"!!!Warning: Maximum iteration in VariationalMagFieldEvolve reached - not converged:\n {err = } \n {tol**2 = }", + f"!!!Warning: Maximum iteration in VariationalMagFieldEvolve reached - not converged:\n {err =} \n {tol**2 =}", ) self._tmp_un_diff = un1 - un @@ -4813,7 +4821,7 @@ def __call_picard(self, dt): if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( - f"!!!Warning: Maximum iteration in VariationalPBEvolve reached - not converged:\n {err = } \n {tol**2 = }", + f"!!!Warning: Maximum iteration in VariationalPBEvolve reached - not converged:\n {err =} \n {tol**2 =}", ) self._tmp_un_diff = un1 - un @@ -5404,7 +5412,7 @@ def __call_picard(self, dt): if it == self._nonlin_solver.maxiter - 1 or xp.isnan(err): print( - f"!!!Warning: Maximum iteration in VariationalPBEvolve reached - not converged:\n {err = } \n {tol**2 = }", + f"!!!Warning: Maximum iteration in VariationalPBEvolve reached - not converged:\n {err =} \n {tol**2 =}", ) self._tmp_un_diff = un1 - un @@ -5996,7 +6004,7 @@ def __call_newton(self, dt): if it == self._nonlin_solver["maxiter"] - 1 or xp.isnan(err): print( - f"!!!Warning: Maximum iteration in VariationalViscosity reached - not converged:\n {err = } \n {tol**2 = }", + f"!!!Warning: Maximum iteration in VariationalViscosity reached - not converged:\n {err =} \n {tol**2 =}", ) self.update_feec_variables(s=sn1, u=un1) @@ -6760,7 +6768,7 @@ def __call_newton(self, dt): if it == self._nonlin_solver["maxiter"] - 1 or xp.isnan(err): print( - f"!!!Warning: Maximum iteration in VariationalResistivity reached - not converged:\n {err = } \n {tol**2 = }", + f"!!!Warning: Maximum iteration in VariationalResistivity reached - not converged:\n {err =} \n {tol**2 =}", ) self.update_feec_variables(s=sn1, b=bn1) @@ -7222,7 +7230,7 @@ def hfun(t): def hfun(t): return xp.sin(self.options.omega * t) else: - raise NotImplementedError(f"{self.options.hfun = } not implemented.") + raise NotImplementedError(f"{self.options.hfun =} not implemented.") self._hfun = hfun self._c0 = self.variables.source.spline.vector.copy() @@ -7557,7 +7565,7 @@ def allocate(self): if self.options.c_fun == "const": c_fun = lambda e1, e2, e3: 0.0 + 0.0 * e1 else: - raise NotImplementedError(f"{self.options.c_fun = } is not available.") + raise NotImplementedError(f"{self.options.c_fun =} is not available.") # expose equation parameters self._kappa = self.options.kappa @@ -7939,16 +7947,32 @@ def allocate(self): # pullback callable funx = TransformedPformComponent( - forceterm_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, ) funy = TransformedPformComponent( - forceterm_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, ) fun_electronsx = TransformedPformComponent( - forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, ) fun_electronsy = TransformedPformComponent( - forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, ) l2_proj = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) self._F1 = l2_proj([funx, funy, _forceterm_logical]) @@ -7983,22 +8007,46 @@ def allocate(self): # pullback callable fun_pb_1 = TransformedPformComponent( - forceterm_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, ) fun_pb_2 = TransformedPformComponent( - forceterm_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, ) fun_pb_3 = TransformedPformComponent( - forceterm_class, given_in_basis="physical", out_form="2", comp=2, domain=self.domain + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=2, + domain=self.domain, ) fun_electrons_pb_1 = TransformedPformComponent( - forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, ) fun_electrons_pb_2 = TransformedPformComponent( - forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, ) fun_electrons_pb_3 = TransformedPformComponent( - forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=2, domain=self.domain + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=2, + domain=self.domain, ) if self._lifting: l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) @@ -8104,22 +8152,46 @@ def allocate(self): # pullback callable fun_pb_1 = TransformedPformComponent( - forceterm_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, ) fun_pb_2 = TransformedPformComponent( - forceterm_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, ) fun_pb_3 = TransformedPformComponent( - forceterm_class, given_in_basis="physical", out_form="2", comp=2, domain=self.domain + forceterm_class, + given_in_basis="physical", + out_form="2", + comp=2, + domain=self.domain, ) fun_electrons_pb_1 = TransformedPformComponent( - forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=0, domain=self.domain + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=0, + domain=self.domain, ) fun_electrons_pb_2 = TransformedPformComponent( - forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=1, domain=self.domain + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=1, + domain=self.domain, ) fun_electrons_pb_3 = TransformedPformComponent( - forcetermelectrons_class, given_in_basis="physical", out_form="2", comp=2, domain=self.domain + forcetermelectrons_class, + given_in_basis="physical", + out_form="2", + comp=2, + domain=self.domain, ) if self._lifting: l2_proj = L2Projector(space_id="Hdiv", mass_ops=self._mass_opsv0) @@ -8200,7 +8272,10 @@ def allocate(self): self._S21 = None if self.derhamv0.with_local_projectors: self._S21 = BasisProjectionOperatorLocal( - self.derhamv0._Ploc["1"], self.derhamv0.Vh_fem["2"], fun, transposed=False + self.derhamv0._Ploc["1"], + self.derhamv0.Vh_fem["2"], + fun, + transposed=False, ) if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): @@ -8310,7 +8385,10 @@ def allocate(self): self._S21 = None if self.derham.with_local_projectors: self._S21 = BasisProjectionOperatorLocal( - self.derham._Ploc["1"], self.derham.Vh_fem["2"], fun, transposed=False + self.derham._Ploc["1"], + self.derham.Vh_fem["2"], + fun, + transposed=False, ) if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): @@ -8603,14 +8681,18 @@ def __call__(self, dt): self._solver_UzawaNumpy.F = _Fnp if self._lifting: un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( - u0.vector, ue0.vector, phinfeec + u0.vector, + ue0.vector, + phinfeec, ) un += u_prime.vector.toarray() uen += ue_prime.vector.toarray() else: un, uen, phin, info, residual_norms, spectralresult = self._solver_UzawaNumpy( - unfeec, uenfeec, phinfeec + unfeec, + uenfeec, + phinfeec, ) dimlist = [[shp - 2 * pi for shp, pi in zip(unfeec[i][:].shape, self.derham.p)] for i in range(3)] diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 69aa86159..9f9516397 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -220,7 +220,7 @@ def allocate(self): elif self.options.algo == "implicit": kernel = Pyccelkernel(pusher_kernels.push_vxb_implicit) else: - raise ValueError(f"{self.options.algo = } not supported.") + raise ValueError(f"{self.options.algo =} not supported.") # instantiate Pusher args_kernel = ( @@ -431,7 +431,7 @@ def __init__( kernel = Pyccelkernel(pusher_kernels.push_pc_eta_rk4_H1vec) else: raise ValueError( - f'{u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', + f'{u_space =} not valid, choose from "Hcurl", "Hdiv" or "H1vec.', ) else: if u_space == "Hcurl": @@ -442,7 +442,7 @@ def __init__( kernel = Pyccelkernel(pusher_kernels.push_pc_eta_rk4_H1vec_full) else: raise ValueError( - f'{u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', + f'{u_space =} not valid, choose from "Hcurl", "Hdiv" or "H1vec.', ) args_kernel = ( diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index ce29c8141..68ba44bcd 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -56,9 +56,9 @@ def test_poisson_M1perp_1d(direction, bc_type, mapping, projected_rhs, show_plot errors = [] h_vec = [] if show_plot: - plt.figure(f"degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }", figsize=(24, 16)) - plt.figure(f"degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }", figsize=(24, 16)) - plt.figure(f"degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }", figsize=(24, 16)) + plt.figure(f"degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}", figsize=(24, 16)) + plt.figure(f"degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}", figsize=(24, 16)) + plt.figure(f"degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}", figsize=(24, 16)) for n, Neli in enumerate(Nels): # boundary conditions (overwritten below) @@ -122,7 +122,7 @@ def rho1_xyz(x, y, z): print("Direction should be either 0 or 1") # create derham object - print(f"{dirichlet_bc = }") + print(f"{dirichlet_bc =}") derham = Derham(Nel, p, spl_kind, dirichlet_bc=dirichlet_bc, comm=comm) # mass matrices @@ -183,7 +183,7 @@ def rho_pulled(e1, e2, e3): analytic_value1 = sol1_xyz(x, y, z) if show_plot: - plt.figure(f"degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }") + plt.figure(f"degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}") plt.subplot(2, 3, n + 1) if direction == 0: plt.plot(x[:, 0, 0], sol_val1[:, 0, 0], "ob", label="numerical") @@ -193,24 +193,25 @@ def rho_pulled(e1, e2, e3): plt.plot(y[0, :, 0], sol_val1[0, :, 0], "ob", label="numerical") plt.plot(y[0, :, 0], analytic_value1[0, :, 0], "r--", label="exact") plt.xlabel("y") - plt.title(f"{Nel = }") + plt.title(f"{Nel =}") plt.legend() error = xp.max(xp.abs(analytic_value1 - sol_val1)) - print(f"{direction = }, {pi = }, {Neli = }, {error=}") + print(f"{direction =}, {pi =}, {Neli =}, {error=}") errors.append(error) h = 1 / (Neli) h_vec.append(h) m, _ = xp.polyfit(xp.log(Nels), xp.log(errors), deg=1) - print(f"For {pi = }, solution converges in {direction=} with rate {-m = } ") + print(f"For {pi =}, solution converges in {direction=} with rate {-m =} ") assert -m > (pi + 1 - 0.07) # Plot convergence in 1D if show_plot: plt.figure( - f"Convergence for degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }", figsize=(12, 8) + f"Convergence for degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}", + figsize=(12, 8), ) plt.plot(h_vec, errors, "o", label=f"p={p[direction]}") plt.plot( @@ -287,7 +288,7 @@ def rho2_xyz(x, y, z): spl_kind = [False, True, True] dirichlet_bc = [(not kd,) * 2 for kd in spl_kind] dirichlet_bc = tuple(dirichlet_bc) - print(f"{dirichlet_bc = }") + print(f"{dirichlet_bc =}") # manufactured solution in 2D def sol2_xyz(x, y, z): @@ -418,9 +419,9 @@ def rho2_pulled(e1, e2, e3): error1 = xp.max(xp.abs(analytic_value1 - sol_val1)) error2 = xp.max(xp.abs(analytic_value2 - sol_val2)) - print(f"{p = }, {bc_type = }, {mapping = }") - print(f"{error1 = }") - print(f"{error2 = }") + print(f"{p =}, {bc_type =}, {mapping =}") + print(f"{error1 =}") + print(f"{error2 =}") print("") if show_plot and rank == 0: @@ -503,7 +504,7 @@ def rho(e1, e2, e3): l2_proj = L2Projector("H1", mass_ops) rho_vec = l2_proj.get_dofs(rho, apply_bc=True) - print(f"{rho_vec[:].shape = }") + print(f"{rho_vec[:].shape =}") # Create 3d Poisson solver solver_params = SolverParameters( diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 78567f8d1..bd425170a 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -73,9 +73,9 @@ def test_poisson_1d( errors = [] h_vec = [] if show_plot: - plt.figure(f"degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }", figsize=(24, 16)) - plt.figure(f"degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }", figsize=(24, 16)) - plt.figure(f"degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }", figsize=(24, 16)) + plt.figure(f"degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}", figsize=(24, 16)) + plt.figure(f"degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}", figsize=(24, 16)) + plt.figure(f"degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}", figsize=(24, 16)) for n, Neli in enumerate(Nels): # boundary conditions (overwritten below) @@ -222,7 +222,7 @@ def rho_pulled(e1, e2, e3): analytic_value1 = sol1_xyz(x, y, z) if show_plot: - plt.figure(f"degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }") + plt.figure(f"degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}") plt.subplot(2, 3, n + 1) if direction == 0: plt.plot(x[:, 0, 0], sol_val1[:, 0, 0], "ob", label="numerical") @@ -236,24 +236,25 @@ def rho_pulled(e1, e2, e3): plt.plot(z[0, 0, :], sol_val1[0, 0, :], "ob", label="numerical") plt.plot(z[0, 0, :], analytic_value1[0, 0, :], "r--", label="exact") plt.xlabel("z") - plt.title(f"{Nel = }") + plt.title(f"{Nel =}") plt.legend() error = xp.max(xp.abs(analytic_value1 - sol_val1)) - print(f"{direction = }, {pi = }, {Neli = }, {error=}") + print(f"{direction =}, {pi =}, {Neli =}, {error=}") errors.append(error) h = 1 / (Neli) h_vec.append(h) m, _ = xp.polyfit(xp.log(Nels), xp.log(errors), deg=1) - print(f"For {pi = }, solution converges in {direction=} with rate {-m = } ") + print(f"For {pi =}, solution converges in {direction=} with rate {-m =} ") assert -m > (pi + 1 - 0.07) # Plot convergence in 1D if show_plot: plt.figure( - f"Convergence for degree {pi = }, {direction + 1 = }, {bc_type = }, {mapping[0] = }", figsize=(12, 8) + f"Convergence for degree {pi =}, {direction + 1 =}, {bc_type =}, {mapping[0] =}", + figsize=(12, 8), ) plt.plot(h_vec, errors, "o", label=f"p={p[direction]}") plt.plot( @@ -485,7 +486,7 @@ def rho2_xyz(x, y, z): spl_kind = [False, True, True] dirichlet_bc = [(not kd,) * 2 for kd in spl_kind] dirichlet_bc = tuple(dirichlet_bc) - print(f"{dirichlet_bc = }") + print(f"{dirichlet_bc =}") # manufactured solution in 2D def sol2_xyz(x, y, z): @@ -628,9 +629,9 @@ def rho2_pulled(e1, e2, e3): error1 = xp.max(xp.abs(analytic_value1 - sol_val1)) error2 = xp.max(xp.abs(analytic_value2 - sol_val2)) - print(f"{p = }, {bc_type = }, {mapping = }") - print(f"{error1 = }") - print(f"{error2 = }") + print(f"{p =}, {bc_type =}, {mapping =}") + print(f"{error1 =}") + print(f"{error2 =}") print("") if show_plot and rank == 0: diff --git a/src/struphy/utils/clone_config.py b/src/struphy/utils/clone_config.py index d23bee47c..b6e14bc8d 100644 --- a/src/struphy/utils/clone_config.py +++ b/src/struphy/utils/clone_config.py @@ -230,7 +230,7 @@ def print_particle_config(self): if marker_key in self.params["kinetic"][species_name]["markers"].keys(): params_value = self.params["kinetic"][species_name]["markers"][marker_key] if params_value is not None: - assert sum_value == params_value, f"{sum_value = } and {params_value = }" + assert sum_value == params_value, f"{sum_value =} and {params_value =}" sum_row += f"| {str(sum_value):30} " # Print the final message diff --git a/src/struphy/utils/test_clone_config.py b/src/struphy/utils/test_clone_config.py index a9aa2c5c7..b1c84139b 100644 --- a/src/struphy/utils/test_clone_config.py +++ b/src/struphy/utils/test_clone_config.py @@ -24,8 +24,8 @@ def test_clone_config(Nel, Np, num_clones): species: { "markers": { "Np": Np, - } - } + }, + }, }, } @@ -37,7 +37,7 @@ def test_clone_config(Nel, Np, num_clones): # Print outputs pconf.print_clone_config() pconf.print_particle_config() - print(f"{pconf.get_Np_clone(Np) = }") + print(f"{pconf.get_Np_clone(Np) =}") if __name__ == "__main__": diff --git a/src/struphy/utils/utils.py b/src/struphy/utils/utils.py index 1707f0c07..171b61c23 100644 --- a/src/struphy/utils/utils.py +++ b/src/struphy/utils/utils.py @@ -203,6 +203,6 @@ def subp_run(cmd, cwd="libpath", check=True): for k, val in state.items(): print(k, val) i_path, o_path, b_path = get_paths(state) - print(f"{i_path = }") - print(f"{o_path = }") - print(f"{b_path = }") + print(f"{i_path =}") + print(f"{o_path =}") + print(f"{b_path =}") From 1b990b42c32af88ed99b92aef0afd6338fbe179c Mon Sep 17 00:00:00 2001 From: Byung Kyu Na Date: Fri, 24 Oct 2025 05:49:08 +0000 Subject: [PATCH 181/292] Porting the moel LinearMHDVlasovPC --- src/struphy/models/hybrid.py | 325 ++++++----- src/struphy/models/tests/test_models.py | 12 +- src/struphy/pic/accumulation/accum_kernels.py | 24 +- src/struphy/pic/pushing/pusher_kernels.py | 523 ++---------------- src/struphy/pic/tests/test_accumulation.py | 4 +- .../propagators/propagators_coupling.py | 258 ++++----- .../propagators/propagators_markers.py | 130 +++-- 7 files changed, 435 insertions(+), 841 deletions(-) diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index b5e7c2547..6994534f9 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -297,209 +297,192 @@ class LinearMHDVlasovPC(StruphyModel): :ref:`Model info `: """ - @staticmethod - def species(): - dct = {"em_fields": {}, "fluid": {}, "kinetic": {}} - - dct["em_fields"]["b_field"] = "Hdiv" - dct["fluid"]["mhd"] = { - "density": "L2", - "velocity": "Hdiv", - "pressure": "L2", - } - dct["kinetic"]["energetic_ions"] = "Particles6D" - return dct - - @staticmethod - def bulk_species(): - return "mhd" - - @staticmethod - def velocity_scale(): - return "alfvén" + ## species + class EnergeticIons(ParticleSpecies): + def __init__(self): + self.var = PICVariable(space="Particles6D") + self.init_variables() - @staticmethod - def propagators_dct(): - return { - propagators_markers.PushEtaPC: ["energetic_ions"], - propagators_markers.PushVxB: ["energetic_ions"], - propagators_coupling.PressureCoupling6D: ["energetic_ions", "mhd_velocity"], - propagators_fields.ShearAlfven: ["mhd_velocity", "b_field"], - propagators_fields.Magnetosonic: ["mhd_density", "mhd_velocity", "mhd_pressure"], - } - - __em_fields__ = species()["em_fields"] - __fluid_species__ = species()["fluid"] - __kinetic_species__ = species()["kinetic"] - __bulk_species__ = bulk_species() - __velocity_scale__ = velocity_scale() - __propagators__ = [prop.__name__ for prop in propagators_dct()] - - # add special options - @classmethod - def options(cls): - dct = super().options() - cls.add_option( - species=["fluid", "mhd"], - key="u_space", - option="Hdiv", - dct=dct, - ) - return dct + class EMFields(FieldSpecies): + def __init__(self): + self.b_field = FEECVariable(space="Hdiv") + self.init_variables() - def __init__(self, params, comm, clone_config=None): - # initialize base class - super().__init__(params, comm=comm, clone_config=clone_config) + class MHD(FluidSpecies): + def __init__(self): + self.density = FEECVariable(space="L2") + self.pressure = FEECVariable(space="L2") + self.velocity = FEECVariable(space="Hdiv") + self.init_variables() - from struphy.polar.basic import PolarVector + ## propagators - # extract necessary parameters - u_space = params["fluid"]["mhd"]["options"]["u_space"] - params_alfven = params["fluid"]["mhd"]["options"]["ShearAlfven"] - params_sonic = params["fluid"]["mhd"]["options"]["Magnetosonic"] - params_vxb = params["kinetic"]["energetic_ions"]["options"]["PushVxB"] - params_pressure = params["kinetic"]["energetic_ions"]["options"]["PressureCoupling6D"] + class Propagators: + def __init__(self, turn_off: tuple[str, ...] = (None,)): + if not "PushEtaPC" in turn_off: + self.push_eta_pc = propagators_markers.PushEtaPC() + if not "PushVxB" in turn_off: + self.push_vxb = propagators_markers.PushVxB() + if not "PressureCoupling6D" in turn_off: + self.pc6d = propagators_coupling.PressureCoupling6D() + if not "ShearAlfven" in turn_off: + self.shearalfven = propagators_fields.ShearAlfven() + if not "Magnetosonic" in turn_off: + self.magnetosonic = propagators_fields.Magnetosonic() - # use perp model - assert ( - params["kinetic"]["energetic_ions"]["options"]["PressureCoupling6D"]["use_perp_model"] - == params["kinetic"]["energetic_ions"]["options"]["PressureCoupling6D"]["use_perp_model"] - ) - use_perp_model = params["kinetic"]["energetic_ions"]["options"]["PressureCoupling6D"]["use_perp_model"] + def __init__(self, turn_off: tuple[str, ...] = (None,)): + if rank == 0: + print(f"\n*** Creating light-weight instance of model '{self.__class__.__name__}':") - # compute coupling parameters - Ab = params["fluid"]["mhd"]["phys_params"]["A"] - Ah = params["kinetic"]["energetic_ions"]["phys_params"]["A"] - epsilon = self.equation_params["energetic_ions"]["epsilon"] + # 1. instantiate all species + self.em_fields = self.EMFields() + self.mhd = self.MHD() + self.energetic_ions = self.EnergeticIons() - if abs(epsilon - 1) < 1e-6: - epsilon = 1.0 + # 2. instantiate all propagators + self.propagators = self.Propagators(turn_off) - self._coupling_params = {} - self._coupling_params["Ab"] = Ab - self._coupling_params["Ah"] = Ah - self._coupling_params["epsilon"] = epsilon + # 3. assign variables to propagators + if not "ShearAlfven" in turn_off: + self.propagators.shearalfven.variables.u = self.mhd.velocity + self.propagators.shearalfven.variables.b = self.em_fields.b_field + if not "Magnetosonic" in turn_off: + self.propagators.magnetosonic.variables.n = self.mhd.density + self.propagators.magnetosonic.variables.u = self.mhd.velocity + self.propagators.magnetosonic.variables.p = self.mhd.pressure + if not "PressureCoupling6D" in turn_off: + self.propagators.pc6d.variables.u = self.mhd.velocity + self.propagators.pc6d.variables.energetic_ions = self.energetic_ions.var + if not "PushEtaPC" in turn_off: + self.propagators.push_eta_pc.variables.var = self.energetic_ions.var + if not "PushVxB" in turn_off: + self.propagators.push_vxb.variables.ions = self.energetic_ions.var - # add control variate to mass_ops object - if self.pointer["energetic_ions"].control_variate: - self.mass_ops.weights["f0"] = self.pointer["energetic_ions"].f0 - - # Project magnetic field - self._b_eq = self.derham.P["2"]( - [ - self.equil.b2_1, - self.equil.b2_2, - self.equil.b2_3, - ] + # define scalars for update_scalar_quantities + self.add_scalar("en_U") + self.add_scalar("en_p") + self.add_scalar("en_B") + self.add_scalar("en_f", compute="from_particles", variable=self.energetic_ions.var) + self.add_scalar( + "en_tot", + summands=[ + "en_U", + "en_p", + "en_B", + "en_f", + ], ) - self._p_eq = self.derham.P["3"](self.equil.p3) - self._ones = self._p_eq.space.zeros() + @property + def bulk_species(self): + return self.mhd + + @property + def velocity_scale(self): + return "alfvén" + + def allocate_helpers(self): + self._ones = self.projected_equil.p3.space.zeros() if isinstance(self._ones, PolarVector): self._ones.tp[:] = 1.0 else: self._ones[:] = 1.0 - # set keyword arguments for propagators - self._kwargs[propagators_markers.PushEtaPC] = { - "u": self.pointer["mhd_velocity"], - "use_perp_model": use_perp_model, - "u_space": u_space, - } - - self._kwargs[propagators_markers.PushVxB] = { - "algo": params_vxb["algo"], - "kappa": epsilon, - "b2": self.pointer["b_field"], - "b2_add": self._b_eq, - } - - if params_pressure["turn_off"]: - self._kwargs[propagators_coupling.PressureCoupling6D] = None - else: - self._kwargs[propagators_coupling.PressureCoupling6D] = { - "use_perp_model": use_perp_model, - "u_space": u_space, - "solver": params_pressure["solver"], - "coupling_params": self._coupling_params, - "filter": params_pressure["filter"], - "boundary_cut": params_pressure["boundary_cut"], - } - - if params_alfven["turn_off"]: - self._kwargs[propagators_fields.ShearAlfven] = None - else: - self._kwargs[propagators_fields.ShearAlfven] = { - "u_space": u_space, - "solver": params_alfven["solver"], - } - - if params_sonic["turn_off"]: - self._kwargs[propagators_fields.Magnetosonic] = None - else: - self._kwargs[propagators_fields.Magnetosonic] = { - "b": self.pointer["b_field"], - "u_space": u_space, - "solver": params_sonic["solver"], - } - - # Initialize propagators used in splitting substeps - self.init_propagators() - - # Scalar variables to be saved during simulation: - self.add_scalar("en_U", compute="from_field") - self.add_scalar("en_p", compute="from_field") - self.add_scalar("en_B", compute="from_field") - self.add_scalar("en_f", compute="from_particles", species="energetic_ions") - self.add_scalar("en_tot", summands=["en_U", "en_p", "en_B", "en_f"]) - self.add_scalar("n_lost_particles", compute="from_particles", species="energetic_ions") - - # temporary vectors for scalar quantities - self._tmp_u = self.derham.Vh["2"].zeros() - self._tmp_b1 = self.derham.Vh["2"].zeros() - self._tmp = xp.empty(1, dtype=float) + self._en_f = xp.empty(1, dtype=float) self._n_lost_particles = xp.empty(1, dtype=float) def update_scalar_quantities(self): + # scaling factor + Ab = self.mhd.mass_number + Ah = self.energetic_ions.var.species.mass_number + # perturbed fields - if "Hdiv" == "Hdiv": - en_U = 0.5 * self.mass_ops.M2n.dot_inner(self.pointer["mhd_velocity"], self.pointer["mhd_velocity"]) - else: - en_U = 0.5 * self.mass_ops.Mvn.dot_inner(self.pointer["mhd_velocity"], self.pointer["mhd_velocity"]) - en_B = 0.5 * self.mass_ops.M2.dot_inner(self.pointer["b_field"], self.pointer["b_field"]) - en_p = self.pointer["mhd_pressure"].inner(self._ones) / (5 / 3 - 1) + en_U = 0.5 * self.mass_ops.M2n.dot_inner( + self.mhd.velocity.spline.vector, + self.mhd.velocity.spline.vector, + ) + en_B = 0.5 * self.mass_ops.M2.dot_inner( + self.em_fields.b_field.spline.vector, + self.em_fields.b_field.spline.vector, + ) + en_p = self.mhd.pressure.spline.vector.inner(self._ones) / (5 / 3 - 1) self.update_scalar("en_U", en_U) self.update_scalar("en_B", en_B) self.update_scalar("en_p", en_p) - # particles - self._tmp[0] = ( - self._coupling_params["Ah"] - / self._coupling_params["Ab"] - * self.pointer["energetic_ions"] - .markers_wo_holes[:, 6] - .dot( - self.pointer["energetic_ions"].markers_wo_holes[:, 3] ** 2 - + self.pointer["energetic_ions"].markers_wo_holes[:, 4] ** 2 - + self.pointer["energetic_ions"].markers_wo_holes[:, 5] ** 2, + # particles' energy + particles = self.energetic_ions.var.particles + + self._en_f[0] = ( + particles.markers[~particles.holes, 6].dot( + particles.markers[~particles.holes, 3] ** 2 + + particles.markers[~particles.holes, 4] ** 2 + + particles.markers[~particles.holes, 5] ** 2 ) - / (2.0) + / 2.0 + * Ah + / Ab ) - self.update_scalar("en_f", self._tmp[0]) - self.update_scalar("en_tot", en_U + en_B + en_p + self._tmp[0]) + self.update_scalar("en_f", self._en_f[0]) + self.update_scalar("en_tot") - # Print number of lost ions - self._n_lost_particles[0] = self.pointer["energetic_ions"].n_lost_markers - self.update_scalar("n_lost_particles", self._n_lost_particles[0]) - if self.rank_world == 0: + # print number of lost particles + n_lost_markers = xp.array(particles.n_lost_markers) + + if self.derham.comm is not None: + self.derham.comm.Allreduce( + MPI.IN_PLACE, + n_lost_markers, + op=MPI.SUM, + ) + + if self.clone_config is not None: + self.clone_config.inter_comm.Allreduce( + MPI.IN_PLACE, + n_lost_markers, + op=MPI.SUM, + ) + + if rank == 0: print( - "ratio of lost particles: ", - self._n_lost_particles[0] / self.pointer["energetic_ions"].Np * 100, - "%", + "Lost particle ratio: ", + n_lost_markers / particles.Np * 100, + "% \n", ) + ## default parameters + def generate_default_parameter_file(self, path=None, prompt=True): + params_path = super().generate_default_parameter_file(path=path, prompt=prompt) + new_file = [] + with open(params_path, "r") as f: + for line in f: + if "magnetosonic.Options" in line: + new_file += [ + """model.propagators.magnetosonic.options = model.propagators.magnetosonic.Options( + b_field=model.em_fields.b_field,)\n""" + ] + + elif "push_eta_pc.Options" in line: + new_file += [ + """model.propagators.push_eta_pc.options = model.propagators.push_eta_pc.Options( + u_tilde = model.mhd.velocity,)\n""" + ] + + elif "push_vxb.Options" in line: + new_file += [ + """model.propagators.push_vxb.options = model.propagators.push_vxb.Options( + b2_var = model.em_fields.b_field,)\n""" + ] + + else: + new_file += [line] + + with open(params_path, "w") as f: + for line in new_file: + f.write(line) + class LinearMHDDriftkineticCC(StruphyModel): r"""Hybrid linear ideal MHD + energetic ions (5D Driftkinetic) with **current coupling scheme**. diff --git a/src/struphy/models/tests/test_models.py b/src/struphy/models/tests/test_models.py index aa2af4013..36a9ea01b 100644 --- a/src/struphy/models/tests/test_models.py +++ b/src/struphy/models/tests/test_models.py @@ -35,14 +35,10 @@ if rank == 0: print(f"\n{kinetic_models = }") -hybrid_models = [ - "LinearMHDDriftkineticCC", - "LinearMHDVlasovCC", - "ColdPlasmaVlasov", -] -# for name, obj in inspect.getmembers(hybrid): -# if inspect.isclass(obj) and "models.hybrid" in obj.__module__: -# hybrid_models += [name] +hybrid_models = [] +for name, obj in inspect.getmembers(hybrid): + if inspect.isclass(obj) and "models.hybrid" in obj.__module__: + hybrid_models += [name] if rank == 0: print(f"\n{hybrid_models = }") diff --git a/src/struphy/pic/accumulation/accum_kernels.py b/src/struphy/pic/accumulation/accum_kernels.py index ab0a49469..2a82a9bcf 100644 --- a/src/struphy/pic/accumulation/accum_kernels.py +++ b/src/struphy/pic/accumulation/accum_kernels.py @@ -1112,9 +1112,7 @@ def pc_lin_mhd_6d_full( vec1_3: "float[:,:,:]", vec2_3: "float[:,:,:]", vec3_3: "float[:,:,:]", - scale_mat: "float", - scale_vec: "float", - boundary_cut: "float", + ep_scale: "float", ): r"""Accumulates into V1 with the filling functions @@ -1158,10 +1156,6 @@ def pc_lin_mhd_6d_full( if markers[ip, 0] == -1.0: continue - # boundary cut - if markers[ip, 0] < boundary_cut or markers[ip, 0] > 1.0 - boundary_cut: - continue - # marker positions eta1 = markers[ip, 0] eta2 = markers[ip, 1] @@ -1192,8 +1186,8 @@ def pc_lin_mhd_6d_full( weight = markers[ip, 8] - filling_m[:, :] = weight * tmp1 / Np * scale_mat - filling_v[:] = weight * tmp_v / Np * scale_vec + filling_m[:, :] = weight * tmp1 / Np * ep_scale + filling_v[:] = weight * tmp_v / Np * ep_scale # call the appropriate matvec filler particle_to_mat_kernels.m_v_fill_v1_pressure_full( @@ -1311,9 +1305,7 @@ def pc_lin_mhd_6d( vec1_3: "float[:,:,:]", vec2_3: "float[:,:,:]", vec3_3: "float[:,:,:]", - scale_mat: "float", - scale_vec: "float", - boundary_cut: "float", + ep_scale: "float", ): r"""Accumulates into V1 with the filling functions @@ -1356,10 +1348,6 @@ def pc_lin_mhd_6d( if markers[ip, 0] == -1.0: continue - # boundary cut - if markers[ip, 0] < boundary_cut or markers[ip, 0] > 1.0 - boundary_cut: - continue - # marker positions eta1 = markers[ip, 0] eta2 = markers[ip, 1] @@ -1390,8 +1378,8 @@ def pc_lin_mhd_6d( linalg_kernels.matrix_matrix(df_inv, df_inv_t, tmp1) linalg_kernels.matrix_vector(df_inv, v, tmp_v) - filling_m[:, :] = weight * tmp1 * scale_mat - filling_v[:] = weight * tmp_v * scale_vec + filling_m[:, :] = weight * tmp1 * ep_scale + filling_v[:] = weight * tmp_v * ep_scale # call the appropriate matvec filler particle_to_mat_kernels.m_v_fill_v1_pressure( diff --git a/src/struphy/pic/pushing/pusher_kernels.py b/src/struphy/pic/pushing/pusher_kernels.py index cdaa9059b..47b6a71ba 100644 --- a/src/struphy/pic/pushing/pusher_kernels.py +++ b/src/struphy/pic/pushing/pusher_kernels.py @@ -1615,6 +1615,15 @@ def push_bxu_Hdiv_pauli( # -- removed omp: #$ omp end parallel +@stack_array( + "dfm", + "dfinv", + "dfinv_t", + "e", + "e_cart", + "GXu", + "v", +) def push_pc_GXu_full( dt: float, stage: int, @@ -1630,7 +1639,6 @@ def push_pc_GXu_full( GXu_31: "float[:,:,:]", GXu_32: "float[:,:,:]", GXu_33: "float[:,:,:]", - boundary_cut: "float", ): r"""Updates @@ -1671,10 +1679,6 @@ def push_pc_GXu_full( if markers[ip, 0] == -1.0: continue - # boundary cut - if markers[ip, 0] < boundary_cut or markers[ip, 0] > 1.0 - boundary_cut: - continue - eta1 = markers[ip, 0] eta2 = markers[ip, 1] eta3 = markers[ip, 2] @@ -1740,6 +1744,15 @@ def push_pc_GXu_full( markers[ip, 3:6] -= dt * e_cart / 2.0 +@stack_array( + "dfm", + "dfinv", + "dfinv_t", + "e", + "e_cart", + "GXu", + "v", +) def push_pc_GXu( dt: float, stage: int, @@ -1755,7 +1768,6 @@ def push_pc_GXu( GXu_31: "float[:,:,:]", GXu_32: "float[:,:,:]", GXu_33: "float[:,:,:]", - boundary_cut: "float", ): r"""Updates @@ -1783,7 +1795,6 @@ def push_pc_GXu( e = empty(3, dtype=float) e_cart = empty(3, dtype=float) GXu = empty((3, 3), dtype=float) - GXu_t = empty((3, 3), dtype=float) # particle velocity v = empty(3, dtype=float) @@ -1797,10 +1808,6 @@ def push_pc_GXu( if markers[ip, 0] == -1.0: continue - # boundary cut - if markers[ip, 0] < boundary_cut or markers[ip, 0] > 1.0 - boundary_cut: - continue - eta1 = markers[ip, 0] eta2 = markers[ip, 1] eta3 = markers[ip, 2] @@ -1939,7 +1946,7 @@ def push_eta_stage( @stack_array("dfm", "dfinv", "dfinv_t", "ginv", "v", "u", "k", "k_v", "k_u") -def push_pc_eta_rk4_Hcurl_full( +def push_pc_eta_stage_Hcurl( dt: float, stage: int, args_markers: "MarkerArguments", @@ -1948,6 +1955,10 @@ def push_pc_eta_rk4_Hcurl_full( u_1: "float[:,:,:]", u_2: "float[:,:,:]", u_3: "float[:,:,:]", + use_perp_model: "bool", + a: "float[:]", + b: "float[:]", + c: "float[:]", ): r"""Fourth order Runge-Kutta solve of @@ -1960,14 +1971,6 @@ def push_pc_eta_rk4_Hcurl_full( .. math:: \textnormal{vec}( \hat{\mathbf U}^{1}) = G^{-1}\hat{\mathbf U}^{1}\,,\qquad \textnormal{vec}( \hat{\mathbf U}^{2}) = \frac{\hat{\mathbf U}^{2}}{\sqrt g}\,. - - Parameters - ---------- - u_1, u_2, u_3: array[float] - 3d array of FE coeffs of U-field, either as 1-form or as 2-form. - - u_basis : int - U is 1-form (u_basis=1) or a 2-form (u_basis=2). """ # allocate metric coeffs @@ -1993,22 +1996,13 @@ def push_pc_eta_rk4_Hcurl_full( first_init_idx = args_markers.first_init_idx first_free_idx = args_markers.first_free_idx - # assign factor of k for each stage - if stage == 0 or stage == 3: - nk = 1.0 - else: - nk = 2.0 + # get number of stages + n_stages = shape(b)[0] - # which stage - if stage == 3: + if stage == n_stages - 1: last = 1.0 - cont = 0.0 - elif stage == 2: - last = 0.0 - cont = 2.0 else: last = 0.0 - cont = 1.0 for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) @@ -2053,396 +2047,8 @@ def push_pc_eta_rk4_Hcurl_full( u, ) - # transform to vector field - linalg_kernels.matrix_vector(ginv, u, k_u) - - # sum contribs - k[:] = k_v + k_u - - # accum k - markers[ip, first_free_idx : first_free_idx + 3] += k * nk / 6.0 - - # update markers for the next stage - markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] - + dt * k / 2 * cont - + dt * markers[ip, first_free_idx : first_free_idx + 3] * last - ) - - -@stack_array("dfm", "dfinv", "dfinv_t", "ginv", "v", "u", "k", "k_v", "k_u") -def push_pc_eta_rk4_Hdiv_full( - dt: float, - stage: int, - args_markers: "MarkerArguments", - args_domain: "DomainArguments", - args_derham: "DerhamArguments", - u_1: "float[:,:,:]", - u_2: "float[:,:,:]", - u_3: "float[:,:,:]", -): - r"""Fourth order Runge-Kutta solve of - - .. math:: - - \frac{\textnormal d \boldsymbol \eta_p(t)}{\textnormal d t} = DF^{-1}(\boldsymbol \eta_p(t)) \mathbf v + \textnormal{vec}( \hat{\mathbf U}^{1(2)}) - - for each marker :math:`p` in markers array, where :math:`\mathbf v` is constant and - - .. math:: - - \textnormal{vec}( \hat{\mathbf U}^{1}) = G^{-1}\hat{\mathbf U}^{1}\,,\qquad \textnormal{vec}( \hat{\mathbf U}^{2}) = \frac{\hat{\mathbf U}^{2}}{\sqrt g}\,. - - Parameters - ---------- - u_1, u_2, u_3: array[float] - 3d array of FE coeffs of U-field, either as 1-form or as 2-form. - - u_basis : int - U is 1-form (u_basis=1) or a 2-form (u_basis=2). - """ - - # allocate metric coeffs - dfm = empty((3, 3), dtype=float) - dfinv = empty((3, 3), dtype=float) - dfinv_t = empty((3, 3), dtype=float) - ginv = empty((3, 3), dtype=float) - - # marker velocity - v = empty(3, dtype=float) - - # U-fiels - u = empty(3, dtype=float) - - # intermediate stages in RK4 - k = empty(3, dtype=float) - k_v = empty(3, dtype=float) - k_u = empty(3, dtype=float) - - # get marker arguments - markers = args_markers.markers - n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx - first_free_idx = args_markers.first_free_idx - - # assign factor of k for each stage - if stage == 0 or stage == 3: - nk = 1.0 - else: - nk = 2.0 - - # is it the last stage? - if stage == 3: - last = 1.0 - cont = 0.0 - else: - last = 0.0 - cont = 1.0 - - for ip in range(n_markers): - # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: - continue - - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] - v[:] = markers[ip, 3:6] - - # ----------------- stage n in Runge-Kutta method ------------------- - # evaluate Jacobian, result in dfm - evaluation_kernels.df( - e1, - e2, - e3, - args_domain, - dfm, - ) - - # metric coeffs - det_df = linalg_kernels.det(dfm) - linalg_kernels.matrix_inv(dfm, dfinv) - linalg_kernels.transpose(dfinv, dfinv_t) - linalg_kernels.matrix_matrix(dfinv, dfinv_t, ginv) - - # pull-back of velocity - linalg_kernels.matrix_vector(dfinv, v, k_v) - - # spline evaluation - span1, span2, span3 = get_spans(e1, e2, e3, args_derham) - - # U-field - eval_2form_spline_mpi( - span1, - span2, - span3, - args_derham, - u_1, - u_2, - u_3, - u, - ) - - # transform to vector field - k_u[:] = u / det_df - - # sum contribs - k[:] = k_v + k_u - - # accum k - markers[ip, first_free_idx : first_free_idx + 3] += k * nk / 6.0 - - # update markers for the next stage - markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] - + dt * k / 2 * cont - + dt * markers[ip, first_free_idx : first_free_idx + 3] * last - ) - - -@stack_array("dfm", "dfinv", "dfinv_t", "ginv", "v", "u", "k", "k_v") -def push_pc_eta_rk4_H1vec_full( - dt: float, - stage: int, - args_markers: "MarkerArguments", - args_domain: "DomainArguments", - args_derham: "DerhamArguments", - u_1: "float[:,:,:]", - u_2: "float[:,:,:]", - u_3: "float[:,:,:]", -): - r"""Fourth order Runge-Kutta solve of - - .. math:: - - \frac{\textnormal d \boldsymbol \eta_p(t)}{\textnormal d t} = DF^{-1}(\boldsymbol \eta_p(t)) \mathbf v + \textnormal{vec}( \hat{\mathbf U}^{1(2)}) - - for each marker :math:`p` in markers array, where :math:`\mathbf v` is constant and - - .. math:: - - \textnormal{vec}( \hat{\mathbf U}^{1}) = G^{-1}\hat{\mathbf U}^{1}\,,\qquad \textnormal{vec}( \hat{\mathbf U}^{2}) = \frac{\hat{\mathbf U}^{2}}{\sqrt g}\,. - - Parameters - ---------- - u_1, u_2, u_3 : array[float] - 3d array of FE coeffs of U-field, either as 1-form or as 2-form. - - u_basis : int - U is 1-form (u_basis=1) or a 2-form (u_basis=2). - """ - - # allocate metric coeffs - dfm = empty((3, 3), dtype=float) - dfinv = empty((3, 3), dtype=float) - dfinv_t = empty((3, 3), dtype=float) - ginv = empty((3, 3), dtype=float) - - # marker and velocity - v = empty(3, dtype=float) - - # U-fiels - u = empty(3, dtype=float) - - # intermediate stages in RK4 - k = empty(3, dtype=float) - k_v = empty(3, dtype=float) - - # get marker arguments - markers = args_markers.markers - n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx - first_free_idx = args_markers.first_free_idx - - # assign factor of k for each stage - if stage == 0 or stage == 3: - nk = 1.0 - else: - nk = 2.0 - - # which stage - if stage == 3: - last = 1.0 - cont = 0.0 - elif stage == 2: - last = 0.0 - cont = 2.0 - else: - last = 0.0 - cont = 1.0 - - for ip in range(n_markers): - # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: - continue - - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] - v[:] = markers[ip, 3:6] - - # ----------------- stage n in Runge-Kutta method ------------------- - # evaluate Jacobian, result in dfm - evaluation_kernels.df( - e1, - e2, - e3, - args_domain, - dfm, - ) - - # metric coeffs - linalg_kernels.matrix_inv(dfm, dfinv) - linalg_kernels.transpose(dfinv, dfinv_t) - linalg_kernels.matrix_matrix(dfinv, dfinv_t, ginv) - - # pull-back of velocity - linalg_kernels.matrix_vector(dfinv, v, k_v) - - # spline evaluation - span1, span2, span3 = get_spans(e1, e2, e3, args_derham) - - # U-field - eval_vectorfield_spline_mpi( - span1, - span2, - span3, - args_derham, - u_1, - u_2, - u_3, - u, - ) - - # sum contribs - k[:] = k_v + u - - # accum k - markers[ip, first_free_idx : first_free_idx + 3] += k * nk / 6.0 - - # update markers for the next stage - markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] - + dt * k / 2 * cont - + dt * markers[ip, first_free_idx : first_free_idx + 3] * last - ) - - -@stack_array("dfm", "dfinv", "dfinv_t", "ginv", "v", "u", "k", "k_v", "k_u") -def push_pc_eta_rk4_Hcurl( - dt: float, - stage: int, - args_markers: "MarkerArguments", - args_domain: "DomainArguments", - args_derham: "DerhamArguments", - u_1: "float[:,:,:]", - u_2: "float[:,:,:]", - u_3: "float[:,:,:]", -): - r"""Fourth order Runge-Kutta solve of - - .. math:: - - \frac{\textnormal d \boldsymbol \eta_p(t)}{\textnormal d t} = DF^{-1}(\boldsymbol \eta_p(t)) \mathbf v + \textnormal{vec}( \hat{\mathbf U}^{1(2)}) - - for each marker :math:`p` in markers array, where :math:`\mathbf v` is constant and - - .. math:: - - \textnormal{vec}( \hat{\mathbf U}^{1}) = G^{-1}\hat{\mathbf U}^{1}\,,\qquad \textnormal{vec}( \hat{\mathbf U}^{2}) = \frac{\hat{\mathbf U}^{2}}{\sqrt g}\,. - - Parameters - ---------- - u_1, u_2, u_3 : array[float] - 3d array of FE coeffs of U-field, either as 1-form or as 2-form. - - u_basis : int - U is 1-form (u_basis=1) or a 2-form (u_basis=2). - """ - - # allocate metric coeffs - dfm = empty((3, 3), dtype=float) - dfinv = empty((3, 3), dtype=float) - dfinv_t = empty((3, 3), dtype=float) - ginv = empty((3, 3), dtype=float) - - # marker velocity - v = empty(3, dtype=float) - - # U-fiels - u = empty(3, dtype=float) - - # intermediate stages in RK4 - k = empty(3, dtype=float) - k_v = empty(3, dtype=float) - k_u = empty(3, dtype=float) - - # get marker arguments - markers = args_markers.markers - n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx - first_free_idx = args_markers.first_free_idx - - # assign factor of k for each stage - if stage == 0 or stage == 3: - nk = 1.0 - else: - nk = 2.0 - - # which stage - if stage == 3: - last = 1.0 - cont = 0.0 - elif stage == 2: - last = 0.0 - cont = 2.0 - else: - last = 0.0 - cont = 1.0 - - for ip in range(n_markers): - # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: - continue - - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] - v[:] = markers[ip, 3:6] - - # ----------------- stage n in Runge-Kutta method ------------------- - # evaluate Jacobian, result in dfm - evaluation_kernels.df( - e1, - e2, - e3, - args_domain, - dfm, - ) - - # metric coeffs - linalg_kernels.matrix_inv(dfm, dfinv) - linalg_kernels.transpose(dfinv, dfinv_t) - linalg_kernels.matrix_matrix(dfinv, dfinv_t, ginv) - - # pull-back of velocity - linalg_kernels.matrix_vector(dfinv, v, k_v) - - # spline evaluation - span1, span2, span3 = get_spans(e1, e2, e3, args_derham) - - # U-field - eval_1form_spline_mpi( - span1, - span2, - span3, - args_derham, - u_1, - u_2, - u_3, - u, - ) - u[2] = 0.0 + if use_perp_model: + u[2] = 0.0 # transform to vector field linalg_kernels.matrix_vector(ginv, u, k_u) @@ -2451,18 +2057,18 @@ def push_pc_eta_rk4_Hcurl( k[:] = k_v + k_u # accum k - markers[ip, first_free_idx : first_free_idx + 3] += k * nk / 6.0 + markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k # update markers for the next stage markers[ip, 0:3] = ( markers[ip, first_init_idx : first_init_idx + 3] - + dt * k / 2 * cont - + dt * markers[ip, first_free_idx : first_free_idx + 3] * last + + dt * k * a[stage] + + last * markers[ip, first_free_idx : first_free_idx + 3] ) @stack_array("dfm", "dfinv", "dfinv_t", "ginv", "v", "u", "k", "k_v", "k_u") -def push_pc_eta_rk4_Hdiv( +def push_pc_eta_stage_Hdiv( dt: float, stage: int, args_markers: "MarkerArguments", @@ -2471,6 +2077,10 @@ def push_pc_eta_rk4_Hdiv( u_1: "float[:,:,:]", u_2: "float[:,:,:]", u_3: "float[:,:,:]", + use_perp_model: "bool", + a: "float[:]", + b: "float[:]", + c: "float[:]", ): r"""Fourth order Runge-Kutta solve of @@ -2483,14 +2093,6 @@ def push_pc_eta_rk4_Hdiv( .. math:: \textnormal{vec}( \hat{\mathbf U}^{1}) = G^{-1}\hat{\mathbf U}^{1}\,,\qquad \textnormal{vec}( \hat{\mathbf U}^{2}) = \frac{\hat{\mathbf U}^{2}}{\sqrt g}\,. - - Parameters - ---------- - u_1, u_2, u_3 : array[float] - 3d array of FE coeffs of U-field, either as 1-form or as 2-form. - - u_basis : int - U is 1-form (u_basis=1) or a 2-form (u_basis=2). """ # allocate metric coeffs @@ -2516,19 +2118,13 @@ def push_pc_eta_rk4_Hdiv( first_init_idx = args_markers.first_init_idx first_free_idx = args_markers.first_free_idx - # assign factor of k for each stage - if stage == 0 or stage == 3: - nk = 1.0 - else: - nk = 2.0 + # get number of stages + n_stages = shape(b)[0] - # is it the last stage? - if stage == 3: + if stage == n_stages - 1: last = 1.0 - cont = 0.0 else: last = 0.0 - cont = 1.0 for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) @@ -2573,7 +2169,9 @@ def push_pc_eta_rk4_Hdiv( u_3, u, ) - u[2] = 0.0 + + if use_perp_model: + u[2] = 0.0 # transform to vector field k_u[:] = u / det_df @@ -2582,18 +2180,18 @@ def push_pc_eta_rk4_Hdiv( k[:] = k_v + k_u # accum k - markers[ip, first_free_idx : first_free_idx + 3] += k * nk / 6.0 + markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k # update markers for the next stage markers[ip, 0:3] = ( markers[ip, first_init_idx : first_init_idx + 3] - + dt * k / 2 * cont - + dt * markers[ip, first_free_idx : first_free_idx + 3] * last + + dt * k * a[stage] + + last * markers[ip, first_free_idx : first_free_idx + 3] ) @stack_array("dfm", "dfinv", "dfinv_t", "ginv", "v", "u", "k", "k_v") -def push_pc_eta_rk4_H1vec( +def push_pc_eta_stage_H1vec( dt: float, stage: int, args_markers: "MarkerArguments", @@ -2602,6 +2200,10 @@ def push_pc_eta_rk4_H1vec( u_1: "float[:,:,:]", u_2: "float[:,:,:]", u_3: "float[:,:,:]", + use_perp_model: "bool", + a: "float[:]", + b: "float[:]", + c: "float[:]", ): r"""Fourth order Runge-Kutta solve of @@ -2646,22 +2248,13 @@ def push_pc_eta_rk4_H1vec( first_init_idx = args_markers.first_init_idx first_free_idx = args_markers.first_free_idx - # assign factor of k for each stage - if stage == 0 or stage == 3: - nk = 1.0 - else: - nk = 2.0 + # get number of stages + n_stages = shape(b)[0] - # which stage - if stage == 3: + if stage == n_stages - 1: last = 1.0 - cont = 0.0 - elif stage == 2: - last = 0.0 - cont = 2.0 else: last = 0.0 - cont = 1.0 for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) @@ -2705,19 +2298,21 @@ def push_pc_eta_rk4_H1vec( u_3, u, ) - u[2] = 0.0 + + if use_perp_model: + u[2] = 0.0 # sum contribs k[:] = k_v + u # accum k - markers[ip, first_free_idx : first_free_idx + 3] += k * nk / 6.0 + markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k # update markers for the next stage markers[ip, 0:3] = ( markers[ip, first_init_idx : first_init_idx + 3] - + dt * k / 2 * cont - + dt * markers[ip, first_free_idx : first_free_idx + 3] * last + + dt * k * a[stage] + + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/tests/test_accumulation.py b/src/struphy/pic/tests/test_accumulation.py index 012c73e33..805c578d5 100644 --- a/src/struphy/pic/tests/test_accumulation.py +++ b/src/struphy/pic/tests/test_accumulation.py @@ -247,7 +247,9 @@ def pc_lin_mhd_6d_step_ph_full(Nel, p, spl_kind, mapping, Np, verbose=False): ) start_time = time() - ACC(1.0, 1.0, 0.0) + ACC( + 1.0, + ) end_time = time() tot_time = xp.round(end_time - start_time, 3) diff --git a/src/struphy/propagators/propagators_coupling.py b/src/struphy/propagators/propagators_coupling.py index 76c2396ef..80d483568 100644 --- a/src/struphy/propagators/propagators_coupling.py +++ b/src/struphy/propagators/propagators_coupling.py @@ -551,119 +551,133 @@ class PressureCoupling6D(Propagator): \begin{bmatrix} {\mathbb M^n}(u^{n+1} + u^n) \\ \bar W (V^{n+1} + V^{n} \end{bmatrix} \,. """ - @staticmethod - def options(default=False): - dct = {} - dct["use_perp_model"] = [True, False] - dct["solver"] = { - "type": [ - ("pcg", "MassMatrixPreconditioner"), - ("cg", None), - ], - "tol": 1.0e-8, - "maxiter": 3000, - "info": False, - "verbose": False, - "recycle": True, - } - dct["filter"] = { - "use_filter": None, - "modes": (1), - "repeat": 1, - "alpha": 0.5, - } - dct["boundary_cut"] = { - "e1": 0.0, - "e2": 0.0, - "e3": 0.0, - } - dct["turn_off"] = False - - if default: - dct = descend_options_dict(dct, []) - - return dct - - def __init__( - self, - particles: Particles5D, - u: BlockVector | PolarVector, - *, - use_perp_model: bool = options(default=True)["use_perp_model"], - u_space: str, - solver: dict = options(default=True)["solver"], - coupling_params: dict, - filter: dict = options(default=True)["filter"], - boundary_cut: dict = options(default=True)["boundary_cut"], - ): - super().__init__(particles, u) - - self._G = self.derham.grad - self._GT = self.derham.grad.transpose() - - self._info = solver["info"] - if self.derham.comm is None: - self._rank = 0 - else: - self._rank = self.derham.comm.Get_rank() + class Variables: + def __init__(self): + self._u: FEECVariable = None + self._energetic_ions: PICVariable = None + + @property + def u(self) -> FEECVariable: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space in ("Hcurl", "Hdiv", "H1vec") + self._u = new + + @property + def energetic_ions(self) -> PICVariable: + return self._energetic_ions + + @energetic_ions.setter + def energetic_ions(self, new): + assert isinstance(new, PICVariable) + assert new.space == "Particles6D" + self._energetic_ions = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass + class Options: + # propagator options + ep_scale: float = 1.0 + u_space: OptsVecSpace = "Hdiv" + solver: OptsSymmSolver = "pcg" + precond: OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + filter_params: FilterParameters = None + use_perp_model: bool = True + + def __post_init__(self): + # checks + check_option(self.u_space, OptsVecSpace) + check_option(self.solver, OptsSymmSolver) + check_option(self.precond, OptsMassPrecond) + assert isinstance(self.ep_scale, float) + assert isinstance(self.use_perp_model, bool) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + + if self.filter_params is None: + self.filter_params = FilterParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new - assert u_space in {"Hcurl", "Hdiv", "H1vec"} + @profile + def allocate(self): + if self.options.u_space == "H1vec": + self._u_form_int = 0 + else: + self._u_form_int = int(self.derham.space_to_form[self.options.u_space]) - if u_space == "Hcurl": + if self.options.u_space == "Hcurl": id_Mn = "M1n" id_X = "X1" - elif u_space == "Hdiv": + elif self.options.u_space == "Hdiv": id_Mn = "M2n" id_X = "X2" - elif u_space == "H1vec": + elif self.options.u_space == "H1vec": id_Mn = "Mvn" id_X = "Xv" - if u_space == "H1vec": - self._space_key_int = 0 - else: - self._space_key_int = int( - self.derham.space_to_form[u_space], - ) + # call operatros + id_M = "M" + self.derham.space_to_form[self.options.u_space] + "n" + _A = getattr(self.mass_ops, id_M) + self._X = getattr(self.basis_ops, id_X) + self._XT = self._X.transpose() + grad = self.derham.grad + gradT = grad.transpose() # Preconditioner - if solver["type"][1] is None: + if self.options.precond is None: pc = None else: - pc_class = getattr(preconditioner, solver["type"][1]) - pc = pc_class(getattr(self.mass_ops, id_Mn)) + pc_class = getattr(preconditioner, self.options.precond) + pc = pc_class(getattr(self.mass_ops, id_M)) # Call the accumulation and Pusher class - if use_perp_model: + if self.options.use_perp_model: accum_ker = Pyccelkernel(accum_kernels.pc_lin_mhd_6d) pusher_ker = Pyccelkernel(pusher_kernels.push_pc_GXu) else: accum_ker = Pyccelkernel(accum_kernels.pc_lin_mhd_6d_full) pusher_ker = Pyccelkernel(pusher_kernels.push_pc_GXu_full) - self._coupling_mat = coupling_params["Ah"] / coupling_params["Ab"] - self._coupling_vec = coupling_params["Ah"] / coupling_params["Ab"] - self._scale_push = 1 - - self._boundary_cut_e1 = boundary_cut["e1"] - + # define Accumulator and arguments self._ACC = Accumulator( - particles, - "Hcurl", + self.variables.energetic_ions.particles, + "Hcurl", # TODO:check accum_ker, self.mass_ops, self.domain.args_domain, add_vector=True, symmetry="pressure", - filter_params=filter, + filter_params=self.options.filter_params, ) - self._tmp_g1 = self._G.codomain.zeros() - self._tmp_g2 = self._G.codomain.zeros() - self._tmp_g3 = self._G.codomain.zeros() + self._tmp_g1 = grad.codomain.zeros() + self._tmp_g2 = grad.codomain.zeros() + self._tmp_g3 = grad.codomain.zeros() # instantiate Pusher - args_kernel = ( + args_pusher_kernel = ( self.derham.args_derham, self._tmp_g1[0]._data, self._tmp_g1[1]._data, @@ -674,38 +688,20 @@ def __init__( self._tmp_g3[0]._data, self._tmp_g3[1]._data, self._tmp_g3[2]._data, - self._boundary_cut_e1, ) self._pusher = Pusher( - particles, + self.variables.energetic_ions.particles, pusher_ker, - args_kernel, + args_pusher_kernel, self.domain.args_domain, alpha_in_kernel=1.0, ) - # Define operators - self._A = getattr(self.mass_ops, id_Mn) - self._X = getattr(self.basis_ops, id_X) - self._XT = self._X.transpose() - - # Instantiate schur solver with dummy BC - self._schur_solver = SchurSolver( - self._A, - self._XT @ self._X, - solver["type"][0], - pc=pc, - tol=solver["tol"], - maxiter=solver["maxiter"], - verbose=solver["verbose"], - recycle=solver["recycle"], - ) - - self.u_temp = u.space.zeros() - self.u_temp2 = u.space.zeros() + self.u_temp = self.variables.u.spline.vector.space.zeros() + self.u_temp2 = self.variables.u.spline.vector.space.zeros() self._tmp = self._X.codomain.zeros() - self._BV = u.space.zeros() + self._BV = self.variables.u.spline.vector.space.zeros() self._MAT = [ [self._ACC.operators[0], self._ACC.operators[1], self._ACC.operators[2]], @@ -715,20 +711,32 @@ def __init__( self._GT_VEC = BlockVector(self.derham.Vh["v"]) + _BC = -1 / 4 * self._XT @ self.GT_MAT_G(self.derham, self._MAT) @ self._X + + self._schur_solver = SchurSolver( + _A, + _BC, + self.options.solver, + precond=pc, + solver_params=self.options.solver_params, + ) + def __call__(self, dt): - # current u - un = self.feec_vars[0] - un.update_ghost_regions() + # current FE coeffs + un = self.variables.u.spline.vector + + # operators + grad = self.derham.grad + gradT = grad.transpose() # acuumulate MAT and VEC - self._ACC(self._coupling_mat, self._coupling_vec, self._boundary_cut_e1) + self._ACC( + self.options.ep_scale, + ) # update GT_VEC for i in range(3): - self._GT_VEC[i] = self._GT.dot(self._ACC.vectors[i]) - - # define BC and B dot V of the Schur block matrix [[A, B], [C, I]] - self._schur_solver.BC = -1 / 4 * self._XT @ self.GT_MAT_G(self.derham, self._MAT) @ self._X + self._GT_VEC[i] = gradT.dot(self._ACC.vectors[i]) self._BV = self._XT.dot(self._GT_VEC) * (-1 / 2) @@ -741,9 +749,9 @@ def __call__(self, dt): # calculate GXu Xu = self._X.dot(_u, out=self._tmp) - GXu_1 = self._G.dot(Xu[0], out=self._tmp_g1) - GXu_2 = self._G.dot(Xu[1], out=self._tmp_g2) - GXu_3 = self._G.dot(Xu[2], out=self._tmp_g3) + GXu_1 = grad.dot(Xu[0], out=self._tmp_g1) + GXu_2 = grad.dot(Xu[1], out=self._tmp_g2) + GXu_3 = grad.dot(Xu[2], out=self._tmp_g3) GXu_1.update_ghost_regions() GXu_2.update_ghost_regions() @@ -753,16 +761,16 @@ def __call__(self, dt): self._pusher(dt) # write new coeffs into Propagator.variables - (max_du,) = self.feec_vars_update(un1) + diffs = self.update_feec_variables(u=un1) # update weights in case of control variate - if self.particles[0].control_variate: - self.particles[0].update_weights() + if self.variables.energetic_ions.species.weights_params.control_variate: + self.variables.energetic_ions.particles.update_weights() - if self._info and self._rank == 0: + if self.options.solver_params.info and MPI.COMM_WORLD.Get_rank() == 0: print("Status for StepPressurecoupling:", info["success"]) print("Iterations for StepPressurecoupling:", info["niter"]) - print("Maxdiff u1 for StepPressurecoupling:", max_du) + print("Maxdiff u1 for StepPressurecoupling:", diffs["u"]) print() class GT_MAT_G(LinOpWithTransp): @@ -781,8 +789,8 @@ class GT_MAT_G(LinOpWithTransp): def __init__(self, derham, MAT, transposed=False): self._derham = derham - self._G = derham.grad - self._GT = derham.grad.transpose() + self._grad = derham.grad + self._gradT = derham.grad.transpose() self._domain = derham.Vh["v"] self._codomain = derham.Vh["v"] @@ -836,9 +844,9 @@ def dot(self, v, out=None): for i in range(3): for j in range(3): - self._temp += self._MAT[i][j].dot(self._G.dot(v[j])) + self._temp += self._MAT[i][j].dot(self._grad.dot(v[j])) - self._vector[i] = self._GT.dot(self._temp) + self._vector[i] = self._gradT.dot(self._temp) self._temp *= 0.0 self._vector.update_ghost_regions() diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 69aa86159..0563c9dd1 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -18,6 +18,7 @@ from struphy.io.options import ( OptsKernel, OptsMPIsort, + OptsVecSpace, check_option, ) from struphy.io.setup import descend_options_dict @@ -397,85 +398,106 @@ class PushEtaPC(Propagator): * ``heun3`` (3rd order) """ - @staticmethod - def options(default=False): - dct = {} - dct["use_perp_model"] = [True, False] + class Variables: + def __init__(self): + self._var: PICVariable | SPHVariable = None - if default: - dct = descend_options_dict(dct, []) + @property + def var(self) -> PICVariable | SPHVariable: + return self._var - return dct + @var.setter + def var(self, new): + assert isinstance(new, PICVariable | SPHVariable) + self._var = new - def __init__( - self, - particles: Particles, - *, - u: BlockVector | PolarVector, - use_perp_model: bool = options(default=True)["use_perp_model"], - u_space: str, - ): - super().__init__(particles) + def __init__(self): + self.variables = self.Variables() - assert isinstance(u, (BlockVector, PolarVector)) + @dataclass + class Options: + butcher: ButcherTableau = None + use_perp_model: bool = True + u_tilde: FEECVariable = None + u_space: OptsVecSpace = "Hdiv" - self._u = u + def __post_init__(self): + # checks + check_option(self.u_space, OptsVecSpace) + assert isinstance(self.u_tilde, FEECVariable) - # call Pusher class - if use_perp_model: - if u_space == "Hcurl": - kernel = Pyccelkernel(pusher_kernels.push_pc_eta_rk4_Hcurl) - elif u_space == "Hdiv": - kernel = Pyccelkernel(pusher_kernels.push_pc_eta_rk4_Hdiv) - elif u_space == "H1vec": - kernel = Pyccelkernel(pusher_kernels.push_pc_eta_rk4_H1vec) - else: - raise ValueError( - f'{u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', - ) + # defaults + if self.butcher is None: + self.butcher = ButcherTableau() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + if MPI.COMM_WORLD.Get_rank() == 0: + print(f"\nNew options for propagator '{self.__class__.__name__}':") + for k, v in new.__dict__.items(): + print(f" {k}: {v}") + self._options = new + + @profile + def allocate(self): + self._u_tilde = self.options.u_tilde.spline.vector + + # get kernell: + if self.options.u_space == "Hcurl": + kernel = Pyccelkernel(pusher_kernels.push_pc_eta_stage_Hcurl) + elif self.options.u_space == "Hdiv": + kernel = Pyccelkernel(pusher_kernels.push_pc_eta_stage_Hdiv) + elif self.options.u_space == "H1vec": + kernel = Pyccelkernel(pusher_kernels.push_pc_eta_stage_H1vec) else: - if u_space == "Hcurl": - kernel = Pyccelkernel(pusher_kernels.push_pc_eta_rk4_Hcurl_full) - elif u_space == "Hdiv": - kernel = Pyccelkernel(pusher_kernels.push_pc_eta_rk4_Hdiv_full) - elif u_space == "H1vec": - kernel = Pyccelkernel(pusher_kernels.push_pc_eta_rk4_H1vec_full) - else: - raise ValueError( - f'{u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', - ) + raise ValueError( + f'{self.options.u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', + ) + + # define algorithm + butcher = self.options.butcher + # temp fix due to refactoring of ButcherTableau: + import cunumpy as xp + + butcher._a = xp.diag(butcher.a, k=-1) + butcher._a = xp.array(list(butcher.a) + [0.0]) args_kernel = ( self.derham.args_derham, - self._u[0]._data, - self._u[1]._data, - self._u[2]._data, + self._u_tilde[0]._data, + self._u_tilde[1]._data, + self._u_tilde[2]._data, + self.options.use_perp_model, + butcher.a, + butcher.b, + butcher.c, ) self._pusher = Pusher( - particles, + self.variables.var.particles, kernel, args_kernel, self.domain.args_domain, alpha_in_kernel=1.0, - n_stages=4, + n_stages=butcher.n_stages, mpi_sort="each", ) def __call__(self, dt): - # check if ghost regions are synchronized - if not self._u[0].ghost_regions_in_sync: - self._u[0].update_ghost_regions() - if not self._u[1].ghost_regions_in_sync: - self._u[1].update_ghost_regions() - if not self._u[2].ghost_regions_in_sync: - self._u[2].update_ghost_regions() + self._u_tilde.update_ghost_regions() self._pusher(dt) # update_weights - if self.particles[0].control_variate: - self.particles[0].update_weights() + if self.variables.var.particles.control_variate: + self.variables.var.particles.update_weights() class PushGuidingCenterBxEstar(Propagator): From 050af73d4fa2734c5761734a7601d1ff2b0f521f Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 08:55:18 +0200 Subject: [PATCH 182/292] Update environment variables in GITHUB_ENV --- .github/actions/install/macos-latest/action.yml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index 15c5e503e..050adb02c 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -22,7 +22,8 @@ runs: cmake --version make -v system_profiler SPHardwareDataType - export FC=`which gfortran` # for gvec - export CC=`which gcc` # for gvec - export CXX=`which g++` # for gvec + # Update environment variables + echo "FC=$(which gfortran)" >> $GITHUB_ENV # for gvec + echo "CC=$(which gcc)" >> $GITHUB_ENV # for gvec + echo "CXX=$(which g++)" >> $GITHUB_ENV # for gvec brew install netcdf-fortran From cc52ce24161115ff390a6f7a140e955774f521d4 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 09:04:18 +0200 Subject: [PATCH 183/292] Added env to install struphy stage --- .github/workflows/testing.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index f410539f1..856856213 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -72,6 +72,10 @@ jobs: # Clone struphy-ci-testing - name: Install struphy uses: ./.github/actions/install/clone-and-install-struphy + env: + FC: ${{ env.FC }} + CC: ${{ env.CC }} + CXX: ${{ env.CXX }} # Compile - name: Compile kernels From 9546472f933298b5e02358db3b94e299f4fccccf Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 09:10:29 +0200 Subject: [PATCH 184/292] Added prints --- .../actions/install/clone-and-install-struphy/action.yml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/.github/actions/install/clone-and-install-struphy/action.yml b/.github/actions/install/clone-and-install-struphy/action.yml index c28e6e688..bce677589 100644 --- a/.github/actions/install/clone-and-install-struphy/action.yml +++ b/.github/actions/install/clone-and-install-struphy/action.yml @@ -7,6 +7,13 @@ runs: - name: Install struphy shell: bash run: | + echo $FC + echo $CC + echo $CXX + echo "----------------" + which gfortran + which gcc + which g++ pip install --upgrade pip pip uninstall -y gvec pip install ".[phys,mpi]" From 77302ebb268ac54a0810e4b585e2606c8a77cfe5 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 09:14:53 +0200 Subject: [PATCH 185/292] Show FC in macos install --- .github/actions/install/macos-latest/action.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index 050adb02c..ca2806a48 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -17,6 +17,7 @@ runs: brew install git brew install pandoc brew install hdf5 + brew install netcdf-fortran brew install cmake brew link --overwrite cmake cmake --version @@ -26,4 +27,7 @@ runs: echo "FC=$(which gfortran)" >> $GITHUB_ENV # for gvec echo "CC=$(which gcc)" >> $GITHUB_ENV # for gvec echo "CXX=$(which g++)" >> $GITHUB_ENV # for gvec - brew install netcdf-fortran + echo "----------------" + which gfortran + which gcc + which g++ From ea4f0954372387df3b0ad8b9370cf6b8c77c9506 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 09:27:04 +0200 Subject: [PATCH 186/292] Don't install python3 manually on mac --- .github/actions/install/macos-latest/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index ca2806a48..9d7acd505 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -7,7 +7,7 @@ runs: shell: bash run: | brew update - brew install python3 + # brew install python3 brew install gcc brew install openblas brew install lapack From a3a1c39d38e09d0b3953ec7a1c09ae024af8fab6 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 09:31:52 +0200 Subject: [PATCH 187/292] Remove system profiler print --- .github/actions/install/macos-latest/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index 9d7acd505..2de95ff29 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -22,7 +22,7 @@ runs: brew link --overwrite cmake cmake --version make -v - system_profiler SPHardwareDataType + # system_profiler SPHardwareDataType # Update environment variables echo "FC=$(which gfortran)" >> $GITHUB_ENV # for gvec echo "CC=$(which gcc)" >> $GITHUB_ENV # for gvec From 26343480d23d39d292b21f6cecce16e4268d5ebe Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 09:34:31 +0200 Subject: [PATCH 188/292] changed order --- .github/actions/install/macos-latest/action.yml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index 2de95ff29..060379fd1 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -22,12 +22,13 @@ runs: brew link --overwrite cmake cmake --version make -v - # system_profiler SPHardwareDataType + system_profiler SPHardwareDataType + which gfortran + which gcc + which g++ # Update environment variables echo "FC=$(which gfortran)" >> $GITHUB_ENV # for gvec echo "CC=$(which gcc)" >> $GITHUB_ENV # for gvec echo "CXX=$(which g++)" >> $GITHUB_ENV # for gvec echo "----------------" - which gfortran - which gcc - which g++ + From 1e864e18ddaaa25079ca8a6aa31502f2f1763fc5 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 09:38:16 +0200 Subject: [PATCH 189/292] Added prints --- .github/actions/install/macos-latest/action.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index 060379fd1..85a02c3ee 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -23,9 +23,9 @@ runs: cmake --version make -v system_profiler SPHardwareDataType - which gfortran - which gcc - which g++ + which gfortran || echo "could not find gfortran" + which gcc || echo "could not find gcc" + which g++ || echo "could not find g++" # Update environment variables echo "FC=$(which gfortran)" >> $GITHUB_ENV # for gvec echo "CC=$(which gcc)" >> $GITHUB_ENV # for gvec From b73b72b43f890b45f0e974974cd9a7b841227c71 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 09:52:59 +0200 Subject: [PATCH 190/292] Update PATH --- .github/actions/install/macos-latest/action.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index 85a02c3ee..3bb754917 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -6,6 +6,9 @@ runs: - name: Install dependencies shell: bash run: | + set -euo pipefail + # ensure Homebrew bin is first + export PATH="/opt/homebrew/bin:$PATH" brew update # brew install python3 brew install gcc From 9d2f33890be203e799a81021bdbc9d19b6c9666e Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 09:56:21 +0200 Subject: [PATCH 191/292] uncommented brew install gcc --- .github/actions/install/macos-latest/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/install/macos-latest/action.yml b/.github/actions/install/macos-latest/action.yml index 3bb754917..6e0593685 100644 --- a/.github/actions/install/macos-latest/action.yml +++ b/.github/actions/install/macos-latest/action.yml @@ -11,7 +11,7 @@ runs: export PATH="/opt/homebrew/bin:$PATH" brew update # brew install python3 - brew install gcc + # brew install gcc brew install openblas brew install lapack brew install open-mpi From 482adbfa7767e7dcec9d384a488aae94e02e71df Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 24 Oct 2025 10:24:11 +0200 Subject: [PATCH 192/292] Add Open Source at Microsoft image to README Added an image to the README for Open Source at Microsoft. --- README.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 6a25a91a7..8f6129e5b 100755 --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ +![Open Source at Microsoft](https://private-user-images.githubusercontent.com/181350288/505158519-f7d9fbd6-99a1-4fa5-85c2-8eb8a678888d.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NjEyOTQ0MTAsIm5iZiI6MTc2MTI5NDExMCwicGF0aCI6Ii8xODEzNTAyODgvNTA1MTU4NTE5LWY3ZDlmYmQ2LTk5YTEtNGZhNS04NWMyLThlYjhhNjc4ODg4ZC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjUxMDI0JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MTAyNFQwODIxNTBaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0xNzExNjhlMWM2YmIwYmZkMjhhOGEzY2NiYzNiOWJhMTJmOGI1NzQwNTlmMTNmOTdmY2YyNGFkY2UyYzdmYTdiJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.eTzu_ulrMzJJBUy82FWxbf6tgjdkZdpKYHh1aS0b6SU) + # Struphy - Structure-preserving hybrid codes A Python package for plasma physics PDEs. @@ -52,4 +54,4 @@ Struphy tutorials are available in the form of [Jupyter notebooks](https://gitla * Stefan Possanner [stefan.possanner@ipp.mpg.de](mailto:spossann@ipp.mpg.de) * Eric Sonnendrücker [eric.sonnendruecker@ipp.mpg.de](mailto:eric.sonnendruecker@ipp.mpg.de) -* Xin Wang [xin.wang@ipp.mpg.de](mailto:xin.wang@ipp.mpg.de) \ No newline at end of file +* Xin Wang [xin.wang@ipp.mpg.de](mailto:xin.wang@ipp.mpg.de) From 3e4e1e2cbb5fd37a98917e8778469b3f7e0e9f04 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 24 Oct 2025 10:25:12 +0200 Subject: [PATCH 193/292] Update header image in README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 8f6129e5b..962093256 100755 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -![Open Source at Microsoft](https://private-user-images.githubusercontent.com/181350288/505158519-f7d9fbd6-99a1-4fa5-85c2-8eb8a678888d.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NjEyOTQ0MTAsIm5iZiI6MTc2MTI5NDExMCwicGF0aCI6Ii8xODEzNTAyODgvNTA1MTU4NTE5LWY3ZDlmYmQ2LTk5YTEtNGZhNS04NWMyLThlYjhhNjc4ODg4ZC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjUxMDI0JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MTAyNFQwODIxNTBaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0xNzExNjhlMWM2YmIwYmZkMjhhOGEzY2NiYzNiOWJhMTJmOGI1NzQwNTlmMTNmOTdmY2YyNGFkY2UyYzdmYTdiJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.eTzu_ulrMzJJBUy82FWxbf6tgjdkZdpKYHh1aS0b6SU) +![Struphy header](https://private-user-images.githubusercontent.com/181350288/505158519-f7d9fbd6-99a1-4fa5-85c2-8eb8a678888d.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NjEyOTQ0MTAsIm5iZiI6MTc2MTI5NDExMCwicGF0aCI6Ii8xODEzNTAyODgvNTA1MTU4NTE5LWY3ZDlmYmQ2LTk5YTEtNGZhNS04NWMyLThlYjhhNjc4ODg4ZC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjUxMDI0JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MTAyNFQwODIxNTBaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0xNzExNjhlMWM2YmIwYmZkMjhhOGEzY2NiYzNiOWJhMTJmOGI1NzQwNTlmMTNmOTdmY2YyNGFkY2UyYzdmYTdiJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.eTzu_ulrMzJJBUy82FWxbf6tgjdkZdpKYHh1aS0b6SU) # Struphy - Structure-preserving hybrid codes From 242fa52e8f8710db906241a4951a86e910cd63ee Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 24 Oct 2025 10:34:54 +0200 Subject: [PATCH 194/292] Revise README header and introduction Updated header image and modified introductory text for clarity. --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 962093256..1730f5e00 100755 --- a/README.md +++ b/README.md @@ -1,13 +1,13 @@ -![Struphy header](https://private-user-images.githubusercontent.com/181350288/505158519-f7d9fbd6-99a1-4fa5-85c2-8eb8a678888d.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NjEyOTQ0MTAsIm5iZiI6MTc2MTI5NDExMCwicGF0aCI6Ii8xODEzNTAyODgvNTA1MTU4NTE5LWY3ZDlmYmQ2LTk5YTEtNGZhNS04NWMyLThlYjhhNjc4ODg4ZC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjUxMDI0JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MTAyNFQwODIxNTBaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0xNzExNjhlMWM2YmIwYmZkMjhhOGEzY2NiYzNiOWJhMTJmOGI1NzQwNTlmMTNmOTdmY2YyNGFkY2UyYzdmYTdiJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.eTzu_ulrMzJJBUy82FWxbf6tgjdkZdpKYHh1aS0b6SU) +![Struphy header](https://raw.githubusercontent.com/struphy-hub/.github/refs/heads/main/profile/struphy_header_with_subs.png) -# Struphy - Structure-preserving hybrid codes +# Welcome! -A Python package for plasma physics PDEs. +Struphy is a Python package for plasma physics PDEs. Join the [Struphy mailing list](https://listserv.gwdg.de/mailman/listinfo/struphy) and stay informed on updates. ## Documentation -See the [Struphy pages](https://struphy.pages.mpcdf.de/struphy/index.html) for details regarding installation, tutorials, use and development. +See the [Struphy pages](https://struphy.pages.mpcdf.de/struphy/index.html) for details regarding installation, tutorials, use, and development. ## Quick install From 2fd7bc2237b82334e91bad65f338753515affd7d Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 11:00:03 +0200 Subject: [PATCH 195/292] Renamed install actions --- .../action.yml | 0 .../{clone-and-install-struphy => install-struphy}/action.yml | 0 .github/workflows/static_analysis.yml | 2 +- .github/workflows/testing.yml | 2 +- 4 files changed, 2 insertions(+), 2 deletions(-) rename .github/actions/install/{clone-and-install-struphy-editable => install-struphy-editable}/action.yml (100%) rename .github/actions/install/{clone-and-install-struphy => install-struphy}/action.yml (100%) diff --git a/.github/actions/install/clone-and-install-struphy-editable/action.yml b/.github/actions/install/install-struphy-editable/action.yml similarity index 100% rename from .github/actions/install/clone-and-install-struphy-editable/action.yml rename to .github/actions/install/install-struphy-editable/action.yml diff --git a/.github/actions/install/clone-and-install-struphy/action.yml b/.github/actions/install/install-struphy/action.yml similarity index 100% rename from .github/actions/install/clone-and-install-struphy/action.yml rename to .github/actions/install/install-struphy/action.yml diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index c2f485b78..cfea26997 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -56,7 +56,7 @@ jobs: uses: ./.github/actions/install/ubuntu-latest - name: Install struphy - uses: ./.github/actions/install/clone-and-install-struphy-editable + uses: ./.github/actions/install/install-struphy-editable - name: Lint the repo run: | diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 856856213..a34db4767 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -71,7 +71,7 @@ jobs: # Clone struphy-ci-testing - name: Install struphy - uses: ./.github/actions/install/clone-and-install-struphy + uses: ./.github/actions/install/install-struphy env: FC: ${{ env.FC }} CC: ${{ env.CC }} From ee02500bc03d9dde1058356359be05348a7a1b8a Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 11:01:22 +0200 Subject: [PATCH 196/292] Removed macos-latest --- .github/workflows/testing.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index a34db4767..3404f6348 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -20,7 +20,7 @@ jobs: fail-fast: false matrix: python-version: ["3.12"] - os: [ubuntu-latest, "macos-latest"] + os: ["ubuntu-latest"] #, "macos-latest"] compile-language: ["fortran", "c"] test-type: ["unit", "model"] #, "quickstart"] From 1e551a9cb85a5c48229a0432f8b10718a0ec0d3d Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 11:07:03 +0200 Subject: [PATCH 197/292] Bugfix in the static analysus --- .github/workflows/static_analysis.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index cfea26997..5dc140173 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -28,7 +28,7 @@ jobs: ./cloc --version ./cloc $(git ls-files) - struphy lint all: + struphy_lint_all: runs-on: ubuntu-latest steps: - name: Checkout code From b0edcdeb1b74b5f56fed89b92daf96c40ae6ec16 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 24 Oct 2025 10:24:11 +0200 Subject: [PATCH 198/292] Add Open Source at Microsoft image to README Added an image to the README for Open Source at Microsoft. --- README.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 6a25a91a7..8f6129e5b 100755 --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ +![Open Source at Microsoft](https://private-user-images.githubusercontent.com/181350288/505158519-f7d9fbd6-99a1-4fa5-85c2-8eb8a678888d.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NjEyOTQ0MTAsIm5iZiI6MTc2MTI5NDExMCwicGF0aCI6Ii8xODEzNTAyODgvNTA1MTU4NTE5LWY3ZDlmYmQ2LTk5YTEtNGZhNS04NWMyLThlYjhhNjc4ODg4ZC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjUxMDI0JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MTAyNFQwODIxNTBaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0xNzExNjhlMWM2YmIwYmZkMjhhOGEzY2NiYzNiOWJhMTJmOGI1NzQwNTlmMTNmOTdmY2YyNGFkY2UyYzdmYTdiJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.eTzu_ulrMzJJBUy82FWxbf6tgjdkZdpKYHh1aS0b6SU) + # Struphy - Structure-preserving hybrid codes A Python package for plasma physics PDEs. @@ -52,4 +54,4 @@ Struphy tutorials are available in the form of [Jupyter notebooks](https://gitla * Stefan Possanner [stefan.possanner@ipp.mpg.de](mailto:spossann@ipp.mpg.de) * Eric Sonnendrücker [eric.sonnendruecker@ipp.mpg.de](mailto:eric.sonnendruecker@ipp.mpg.de) -* Xin Wang [xin.wang@ipp.mpg.de](mailto:xin.wang@ipp.mpg.de) \ No newline at end of file +* Xin Wang [xin.wang@ipp.mpg.de](mailto:xin.wang@ipp.mpg.de) From 6f3b748e58fe10524f72b63d67de150ea7885704 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 24 Oct 2025 10:25:12 +0200 Subject: [PATCH 199/292] Update header image in README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 8f6129e5b..962093256 100755 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -![Open Source at Microsoft](https://private-user-images.githubusercontent.com/181350288/505158519-f7d9fbd6-99a1-4fa5-85c2-8eb8a678888d.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NjEyOTQ0MTAsIm5iZiI6MTc2MTI5NDExMCwicGF0aCI6Ii8xODEzNTAyODgvNTA1MTU4NTE5LWY3ZDlmYmQ2LTk5YTEtNGZhNS04NWMyLThlYjhhNjc4ODg4ZC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjUxMDI0JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MTAyNFQwODIxNTBaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0xNzExNjhlMWM2YmIwYmZkMjhhOGEzY2NiYzNiOWJhMTJmOGI1NzQwNTlmMTNmOTdmY2YyNGFkY2UyYzdmYTdiJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.eTzu_ulrMzJJBUy82FWxbf6tgjdkZdpKYHh1aS0b6SU) +![Struphy header](https://private-user-images.githubusercontent.com/181350288/505158519-f7d9fbd6-99a1-4fa5-85c2-8eb8a678888d.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NjEyOTQ0MTAsIm5iZiI6MTc2MTI5NDExMCwicGF0aCI6Ii8xODEzNTAyODgvNTA1MTU4NTE5LWY3ZDlmYmQ2LTk5YTEtNGZhNS04NWMyLThlYjhhNjc4ODg4ZC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjUxMDI0JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MTAyNFQwODIxNTBaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0xNzExNjhlMWM2YmIwYmZkMjhhOGEzY2NiYzNiOWJhMTJmOGI1NzQwNTlmMTNmOTdmY2YyNGFkY2UyYzdmYTdiJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.eTzu_ulrMzJJBUy82FWxbf6tgjdkZdpKYHh1aS0b6SU) # Struphy - Structure-preserving hybrid codes From 5736beea51e84eb922c0af32b9c5d9b88b2f5cc4 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 24 Oct 2025 10:34:54 +0200 Subject: [PATCH 200/292] Revise README header and introduction Updated header image and modified introductory text for clarity. --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 962093256..1730f5e00 100755 --- a/README.md +++ b/README.md @@ -1,13 +1,13 @@ -![Struphy header](https://private-user-images.githubusercontent.com/181350288/505158519-f7d9fbd6-99a1-4fa5-85c2-8eb8a678888d.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NjEyOTQ0MTAsIm5iZiI6MTc2MTI5NDExMCwicGF0aCI6Ii8xODEzNTAyODgvNTA1MTU4NTE5LWY3ZDlmYmQ2LTk5YTEtNGZhNS04NWMyLThlYjhhNjc4ODg4ZC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjUxMDI0JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MTAyNFQwODIxNTBaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT0xNzExNjhlMWM2YmIwYmZkMjhhOGEzY2NiYzNiOWJhMTJmOGI1NzQwNTlmMTNmOTdmY2YyNGFkY2UyYzdmYTdiJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.eTzu_ulrMzJJBUy82FWxbf6tgjdkZdpKYHh1aS0b6SU) +![Struphy header](https://raw.githubusercontent.com/struphy-hub/.github/refs/heads/main/profile/struphy_header_with_subs.png) -# Struphy - Structure-preserving hybrid codes +# Welcome! -A Python package for plasma physics PDEs. +Struphy is a Python package for plasma physics PDEs. Join the [Struphy mailing list](https://listserv.gwdg.de/mailman/listinfo/struphy) and stay informed on updates. ## Documentation -See the [Struphy pages](https://struphy.pages.mpcdf.de/struphy/index.html) for details regarding installation, tutorials, use and development. +See the [Struphy pages](https://struphy.pages.mpcdf.de/struphy/index.html) for details regarding installation, tutorials, use, and development. ## Quick install From 68fee4ad5626050f2a70560da0b738cb70470bb4 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 11:51:01 +0200 Subject: [PATCH 201/292] Added struphy --refresh-models --- .github/actions/install/install-struphy/action.yml | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/.github/actions/install/install-struphy/action.yml b/.github/actions/install/install-struphy/action.yml index bce677589..007c4233f 100644 --- a/.github/actions/install/install-struphy/action.yml +++ b/.github/actions/install/install-struphy/action.yml @@ -7,15 +7,9 @@ runs: - name: Install struphy shell: bash run: | - echo $FC - echo $CC - echo $CXX - echo "----------------" - which gfortran - which gcc - which g++ pip install --upgrade pip pip uninstall -y gvec pip install ".[phys,mpi]" pip list struphy -h + struphy --refresh-models From b6b35a4bb27c208d0d808cbcd44ccd85d7c38ef4 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 11:52:50 +0200 Subject: [PATCH 202/292] Updated model tests --- .github/actions/tests/models/action.yml | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/.github/actions/tests/models/action.yml b/.github/actions/tests/models/action.yml index 0a9fafc16..435276caa 100644 --- a/.github/actions/tests/models/action.yml +++ b/.github/actions/tests/models/action.yml @@ -3,14 +3,21 @@ name: "Run model tests" runs: using: composite steps: - - name: Install dependencies + - name: Model tests shell: bash run: | struphy compile --status + struphy test LinearMHD + struphy test toy + struphy test models + struphy test verification + - name: Model tests with MPI + shell: bash + run: | struphy compile --status - struphy test models --fast - struphy test models --fast --mpi 2 - struphy test models --fast --verification --mpi 1 - struphy test models --fast --verification --mpi 4 - struphy test models --fast --verification --mpi 4 --nclones 2 - struphy test DriftKineticElectrostaticAdiabatic --mpi 2 --nclones 2 \ No newline at end of file + struphy test models + struphy test models --mpi 2 + struphy test verification --mpi 1 + struphy test verification --mpi 4 + struphy test verification --mpi 4 --nclones 2 + struphy test VlasovAmpereOneSpecies --mpi 2 --nclones 2 From d6d47894dbfd0e37f0af5ace81ef50acd5f2679a Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 11:53:54 +0200 Subject: [PATCH 203/292] Updated unit tests --- .github/actions/tests/unit/action.yml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/.github/actions/tests/unit/action.yml b/.github/actions/tests/unit/action.yml index 2d0e70a2d..73edd4d03 100644 --- a/.github/actions/tests/unit/action.yml +++ b/.github/actions/tests/unit/action.yml @@ -6,4 +6,13 @@ runs: - name: Run unit tests shell: bash run: | + struphy compile --status + struphy --refresh-models struphy test unit + + - name: Run unit tests with MPI + shell: bash + run: | + struphy compile --status + struphy --refresh-models + struphy test unit --mpi 2 From 0baca29c76ef5966915b8c309819f6019725aa75 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 11:55:22 +0200 Subject: [PATCH 204/292] Added back the quickstart tests --- .github/actions/tests/quickstart/action.yml | 2 +- .github/workflows/testing.yml | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/actions/tests/quickstart/action.yml b/.github/actions/tests/quickstart/action.yml index b78ade5e7..3845360d3 100644 --- a/.github/actions/tests/quickstart/action.yml +++ b/.github/actions/tests/quickstart/action.yml @@ -8,7 +8,7 @@ runs: run: | struphy -p struphy -h - struphy params VlasovAmpereOneSpecies -y + struphy params VlasovAmpereOneSpecies ls -1a mv params_VlasovAmpereOneSpecies.py test.py python3 test.py diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 3404f6348..fab933fb0 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -22,7 +22,7 @@ jobs: python-version: ["3.12"] os: ["ubuntu-latest"] #, "macos-latest"] compile-language: ["fortran", "c"] - test-type: ["unit", "model"] #, "quickstart"] + test-type: ["unit", "model", "quickstart"] steps: # Checkout the repository @@ -90,6 +90,6 @@ jobs: if: matrix.test-type == 'model' uses: ./.github/actions/tests/models - #- name: Run quickstart tests - # if: matrix.test-type == 'quickstart' - # uses: ./.github/actions/tests/quickstart + - name: Run quickstart tests + if: matrix.test-type == 'quickstart' + uses: ./.github/actions/tests/quickstart From bf39ab06311f4649ffd41bb16b3976e763eccbc0 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 12:00:11 +0200 Subject: [PATCH 205/292] Uninstall mpi4py for unit tests --- .github/actions/tests/unit/action.yml | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/.github/actions/tests/unit/action.yml b/.github/actions/tests/unit/action.yml index 73edd4d03..385d6f419 100644 --- a/.github/actions/tests/unit/action.yml +++ b/.github/actions/tests/unit/action.yml @@ -3,16 +3,19 @@ name: "Run unit tests" runs: using: composite steps: - - name: Run unit tests + - name: Run unit tests with MPI shell: bash run: | struphy compile --status struphy --refresh-models - struphy test unit + struphy test unit --mpi 2 - - name: Run unit tests with MPI + - name: Run unit tests shell: bash run: | struphy compile --status struphy --refresh-models - struphy test unit --mpi 2 + pip show mpi4py + pip uninstall -y mpi4py + pip list + struphy test unit From 746b2fbbb1d6a1734380b4027d01ffd1087e129b Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 13:12:42 +0200 Subject: [PATCH 206/292] Added tutorial tests --- .github/actions/tests/tutorials/action.yml | 12 ++++++++++++ .github/workflows/testing.yml | 6 +++++- 2 files changed, 17 insertions(+), 1 deletion(-) create mode 100644 .github/actions/tests/tutorials/action.yml diff --git a/.github/actions/tests/tutorials/action.yml b/.github/actions/tests/tutorials/action.yml new file mode 100644 index 000000000..17d41603f --- /dev/null +++ b/.github/actions/tests/tutorials/action.yml @@ -0,0 +1,12 @@ +name: "Run tutorial tests" + +runs: + using: composite + steps: + - name: Run turorial tests + shell: bash + run: | + pwd + ls -1a + which python + jupyter nbconvert --to notebook --execute tutorials/*.ipynb diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index fab933fb0..f0950702a 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -22,7 +22,7 @@ jobs: python-version: ["3.12"] os: ["ubuntu-latest"] #, "macos-latest"] compile-language: ["fortran", "c"] - test-type: ["unit", "model", "quickstart"] + test-type: ["unit", "unit-mpi", "model", "model-mpi", "quickstart", "tutorials"] steps: # Checkout the repository @@ -93,3 +93,7 @@ jobs: - name: Run quickstart tests if: matrix.test-type == 'quickstart' uses: ./.github/actions/tests/quickstart + + - name: Run tutorials + if: matrix.test-type == 'tutorials' + uses: ./.github/actions/tests/tutorials From d68bec01868608fd20723901fbb9607cb286fe62 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 13:13:36 +0200 Subject: [PATCH 207/292] Updated isort job --- .github/workflows/static_analysis.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index 5dc140173..7fff5f1c3 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -86,7 +86,7 @@ jobs: run: | pip install isort isort --check src/ - isort --check doc/tutorials/ + isort --check tutorials/ # mypy: # runs-on: ubuntu-latest From ea215c08676988ded0bb42f441888ad5c421a5bb Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 13:27:33 +0200 Subject: [PATCH 208/292] Install with doc --- .github/actions/install/install-struphy/action.yml | 2 +- .github/workflows/testing.yml | 2 +- pyproject.toml | 3 +++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/actions/install/install-struphy/action.yml b/.github/actions/install/install-struphy/action.yml index 007c4233f..471c9822a 100644 --- a/.github/actions/install/install-struphy/action.yml +++ b/.github/actions/install/install-struphy/action.yml @@ -9,7 +9,7 @@ runs: run: | pip install --upgrade pip pip uninstall -y gvec - pip install ".[phys,mpi]" + pip install ".[phys,mpi,doc]" pip list struphy -h struphy --refresh-models diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index f0950702a..654fb0fc2 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -22,7 +22,7 @@ jobs: python-version: ["3.12"] os: ["ubuntu-latest"] #, "macos-latest"] compile-language: ["fortran", "c"] - test-type: ["unit", "unit-mpi", "model", "model-mpi", "quickstart", "tutorials"] + test-type: ["unit", "model", "quickstart", "tutorials"] steps: # Checkout the repository diff --git a/pyproject.toml b/pyproject.toml index c9e68a74f..c4a2c9d27 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,9 @@ mpi = [ ] doc = [ "struphy[phys]", + "jupyter", + "nbconvert", + "ipykernel", "sphinx", "sphinx-design", "lxml_html_clean", From 4f7db70ccee9a554e6c0f0a0e82e209ee65e1bd0 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 13:47:34 +0200 Subject: [PATCH 209/292] Added pull request template --- .../pull_request_template.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 .github/PULL_REQUEST_TEMPLATE/pull_request_template.md diff --git a/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md b/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md new file mode 100644 index 000000000..87f9ff3cc --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md @@ -0,0 +1,15 @@ +**Solves the following issue(s):** + +Closes #... + +**Core changes:** + +None + +**Model-specific changes:** + +None + +**Documentation changes:** + +None \ No newline at end of file From d838f36341fd998947946239883d5359056a6270 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 14:27:38 +0200 Subject: [PATCH 210/292] fixes --- .github/actions/install/install-struphy-editable/action.yml | 2 +- .github/actions/install/install-struphy/action.yml | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/actions/install/install-struphy-editable/action.yml b/.github/actions/install/install-struphy-editable/action.yml index 49d871e15..f395d5caa 100644 --- a/.github/actions/install/install-struphy-editable/action.yml +++ b/.github/actions/install/install-struphy-editable/action.yml @@ -9,6 +9,6 @@ runs: run: | pip install --upgrade pip pip uninstall -y gvec - pip install -e ".[dev,mpi]" + pip install -e ".[dev,mpi,doc]" pip list struphy -h diff --git a/.github/actions/install/install-struphy/action.yml b/.github/actions/install/install-struphy/action.yml index 471c9822a..a9589c4c4 100644 --- a/.github/actions/install/install-struphy/action.yml +++ b/.github/actions/install/install-struphy/action.yml @@ -8,7 +8,6 @@ runs: shell: bash run: | pip install --upgrade pip - pip uninstall -y gvec pip install ".[phys,mpi,doc]" pip list struphy -h From b3b83aa7d54400d2f479c7d746ee60f5ff8aee70 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 14:43:02 +0200 Subject: [PATCH 211/292] Moved tutorial 7 out of the tutorials/ dir so it's not executed in the CI --- ..._07_data_structures.ipynb => tutorial_07_data_structures.ipynb | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tutorials/tutorial_07_data_structures.ipynb => tutorial_07_data_structures.ipynb (100%) diff --git a/tutorials/tutorial_07_data_structures.ipynb b/tutorial_07_data_structures.ipynb similarity index 100% rename from tutorials/tutorial_07_data_structures.ipynb rename to tutorial_07_data_structures.ipynb From ba2bab1aaffcd961ea665c55627a2032a0ede4e7 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 14:50:22 +0200 Subject: [PATCH 212/292] Moved templates --- .github/issue_template.md | 13 +++++++++++++ .github/pull_request_template.md | 15 +++++++++++++++ 2 files changed, 28 insertions(+) create mode 100644 .github/issue_template.md create mode 100644 .github/pull_request_template.md diff --git a/.github/issue_template.md b/.github/issue_template.md new file mode 100644 index 000000000..b2c4eb7bc --- /dev/null +++ b/.github/issue_template.md @@ -0,0 +1,13 @@ +**Bug description / feature request:** + +... + +**Expected behavior:** + +... + +**Proposed solution:** + +... + + diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 000000000..87f9ff3cc --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,15 @@ +**Solves the following issue(s):** + +Closes #... + +**Core changes:** + +None + +**Model-specific changes:** + +None + +**Documentation changes:** + +None \ No newline at end of file From fb0d0e886eaa9bf09eba1699312b1548740da77d Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:03:35 +0200 Subject: [PATCH 213/292] Added badges --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 1730f5e00..9d4630a74 100755 --- a/README.md +++ b/README.md @@ -1,4 +1,5 @@ -![Struphy header](https://raw.githubusercontent.com/struphy-hub/.github/refs/heads/main/profile/struphy_header_with_subs.png) +[![Testing](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml) +![PyPI](https://img.shields.io/pypi/v/struphy?label=pypi%20package) # Welcome! From 6524c9bd18b808afc71cc82213b70c93499a4f33 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:05:50 +0200 Subject: [PATCH 214/292] Added stars badge --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 9d4630a74..7f2d80d49 100755 --- a/README.md +++ b/README.md @@ -1,5 +1,6 @@ [![Testing](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml) ![PyPI](https://img.shields.io/pypi/v/struphy?label=pypi%20package) + # Welcome! From 02da58cf5b35d2cea467d17c1f3bd667cfc1ecef Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:07:23 +0200 Subject: [PATCH 215/292] Added license badge --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 7f2d80d49..6be8f68a3 100755 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ [![Testing](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml) ![PyPI](https://img.shields.io/pypi/v/struphy?label=pypi%20package) - +![License](https://img.shields.io/badge/License-MIT-violet) # Welcome! From e063917bf3a0c50a414d86d5838d2afd3191c627 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:09:18 +0200 Subject: [PATCH 216/292] Added release --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 6be8f68a3..73e7b24fd 100755 --- a/README.md +++ b/README.md @@ -1,5 +1,6 @@ [![Testing](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml) ![PyPI](https://img.shields.io/pypi/v/struphy?label=pypi%20package) +![Release](https://img.shields.io/github/v/release/struphy-hub/struphy) ![License](https://img.shields.io/badge/License-MIT-violet) # Welcome! From 4efb4690b66a5ded095a774578299f0e94c7ada2 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:10:28 +0200 Subject: [PATCH 217/292] Added static analysis --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 73e7b24fd..c1432c460 100755 --- a/README.md +++ b/README.md @@ -1,4 +1,5 @@ [![Testing](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml) +[![Testing](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml) ![PyPI](https://img.shields.io/pypi/v/struphy?label=pypi%20package) ![Release](https://img.shields.io/github/v/release/struphy-hub/struphy) ![License](https://img.shields.io/badge/License-MIT-violet) From e534953986759ba1c9dca2e6a6defee3cb87622c Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:11:41 +0200 Subject: [PATCH 218/292] Added pypi downloads --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index c1432c460..6cd28a6eb 100755 --- a/README.md +++ b/README.md @@ -1,6 +1,8 @@ [![Testing](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml) [![Testing](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml) ![PyPI](https://img.shields.io/pypi/v/struphy?label=pypi%20package) +[![PyPI Downloads](https://img.shields.io/pypi/dm/struphy.svg?label=PyPI%20downloads)]( +https://pypi.org/project/struphy/) ![Release](https://img.shields.io/github/v/release/struphy-hub/struphy) ![License](https://img.shields.io/badge/License-MIT-violet) From c456d99443086f515a9b2278998ed611c563a604 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:22:02 +0200 Subject: [PATCH 219/292] Reusable workflow --- .github/workflows/macos-latest.yml | 8 ++++++++ .github/workflows/main.yml | 23 +++++++++++++++++++++++ .github/workflows/ubuntu-latest.yml | 8 ++++++++ 3 files changed, 39 insertions(+) create mode 100644 .github/workflows/macos-latest.yml create mode 100644 .github/workflows/main.yml create mode 100644 .github/workflows/ubuntu-latest.yml diff --git a/.github/workflows/macos-latest.yml b/.github/workflows/macos-latest.yml new file mode 100644 index 000000000..e4028e472 --- /dev/null +++ b/.github/workflows/macos-latest.yml @@ -0,0 +1,8 @@ +name: macos-latest +on: push + +jobs: + macos-latest-build: + uses: ./.github/workflows/main.yml + with: + os: macos-latest \ No newline at end of file diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml new file mode 100644 index 000000000..b1ba9d2fc --- /dev/null +++ b/.github/workflows/main.yml @@ -0,0 +1,23 @@ +name: CI +on: + # workflow_call is a required field of the "on" property, + # when intending to use the workflow as a reusable workflow + workflow_call: + inputs: + os: + required: true + type: string + +jobs: + build: + runs-on: ${{ inputs.os }} + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + - name: Set up Go 1.22 + uses: actions/setup-go@v5 + with: + go-version: 1.22 + - name: Build application + run: echo "Hello from ${{ inputs.os }}" diff --git a/.github/workflows/ubuntu-latest.yml b/.github/workflows/ubuntu-latest.yml new file mode 100644 index 000000000..f7be8f138 --- /dev/null +++ b/.github/workflows/ubuntu-latest.yml @@ -0,0 +1,8 @@ +name: ubuntu-latest +on: push + +jobs: + ubuntu-latest-build: + uses: ./.github/workflows/main.yml + with: + os: ubuntu-latest \ No newline at end of file From 3f6485458e06bf06a539882fdf2dcbfc67de8a93 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:24:22 +0200 Subject: [PATCH 220/292] added badges --- README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/README.md b/README.md index 6cd28a6eb..d493f2ed3 100755 --- a/README.md +++ b/README.md @@ -1,5 +1,9 @@ [![Testing](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml) [![Testing](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml) + +[![Ubuntu latest](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml) +[![MacOS latest](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml) + ![PyPI](https://img.shields.io/pypi/v/struphy?label=pypi%20package) [![PyPI Downloads](https://img.shields.io/pypi/dm/struphy.svg?label=PyPI%20downloads)]( https://pypi.org/project/struphy/) From 8f1038be1035fa1f473bfd547af457ed730df092 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:26:43 +0200 Subject: [PATCH 221/292] Removed the old testing badge --- README.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/README.md b/README.md index d493f2ed3..7b1ab39e9 100755 --- a/README.md +++ b/README.md @@ -1,6 +1,4 @@ -[![Testing](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/testing.yml) [![Testing](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml) - [![Ubuntu latest](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml) [![MacOS latest](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml) From 6f2732d0dc84689b5ab1fd1d4f9e98e52b6024ce Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:27:21 +0200 Subject: [PATCH 222/292] Removed the old testing badge --- README.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/README.md b/README.md index 7b1ab39e9..e9fc36925 100755 --- a/README.md +++ b/README.md @@ -1,7 +1,5 @@ -[![Testing](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml) [![Ubuntu latest](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml) [![MacOS latest](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml) - ![PyPI](https://img.shields.io/pypi/v/struphy?label=pypi%20package) [![PyPI Downloads](https://img.shields.io/pypi/dm/struphy.svg?label=PyPI%20downloads)]( https://pypi.org/project/struphy/) From 4c4f967a823769c1666e3a227b1fd5a733027161 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:41:50 +0200 Subject: [PATCH 223/292] Renamed static analysis --- .github/workflows/static_analysis.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index 5dc140173..f8a89bf9b 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -1,4 +1,4 @@ -name: Static analysis +name: isort & ruff on: push: From d3a29d389a32e8da0b7fb3b7c4f2b4956c457fe2 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:45:43 +0200 Subject: [PATCH 224/292] Renaming --- .github/workflows/macos-latest.yml | 2 +- .github/workflows/ubuntu-latest.yml | 2 +- README.md | 5 +++-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/.github/workflows/macos-latest.yml b/.github/workflows/macos-latest.yml index e4028e472..706d1ff1e 100644 --- a/.github/workflows/macos-latest.yml +++ b/.github/workflows/macos-latest.yml @@ -1,4 +1,4 @@ -name: macos-latest +name: Mac OS on: push jobs: diff --git a/.github/workflows/ubuntu-latest.yml b/.github/workflows/ubuntu-latest.yml index f7be8f138..abe4bde57 100644 --- a/.github/workflows/ubuntu-latest.yml +++ b/.github/workflows/ubuntu-latest.yml @@ -1,4 +1,4 @@ -name: ubuntu-latest +name: Ubuntu on: push jobs: diff --git a/README.md b/README.md index e9fc36925..172b7b3fc 100755 --- a/README.md +++ b/README.md @@ -1,9 +1,10 @@ [![Ubuntu latest](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml) [![MacOS latest](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml) -![PyPI](https://img.shields.io/pypi/v/struphy?label=pypi%20package) +[![isort and ruff](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml) +![PyPI](https://img.shields.io/pypi/v/struphy?label=PyPI) [![PyPI Downloads](https://img.shields.io/pypi/dm/struphy.svg?label=PyPI%20downloads)]( https://pypi.org/project/struphy/) -![Release](https://img.shields.io/github/v/release/struphy-hub/struphy) +![Release](https://img.shields.io/github/v/release/struphy-hub/struphy?label=Release) ![License](https://img.shields.io/badge/License-MIT-violet) # Welcome! From 2652ebe122a5ae60b26f571d9e462f1cd6d642e9 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:49:43 +0200 Subject: [PATCH 225/292] MacOS --- .github/workflows/macos-latest.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/macos-latest.yml b/.github/workflows/macos-latest.yml index 706d1ff1e..00017f112 100644 --- a/.github/workflows/macos-latest.yml +++ b/.github/workflows/macos-latest.yml @@ -1,4 +1,4 @@ -name: Mac OS +name: MacOS on: push jobs: From fb00f7fa9db2263ce2371fa793e72633e6bb1ab0 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:52:55 +0200 Subject: [PATCH 226/292] Added concurrency --- .github/workflows/docs.yml | 2 +- .github/workflows/macos-latest.yml | 4 ++++ .github/workflows/static_analysis.yml | 3 +++ .github/workflows/testing.yml | 4 ++++ .github/workflows/ubuntu-latest.yml | 4 ++++ 5 files changed, 16 insertions(+), 1 deletion(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 79caa16e7..b4827d372 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -2,7 +2,7 @@ name: Deploy docs to GitHub Pages on: push: - branches: ["devel", "main"] # TODO: Set to main only after release + branches: ["devel", "main"] workflow_dispatch: defaults: diff --git a/.github/workflows/macos-latest.yml b/.github/workflows/macos-latest.yml index 00017f112..fd95c2658 100644 --- a/.github/workflows/macos-latest.yml +++ b/.github/workflows/macos-latest.yml @@ -1,6 +1,10 @@ name: MacOS on: push +concurrency: + group: ${{ github.ref }} + cancel-in-progress: true + jobs: macos-latest-build: uses: ./.github/workflows/main.yml diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index e06cb1d4d..81b248754 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -10,6 +10,9 @@ on: - main - devel +concurrency: + group: ${{ github.ref }} + cancel-in-progress: true defaults: run: shell: bash diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 654fb0fc2..a95e3c365 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -10,6 +10,10 @@ on: - main - devel +concurrency: + group: ${{ github.ref }} + cancel-in-progress: true + jobs: test: runs-on: ${{ matrix.os }} diff --git a/.github/workflows/ubuntu-latest.yml b/.github/workflows/ubuntu-latest.yml index abe4bde57..783f03505 100644 --- a/.github/workflows/ubuntu-latest.yml +++ b/.github/workflows/ubuntu-latest.yml @@ -1,6 +1,10 @@ name: Ubuntu on: push +concurrency: + group: ${{ github.ref }} + cancel-in-progress: true + jobs: ubuntu-latest-build: uses: ./.github/workflows/main.yml From b5b1a60309dfaa4884bdf6ad78046282cade636d Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 15:58:41 +0200 Subject: [PATCH 227/292] Use testing.yml as reusable workflow --- .github/workflows/macos-latest.yml | 12 ++++++++++-- .github/workflows/main.yml | 23 ----------------------- .github/workflows/testing.yml | 23 +++++++++-------------- .github/workflows/ubuntu-latest.yml | 12 ++++++++++-- 4 files changed, 29 insertions(+), 41 deletions(-) delete mode 100644 .github/workflows/main.yml diff --git a/.github/workflows/macos-latest.yml b/.github/workflows/macos-latest.yml index fd95c2658..48aacf8b7 100644 --- a/.github/workflows/macos-latest.yml +++ b/.github/workflows/macos-latest.yml @@ -1,5 +1,13 @@ name: MacOS -on: push +on: + push: + branches: + - main + - devel + pull_request: + branches: + - main + - devel concurrency: group: ${{ github.ref }} @@ -7,6 +15,6 @@ concurrency: jobs: macos-latest-build: - uses: ./.github/workflows/main.yml + uses: ./.github/workflows/testing.yml with: os: macos-latest \ No newline at end of file diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml deleted file mode 100644 index b1ba9d2fc..000000000 --- a/.github/workflows/main.yml +++ /dev/null @@ -1,23 +0,0 @@ -name: CI -on: - # workflow_call is a required field of the "on" property, - # when intending to use the workflow as a reusable workflow - workflow_call: - inputs: - os: - required: true - type: string - -jobs: - build: - runs-on: ${{ inputs.os }} - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - - name: Set up Go 1.22 - uses: actions/setup-go@v5 - with: - go-version: 1.22 - - name: Build application - run: echo "Hello from ${{ inputs.os }}" diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index a95e3c365..1ac7721a7 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -1,22 +1,18 @@ name: Testing on: - push: - branches: - - main - - devel - pull_request: - branches: - - main - - devel - + workflow_call: + inputs: + os: + required: true + type: string concurrency: group: ${{ github.ref }} cancel-in-progress: true jobs: test: - runs-on: ${{ matrix.os }} + runs-on: ${{ inputs.os }} env: OMPI_MCA_rmaps_base_oversubscribe: 1 # Linux PRRTE_MCA_rmaps_base_oversubscribe: 1 # MacOS @@ -24,7 +20,6 @@ jobs: fail-fast: false matrix: python-version: ["3.12"] - os: ["ubuntu-latest"] #, "macos-latest"] compile-language: ["fortran", "c"] test-type: ["unit", "model", "quickstart", "tutorials"] @@ -54,13 +49,13 @@ jobs: # Install prereqs # I don't think it's possible to use a single action for this because - # we can't use ${matrix.os} in an if statement, so we have to use two different actions. + # we can't use ${inputs.os} in an if statement, so we have to use two different actions. - name: Install prerequisites (Ubuntu) - if: matrix.os == 'ubuntu-latest' + if: inputs.os == 'ubuntu-latest' uses: ./.github/actions/install/ubuntu-latest - name: Install prerequisites (macOS) - if: matrix.os == 'macos-latest' + if: inputs.os == 'macos-latest' uses: ./.github/actions/install/macos-latest # Check that mpirun oversubscribing works, doesn't work unless OMPI_MCA_rmaps_base_oversubscribe==1 diff --git a/.github/workflows/ubuntu-latest.yml b/.github/workflows/ubuntu-latest.yml index 783f03505..864d9861d 100644 --- a/.github/workflows/ubuntu-latest.yml +++ b/.github/workflows/ubuntu-latest.yml @@ -1,5 +1,13 @@ name: Ubuntu -on: push +on: + push: + branches: + - main + - devel + pull_request: + branches: + - main + - devel concurrency: group: ${{ github.ref }} @@ -7,6 +15,6 @@ concurrency: jobs: ubuntu-latest-build: - uses: ./.github/workflows/main.yml + uses: ./.github/workflows/testing.yml with: os: ubuntu-latest \ No newline at end of file From 7b061899456dd013aecd75445bd30f1b9f2ddf62 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 16:00:11 +0200 Subject: [PATCH 228/292] trigger CI From b4c483732f6c565328f64b538ae4bcf3fe334e8a Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 16:02:55 +0200 Subject: [PATCH 229/292] Change concurrency --- .github/workflows/testing.yml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 1ac7721a7..c27af6914 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -6,9 +6,6 @@ on: os: required: true type: string -concurrency: - group: ${{ github.ref }} - cancel-in-progress: true jobs: test: From 9790beb252cb130eef65bba9adde4589b8318ea7 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 16:05:34 +0200 Subject: [PATCH 230/292] Removed concurrency --- .github/workflows/macos-latest.yml | 6 +++--- .github/workflows/static_analysis.yml | 6 +++--- .github/workflows/ubuntu-latest.yml | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/.github/workflows/macos-latest.yml b/.github/workflows/macos-latest.yml index 48aacf8b7..8ae847289 100644 --- a/.github/workflows/macos-latest.yml +++ b/.github/workflows/macos-latest.yml @@ -9,9 +9,9 @@ on: - main - devel -concurrency: - group: ${{ github.ref }} - cancel-in-progress: true +# concurrency: +# group: ${{ github.ref }} +# cancel-in-progress: true jobs: macos-latest-build: diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index 81b248754..2fd4c2051 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -10,9 +10,9 @@ on: - main - devel -concurrency: - group: ${{ github.ref }} - cancel-in-progress: true +# concurrency: +# group: ${{ github.ref }} +# cancel-in-progress: true defaults: run: shell: bash diff --git a/.github/workflows/ubuntu-latest.yml b/.github/workflows/ubuntu-latest.yml index 864d9861d..e66aeb3b3 100644 --- a/.github/workflows/ubuntu-latest.yml +++ b/.github/workflows/ubuntu-latest.yml @@ -9,9 +9,9 @@ on: - main - devel -concurrency: - group: ${{ github.ref }} - cancel-in-progress: true +# concurrency: +# group: ${{ github.ref }} +# cancel-in-progress: true jobs: ubuntu-latest-build: From 27da0de95664b2dd9a2b5a4891d0c71051de3968 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 16:07:27 +0200 Subject: [PATCH 231/292] Added header --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 172b7b3fc..db8f90368 100755 --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ +![Struphy header](https://raw.githubusercontent.com/struphy-hub/.github/refs/heads/main/profile/struphy_header_with_subs.png) + [![Ubuntu latest](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml) [![MacOS latest](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml) [![isort and ruff](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml) From 8d449b406c3e4a604489ac6450563b85e70cefed Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 16:14:24 +0200 Subject: [PATCH 232/292] New banner --- README.md | 4 +++- src/struphy/models/base.py | 3 +++ src/struphy/models/toy.py | 3 +++ 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index db8f90368..8c6c734ba 100755 --- a/README.md +++ b/README.md @@ -1,4 +1,6 @@ -![Struphy header](https://raw.githubusercontent.com/struphy-hub/.github/refs/heads/main/profile/struphy_header_with_subs.png) +

+ +


[![Ubuntu latest](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml) [![MacOS latest](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index d18942154..317121ccc 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -44,6 +44,9 @@ class StruphyModel(metaclass=ABCMeta): """ Base class for all Struphy models. +

x2 + y2 = z2

+ Note ---- All Struphy models are subclasses of ``StruphyModel`` and should be added to ``struphy/models/`` diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index bad2b7916..aa295c513 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -14,6 +14,9 @@ class Maxwell(StruphyModel): r"""Maxwell's equations in vacuum. +

x2 + y2 = z2

+ :ref:`normalization`: .. math:: From da107313b5611b118eb2fe2daee9f34d61fc7c24 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 16:26:10 +0200 Subject: [PATCH 233/292] Commented out Mac OS latest --- .github/workflows/macos-latest.yml | 36 +++++++++++++++--------------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/.github/workflows/macos-latest.yml b/.github/workflows/macos-latest.yml index 8ae847289..2b5a1b0e5 100644 --- a/.github/workflows/macos-latest.yml +++ b/.github/workflows/macos-latest.yml @@ -1,20 +1,20 @@ -name: MacOS -on: - push: - branches: - - main - - devel - pull_request: - branches: - - main - - devel +# name: MacOS +# on: +# push: +# branches: +# - main +# - devel +# pull_request: +# branches: +# - main +# - devel -# concurrency: -# group: ${{ github.ref }} -# cancel-in-progress: true +# # concurrency: +# # group: ${{ github.ref }} +# # cancel-in-progress: true -jobs: - macos-latest-build: - uses: ./.github/workflows/testing.yml - with: - os: macos-latest \ No newline at end of file +# jobs: +# macos-latest-build: +# uses: ./.github/workflows/testing.yml +# with: +# os: macos-latest \ No newline at end of file From 1a753f38cc84807a37d1b8daf7cfb3dbb4d245d4 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 17:16:16 +0200 Subject: [PATCH 234/292] Reverted toy.py and base.py --- src/struphy/models/base.py | 3 --- src/struphy/models/toy.py | 3 --- 2 files changed, 6 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 317121ccc..d18942154 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -44,9 +44,6 @@ class StruphyModel(metaclass=ABCMeta): """ Base class for all Struphy models. -

x2 + y2 = z2

- Note ---- All Struphy models are subclasses of ``StruphyModel`` and should be added to ``struphy/models/`` diff --git a/src/struphy/models/toy.py b/src/struphy/models/toy.py index aa295c513..bad2b7916 100644 --- a/src/struphy/models/toy.py +++ b/src/struphy/models/toy.py @@ -14,9 +14,6 @@ class Maxwell(StruphyModel): r"""Maxwell's equations in vacuum. -

x2 + y2 = z2

- :ref:`normalization`: .. math:: From 72e1e5dca00bd99e5867a4439d3accc9f12369d2 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 17:18:54 +0200 Subject: [PATCH 235/292] Added trailing commas --- src/struphy/models/hybrid.py | 30 +++++++++---------- .../propagators/propagators_markers.py | 4 +-- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index 6994534f9..bcc9f6492 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -211,18 +211,18 @@ def generate_default_parameter_file(self, path=None, prompt=True): for line in f: if "mag_sonic.Options" in line: new_file += [ - "model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n" + "model.propagators.mag_sonic.options = model.propagators.mag_sonic.Options(b_field=model.em_fields.b_field)\n", ] elif "couple_dens.Options" in line: new_file += [ - "model.propagators.couple_dens.options = model.propagators.couple_dens.Options(energetic_ions=model.energetic_ions.var,\n" + "model.propagators.couple_dens.options = model.propagators.couple_dens.Options(energetic_ions=model.energetic_ions.var,\n", ] new_file += [ - " b_tilde=model.em_fields.b_field)\n" + " b_tilde=model.em_fields.b_field)\n", ] elif "couple_curr.Options" in line: new_file += [ - "model.propagators.couple_curr.options = model.propagators.couple_curr.Options(b_tilde=model.em_fields.b_field)\n" + "model.propagators.couple_curr.options = model.propagators.couple_curr.Options(b_tilde=model.em_fields.b_field)\n", ] elif "set_save_data" in line: new_file += ["\nbinplot = BinningPlot(slice='e1', n_bins=128, ranges=(0.0, 1.0))\n"] @@ -418,7 +418,7 @@ def update_scalar_quantities(self): particles.markers[~particles.holes, 6].dot( particles.markers[~particles.holes, 3] ** 2 + particles.markers[~particles.holes, 4] ** 2 - + particles.markers[~particles.holes, 5] ** 2 + + particles.markers[~particles.holes, 5] ** 2, ) / 2.0 * Ah @@ -461,19 +461,19 @@ def generate_default_parameter_file(self, path=None, prompt=True): if "magnetosonic.Options" in line: new_file += [ """model.propagators.magnetosonic.options = model.propagators.magnetosonic.Options( - b_field=model.em_fields.b_field,)\n""" + b_field=model.em_fields.b_field,)\n""", ] elif "push_eta_pc.Options" in line: new_file += [ """model.propagators.push_eta_pc.options = model.propagators.push_eta_pc.Options( - u_tilde = model.mhd.velocity,)\n""" + u_tilde = model.mhd.velocity,)\n""", ] elif "push_vxb.Options" in line: new_file += [ """model.propagators.push_vxb.options = model.propagators.push_vxb.Options( - b2_var = model.em_fields.b_field,)\n""" + b2_var = model.em_fields.b_field,)\n""", ] else: @@ -743,44 +743,44 @@ def generate_default_parameter_file(self, path=None, prompt=True): if "shearalfen_cc5d.Options" in line: new_file += [ """model.propagators.shearalfen_cc5d.options = model.propagators.shearalfen_cc5d.Options( - energetic_ions = model.energetic_ions.var,)\n""" + energetic_ions = model.energetic_ions.var,)\n""", ] elif "magnetosonic.Options" in line: new_file += [ """model.propagators.magnetosonic.options = model.propagators.magnetosonic.Options( - b_field=model.em_fields.b_field,)\n""" + b_field=model.em_fields.b_field,)\n""", ] elif "cc5d_density.Options" in line: new_file += [ """model.propagators.cc5d_density.options = model.propagators.cc5d_density.Options( energetic_ions = model.energetic_ions.var, - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] elif "cc5d_curlb.Options" in line: new_file += [ """model.propagators.cc5d_curlb.options = model.propagators.cc5d_curlb.Options( - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] elif "cc5d_gradb.Options" in line: new_file += [ """model.propagators.cc5d_gradb.options = model.propagators.cc5d_gradb.Options( - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] elif "push_bxe.Options" in line: new_file += [ """model.propagators.push_bxe.options = model.propagators.push_bxe.Options( - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] elif "push_parallel.Options" in line: new_file += [ """model.propagators.push_parallel.options = model.propagators.push_parallel.Options( - b_tilde = model.em_fields.b_field,)\n""" + b_tilde = model.em_fields.b_field,)\n""", ] else: diff --git a/src/struphy/propagators/propagators_markers.py b/src/struphy/propagators/propagators_markers.py index 0563c9dd1..f1dbbe5f6 100644 --- a/src/struphy/propagators/propagators_markers.py +++ b/src/struphy/propagators/propagators_markers.py @@ -221,7 +221,7 @@ def allocate(self): elif self.options.algo == "implicit": kernel = Pyccelkernel(pusher_kernels.push_vxb_implicit) else: - raise ValueError(f"{self.options.algo = } not supported.") + raise ValueError(f"{self.options.algo =} not supported.") # instantiate Pusher args_kernel = ( @@ -458,7 +458,7 @@ def allocate(self): kernel = Pyccelkernel(pusher_kernels.push_pc_eta_stage_H1vec) else: raise ValueError( - f'{self.options.u_space = } not valid, choose from "Hcurl", "Hdiv" or "H1vec.', + f'{self.options.u_space =} not valid, choose from "Hcurl", "Hdiv" or "H1vec.', ) # define algorithm From b55a3cfd0f93a7409b5cba52c26f136aec5a0700 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 18:24:00 +0200 Subject: [PATCH 236/292] Updated psydac dependency to https://github.com/struphy-hub/psydac-for-struphy.git --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c4a2c9d27..82c2e2605 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ 'numpy', 'cunumpy', 'pyccel>=2.0', - 'psydac @ git+https://github.com/max-models/psydac-for-struphy.git@devel-tiny', + 'psydac @ git+https://github.com/struphy-hub/psydac-for-struphy.git@devel-tiny', 'scipy', 'h5py', 'matplotlib', From 4631c20bbcf4ad21e209695c5c0801a2fcf0d18e Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 18:40:21 +0200 Subject: [PATCH 237/292] Added pytest-testmon --- pyproject.toml | 1 + src/struphy/console/test.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 82c2e2605..5797cfd8e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ 'argcomplete', 'pytest', 'pytest-mpi', + 'pytest-testmon', 'line_profiler', ] diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index ebcff34d4..9aacbce87 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -41,6 +41,7 @@ def struphy_test( "-n", str(mpi), "pytest", + "--testmon", "-k", "not _models and not _tutorial and not pproc", "--with-mpi", @@ -48,6 +49,7 @@ def struphy_test( else: cmd = [ "pytest", + "--testmon", "-k", "not _models and not _tutorial and not pproc", ] From 6775cc43929619a022a7d3b912a67f2a4350ced7 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 18:44:48 +0200 Subject: [PATCH 238/292] Run unit tests twice --- .github/actions/tests/unit/action.yml | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/.github/actions/tests/unit/action.yml b/.github/actions/tests/unit/action.yml index 385d6f419..e836e5a52 100644 --- a/.github/actions/tests/unit/action.yml +++ b/.github/actions/tests/unit/action.yml @@ -10,7 +10,17 @@ runs: struphy --refresh-models struphy test unit --mpi 2 - - name: Run unit tests + - name: Run unit tests without MPI + shell: bash + run: | + struphy compile --status + struphy --refresh-models + pip show mpi4py + pip uninstall -y mpi4py + pip list + struphy test unit + + - name: Run unit tests without MPI again shell: bash run: | struphy compile --status From 4d59bcaf5faa35239779510e55764bb0a1833c5e Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 24 Oct 2025 18:48:07 +0200 Subject: [PATCH 239/292] Reverted accidental changes --- .github/actions/tests/unit/action.yml | 12 +----------- pyproject.toml | 1 - src/struphy/console/test.py | 2 -- 3 files changed, 1 insertion(+), 14 deletions(-) diff --git a/.github/actions/tests/unit/action.yml b/.github/actions/tests/unit/action.yml index e836e5a52..385d6f419 100644 --- a/.github/actions/tests/unit/action.yml +++ b/.github/actions/tests/unit/action.yml @@ -10,17 +10,7 @@ runs: struphy --refresh-models struphy test unit --mpi 2 - - name: Run unit tests without MPI - shell: bash - run: | - struphy compile --status - struphy --refresh-models - pip show mpi4py - pip uninstall -y mpi4py - pip list - struphy test unit - - - name: Run unit tests without MPI again + - name: Run unit tests shell: bash run: | struphy compile --status diff --git a/pyproject.toml b/pyproject.toml index 5797cfd8e..82c2e2605 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,6 @@ dependencies = [ 'argcomplete', 'pytest', 'pytest-mpi', - 'pytest-testmon', 'line_profiler', ] diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 9aacbce87..ebcff34d4 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -41,7 +41,6 @@ def struphy_test( "-n", str(mpi), "pytest", - "--testmon", "-k", "not _models and not _tutorial and not pproc", "--with-mpi", @@ -49,7 +48,6 @@ def struphy_test( else: cmd = [ "pytest", - "--testmon", "-k", "not _models and not _tutorial and not pproc", ] From c89866dc7c01a2d9dff425b027289fd5fc04d183 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 24 Oct 2025 20:27:42 +0200 Subject: [PATCH 240/292] Added links to README.md --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 8c6c734ba..62263ed54 100755 --- a/README.md +++ b/README.md @@ -5,11 +5,11 @@ [![Ubuntu latest](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/ubuntu-latest.yml) [![MacOS latest](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/macos-latest.yml) [![isort and ruff](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml/badge.svg)](https://github.com/struphy-hub/struphy/actions/workflows/static_analysis.yml) -![PyPI](https://img.shields.io/pypi/v/struphy?label=PyPI) +[![PyPI](https://img.shields.io/pypi/v/struphy?label=PyPI)](https://pypi.org/project/struphy/) [![PyPI Downloads](https://img.shields.io/pypi/dm/struphy.svg?label=PyPI%20downloads)]( https://pypi.org/project/struphy/) -![Release](https://img.shields.io/github/v/release/struphy-hub/struphy?label=Release) -![License](https://img.shields.io/badge/License-MIT-violet) +[![Release](https://img.shields.io/github/v/release/struphy-hub/struphy?label=Release)](https://github.com/struphy-hub/struphy/releases) +[![License](https://img.shields.io/badge/License-MIT-violet)](https://github.com/struphy-hub/struphy/blob/devel/LICENSE) # Welcome! From 1e16093ada1df884bbe34aab55199d6640348266 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 27 Oct 2025 18:14:12 +0100 Subject: [PATCH 241/292] ruff check --select E713 --fix --- src/struphy/models/base.py | 4 +-- src/struphy/models/hybrid.py | 48 ++++++++++++++++++------------------ 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index e6eb5b665..fb77bcb65 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -428,13 +428,13 @@ def getFromDict(dataDict, mapList): def setInDict(dataDict, mapList, value): # Loop over dicitionary and creaty empty dicts where the path does not exist for k in range(len(mapList)): - if not mapList[k] in getFromDict(dataDict, mapList[:k]).keys(): + if mapList[k] not in getFromDict(dataDict, mapList[:k]).keys(): getFromDict(dataDict, mapList[:k])[mapList[k]] = {} getFromDict(dataDict, mapList[:-1])[mapList[-1]] = value # make sure that the base keys are top-level keys for base_key in ["em_fields", "fluid", "kinetic"]: - if not base_key in dct.keys(): + if base_key not in dct.keys(): dct[base_key] = {} if isinstance(species, str): diff --git a/src/struphy/models/hybrid.py b/src/struphy/models/hybrid.py index bcc9f6492..c1952f59c 100644 --- a/src/struphy/models/hybrid.py +++ b/src/struphy/models/hybrid.py @@ -319,15 +319,15 @@ def __init__(self): class Propagators: def __init__(self, turn_off: tuple[str, ...] = (None,)): - if not "PushEtaPC" in turn_off: + if "PushEtaPC" not in turn_off: self.push_eta_pc = propagators_markers.PushEtaPC() - if not "PushVxB" in turn_off: + if "PushVxB" not in turn_off: self.push_vxb = propagators_markers.PushVxB() - if not "PressureCoupling6D" in turn_off: + if "PressureCoupling6D" not in turn_off: self.pc6d = propagators_coupling.PressureCoupling6D() - if not "ShearAlfven" in turn_off: + if "ShearAlfven" not in turn_off: self.shearalfven = propagators_fields.ShearAlfven() - if not "Magnetosonic" in turn_off: + if "Magnetosonic" not in turn_off: self.magnetosonic = propagators_fields.Magnetosonic() def __init__(self, turn_off: tuple[str, ...] = (None,)): @@ -343,19 +343,19 @@ def __init__(self, turn_off: tuple[str, ...] = (None,)): self.propagators = self.Propagators(turn_off) # 3. assign variables to propagators - if not "ShearAlfven" in turn_off: + if "ShearAlfven" not in turn_off: self.propagators.shearalfven.variables.u = self.mhd.velocity self.propagators.shearalfven.variables.b = self.em_fields.b_field - if not "Magnetosonic" in turn_off: + if "Magnetosonic" not in turn_off: self.propagators.magnetosonic.variables.n = self.mhd.density self.propagators.magnetosonic.variables.u = self.mhd.velocity self.propagators.magnetosonic.variables.p = self.mhd.pressure - if not "PressureCoupling6D" in turn_off: + if "PressureCoupling6D" not in turn_off: self.propagators.pc6d.variables.u = self.mhd.velocity self.propagators.pc6d.variables.energetic_ions = self.energetic_ions.var - if not "PushEtaPC" in turn_off: + if "PushEtaPC" not in turn_off: self.propagators.push_eta_pc.variables.var = self.energetic_ions.var - if not "PushVxB" in turn_off: + if "PushVxB" not in turn_off: self.propagators.push_vxb.variables.ions = self.energetic_ions.var # define scalars for update_scalar_quantities @@ -584,19 +584,19 @@ def __init__(self): class Propagators: def __init__(self, turn_off: tuple[str, ...] = (None,)): - if not "PushGuidingCenterBxEstar" in turn_off: + if "PushGuidingCenterBxEstar" not in turn_off: self.push_bxe = propagators_markers.PushGuidingCenterBxEstar() - if not "PushGuidingCenterParallel" in turn_off: + if "PushGuidingCenterParallel" not in turn_off: self.push_parallel = propagators_markers.PushGuidingCenterParallel() - if not "ShearAlfvenCurrentCoupling5D" in turn_off: + if "ShearAlfvenCurrentCoupling5D" not in turn_off: self.shearalfen_cc5d = propagators_fields.ShearAlfvenCurrentCoupling5D() - if not "Magnetosonic" in turn_off: + if "Magnetosonic" not in turn_off: self.magnetosonic = propagators_fields.Magnetosonic() - if not "CurrentCoupling5DDensity" in turn_off: + if "CurrentCoupling5DDensity" not in turn_off: self.cc5d_density = propagators_fields.CurrentCoupling5DDensity() - if not "CurrentCoupling5DGradB" in turn_off: + if "CurrentCoupling5DGradB" not in turn_off: self.cc5d_gradb = propagators_coupling.CurrentCoupling5DGradB() - if not "CurrentCoupling5DCurlb" in turn_off: + if "CurrentCoupling5DCurlb" not in turn_off: self.cc5d_curlb = propagators_coupling.CurrentCoupling5DCurlb() def __init__(self, turn_off: tuple[str, ...] = (None,)): @@ -612,24 +612,24 @@ def __init__(self, turn_off: tuple[str, ...] = (None,)): self.propagators = self.Propagators(turn_off) # 3. assign variables to propagators - if not "ShearAlfvenCurrentCoupling5D" in turn_off: + if "ShearAlfvenCurrentCoupling5D" not in turn_off: self.propagators.shearalfen_cc5d.variables.u = self.mhd.velocity self.propagators.shearalfen_cc5d.variables.b = self.em_fields.b_field - if not "Magnetosonic" in turn_off: + if "Magnetosonic" not in turn_off: self.propagators.magnetosonic.variables.n = self.mhd.density self.propagators.magnetosonic.variables.u = self.mhd.velocity self.propagators.magnetosonic.variables.p = self.mhd.pressure - if not "CurrentCoupling5DDensity" in turn_off: + if "CurrentCoupling5DDensity" not in turn_off: self.propagators.cc5d_density.variables.u = self.mhd.velocity - if not "CurrentCoupling5DGradB" in turn_off: + if "CurrentCoupling5DGradB" not in turn_off: self.propagators.cc5d_gradb.variables.u = self.mhd.velocity self.propagators.cc5d_gradb.variables.energetic_ions = self.energetic_ions.var - if not "CurrentCoupling5DCurlb" in turn_off: + if "CurrentCoupling5DCurlb" not in turn_off: self.propagators.cc5d_curlb.variables.u = self.mhd.velocity self.propagators.cc5d_curlb.variables.energetic_ions = self.energetic_ions.var - if not "PushGuidingCenterBxEstar" in turn_off: + if "PushGuidingCenterBxEstar" not in turn_off: self.propagators.push_bxe.variables.ions = self.energetic_ions.var - if not "PushGuidingCenterParallel" in turn_off: + if "PushGuidingCenterParallel" not in turn_off: self.propagators.push_parallel.variables.ions = self.energetic_ions.var # define scalars for update_scalar_quantities From d61bb3199c56c83292d990a1e219be89a86c5173 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 27 Oct 2025 18:17:31 +0100 Subject: [PATCH 242/292] ruff check --select E712 --unsafe-fixes --fix --- src/struphy/bsplines/bsplines.py | 4 +- .../legacy/massless_operators/fB_arrays.py | 4 +- .../pro_local/mhd_operators_3d_local.py | 6 +-- .../pro_local/projectors_local.py | 12 ++--- .../shape_function_projectors_L2.py | 8 ++-- .../shape_function_projectors_local.py | 12 ++--- .../eigenvalue_solvers/projectors_global.py | 8 ++-- src/struphy/feec/basis_projection_ops.py | 2 +- src/struphy/feec/linear_operators.py | 8 ++-- src/struphy/feec/local_projectors_kernels.py | 46 +++++++++---------- src/struphy/feec/mass.py | 10 ++-- src/struphy/feec/projectors.py | 16 +++---- src/struphy/feec/psydac_derham.py | 6 +-- src/struphy/feec/variational_utilities.py | 2 +- src/struphy/linear_algebra/saddle_point.py | 14 +++--- .../tests/test_saddlepoint_massmatrices.py | 6 +-- src/struphy/polar/basic.py | 2 +- src/struphy/propagators/propagators_fields.py | 6 +-- 18 files changed, 86 insertions(+), 86 deletions(-) diff --git a/src/struphy/bsplines/bsplines.py b/src/struphy/bsplines/bsplines.py index a04ee4851..9974a9ff2 100644 --- a/src/struphy/bsplines/bsplines.py +++ b/src/struphy/bsplines/bsplines.py @@ -164,7 +164,7 @@ def basis_funs(knots, degree, x, span, normalize=False): saved = left[j - r] * temp values[j + 1] = saved - if normalize == True: + if normalize: values = values * scaling_vector(knots, degree, span) return values @@ -735,7 +735,7 @@ def basis_ders_on_quad_grid(knots, degree, quad_grid, nders, normalize=False): span = find_span(knots, degree, xq) ders = basis_funs_all_ders(knots, degree, xq, span, nders) - if normalize == True: + if normalize: ders = ders * scaling_vector(knots, degree, span) basis[ie, :, :, iq] = ders.transpose() diff --git a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py index e74302878..65faf9209 100644 --- a/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py +++ b/src/struphy/eigenvalue_solvers/legacy/massless_operators/fB_arrays.py @@ -225,7 +225,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): dtype=float, ) # when using delta f method, the values of current equilibrium at all quadrature points - if control == True: + if control: self.Jeqx = xp.empty( ( self.Nel[0], @@ -761,7 +761,7 @@ def __init__(self, TENSOR_SPACE_FEM, DOMAIN, control, mpi_comm): self.df_det[ie1, ie2, ie3, q1, q2, q3] = det_number - if control == True: + if control: x1 = mapping3d.f( TENSOR_SPACE_FEM.pts[0][ie1, q1], TENSOR_SPACE_FEM.pts[1][ie2, q2], diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py index 49464aa58..6734a11b0 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/mhd_operators_3d_local.py @@ -52,7 +52,7 @@ def __init__(self, tensor_space, n_quad): self.coeff_h = [0, 0, 0] for a in range(3): - if self.bc[a] == True: + if self.bc[a]: self.coeff_i[a] = xp.zeros((1, 2 * self.p[a] - 1), dtype=float) self.coeff_h[a] = xp.zeros((1, 2 * self.p[a]), dtype=float) @@ -186,7 +186,7 @@ def __init__(self, tensor_space, n_quad): self.int_shift_N = [0, 0, 0] for a in range(3): - if self.bc[a] == False: + if not self.bc[a]: # maximum number of non-vanishing coefficients if self.p[a] == 1: self.n_int_nvcof_D[a] = 2 @@ -405,7 +405,7 @@ def __init__(self, tensor_space, n_quad): self.his_shift_N = [0, 0, 0] for a in range(3): - if self.bc[a] == False: + if not self.bc[a]: # maximum number of non-vanishing coefficients self.n_his_nvcof_D[a] = 3 * self.p[a] - 2 self.n_his_nvcof_N[a] = 3 * self.p[a] - 1 diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py index 3b27b1b5f..9ede3f608 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/pro_local/projectors_local.py @@ -45,7 +45,7 @@ def __init__(self, spline_space, n_quad): self.wts_loc = xp.polynomial.legendre.leggauss(self.n_quad)[1] # set interpolation and histopolation coefficients - if self.bc == True: + if self.bc: self.coeff_i = xp.zeros((1, 2 * self.p - 1), dtype=float) self.coeff_h = xp.zeros((1, 2 * self.p), dtype=float) @@ -152,7 +152,7 @@ def __init__(self, spline_space, n_quad): self.coeffi_indices = xp.zeros(n_lambda_int, dtype=int) - if self.bc == False: + if not self.bc: # maximum number of non-vanishing coefficients if self.p == 1: self.n_int_nvcof_D = 2 @@ -318,7 +318,7 @@ def __init__(self, spline_space, n_quad): self.coeffh_indices = xp.zeros(n_lambda_his, dtype=int) - if self.bc == False: + if not self.bc: # maximum number of non-vanishing coefficients self.n_his_nvcof_D = 3 * self.p - 2 self.n_his_nvcof_N = 3 * self.p - 1 @@ -629,7 +629,7 @@ def __init__(self, tensor_space, n_quad): self.coeff_h = [0, 0, 0] for a in range(3): - if self.bc[a] == True: + if self.bc[a]: self.coeff_i[a] = xp.zeros((1, 2 * self.p[a] - 1), dtype=float) self.coeff_h[a] = xp.zeros((1, 2 * self.p[a]), dtype=float) @@ -763,7 +763,7 @@ def __init__(self, tensor_space, n_quad): self.int_shift_N = [0, 0, 0] for a in range(3): - if self.bc[a] == False: + if not self.bc[a]: # maximum number of non-vanishing coefficients if self.p[a] == 1: self.n_int_nvcof_D[a] = 2 @@ -979,7 +979,7 @@ def __init__(self, tensor_space, n_quad): self.his_shift_N = [0, 0, 0] for a in range(3): - if self.bc[a] == False: + if not self.bc[a]: # maximum number of non-vanishing coefficients self.n_his_nvcof_D[a] = 3 * self.p[a] - 2 self.n_his_nvcof_N[a] = 3 * self.p[a] - 1 diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py index 8978e2464..137df7f09 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_L2.py @@ -590,7 +590,7 @@ def potential_pi_0(self, particles_loc, Np, domain, mpi_comm): ------- kernel_0 matrix """ - if self.bc[0] == True and self.bc[1] == True and self.bc[2] == True: + if self.bc[0] and self.bc[1] and self.bc[2]: ker_loc.potential_kernel_0_form( Np, self.p, @@ -637,7 +637,7 @@ def S_pi_0(self, particles_loc, Np, domain): kernel_0 matrix """ self.kernel_0[:, :, :, :, :, :] = 0.0 - if self.bc[0] == True and self.bc[1] == True and self.bc[2] == True: + if self.bc[0] and self.bc[1] and self.bc[2]: ker_loc.kernel_0_form( Np, self.p, @@ -699,7 +699,7 @@ def S_pi_1(self, particles_loc, Np, domain): self.right_loc_2[:, :, :] = 0.0 self.right_loc_3[:, :, :] = 0.0 - if self.bc[0] == True and self.bc[1] == True and self.bc[2] == True: + if self.bc[0] and self.bc[1] and self.bc[2]: ker_loc.kernel_1_form( self.indN[0], self.indN[1], @@ -764,7 +764,7 @@ def S_pi_1(self, particles_loc, Np, domain): print("non-periodic case not implemented!!!") def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): - if self.bc[0] == True and self.bc[1] == True and self.bc[2] == True: + if self.bc[0] and self.bc[1] and self.bc[2]: if index_label == 1: ker_loc.vv_1_form( self.wts[0][0], diff --git a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py index 43c8c8ff9..2ebb497a3 100644 --- a/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py +++ b/src/struphy/eigenvalue_solvers/legacy/projectors_local/shape_pro_local/shape_function_projectors_local.py @@ -304,7 +304,7 @@ def __init__(self, tensor_space, n_quad, p_shape, p_size, NbaseN, NbaseD, mpi_co self.coeff_i = [0, 0, 0] self.coeff_h = [0, 0, 0] for a in range(3): - if self.bc[a] == True: + if self.bc[a]: self.coeff_i[a] = xp.zeros(2 * self.p[a], dtype=float) self.coeff_h[a] = xp.zeros(2 * self.p[a], dtype=float) @@ -686,7 +686,7 @@ def potential_pi_0(self, particles_loc, Np, domain, mpi_comm): ------- kernel_0 matrix """ - if self.bc[0] == True and self.bc[1] == True and self.bc[2] == True: + if self.bc[0] and self.bc[1] and self.bc[2]: ker_loc.potential_kernel_0_form( Np, self.p, @@ -733,7 +733,7 @@ def S_pi_0(self, particles_loc, Np, domain): kernel_0 matrix """ self.kernel_0[:, :, :, :, :, :] = 0.0 - if self.bc[0] == True and self.bc[1] == True and self.bc[2] == True: + if self.bc[0] and self.bc[1] and self.bc[2]: ker_loc.kernel_0_form( Np, self.p, @@ -795,7 +795,7 @@ def S_pi_1(self, particles_loc, Np, domain): self.right_loc_2[:, :, :] = 0.0 self.right_loc_3[:, :, :] = 0.0 - if self.bc[0] == True and self.bc[1] == True and self.bc[2] == True: + if self.bc[0] and self.bc[1] and self.bc[2]: ker_loc.kernel_1_form( self.right_loc_1, self.right_loc_2, @@ -882,7 +882,7 @@ def S_pi_01(self, particles_loc, Np, domain): self.right_loc_2[:, :, :] = 0.0 self.right_loc_3[:, :, :] = 0.0 - if self.bc[0] == True and self.bc[1] == True and self.bc[2] == True: + if self.bc[0] and self.bc[1] and self.bc[2]: ker_loc.kernel_01_form( self.right_loc_1, self.right_loc_2, @@ -933,7 +933,7 @@ def S_pi_01(self, particles_loc, Np, domain): print("non-periodic case not implemented!!!") def vv_S1(self, particles_loc, Np, domain, index_label, accvv, dt, mpi_comm): - if self.bc[0] == True and self.bc[1] == True and self.bc[2] == True: + if self.bc[0] and self.bc[1] and self.bc[2]: if index_label == 1: ker_loc.vv_1_form( self.wts[0][0], diff --git a/src/struphy/eigenvalue_solvers/projectors_global.py b/src/struphy/eigenvalue_solvers/projectors_global.py index ca67c66e6..9d246cdac 100644 --- a/src/struphy/eigenvalue_solvers/projectors_global.py +++ b/src/struphy/eigenvalue_solvers/projectors_global.py @@ -169,7 +169,7 @@ def __init__(self, spline_space, n_quad=6): for i in range(spline_space.NbaseD): for br in spline_space.el_b: # left and right integration boundaries - if spline_space.spl_kind == False: + if not spline_space.spl_kind: xl = self.x_int[i] xr = self.x_int[i + 1] else: @@ -186,7 +186,7 @@ def __init__(self, spline_space, n_quad=6): self.x_his = xp.append(self.x_his, xr) break - if spline_space.spl_kind == True and spline_space.p % 2 == 0: + if spline_space.spl_kind and spline_space.p % 2 == 0: self.x_his = xp.append(self.x_his, spline_space.el_b[-1] + self.x_his[0]) # cumulative number of sub-intervals for conversion local interval --> global interval @@ -198,7 +198,7 @@ def __init__(self, spline_space, n_quad=6): # quadrature points and weights, ignoring subs (less accurate integration for even degree) self.x_hisG = self.x_int - if spline_space.spl_kind == True: + if spline_space.spl_kind: if spline_space.p % 2 == 0: self.x_hisG = xp.append(self.x_hisG, spline_space.el_b[-1] + self.x_hisG[0]) else: @@ -2153,7 +2153,7 @@ def pi_3(self, fun, include_bc=True, eval_kind="meshgrid", with_subs=True): # ======================================== def assemble_approx_inv(self, tol): - if self.approx_Ik_0_inv == False or (self.approx_Ik_0_inv == True and self.approx_Ik_0_tol != tol): + if not self.approx_Ik_0_inv or (self.approx_Ik_0_inv and self.approx_Ik_0_tol != tol): # poloidal plane I0_pol_0_inv_approx = xp.linalg.inv(self.I0_pol_0.toarray()) I1_pol_0_inv_approx = xp.linalg.inv(self.I1_pol_0.toarray()) diff --git a/src/struphy/feec/basis_projection_ops.py b/src/struphy/feec/basis_projection_ops.py index d76dc7ca9..ab0925828 100644 --- a/src/struphy/feec/basis_projection_ops.py +++ b/src/struphy/feec/basis_projection_ops.py @@ -2385,7 +2385,7 @@ def find_relative_col(col, row, Nbasis, periodic): The relative column position of col with respect to the the current row of the StencilMatrix. """ - if periodic == False: + if not periodic: relativecol = col - row # In the periodic case we must account for the possible looping of the basis functions when computing the relative row postion else: diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 7469d68e9..28b4a0805 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -63,7 +63,7 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): rank = comm.Get_rank() size = comm.Get_size() - if is_sparse == False: + if not is_sparse: if out is None: # We declare the matrix form of our linear operator out = xp.zeros([self.codomain.dimension, self.domain.dimension], dtype=self.dtype) @@ -149,7 +149,7 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): # Compute to which column this iteration belongs col = spoint col += xp.ravel_multi_index(i, npts[h]) - if is_sparse == False: + if not is_sparse: result[:, col] = tmp2.toarray() else: aux = tmp2.toarray() @@ -220,7 +220,7 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): self.dot(v, out=tmp2) # Compute to which column this iteration belongs col = xp.ravel_multi_index(i, npts) - if is_sparse == False: + if not is_sparse: result[:, col] = tmp2.toarray() else: aux = tmp2.toarray() @@ -237,7 +237,7 @@ def toarray_struphy(self, out=None, is_sparse=False, format="csr"): # I cannot conceive any situation where this error should be thrown, but I put it here just in case something unexpected happens. raise Exception("Function toarray_struphy() only supports Stencil Vectors or Block Vectors.") - if is_sparse == False: + if not is_sparse: # Use Allreduce to perform addition reduction and give one copy of the result to all ranks. if comm is None or isinstance(comm, MockComm): out[:] = result diff --git a/src/struphy/feec/local_projectors_kernels.py b/src/struphy/feec/local_projectors_kernels.py index f1eb285c9..706b3f78e 100644 --- a/src/struphy/feec/local_projectors_kernels.py +++ b/src/struphy/feec/local_projectors_kernels.py @@ -63,7 +63,7 @@ def get_local_problem_size(periodic: "bool[:]", p: "int[:]", IoH: "bool[:]"): for h in range(3): # Interpolation - if IoH[h] == False: + if not IoH[h]: lenj[h] = 2 * p[h] - 1 # Histopolation else: @@ -734,7 +734,7 @@ def solve_local_main_loop_weighted( if counteri0 >= rows0[i00] and counteri0 <= rowe0[i00]: compute0 = True break - if compute0 == True: + if compute0: counteri1 = 0 for i1 in range(args_solve.starts[1], args_solve.ends[1] + 1): # This bool variable tell us if this row has a non-zero FE coefficient, based on the current basis function we are using on our projection @@ -744,7 +744,7 @@ def solve_local_main_loop_weighted( if counteri1 >= rows1[i11] and counteri1 <= rowe1[i11]: compute1 = True break - if compute1 == True: + if compute1: counteri2 = 0 for i2 in range(args_solve.starts[2], args_solve.ends[2] + 1): # This bool variable tell us if this row has a non-zero FE coefficient, based on the current basis function we are using on our projection @@ -754,7 +754,7 @@ def solve_local_main_loop_weighted( if counteri2 >= rows2[i22] and counteri2 <= rowe2[i22]: compute2 = True break - if compute2 == True: + if compute2: L123 = 0.0 startj1, endj1 = select_quasi_points( i0, @@ -850,7 +850,7 @@ def find_relative_col(col: int, row: int, Nbasis: int, periodic: bool): The relative column position of col with respect to the the current row of the StencilMatrix. """ - if periodic == False: + if not periodic: relativecol = col - row # In the periodic case we must account for the possible looping of the basis functions when computing the relative row postion else: @@ -944,7 +944,7 @@ def assemble_basis_projection_operator_local( compute0 = True break relativecol0 = find_relative_col(col[0], row0, VNbasis[0], periodic[0]) - if relativecol0 >= -p[0] and relativecol0 <= p[0] and compute0 == True: + if relativecol0 >= -p[0] and relativecol0 <= p[0] and compute0: count1 = 0 for row1 in range(starts[1], ends[1] + 1): # This bool variable tell us if this row has a non-zero FE coefficient, based on the current basis function we are using on our projection @@ -955,7 +955,7 @@ def assemble_basis_projection_operator_local( compute1 = True break relativecol1 = find_relative_col(col[1], row1, VNbasis[1], periodic[1]) - if relativecol1 >= -p[1] and relativecol1 <= p[1] and compute1 == True: + if relativecol1 >= -p[1] and relativecol1 <= p[1] and compute1: count2 = 0 for row2 in range(starts[2], ends[2] + 1): # This bool variable tell us if this row has a non-zero FE coefficient, based on the current basis function we are using on our projection @@ -966,7 +966,7 @@ def assemble_basis_projection_operator_local( compute2 = True break relativecol2 = find_relative_col(col[2], row2, VNbasis[2], periodic[2]) - if relativecol2 >= -p[2] and relativecol2 <= p[2] and compute2 == True: + if relativecol2 >= -p[2] and relativecol2 <= p[2] and compute2: mat[ count0 + pds[0], count1 + pds[1], @@ -1002,7 +1002,7 @@ def are_quadrature_points_zero(aux: "int[:]", p: int, basis: "float[:]"): if basis[in_start + ii] != 0.0: all_zero = False break - if all_zero == True: + if all_zero: aux[i] = 0 @@ -1085,33 +1085,33 @@ def get_rows( Array where we put a one if the current row could have a non-zero FE coefficient for the column given by col. """ # Periodic boundary conditions - if periodic == True: + if periodic: # Histopolation - if IoH == True: + if IoH: # D-splines - if BoD == True: + if BoD: get_rows_periodic(starts, ends, -p + 1, p, Nbasis, col, aux) # B-splines - if BoD == False: + if not BoD: get_rows_periodic(starts, ends, -p + 1, p + 1, Nbasis, col, aux) # Interpolation - if IoH == False: + if not IoH: # D-splines - if BoD == True: + if BoD: # Special case p = 1 if p == 1: get_rows_periodic(starts, ends, -1, 1, Nbasis, col, aux) if p != 1: get_rows_periodic(starts, ends, -p + 1, p - 1, Nbasis, col, aux) # B-splines - if BoD == False: + if not BoD: get_rows_periodic(starts, ends, -p + 1, p, Nbasis, col, aux) # Clamped boundary conditions - if periodic == False: + if not periodic: # Histopolation - if IoH == True: + if IoH: # D-splines - if BoD == True: + if BoD: count = 0 for row in range(starts, ends + 1): if row >= 0 and row <= (p - 2) and col >= 0 and col <= row + p - 1: @@ -1124,7 +1124,7 @@ def get_rows( aux[count] = 1 count += 1 # B-splines - if BoD == False: + if not BoD: count = 0 for row in range(starts, ends + 1): if row >= 0 and row <= (p - 2) and col >= 0 and col <= (row + p): @@ -1135,9 +1135,9 @@ def get_rows( aux[count] = 1 count += 1 # Interpolation - if IoH == False: + if not IoH: # D-splines - if BoD == True: + if BoD: count = 0 for row in range(starts, ends + 1): if row == 0 and col <= (p - 1): @@ -1152,7 +1152,7 @@ def get_rows( aux[count] = 1 count += 1 # B-splines - if BoD == False: + if not BoD: count = 0 for row in range(starts, ends + 1): if row == 0 and col <= p: diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 16f0109d9..5964f5f7c 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -905,7 +905,7 @@ def DFinvT(e1, e2, e3): if weights_rank2: # if matrix exits fun = [] - if listinput == True and len(weights_rank2) == 1: + if listinput and len(weights_rank2) == 1: for m in range(3): fun += [[]] for n in range(3): @@ -2518,10 +2518,10 @@ def tosparse(self): if all(op is None for op in (self._W_extraction_op, self._V_extraction_op)): for bl in self._V_boundary_op.bc: for bc in bl: - assert bc == False, print(".tosparse() only works without boundary conditions at the moment") + assert not bc, print(".tosparse() only works without boundary conditions at the moment") for bl in self._W_boundary_op.bc: for bc in bl: - assert bc == False, print(".tosparse() only works without boundary conditions at the moment") + assert not bc, print(".tosparse() only works without boundary conditions at the moment") return self._mat.tosparse() elif all(isinstance(op, IdentityOperator) for op in (self._W_extraction_op, self._V_extraction_op)): @@ -2534,10 +2534,10 @@ def toarray(self): if all(op is None for op in (self._W_extraction_op, self._V_extraction_op)): for bl in self._V_boundary_op.bc: for bc in bl: - assert bc == False, print(".toarray() only works without boundary conditions at the moment") + assert not bc, print(".toarray() only works without boundary conditions at the moment") for bl in self._W_boundary_op.bc: for bc in bl: - assert bc == False, print(".toarray() only works without boundary conditions at the moment") + assert not bc, print(".toarray() only works without boundary conditions at the moment") return self._mat.toarray() elif all(isinstance(op, IdentityOperator) for op in (self._W_extraction_op, self._V_extraction_op)): diff --git a/src/struphy/feec/projectors.py b/src/struphy/feec/projectors.py index 115ed0aa1..be56cc722 100644 --- a/src/struphy/feec/projectors.py +++ b/src/struphy/feec/projectors.py @@ -1481,7 +1481,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non Builds 3D numpy array with the evaluation of the right-hand-side. """ if self._space_key == "0": - if first_go == True: + if first_go: pre_computed_dofs = [fun(*self._meshgrid)] elif self._space_key == "1" or self._space_key == "2": @@ -1491,12 +1491,12 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non f_eval = [] # If this is the first time this rank has to evaluate the weights degrees of freedom we declare the list where to store them. - if first_go == True: + if first_go: pre_computed_dofs = [] for h in range(3): # Evaluation of the function to compute the h component - if first_go == True: + if first_go: pre_computed_dofs.append(fun[h](*self._meshgrid[h])) # Array into which we will write the Dofs. @@ -1547,7 +1547,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non elif self._space_key == "3": f_eval = xp.zeros(tuple(xp.shape(dim)[0] for dim in self._localpts)) # Evaluation of the function at all Gauss-Legendre quadrature points - if first_go == True: + if first_go: pre_computed_dofs = [fun(*self._meshgrid)] get_dofs_local_3_form_weighted( @@ -1578,7 +1578,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non # We should do nothing here self._do_nothing[h] = 1 - if first_go == True: + if first_go: f_eval = [] for h in range(3): f_eval.append(fun[h](*self._meshgrid[h])) @@ -1588,7 +1588,7 @@ def get_dofs_weighted(self, fun, dofs=None, first_go=True, pre_computed_dofs=Non "Uknown space. It must be either H1, Hcurl, Hdiv, L2 or H1vec.", ) - if first_go == True: + if first_go: if self._space_key == "0": return pre_computed_dofs[0], pre_computed_dofs elif self._space_key == "v": @@ -1654,14 +1654,14 @@ def __call__( coeffs : psydac.linalg.basic.vector | xp.array 3D The FEM spline coefficients after projection. """ - if weighted == False: + if not weighted: return self.solve(self.get_dofs(fun, dofs=dofs), out=out) else: # We set B_or_D and basis_indices as attributes of the projectors so we can easily access them in the get_rowstarts, get_rowends and get_values functions, where they are needed. self._B_or_D = B_or_D self._basis_indices = basis_indices - if first_go == True: + if first_go: # rhs contains the evaluation over the degrees of freedom of the weights multiplied by the basis function # rhs_weights contains the evaluation over the degrees of freedom of only the weights rhs, rhs_weights = self.get_dofs_weighted( diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 523a5bb97..bb833e02d 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1270,7 +1270,7 @@ def _get_neighbour_one_component(self, comp): # if only one process: check if comp is neighbour in non-peridic directions, if this is not the case then return the rank as neighbour id if size == 1: - if (comp[kinds == False] == 1).all(): + if (comp[not kinds] == 1).all(): return rank # multiple processes @@ -2055,7 +2055,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): ) if self.derham.comm is not None: - if local == False: + if not local: self.derham.comm.Allreduce( MPI.IN_PLACE, tmp, @@ -2126,7 +2126,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): ) if self.derham.comm is not None: - if local == False: + if not local: self.derham.comm.Allreduce( MPI.IN_PLACE, tmp, diff --git a/src/struphy/feec/variational_utilities.py b/src/struphy/feec/variational_utilities.py index d03a75e3d..8174a1a5b 100644 --- a/src/struphy/feec/variational_utilities.py +++ b/src/struphy/feec/variational_utilities.py @@ -94,7 +94,7 @@ def __init__( self.Pcoord3 = CoordinateProjector(2, derham.Vh_pol["v"], derham.Vh_pol["0"]) @ derham.boundary_ops["v"] # Initialize the BasisProjectionOperators - if derham._with_local_projectors == True: + if derham._with_local_projectors: self.PiuT = BasisProjectionOperatorLocal( P0, V1h, diff --git a/src/struphy/linear_algebra/saddle_point.py b/src/struphy/linear_algebra/saddle_point.py index 3a191cde4..337664754 100644 --- a/src/struphy/linear_algebra/saddle_point.py +++ b/src/struphy/linear_algebra/saddle_point.py @@ -304,7 +304,7 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): elif self._variant == "Uzawa": info = {} - if self._spectralanalysis == True: + if self._spectralanalysis: self._spectralresult = self._spectral_analysis() else: self._spectralresult = [] @@ -333,9 +333,9 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): self._rhs0np -= self._B1np.transpose().dot(self._Pnp) self._rhs0np -= self._Anp.dot(self._Unp) self._rhs0np += self._F[0] - if self._preconditioner == False: + if not self._preconditioner: self._Unp += self._Anpinv.dot(self._rhs0np) - elif self._preconditioner == True: + elif self._preconditioner: self._Unp += self._Anpinv.dot(self._A11npinv @ self._rhs0np) R1 = self._B1np.dot(self._Unp) @@ -344,9 +344,9 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): self._rhs1np -= self._B2np.transpose().dot(self._Pnp) self._rhs1np -= self._Aenp.dot(self._Uenp) self._rhs1np += self._F[1] - if self._preconditioner == False: + if not self._preconditioner: self._Uenp += self._Aenpinv.dot(self._rhs1np) - elif self._preconditioner == True: + elif self._preconditioner: self._Uenp += self._Aenpinv.dot(self._A22npinv @ self._rhs1np) R2 = self._B2np.dot(self._Uenp) @@ -382,7 +382,7 @@ def __call__(self, U_init=None, Ue_init=None, P_init=None, out=None): # Return with info if maximum iterations reached info["success"] = False info["niter"] = iteration + 1 - if self._verbose == True: + if self._verbose: _plot_residual_norms(self._residual_norms) return self._Unp, self._Uenp, self._Pnp, info, self._residual_norms, self._spectralresult @@ -523,7 +523,7 @@ def _spectral_analysis(self): print(f"{specA22_bef_abs =}") print(f"{condA22_before =}") - if self._preconditioner == True: + if self._preconditioner: # A11 after preconditioning with its inverse if self._method_to_solve in ("DirectNPInverse", "InexactNPInverse"): eigvalsA11_after_prec, eigvecs_after = xp.linalg.eig(self._A11npinv @ self._A[0]) # Implement this diff --git a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py index 42d3ae8d3..8de563e8c 100644 --- a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py +++ b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py @@ -107,7 +107,7 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m Cnp = derhamnumpy.curl.toarray() # Dnp = D.toarray() # Cnp = C.toarray() - if derham.with_local_projectors == True: + if derham.with_local_projectors: S21np = S21.toarray else: S21np = S21.toarray_struphy() @@ -121,7 +121,7 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m Cnp = derhamnumpy.curl.tosparse() # Dnp = D.tosparse() # Cnp = C.tosparse() - if derham.with_local_projectors == True: + if derham.with_local_projectors: S21np = S21.tosparse else: S21np = S21.toarray_struphy(is_sparse=True) @@ -270,7 +270,7 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m x_uzawa = {} x_uzawa[0] = x_u x_uzawa[1] = x_ue - if show_plots == True: + if show_plots: _plot_residual_norms(residual_norms) elif method_for_solving == "SaddlePointSolverGMRES": # Wrong initialization to check if changed diff --git a/src/struphy/polar/basic.py b/src/struphy/polar/basic.py index f737c671e..99a95cc47 100644 --- a/src/struphy/polar/basic.py +++ b/src/struphy/polar/basic.py @@ -19,7 +19,7 @@ class PolarDerhamSpace(VectorSpace): """ def __init__(self, derham, space_id): - assert derham.spl_kind[0] == False, "Spline basis in eta1 must be clamped" + assert not derham.spl_kind[0], "Spline basis in eta1 must be clamped" assert derham.spl_kind[1], "Spline basis in eta2 must be periodic" assert (derham.Nel[1] / 3) % 1 == 0.0, "Number of elements in eta2 must be a multiple of 3" diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index c3f3e1381..1c14c2cd9 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -2190,7 +2190,7 @@ def _initialize_projection_operator_TB(self): self._bf = self.derham.create_spline_function("bf", "Hdiv") # Initialize BasisProjectionOperator - if self.derham._with_local_projectors == True: + if self.derham._with_local_projectors: self._TB = BasisProjectionOperatorLocal( P1, Vh, @@ -8638,7 +8638,7 @@ def __call__(self, dt): # _Anp[1] and _Anppre[1] remain unchanged _Anp = [A11np, A22np] - if self._preconditioner == True: + if self._preconditioner: _A11prenp = self._M2np / dt # + self._A11prenp_notimedependency _Anppre = [_A11prenp, _A22prenp] @@ -8675,7 +8675,7 @@ def __call__(self, dt): _Fnp = [_F1np, _F2np] if self.rank == 0: - if self._preconditioner == True: + if self._preconditioner: self._solver_UzawaNumpy.Apre = _Anppre self._solver_UzawaNumpy.A = _Anp self._solver_UzawaNumpy.F = _Fnp From 87dc0e7d869a645d9f128cb3114b2449bbfb4a6d Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 27 Oct 2025 18:22:02 +0100 Subject: [PATCH 243/292] ruff check --select F541 --fix --- src/struphy/console/compile.py | 2 +- src/struphy/console/main.py | 2 +- src/struphy/diagnostics/diagn_tools.py | 2 +- .../mhd_axisymmetric_main.py | 14 ++-- .../feec/tests/test_lowdim_nel_is_1.py | 2 +- src/struphy/geometry/domains.py | 2 +- .../initial/tests/test_init_perturbations.py | 4 +- src/struphy/io/setup.py | 12 ++-- .../tests/test_maxwellians.py | 20 +++--- .../tests/test_saddlepoint_massmatrices.py | 2 +- src/struphy/main.py | 10 +-- src/struphy/models/base.py | 66 +++++++++---------- src/struphy/models/variables.py | 4 +- src/struphy/ode/tests/test_ode_feec.py | 8 +-- src/struphy/pic/particles.py | 2 +- src/struphy/pic/tests/test_mat_vec_filler.py | 16 ++--- .../likwid/plot_likwidproject.py | 2 +- .../post_processing/post_processing_tools.py | 2 +- src/struphy/propagators/propagators_fields.py | 2 +- .../tests/test_gyrokinetic_poisson.py | 2 +- src/struphy/propagators/tests/test_poisson.py | 2 +- ...l_03_smoothed_particle_hydrodynamics.ipynb | 4 +- .../tutorial_02_fluid_particles.ipynb | 12 ++-- 23 files changed, 97 insertions(+), 97 deletions(-) diff --git a/src/struphy/console/compile.py b/src/struphy/console/compile.py index 432e4fa1f..d92f31285 100644 --- a/src/struphy/console/compile.py +++ b/src/struphy/console/compile.py @@ -272,7 +272,7 @@ def struphy_compile( ) sys.exit(1) else: - print(f"Psydac is not installed. To install it, please re-install struphy (e.g. pip install .)\n") + print("Psydac is not installed. To install it, please re-install struphy (e.g. pip install .)\n") sys.exit(1) else: diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index f6e03ff16..545e4a24c 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -452,7 +452,7 @@ def add_parser_run(subparsers, list_models, model_message, params_files, batch_f default=None, # fallback if nothing is passed choices=list_models, metavar="MODEL", - help=model_message + f" (default: None)", + help=model_message + " (default: None)", ) parser_run.add_argument( diff --git a/src/struphy/diagnostics/diagn_tools.py b/src/struphy/diagnostics/diagn_tools.py index b9e66dbb6..e7a9d8ee3 100644 --- a/src/struphy/diagnostics/diagn_tools.py +++ b/src/struphy/diagnostics/diagn_tools.py @@ -683,7 +683,7 @@ def plots_videos_2d( df_binned = df_data[tuple(f_slicing)].squeeze() - assert t_grid.ndim == grid_1.ndim == grid_2.ndim == 1, f"Input arrays must be 1D!" + assert t_grid.ndim == grid_1.ndim == grid_2.ndim == 1, "Input arrays must be 1D!" assert df_binned.shape[0] == t_grid.size, f"{df_binned.shape =}, {t_grid.shape =}" assert df_binned.shape[1] == grid_1.size, f"{df_binned.shape =}, {grid_1.shape =}" assert df_binned.shape[2] == grid_2.size, f"{df_binned.shape =}, {grid_2.shape =}" diff --git a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py index b8d4aaf81..04a194c7f 100644 --- a/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py +++ b/src/struphy/eigenvalue_solvers/mhd_axisymmetric_main.py @@ -45,13 +45,13 @@ def solve_mhd_ev_problem_2d(num_params, eq_mhd, n_tor, basis_tor="i", path_out=N # print grid info print("\nGrid parameters:") - print(f"number of elements :", num_params["Nel"]) - print(f"spline degrees :", num_params["p"]) - print(f"periodic bcs :", num_params["spl_kind"]) - print(f"hom. Dirichlet bc :", num_params["bc"]) - print(f"GL quad pts (L2) :", num_params["nq_el"]) - print(f"GL quad pts (hist) :", num_params["nq_pr"]) - print(f"polar Ck :", num_params["polar_ck"]) + print("number of elements :", num_params["Nel"]) + print("spline degrees :", num_params["p"]) + print("periodic bcs :", num_params["spl_kind"]) + print("hom. Dirichlet bc :", num_params["bc"]) + print("GL quad pts (L2) :", num_params["nq_el"]) + print("GL quad pts (hist) :", num_params["nq_pr"]) + print("polar Ck :", num_params["polar_ck"]) print("") # extract numerical parameters diff --git a/src/struphy/feec/tests/test_lowdim_nel_is_1.py b/src/struphy/feec/tests/test_lowdim_nel_is_1.py index cefcddf61..325da31ea 100644 --- a/src/struphy/feec/tests/test_lowdim_nel_is_1.py +++ b/src/struphy/feec/tests/test_lowdim_nel_is_1.py @@ -277,7 +277,7 @@ def div_f(x, y, z): plt.subplot(2, 1, 2) plt.plot(e, div_f(e1, e2, e3), "o") plt.plot(e, field_df2_vals) - plt.title(f"div") + plt.title("div") plt.subplots_adjust(wspace=1.0, hspace=0.4) diff --git a/src/struphy/geometry/domains.py b/src/struphy/geometry/domains.py index 7b2c25064..20f995779 100644 --- a/src/struphy/geometry/domains.py +++ b/src/struphy/geometry/domains.py @@ -747,7 +747,7 @@ def __init__( if sfl: assert pol_period == 1, ( - f"Piece-of-cake is only implemented for torus coordinates, not for straight field line coordinates!" + "Piece-of-cake is only implemented for torus coordinates, not for straight field line coordinates!" ) # periodicity in eta3-direction and pole at eta1=0 diff --git a/src/struphy/initial/tests/test_init_perturbations.py b/src/struphy/initial/tests/test_init_perturbations.py index dd391cf56..5d52e9291 100644 --- a/src/struphy/initial/tests/test_init_perturbations.py +++ b/src/struphy/initial/tests/test_init_perturbations.py @@ -169,7 +169,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False plt.xlabel("x") plt.ylabel("y") plt.colorbar() - plt.title(f"exact function") + plt.title("exact function") ax = plt.gca() ax.set_aspect("equal", adjustable="box") @@ -198,7 +198,7 @@ def test_init_modes(Nel, p, spl_kind, mapping, combine_comps=None, do_plot=False plt.xlabel("x") plt.ylabel("z") plt.colorbar() - plt.title(f"exact function") + plt.title("exact function") ax = plt.gca() ax.set_aspect("equal", adjustable="box") diff --git a/src/struphy/io/setup.py b/src/struphy/io/setup.py index f38654160..4ecd96f47 100644 --- a/src/struphy/io/setup.py +++ b/src/struphy/io/setup.py @@ -152,12 +152,12 @@ def setup_derham( if MPI.COMM_WORLD.Get_rank() == 0 and verbose: print("\nDERHAM:") - print(f"number of elements:".ljust(25), Nel) - print(f"spline degrees:".ljust(25), p) - print(f"periodic bcs:".ljust(25), spl_kind) - print(f"hom. Dirichlet bc:".ljust(25), dirichlet_bc) - print(f"GL quad pts (L2):".ljust(25), nquads) - print(f"GL quad pts (hist):".ljust(25), nq_pr) + print("number of elements:".ljust(25), Nel) + print("spline degrees:".ljust(25), p) + print("periodic bcs:".ljust(25), spl_kind) + print("hom. Dirichlet bc:".ljust(25), dirichlet_bc) + print("GL quad pts (L2):".ljust(25), nquads) + print("GL quad pts (hist):".ljust(25), nq_pr) print( "MPI proc. per dir.:".ljust(25), derham.domain_decomposition.nprocs, diff --git a/src/struphy/kinetic_background/tests/test_maxwellians.py b/src/struphy/kinetic_background/tests/test_maxwellians.py index 710a88262..4aaa0624a 100644 --- a/src/struphy/kinetic_background/tests/test_maxwellians.py +++ b/src/struphy/kinetic_background/tests/test_maxwellians.py @@ -294,7 +294,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): continue if "GVECequilibrium" in key: - print(f"Attention: flat (marker) evaluation not tested for GVEC at the moment.") + print("Attention: flat (marker) evaluation not tested for GVEC at the moment.") mhd_equil = val() assert isinstance(mhd_equil, FluidEquilibrium) @@ -496,7 +496,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): plt.ylabel("y") plt.axis("equal") plt.colorbar() - plt.title(f"Maxwellian thermal velocity $v_t$, top view (e1-e3)") + plt.title("Maxwellian thermal velocity $v_t$, top view (e1-e3)") plt.subplot(2, 5, 10) if "Slab" in key or "Pinch" in key: plt.contourf(x[:, :, 0], y[:, :, 0], vth_cart[:, :, 0], levels=levels) @@ -508,7 +508,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): plt.ylabel("z") plt.axis("equal") plt.colorbar() - plt.title(f"Maxwellian thermal velocity $v_t$, poloidal view (e1-e2)") + plt.title("Maxwellian thermal velocity $v_t$, poloidal view (e1-e2)") plt.show() @@ -678,7 +678,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): plt.ylabel("y") plt.axis("equal") plt.colorbar() - plt.title(f"Maxwellian perturbed thermal velocity $v_t$, top view (e1-e3)") + plt.title("Maxwellian perturbed thermal velocity $v_t$, top view (e1-e3)") plt.subplot(2, 5, 10) if "Slab" in key or "Pinch" in key: plt.contourf(x[:, :, 0], y[:, :, 0], vth_cart[:, :, 0], levels=levels) @@ -690,7 +690,7 @@ def test_maxwellian_3d_mhd(Nel, with_desc, show_plot=False): plt.ylabel("z") plt.axis("equal") plt.colorbar() - plt.title(f"Maxwellian perturbed thermal velocity $v_t$, poloidal view (e1-e2)") + plt.title("Maxwellian perturbed thermal velocity $v_t$, poloidal view (e1-e2)") plt.show() @@ -1090,7 +1090,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): continue if "GVECequilibrium" in key: - print(f"Attention: flat (marker) evaluation not tested for GVEC at the moment.") + print("Attention: flat (marker) evaluation not tested for GVEC at the moment.") mhd_equil = val() if not isinstance(mhd_equil, FluidEquilibriumWithB): @@ -1287,7 +1287,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): plt.ylabel("y") plt.axis("equal") plt.colorbar() - plt.title(f"Maxwellian thermal velocity $v_t$, top view (e1-e3)") + plt.title("Maxwellian thermal velocity $v_t$, top view (e1-e3)") plt.subplot(2, 4, 8) if "Slab" in key or "Pinch" in key: plt.contourf(x[:, :, 0], y[:, :, 0], vth_cart[:, :, 0], levels=levels) @@ -1299,7 +1299,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): plt.ylabel("z") plt.axis("equal") plt.colorbar() - plt.title(f"Maxwellian density $v_t$, poloidal view (e1-e2)") + plt.title("Maxwellian density $v_t$, poloidal view (e1-e2)") plt.show() @@ -1463,7 +1463,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): plt.ylabel("y") plt.axis("equal") plt.colorbar() - plt.title(f"Maxwellian perturbed thermal velocity $v_t$, top view (e1-e3)") + plt.title("Maxwellian perturbed thermal velocity $v_t$, top view (e1-e3)") plt.subplot(2, 4, 8) if "Slab" in key or "Pinch" in key: plt.contourf(x[:, :, 0], y[:, :, 0], vth_cart[:, :, 0], levels=levels) @@ -1475,7 +1475,7 @@ def test_maxwellian_2d_mhd(Nel, with_desc, show_plot=False): plt.ylabel("z") plt.axis("equal") plt.colorbar() - plt.title(f"Maxwellian perturbed density $v_t$, poloidal view (e1-e2)") + plt.title("Maxwellian perturbed density $v_t$, poloidal view (e1-e2)") plt.show() diff --git a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py index 42d3ae8d3..4da884d02 100644 --- a/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py +++ b/src/struphy/linear_algebra/tests/test_saddlepoint_massmatrices.py @@ -242,7 +242,7 @@ def test_saddlepointsolver(method_for_solving, Nel, p, spl_kind, dirichlet_bc, m TestA11dot = TestA11.dot(x1) compare_arrays(TestA11dot, TestA11composeddot, mpi_rank, atol=1e-5) # compare_arrays(TestA11dot, TestA11npdot, mpi_rank, atol=1e-5) - print(f"Comparison numpy to psydac succesfull.") + print("Comparison numpy to psydac succesfull.") M2pre = MassMatrixPreconditioner(mass_mats.M2) diff --git a/src/struphy/main.py b/src/struphy/main.py index 4b7b65645..047abea95 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -802,7 +802,7 @@ def load_data(path: str) -> SimData: raise NotImplementedError print("\nThe following data has been loaded:") - print(f"\ngrids:") + print("\ngrids:") print(f"{simdata.t_grid.shape =}") if simdata.grids_log is not None: print(f"{simdata.grids_log[0].shape =}") @@ -812,22 +812,22 @@ def load_data(path: str) -> SimData: print(f"{simdata.grids_phy[0].shape =}") print(f"{simdata.grids_phy[1].shape =}") print(f"{simdata.grids_phy[2].shape =}") - print(f"\nsimdata.spline_values:") + print("\nsimdata.spline_values:") for k, v in simdata.spline_values.items(): print(f" {k}") for kk, vv in v.items(): print(f" {kk}") - print(f"\nsimdata.orbits:") + print("\nsimdata.orbits:") for k, v in simdata.orbits.items(): print(f" {k}") - print(f"\nsimdata.f:") + print("\nsimdata.f:") for k, v in simdata.f.items(): print(f" {k}") for kk, vv in v.items(): print(f" {kk}") for kkk, vvv in vv.items(): print(f" {kkk}") - print(f"\nsimdata.n_sph:") + print("\nsimdata.n_sph:") for k, v in simdata.n_sph.items(): print(f" {k}") for kk, vv in v.items(): diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index e6eb5b665..071fb8093 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -106,7 +106,7 @@ def setup_domain_and_equil(self, domain: Domain, equil: FluidEquilibrium): if MPI.COMM_WORLD.Get_rank() == 0 and self.verbose: print("\nDOMAIN:") - print(f"type:".ljust(25), self.domain.__class__.__name__) + print("type:".ljust(25), self.domain.__class__.__name__) for key, val in self.domain.params.items(): if key not in {"cx", "cy", "cz"}: print((key + ":").ljust(25), val) @@ -721,7 +721,7 @@ def update_markers_to_be_saved(self): for name, species in self.particle_species.items(): assert isinstance(species, ParticleSpecies) - assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." + assert len(species.variables) == 1, "More than 1 variable per kinetic species is not allowed." for _, var in species.variables.items(): assert isinstance(var, PICVariable | SPHVariable) obj = var.particles @@ -746,7 +746,7 @@ def update_distr_functions(self): for name, species in self.particle_species.items(): assert isinstance(species, ParticleSpecies) - assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." + assert len(species.variables) == 1, "More than 1 variable per kinetic species is not allowed." for _, var in species.variables.items(): assert isinstance(var, PICVariable | SPHVariable) obj = var.particles @@ -1107,7 +1107,7 @@ def initialize_data_output(self, data: DataContainer, size): # save kinetic data in group 'kinetic/' for name, species in self.particle_species.items(): assert isinstance(species, ParticleSpecies) - assert len(species.variables) == 1, f"More than 1 variable per kinetic species is not allowed." + assert len(species.variables) == 1, "More than 1 variable per kinetic species is not allowed." for varname, var in species.variables.items(): assert isinstance(var, PICVariable | SPHVariable) obj = var.particles @@ -1332,15 +1332,15 @@ def generate_default_parameter_file( has_plasma = True species_params += f"model.{sn}.set_phys_params()\n" if isinstance(species, ParticleSpecies): - particle_params += f"\nloading_params = LoadingParameters()\n" - particle_params += f"weights_params = WeightsParameters()\n" - particle_params += f"boundary_params = BoundaryParameters()\n" + particle_params += "\nloading_params = LoadingParameters()\n" + particle_params += "weights_params = WeightsParameters()\n" + particle_params += "boundary_params = BoundaryParameters()\n" particle_params += f"model.{sn}.set_markers(loading_params=loading_params,\n" - txt = f"weights_params=weights_params,\n" + txt = "weights_params=weights_params,\n" particle_params += indent(txt, " " * len(f"model.{sn}.set_markers(")) - txt = f"boundary_params=boundary_params,\n" + txt = "boundary_params=boundary_params,\n" particle_params += indent(txt, " " * len(f"model.{sn}.set_markers(")) - txt = f")\n" + txt = ")\n" particle_params += indent(txt, " " * len(f"model.{sn}.set_markers(")) particle_params += f"model.{sn}.set_sorting_boxes()\n" particle_params += f"model.{sn}.set_save_data()\n" @@ -1361,38 +1361,38 @@ def generate_default_parameter_file( elif isinstance(var, PICVariable): has_pic = True - init_pert_pic = f"\n# if .add_initial_condition is not called, the background is the kinetic initial condition\n" - init_pert_pic += f"perturbation = perturbations.TorusModesCos()\n" + init_pert_pic = "\n# if .add_initial_condition is not called, the background is the kinetic initial condition\n" + init_pert_pic += "perturbation = perturbations.TorusModesCos()\n" if "6D" in var.space: - init_bckgr_pic = f"maxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, None))\n" - init_bckgr_pic += f"maxwellian_2 = maxwellians.Maxwellian3D(n=(0.1, None))\n" - init_pert_pic += f"maxwellian_1pt = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n" - init_pert_pic += f"init = maxwellian_1pt + maxwellian_2\n" + init_bckgr_pic = "maxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, None))\n" + init_bckgr_pic += "maxwellian_2 = maxwellians.Maxwellian3D(n=(0.1, None))\n" + init_pert_pic += "maxwellian_1pt = maxwellians.Maxwellian3D(n=(1.0, perturbation))\n" + init_pert_pic += "init = maxwellian_1pt + maxwellian_2\n" init_pert_pic += f"model.{sn}.{vn}.add_initial_condition(init)\n" elif "5D" in var.space: - init_bckgr_pic = f"maxwellian_1 = maxwellians.GyroMaxwellian2D(n=(1.0, None), equil=equil)\n" - init_bckgr_pic += f"maxwellian_2 = maxwellians.GyroMaxwellian2D(n=(0.1, None), equil=equil)\n" + init_bckgr_pic = "maxwellian_1 = maxwellians.GyroMaxwellian2D(n=(1.0, None), equil=equil)\n" + init_bckgr_pic += "maxwellian_2 = maxwellians.GyroMaxwellian2D(n=(0.1, None), equil=equil)\n" init_pert_pic += ( - f"maxwellian_1pt = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), equil=equil)\n" + "maxwellian_1pt = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), equil=equil)\n" ) - init_pert_pic += f"init = maxwellian_1pt + maxwellian_2\n" + init_pert_pic += "init = maxwellian_1pt + maxwellian_2\n" init_pert_pic += f"model.{sn}.{vn}.add_initial_condition(init)\n" if "3D" in var.space: - init_bckgr_pic = f"maxwellian_1 = maxwellians.ColdPlasma(n=(1.0, None))\n" - init_bckgr_pic += f"maxwellian_2 = maxwellians.ColdPlasma(n=(0.1, None))\n" - init_pert_pic += f"maxwellian_1pt = maxwellians.ColdPlasma(n=(1.0, perturbation))\n" - init_pert_pic += f"init = maxwellian_1pt + maxwellian_2\n" + init_bckgr_pic = "maxwellian_1 = maxwellians.ColdPlasma(n=(1.0, None))\n" + init_bckgr_pic += "maxwellian_2 = maxwellians.ColdPlasma(n=(0.1, None))\n" + init_pert_pic += "maxwellian_1pt = maxwellians.ColdPlasma(n=(1.0, perturbation))\n" + init_pert_pic += "init = maxwellian_1pt + maxwellian_2\n" init_pert_pic += f"model.{sn}.{vn}.add_initial_condition(init)\n" - init_bckgr_pic += f"background = maxwellian_1 + maxwellian_2\n" + init_bckgr_pic += "background = maxwellian_1 + maxwellian_2\n" init_bckgr_pic += f"model.{sn}.{vn}.add_background(background)\n" - exclude = f"# model.....save_data = False\n" + exclude = "# model.....save_data = False\n" elif isinstance(var, SPHVariable): has_sph = True - init_bckgr_sph = f"background = equils.ConstantVelocity()\n" + init_bckgr_sph = "background = equils.ConstantVelocity()\n" init_bckgr_sph += f"model.{sn}.{vn}.add_background(background)\n" - init_pert_sph = f"perturbation = perturbations.TorusModesCos()\n" + init_pert_sph = "perturbation = perturbations.TorusModesCos()\n" init_pert_sph += f"model.{sn}.{vn}.add_perturbation(del_n=perturbation)\n" exclude = f"# model.{sn}.{vn}.save_data = False\n" @@ -1583,23 +1583,23 @@ def compute_plasma_params(self, verbose=True): if verbose and MPI.COMM_WORLD.Get_rank() == 0: print("\nPLASMA PARAMETERS:") print( - f"Plasma volume:".ljust(25), + "Plasma volume:".ljust(25), "{:4.3e}".format(plasma_volume) + units_affix["plasma volume"], ) print( - f"Transit length:".ljust(25), + "Transit length:".ljust(25), "{:4.3e}".format(transit_length) + units_affix["transit length"], ) print( - f"Avg. magnetic field:".ljust(25), + "Avg. magnetic field:".ljust(25), "{:4.3e}".format(magnetic_field) + units_affix["magnetic field"], ) print( - f"Max magnetic field:".ljust(25), + "Max magnetic field:".ljust(25), "{:4.3e}".format(B_max) + units_affix["magnetic field"], ) print( - f"Min magnetic field:".ljust(25), + "Min magnetic field:".ljust(25), "{:4.3e}".format(B_min) + units_affix["magnetic field"], ) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index e47ab1cd0..e1c310db0 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -197,7 +197,7 @@ def allocate( ): # assert isinstance(self.species, KineticSpecies) assert isinstance(self.backgrounds, KineticBackground), ( - f"List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." + "List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." ) if derham is None: @@ -340,7 +340,7 @@ def allocate( verbose: bool = False, ): assert isinstance(self.backgrounds, FluidEquilibrium), ( - f"List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." + "List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." ) self.backgrounds.domain = domain diff --git a/src/struphy/ode/tests/test_ode_feec.py b/src/struphy/ode/tests/test_ode_feec.py index c0ef51b08..7dfa87a46 100644 --- a/src/struphy/ode/tests/test_ode_feec.py +++ b/src/struphy/ode/tests/test_ode_feec.py @@ -167,10 +167,10 @@ def f(t, y1, y2, y3, out=out): print(f"Convergence check passed on {rank =}.") if rank == 0: - plt.loglog(h_vec, h_vec, "--", label=f"h") - plt.loglog(h_vec, [h**2 for h in h_vec], "--", label=f"h^2") - plt.loglog(h_vec, [h**3 for h in h_vec], "--", label=f"h^3") - plt.loglog(h_vec, [h**4 for h in h_vec], "--", label=f"h^4") + plt.loglog(h_vec, h_vec, "--", label="h") + plt.loglog(h_vec, [h**2 for h in h_vec], "--", label="h^2") + plt.loglog(h_vec, [h**3 for h in h_vec], "--", label="h^3") + plt.loglog(h_vec, [h**4 for h in h_vec], "--", label="h^4") plt.loglog(h_vec, err_vec, "o-k", label=f"{spaces[j]}-space, {algo}") if rank == 0: plt.xlabel("log(h)") diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 79634d13a..6c818b3ee 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -142,7 +142,7 @@ def s0(self, eta1, eta2, eta3, *v, flat_eval=False, remove_holes=True): The 0-form sampling density. ------- """ - assert self.domain, f"self.domain must be set to call the sampling density 0-form." + assert self.domain, "self.domain must be set to call the sampling density 0-form." return self.domain.transform( self.svol(eta1, eta2, eta3, *v), diff --git a/src/struphy/pic/tests/test_mat_vec_filler.py b/src/struphy/pic/tests/test_mat_vec_filler.py index c6bee1faa..073d52ae7 100644 --- a/src/struphy/pic/tests/test_mat_vec_filler.py +++ b/src/struphy/pic/tests/test_mat_vec_filler.py @@ -261,14 +261,14 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): # testing salar spaces if rank == 0: - print(f"\nTesting mat_fill_b_v0 ...") + print("\nTesting mat_fill_b_v0 ...") ptomat.mat_fill_b_v0(DR.args_derham, eta1, eta2, eta3, mat["v0"], fill_mat[0, 0]) assert_mat(mat["v0"], rows, cols, basis["v0"], basis["v0"], rank) # assertion test of mat count += 1 comm.Barrier() if rank == 0: - print(f"\nTesting m_v_fill_b_v0 ...") + print("\nTesting m_v_fill_b_v0 ...") ptomat.m_v_fill_b_v0(DR.args_derham, eta1, eta2, eta3, mat["v0"], fill_mat[0, 0], vec["v0"], fill_vec[0]) assert_mat(mat["v0"], rows, cols, basis["v0"], basis["v0"], rank) # assertion test of mat assert_vec(vec["v0"], rows, basis["v0"], rank) # assertion test of vec @@ -276,14 +276,14 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): comm.Barrier() if rank == 0: - print(f"\nTesting mat_fill_b_v3 ...") + print("\nTesting mat_fill_b_v3 ...") ptomat.mat_fill_b_v3(DR.args_derham, eta1, eta2, eta3, mat["v3"], fill_mat[0, 0]) assert_mat(mat["v3"], rows, cols, basis["v3"], basis["v3"], rank) # assertion test of mat count += 1 comm.Barrier() if rank == 0: - print(f"\nTesting m_v_fill_b_v3 ...") + print("\nTesting m_v_fill_b_v3 ...") ptomat.m_v_fill_b_v3(DR.args_derham, eta1, eta2, eta3, mat["v3"], fill_mat[0, 0], vec["v3"], fill_vec[0]) assert_mat(mat["v3"], rows, cols, basis["v3"], basis["v3"], rank) # assertion test of mat assert_vec(vec["v3"], rows, basis["v3"], rank) # assertion test of vec @@ -291,14 +291,14 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): comm.Barrier() if rank == 0: - print(f"\nTesting mat_fill_v0 ...") + print("\nTesting mat_fill_v0 ...") ptomat.mat_fill_v0(DR.args_derham, span1, span2, span3, mat["v0"], fill_mat[0, 0]) assert_mat(mat["v0"], rows, cols, basis["v0"], basis["v0"], rank) # assertion test of mat count += 1 comm.Barrier() if rank == 0: - print(f"\nTesting m_v_fill_v0 ...") + print("\nTesting m_v_fill_v0 ...") ptomat.m_v_fill_v0(DR.args_derham, span1, span2, span3, mat["v0"], fill_mat[0, 0], vec["v0"], fill_vec[0]) assert_mat(mat["v0"], rows, cols, basis["v0"], basis["v0"], rank) # assertion test of mat assert_vec(vec["v0"], rows, basis["v0"], rank) # assertion test of vec @@ -306,14 +306,14 @@ def test_particle_to_mat_kernels(Nel, p, spl_kind, n_markers=1): comm.Barrier() if rank == 0: - print(f"\nTesting mat_fill_v3 ...") + print("\nTesting mat_fill_v3 ...") ptomat.mat_fill_v3(DR.args_derham, span1, span2, span3, mat["v3"], fill_mat[0, 0]) assert_mat(mat["v3"], rows, cols, basis["v3"], basis["v3"], rank) # assertion test of mat count += 1 comm.Barrier() if rank == 0: - print(f"\nTesting m_v_fill_v3 ...") + print("\nTesting m_v_fill_v3 ...") ptomat.m_v_fill_v3(DR.args_derham, span1, span2, span3, mat["v3"], fill_mat[0, 0], vec["v3"], fill_vec[0]) assert_mat(mat["v3"], rows, cols, basis["v3"], basis["v3"], rank) # assertion test of mat assert_vec(vec["v3"], rows, basis["v3"], rank) # assertion test of vec diff --git a/src/struphy/post_processing/likwid/plot_likwidproject.py b/src/struphy/post_processing/likwid/plot_likwidproject.py index cde2a2b76..f4c3bb442 100644 --- a/src/struphy/post_processing/likwid/plot_likwidproject.py +++ b/src/struphy/post_processing/likwid/plot_likwidproject.py @@ -387,7 +387,7 @@ def plot_speedup( fig.update_layout( # xaxis_title='Job name', - xaxis_title=f"MPI tasks (#)", + xaxis_title="MPI tasks (#)", yaxis_title=re.sub(r"\[.*?\]", "[relative]", metric2), showlegend=True, xaxis_tickformat=".1f", diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 74a6288f6..e0759bb63 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -156,7 +156,7 @@ def create_femfields( # get fields names, space IDs and time grid from 0-th rank hdf5 file file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") space_ids = {} - print(f"\nReading hdf5 data of following species:") + print("\nReading hdf5 data of following species:") for species, dset in file["feec"].items(): space_ids[species] = {} print(f"{species}:") diff --git a/src/struphy/propagators/propagators_fields.py b/src/struphy/propagators/propagators_fields.py index c3f3e1381..2a4431952 100644 --- a/src/struphy/propagators/propagators_fields.py +++ b/src/struphy/propagators/propagators_fields.py @@ -8722,7 +8722,7 @@ def __call__(self, dt): e = phi_temp.ends phi_temp[s[0] : e[0] + 1, s[1] : e[1] + 1, s[2] : e[2] + 1] = phin.reshape(*dimphi) else: - print(f"TwoFluidQuasiNeutralFull is only running on one MPI.") + print("TwoFluidQuasiNeutralFull is only running on one MPI.") # write new coeffs into self.feec_vars max_du, max_due, max_dphi = self.update_feec_variables(u=u_temp, ue=ue_temp, phi=phi_temp) diff --git a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py index 68ba44bcd..747ec65c7 100644 --- a/src/struphy/propagators/tests/test_gyrokinetic_poisson.py +++ b/src/struphy/propagators/tests/test_gyrokinetic_poisson.py @@ -224,7 +224,7 @@ def rho_pulled(e1, e2, e3): plt.xscale("log") plt.xlabel("Grid Spacing h") plt.ylabel("Error") - plt.title(f"Poisson solver") + plt.title("Poisson solver") plt.legend() if show_plot and rank == 0: diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index bd425170a..588aa2aa1 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -267,7 +267,7 @@ def rho_pulled(e1, e2, e3): plt.xscale("log") plt.xlabel("Grid Spacing h") plt.ylabel("Error") - plt.title(f"Poisson solver") + plt.title("Poisson solver") plt.legend() if show_plot and rank == 0: diff --git a/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb b/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb index 922282da1..f50bfd8a7 100644 --- a/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb +++ b/tutorials/tutorial_03_smoothed_particle_hydrodynamics.ipynb @@ -913,14 +913,14 @@ "plt.pcolor(ee1[:,:,0], ee2[:,:,0], n_sph[:,:,0])\n", "plt.grid()\n", "plt.axis('square')\n", - "plt.title(f'n_sph initial (random)')\n", + "plt.title('n_sph initial (random)')\n", "plt.colorbar()\n", "\n", "plt.subplot(3, 2, 5)\n", "ax = plt.gca()\n", "plt.pcolor(bc_x, bc_y, f_bin)\n", "plt.axis('square')\n", - "plt.title(f'n_binned initial (random)')\n", + "plt.title('n_binned initial (random)')\n", "plt.colorbar()" ] }, diff --git a/tutorials_old/tutorial_02_fluid_particles.ipynb b/tutorials_old/tutorial_02_fluid_particles.ipynb index 49ccf724b..36ab63d1b 100644 --- a/tutorials_old/tutorial_02_fluid_particles.ipynb +++ b/tutorials_old/tutorial_02_fluid_particles.ipynb @@ -805,7 +805,7 @@ "plt.tick_params(labelbottom = False) \n", "plt.plot(eta1, n_sph_init[:, 0, 0])\n", "plt.grid()\n", - "plt.title(f'n_sph_init')\n", + "plt.title('n_sph_init')\n", "\n", "plt.subplot(2, 2, 4)\n", "ax = plt.gca()\n", @@ -815,7 +815,7 @@ "bc_x = (be_x[:-1] + be_x[1:]) / 2. # centers of binning cells\n", "plt.plot(bc_x, df_bin.T)\n", "#plt.grid()\n", - "plt.title(f'n_binned')" + "plt.title('n_binned')" ] }, { @@ -1209,7 +1209,7 @@ "plt.pcolor(ee1[:,:,0], ee2[:,:,0], n_sph_1[:,:,0])\n", "plt.grid()\n", "plt.axis('square')\n", - "plt.title(f'n_sph (random)')\n", + "plt.title('n_sph (random)')\n", "plt.colorbar()\n", "\n", "plt.subplot(4, 2, 6)\n", @@ -1220,7 +1220,7 @@ "plt.pcolor(ee1[:,:,0], ee2[:,:,0], n_sph_2[:,:,0])\n", "plt.grid()\n", "plt.axis('square')\n", - "plt.title(f'n_sph (tesselation)')\n", + "plt.title('n_sph (tesselation)')\n", "plt.colorbar()\n", "\n", "plt.subplot(4, 2, 7)\n", @@ -1233,7 +1233,7 @@ "plt.pcolor(bc_x, bc_y, f_bin_1)\n", "#plt.grid()\n", "plt.axis('square')\n", - "plt.title(f'n_binned (random)')\n", + "plt.title('n_binned (random)')\n", "plt.colorbar()\n", "\n", "plt.subplot(4, 2, 8)\n", @@ -1246,7 +1246,7 @@ "plt.pcolor(bc_x, bc_y, f_bin_2)\n", "#plt.grid()\n", "plt.axis('square')\n", - "plt.title(f'n_binned (tesselation)')\n", + "plt.title('n_binned (tesselation)')\n", "plt.colorbar()" ] }, From 7eacf1c035ae705d82956ad98345ef71adebe510 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 27 Oct 2025 18:22:32 +0100 Subject: [PATCH 244/292] Formatting --- src/struphy/models/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 071fb8093..27076b6dc 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -1361,7 +1361,9 @@ def generate_default_parameter_file( elif isinstance(var, PICVariable): has_pic = True - init_pert_pic = "\n# if .add_initial_condition is not called, the background is the kinetic initial condition\n" + init_pert_pic = ( + "\n# if .add_initial_condition is not called, the background is the kinetic initial condition\n" + ) init_pert_pic += "perturbation = perturbations.TorusModesCos()\n" if "6D" in var.space: init_bckgr_pic = "maxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, None))\n" From 5c5489520147842040957b43dd871ebabd9b55f4 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 27 Oct 2025 19:14:28 +0100 Subject: [PATCH 245/292] not kinds --> ~kinds --- src/struphy/feec/psydac_derham.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index bb833e02d..e5a8cde32 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1270,7 +1270,7 @@ def _get_neighbour_one_component(self, comp): # if only one process: check if comp is neighbour in non-peridic directions, if this is not the case then return the rank as neighbour id if size == 1: - if (comp[not kinds] == 1).all(): + if (comp[~kinds] == 1).all(): return rank # multiple processes From 37af29ba9ec724e463d5214aaafba577d1746b49 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 08:58:11 +0100 Subject: [PATCH 246/292] updates some docker files --- docker/opensuse-latest.dockerfile | 7 +-- docker/ubuntu-latest-with-struphy.dockerfile | 53 ++++++++++++-------- docker/ubuntu-latest.dockerfile | 37 ++++++++------ 3 files changed, 57 insertions(+), 40 deletions(-) diff --git a/docker/opensuse-latest.dockerfile b/docker/opensuse-latest.dockerfile index 04ecff4e4..97d5a5aa0 100644 --- a/docker/opensuse-latest.dockerfile +++ b/docker/opensuse-latest.dockerfile @@ -1,9 +1,10 @@ -# Here is how to build the image and upload it to the mpcdf gitlab registry: +# Here is how to build the image and upload it to the Github package registry: # # We suppose you are in the struphy repo directory. -# Start the docker engine and run "docker login" with the following token: +# Start the docker engine and login to the Github package registry using a github personal acces token (classic): # -# TOKEN=gldt-CgMRBMtePbSwdWTxKw4Q; echo "$TOKEN" | docker login gitlab-registry.mpcdf.mpg.de -u gitlab+deploy-token-162 --password-stdin +# export CR_PAT=YOUR_TOKEN +# echo $CR_PAT | docker login ghcr.io -u USERNAME --password-stdin # docker info # docker build -t gitlab-registry.mpcdf.mpg.de/struphy/struphy/opensuse-latest --provenance=false -f docker/opensuse-latest.dockerfile . # docker push gitlab-registry.mpcdf.mpg.de/struphy/struphy/opensuse-latest diff --git a/docker/ubuntu-latest-with-struphy.dockerfile b/docker/ubuntu-latest-with-struphy.dockerfile index c7fea9c9b..5b8d8c0fd 100644 --- a/docker/ubuntu-latest-with-struphy.dockerfile +++ b/docker/ubuntu-latest-with-struphy.dockerfile @@ -1,12 +1,13 @@ -# Here is how to build the image and upload it to the mpcdf gitlab registry: +# Here is how to build the image and upload it to the Github package registry: # # We suppose you are in the struphy repo directory. -# Start the docker engine and run "docker login" with the following token: +# Start the docker engine and login to the Github package registry using a github personal acces token (classic): # -# TOKEN=gldt-CgMRBMtePbSwdWTxKw4Q; echo "$TOKEN" | docker login gitlab-registry.mpcdf.mpg.de -u gitlab+deploy-token-162 --password-stdin +# export CR_PAT=YOUR_TOKEN +# echo $CR_PAT | docker login ghcr.io -u USERNAME --password-stdinn # docker info -# docker build -t gitlab-registry.mpcdf.mpg.de/struphy/struphy/ubuntu-latest-with-struphy --provenance=false -f docker/ubuntu-latest-with-struphy.dockerfile . -# docker push gitlab-registry.mpcdf.mpg.de/struphy/struphy/ubuntu-latest-with-struphy +# docker build -t ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest --provenance=false -f docker/ubuntu-latest-with-struphy.dockerfile . +# docker push ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest FROM ubuntu:latest @@ -16,49 +17,57 @@ ARG DEBIAN_FRONTEND=noninteractive RUN apt update -y && apt clean \ && apt install -y software-properties-common \ && add-apt-repository -y ppa:deadsnakes/ppa \ - && apt update -y \ - && apt install -y python3 \ + && apt update -y + +RUN apt install -y python3 \ && apt install -y python3-dev \ && apt install -y python3-pip \ - && apt install -y python3-venv \ - && apt install -y gfortran gcc \ - && apt install -y liblapack-dev libblas-dev \ - && apt install -y libopenmpi-dev openmpi-bin \ - && apt install -y libomp-dev libomp5 \ - && apt install -y git \ + && apt install -y python3-venv + +RUN apt install -y gfortran gcc \ + && apt install -y liblapack-dev libblas-dev + +RUN apt install -y libopenmpi-dev openmpi-bin \ + && apt install -y libomp-dev libomp5 + +RUN apt install -y git \ && apt install -y pandoc graphviz \ - && bash -c "source ~/.bashrc" \ - # for gvec - && apt install -y g++ liblapack3 cmake cmake-curses-gui zlib1g-dev libnetcdf-dev libnetcdff-dev \ + && bash -c "source ~/.bashrc" + +# for gvec +RUN apt install -y g++ liblapack3 cmake cmake-curses-gui zlib1g-dev libnetcdf-dev libnetcdff-dev \ && export FC=`which gfortran` \ && export CC=`which gcc` \ && export CXX=`which g++` # install three versions of struphy -RUN git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git struphy_c_ \ +RUN git clone https://github.com/struphy-hub/struphy.git struphy_c_ \ && cd struphy_c_ \ && python3 -m venv env_c_ \ && . env_c_/bin/activate \ && pip install -U pip \ - && pip install -e .[phys] --no-cache-dir \ + && pip install -e .[phys,mpi,doc] --no-cache-dir \ + && struphy compile --status \ && struphy compile \ && deactivate -RUN git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git struphy_fortran_\ +RUN git clone https://github.com/struphy-hub/struphy.git struphy_fortran_\ && cd struphy_fortran_ \ && python3 -m venv env_fortran_ \ && . env_fortran_/bin/activate \ && pip install -U pip \ - && pip install -e .[phys] --no-cache-dir \ + && pip install -e .[phys,mpi,doc] --no-cache-dir \ + && struphy compile --status \ && struphy compile --language fortran -y \ && deactivate -RUN git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git struphy_fortran_--omp-pic\ +RUN git clone https://github.com/struphy-hub/struphy.git struphy_fortran_--omp-pic\ && cd struphy_fortran_--omp-pic \ && python3 -m venv env_fortran_--omp-pic \ && . env_fortran_--omp-pic/bin/activate \ && pip install -U pip \ - && pip install -e .[phys] --no-cache-dir \ + && pip install -e .[phys,mpi,doc] --no-cache-dir \ + && struphy compile --status \ && struphy compile --language fortran --omp-pic -y \ && deactivate diff --git a/docker/ubuntu-latest.dockerfile b/docker/ubuntu-latest.dockerfile index adcf65609..3f0cb3450 100644 --- a/docker/ubuntu-latest.dockerfile +++ b/docker/ubuntu-latest.dockerfile @@ -1,12 +1,13 @@ -# Here is how to build the image and upload it to the mpcdf gitlab registry: +# Here is how to build the image and upload it to the Github package registry: # # We suppose you are in the struphy repo directory. -# Start the docker engine and run "docker login" with the following token: +# Start the docker engine and login to the Github package registry using a github personal acces token (classic): # -# TOKEN=gldt-CgMRBMtePbSwdWTxKw4Q; echo "$TOKEN" | docker login gitlab-registry.mpcdf.mpg.de -u gitlab+deploy-token-162 --password-stdin +# export CR_PAT=YOUR_TOKEN +# echo $CR_PAT | docker login ghcr.io -u USERNAME --password-stdin # docker info -# docker build -t gitlab-registry.mpcdf.mpg.de/struphy/struphy/ubuntu-latest --provenance=false -f docker/ubuntu-latest.dockerfile . -# docker push gitlab-registry.mpcdf.mpg.de/struphy/struphy/ubuntu-latest +# docker build -t ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest --provenance=false -f docker/ubuntu-latest.dockerfile . +# docker push ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest FROM ubuntu:latest @@ -16,20 +17,26 @@ ARG DEBIAN_FRONTEND=noninteractive RUN apt update -y && apt clean \ && apt install -y software-properties-common \ && add-apt-repository -y ppa:deadsnakes/ppa \ - && apt update -y \ - && apt install -y python3 \ + && apt update -y + +RUN apt install -y python3 \ && apt install -y python3-dev \ && apt install -y python3-pip \ - && apt install -y python3-venv \ - && apt install -y gfortran gcc \ - && apt install -y liblapack-dev libblas-dev \ - && apt install -y libopenmpi-dev openmpi-bin \ - && apt install -y libomp-dev libomp5 \ - && apt install -y git \ + && apt install -y python3-venv + +RUN apt install -y gfortran gcc \ + && apt install -y liblapack-dev libblas-dev + +RUN apt install -y libopenmpi-dev openmpi-bin \ + && apt install -y libomp-dev libomp5 + + +RUN apt install -y git \ && apt install -y pandoc graphviz \ - && bash -c "source ~/.bashrc" \ + && bash -c "source ~/.bashrc" + # for gvec - && apt install -y g++ liblapack3 cmake cmake-curses-gui zlib1g-dev libnetcdf-dev libnetcdff-dev \ +RUN apt install -y g++ liblapack3 cmake cmake-curses-gui zlib1g-dev libnetcdf-dev libnetcdff-dev \ && export FC=`which gfortran` \ && export CC=`which gcc` \ && export CXX=`which g++` From c66da598ecb6aecaaf1ca482baf9c5ce819f600b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 08:58:34 +0100 Subject: [PATCH 247/292] new workflow for testing image --- .github/workflows/test_our_image.yml | 35 ++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 .github/workflows/test_our_image.yml diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml new file mode 100644 index 000000000..a79a7fa75 --- /dev/null +++ b/.github/workflows/test_our_image.yml @@ -0,0 +1,35 @@ +# name: Deploy docs to GitHub Pages + +on: + push: + branches: ["devel", "main"] + workflow_dispatch: + +defaults: + run: + shell: bash + +permissions: + contents: read +# id-token: write + +# concurrency: +# group: "pages" +# cancel-in-progress: false + +jobs: + hello-world: + runs-on: ubuntu-latest + container: + image: struphy/ubuntu-with-reqs + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Inspect Image + run: | + ls + python3 -m venv env + source env/bin/activate + pip install -U pip + pip list + From 3bc004074e26fae94c843bf829f7c64c46b8fdbf Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 09:54:03 +0100 Subject: [PATCH 248/292] test images in workflow --- .github/workflows/test_our_image.yml | 22 +++++++++++++++++++--- .github/workflows/ubuntu-latest.yml | 8 ++++---- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index a79a7fa75..42037100f 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -1,8 +1,10 @@ # name: Deploy docs to GitHub Pages on: - push: - branches: ["devel", "main"] + pull_request: + branches: + - main + - devel workflow_dispatch: defaults: @@ -18,7 +20,7 @@ permissions: # cancel-in-progress: false jobs: - hello-world: + with-reqs: runs-on: ubuntu-latest container: image: struphy/ubuntu-with-reqs @@ -33,3 +35,17 @@ jobs: pip install -U pip pip list + with-struphy: + runs-on: ubuntu-latest + container: + image: struphy/ubuntu-with-struphy + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Inspect Image + run: | + ls + source env_c_/bin/activate + struphy -h + struphy compile + diff --git a/.github/workflows/ubuntu-latest.yml b/.github/workflows/ubuntu-latest.yml index e66aeb3b3..b2ba00399 100644 --- a/.github/workflows/ubuntu-latest.yml +++ b/.github/workflows/ubuntu-latest.yml @@ -4,10 +4,10 @@ on: branches: - main - devel - pull_request: - branches: - - main - - devel + # pull_request: + # branches: + # - main + # - devel # concurrency: # group: ${{ github.ref }} From 485ae42745934a4b8da524830c80832beafd3c5e Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 10:03:22 +0100 Subject: [PATCH 249/292] change container --- .github/workflows/static_analysis.yml | 8 ++++---- .github/workflows/test_our_image.yml | 11 +++++++++-- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index 2fd4c2051..1f18a0d09 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -5,10 +5,10 @@ on: branches: - main - devel - pull_request: - branches: - - main - - devel + # pull_request: + # branches: + # - main + # - devel # concurrency: # group: ${{ github.ref }} diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 42037100f..b3775b296 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -23,7 +23,11 @@ jobs: with-reqs: runs-on: ubuntu-latest container: - image: struphy/ubuntu-with-reqs + image: ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest + credentials: + username: ${{ github.actor }} + password: ${{ secrets.github_token }} + steps: - name: Checkout uses: actions/checkout@v4 @@ -38,7 +42,10 @@ jobs: with-struphy: runs-on: ubuntu-latest container: - image: struphy/ubuntu-with-struphy + image: ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest + credentials: + username: ${{ github.actor }} + password: ${{ secrets.github_token }} steps: - name: Checkout uses: actions/checkout@v4 From 714c3f3e7d8147ee80baa8796e9b4408b300b1cd Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 10:07:24 +0100 Subject: [PATCH 250/292] upper case GITHUB_TOKEN --- .github/workflows/test_our_image.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index b3775b296..58bcb24e5 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -26,7 +26,7 @@ jobs: image: ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest credentials: username: ${{ github.actor }} - password: ${{ secrets.github_token }} + password: ${{ secrets.GITHUB_TOKEN }} steps: - name: Checkout @@ -45,7 +45,7 @@ jobs: image: ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest credentials: username: ${{ github.actor }} - password: ${{ secrets.github_token }} + password: ${{ secrets.GITHUB_TOKEN }} steps: - name: Checkout uses: actions/checkout@v4 From 327afec1755eff440a773dd982c6f9ea5416f629 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 10:21:34 +0100 Subject: [PATCH 251/292] use new secrets.GHCR_TOKEN --- .github/workflows/test_our_image.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 58bcb24e5..df0c70078 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -26,7 +26,7 @@ jobs: image: ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest credentials: username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} + password: ${{ secrets.GHCR_TOKEN }} steps: - name: Checkout @@ -45,7 +45,7 @@ jobs: image: ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest credentials: username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} + password: ${{ secrets.GHCR_TOKEN }} steps: - name: Checkout uses: actions/checkout@v4 From d94b189c2c7c9dbabd7be5fafdca1150142681cf Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 10:29:36 +0100 Subject: [PATCH 252/292] remove checkout --- .github/workflows/test_our_image.yml | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index df0c70078..d1593216b 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -29,11 +29,14 @@ jobs: password: ${{ secrets.GHCR_TOKEN }} steps: - - name: Checkout - uses: actions/checkout@v4 + # - name: Checkout + # uses: actions/checkout@v4 - name: Inspect Image run: | ls + cd .. + ls + uname -a python3 -m venv env source env/bin/activate pip install -U pip @@ -47,11 +50,14 @@ jobs: username: ${{ github.actor }} password: ${{ secrets.GHCR_TOKEN }} steps: - - name: Checkout - uses: actions/checkout@v4 + # - name: Checkout + # uses: actions/checkout@v4 - name: Inspect Image run: | ls + cd .. + ls + uname -a source env_c_/bin/activate struphy -h struphy compile From 42c80af2b030edf59bd37fd83be2c94ca9ef1608 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 10:46:17 +0100 Subject: [PATCH 253/292] check for dockerenv file --- .github/workflows/test_our_image.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index d1593216b..dc816e578 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -31,6 +31,8 @@ jobs: steps: # - name: Checkout # uses: actions/checkout@v4 + - name: Check for dockerenv file + run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - name: Inspect Image run: | ls @@ -52,6 +54,8 @@ jobs: steps: # - name: Checkout # uses: actions/checkout@v4 + - name: Check for dockerenv file + run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - name: Inspect Image run: | ls From 43e8f86a4b8165aff8fe32d193d51302b16b7bcd Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 10:50:45 +0100 Subject: [PATCH 254/292] add shell bach --- .github/workflows/test_our_image.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index dc816e578..745704faa 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -34,6 +34,7 @@ jobs: - name: Check for dockerenv file run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - name: Inspect Image + shell: bash run: | ls cd .. @@ -57,6 +58,7 @@ jobs: - name: Check for dockerenv file run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - name: Inspect Image + shell: bash run: | ls cd .. From 97b9455822f3a1421464f3fb685b19dd1b7ff519 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 10:52:56 +0100 Subject: [PATCH 255/292] add pwd --- .github/workflows/test_our_image.yml | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 745704faa..fe802a1e5 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -36,8 +36,7 @@ jobs: - name: Inspect Image shell: bash run: | - ls - cd .. + pwd ls uname -a python3 -m venv env @@ -60,8 +59,7 @@ jobs: - name: Inspect Image shell: bash run: | - ls - cd .. + pwd ls uname -a source env_c_/bin/activate From 552216cf993ae67950a0bc184fa4f0b3f05c3226 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 13:36:32 +0100 Subject: [PATCH 256/292] try installing struphy --- .github/workflows/test_our_image.yml | 39 ++++++++++++++++------------ docker/ubuntu-latest.dockerfile | 6 +---- 2 files changed, 24 insertions(+), 21 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index fe802a1e5..dd43d372b 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -29,20 +29,26 @@ jobs: password: ${{ secrets.GHCR_TOKEN }} steps: - # - name: Checkout - # uses: actions/checkout@v4 - - name: Check for dockerenv file - run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - - name: Inspect Image - shell: bash - run: | - pwd - ls - uname -a - python3 -m venv env - source env/bin/activate - pip install -U pip - pip list + - name: Inspect Image + run: | + pwd + ls + uname -a + python3 -m venv env + source env/bin/activate + pip install -U pip + pip list + - name: Checkout + uses: actions/checkout@v4 + - name: Check for dockerenv file + run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) + - name: Install Struphy + run: | + ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) + pip list + pip install -U pip + pip install ".[test, dev, docs]" + with-struphy: runs-on: ubuntu-latest @@ -60,9 +66,10 @@ jobs: shell: bash run: | pwd - ls + ls -a uname -a - source env_c_/bin/activate + ls / + ls ~ struphy -h struphy compile diff --git a/docker/ubuntu-latest.dockerfile b/docker/ubuntu-latest.dockerfile index 3f0cb3450..386426c29 100644 --- a/docker/ubuntu-latest.dockerfile +++ b/docker/ubuntu-latest.dockerfile @@ -30,20 +30,16 @@ RUN apt install -y gfortran gcc \ RUN apt install -y libopenmpi-dev openmpi-bin \ && apt install -y libomp-dev libomp5 - RUN apt install -y git \ && apt install -y pandoc graphviz \ && bash -c "source ~/.bashrc" - # for gvec +# for gvec RUN apt install -y g++ liblapack3 cmake cmake-curses-gui zlib1g-dev libnetcdf-dev libnetcdff-dev \ && export FC=`which gfortran` \ && export CC=`which gcc` \ && export CXX=`which g++` -# Create a new working directory -WORKDIR /install_struphy_here/ - # allow mpirun as root ENV OMPI_ALLOW_RUN_AS_ROOT=1 ENV OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 From 90a2dd0fa6ff4879dffb4870c2dd5f4b700f419b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 13:37:54 +0100 Subject: [PATCH 257/292] forgotten ( --- .github/workflows/test_our_image.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index dd43d372b..4f6f7e533 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -44,7 +44,7 @@ jobs: run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - name: Install Struphy run: | - ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) + (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) pip list pip install -U pip pip install ".[test, dev, docs]" From d97b8a82d08c0e3397a87cda3acb4bca06174ff0 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 13:47:12 +0100 Subject: [PATCH 258/292] enter virtual envs --- .github/workflows/test_our_image.yml | 46 +++++++++++++++------------- 1 file changed, 25 insertions(+), 21 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 4f6f7e533..6f16080e0 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -29,27 +29,26 @@ jobs: password: ${{ secrets.GHCR_TOKEN }} steps: - - name: Inspect Image + - name: Check for dockerenv file + run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) + - name: Checkout repo + uses: actions/checkout@v4 + - name: Create virtual environment run: | - pwd - ls - uname -a python3 -m venv env source env/bin/activate pip install -U pip pip list - - name: Checkout - uses: actions/checkout@v4 - - name: Check for dockerenv file - run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) + whch python3 - name: Install Struphy run: | - (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) pip list - pip install -U pip + which python3 + ls -a + source env/bin/activate + which python3 pip install ".[test, dev, docs]" - with-struphy: runs-on: ubuntu-latest container: @@ -58,18 +57,23 @@ jobs: username: ${{ github.actor }} password: ${{ secrets.GHCR_TOKEN }} steps: - # - name: Checkout - # uses: actions/checkout@v4 - name: Check for dockerenv file run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - - name: Inspect Image - shell: bash + - name: Checkout repo + uses: actions/checkout@v4 + - name: Actovate pre-installed virtual environment + run: | + ls / -a + source /env_c_/bin/activate + pip install -U pip + pip list + whch python3 + - name: Install Struphy run: | - pwd + pip list + which python3 ls -a - uname -a - ls / - ls ~ - struphy -h - struphy compile + source /env_c_/bin/activate + which python3 + pip install ".[test, dev, docs]" From d2187352870bfbafd7e6c8e40304c63c0d3392ed Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 13:48:23 +0100 Subject: [PATCH 259/292] spelling error --- .github/workflows/test_our_image.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 6f16080e0..0ed2a5820 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -39,7 +39,7 @@ jobs: source env/bin/activate pip install -U pip pip list - whch python3 + which python3 - name: Install Struphy run: | pip list @@ -67,7 +67,7 @@ jobs: source /env_c_/bin/activate pip install -U pip pip list - whch python3 + which python3 - name: Install Struphy run: | pip list From 50eb15cc9135b2ba2fd45157dc471b69a606879a Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 13:51:18 +0100 Subject: [PATCH 260/292] correct env name --- .github/workflows/test_our_image.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 0ed2a5820..a095d8d17 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -61,10 +61,10 @@ jobs: run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - name: Checkout repo uses: actions/checkout@v4 - - name: Actovate pre-installed virtual environment + - name: Activate pre-installed virtual environment run: | ls / -a - source /env_c_/bin/activate + source /struphy_c_/bin/activate pip install -U pip pip list which python3 @@ -73,7 +73,7 @@ jobs: pip list which python3 ls -a - source /env_c_/bin/activate + source /struphy_c_/bin/activate which python3 pip install ".[test, dev, docs]" From bfa356628cea03c54f5df213ecc2eea299bc7234 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 13:55:19 +0100 Subject: [PATCH 261/292] add some which python3 --- .github/workflows/test_our_image.yml | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index a095d8d17..5bb7277af 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -35,19 +35,20 @@ jobs: uses: actions/checkout@v4 - name: Create virtual environment run: | + which python3 python3 -m venv env source env/bin/activate - pip install -U pip - pip list which python3 + pip install -U pip - name: Install Struphy run: | - pip list which python3 - ls -a source env/bin/activate which python3 pip install ".[test, dev, docs]" + # - name: Compile Struphy + # run: | + # which python3 with-struphy: runs-on: ubuntu-latest @@ -64,15 +65,13 @@ jobs: - name: Activate pre-installed virtual environment run: | ls / -a + which python3 source /struphy_c_/bin/activate - pip install -U pip - pip list which python3 + pip install -U pip - name: Install Struphy run: | - pip list which python3 - ls -a source /struphy_c_/bin/activate which python3 pip install ".[test, dev, docs]" From 3f54c23eb416fe6e6c3aa213eb006127a6c2ad78 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 14:00:56 +0100 Subject: [PATCH 262/292] correct env path --- .github/workflows/test_our_image.yml | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 5bb7277af..1860b6245 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -45,10 +45,13 @@ jobs: which python3 source env/bin/activate which python3 - pip install ".[test, dev, docs]" - # - name: Compile Struphy - # run: | - # which python3 + pip install ".[test,dev,docs]" + - name: Compile Struphy + run: | + which python3 + (struphy compile) || (source env/bin/activate) + which python3 + struphy compile with-struphy: runs-on: ubuntu-latest @@ -66,13 +69,13 @@ jobs: run: | ls / -a which python3 - source /struphy_c_/bin/activate + source /struphy_c_/env_c_/bin/activate which python3 pip install -U pip - name: Install Struphy run: | which python3 - source /struphy_c_/bin/activate + source /struphy_c_/env_c_/bin/activate which python3 - pip install ".[test, dev, docs]" + pip install ".[test,dev,docs]" From 4fedd4adc72f03745bbbb3dda909a380b8c483a6 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 14:05:54 +0100 Subject: [PATCH 263/292] compile Struphy --- .github/workflows/test_our_image.yml | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 1860b6245..9ad152ad1 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -45,11 +45,11 @@ jobs: which python3 source env/bin/activate which python3 - pip install ".[test,dev,docs]" + pip install .[test,dev,docs] - name: Compile Struphy run: | which python3 - (struphy compile) || (source env/bin/activate) + source env/bin/activate which python3 struphy compile @@ -78,4 +78,10 @@ jobs: source /struphy_c_/env_c_/bin/activate which python3 pip install ".[test,dev,docs]" + - name: Compile Struphy + run: | + which python3 + source /struphy_c_/env_c_/bin/activate + which python3 + struphy compile From 0ee68dbd20b3f1f58b90978c123cbdda3e23b14b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 14:54:20 +0100 Subject: [PATCH 264/292] find out why struphy compile just refresshes models and then exits --- src/struphy/console/compile.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/struphy/console/compile.py b/src/struphy/console/compile.py index 432e4fa1f..98847e433 100644 --- a/src/struphy/console/compile.py +++ b/src/struphy/console/compile.py @@ -78,6 +78,8 @@ def struphy_compile( # Read struphy state file state = utils.read_state() + + print(f"before {state = }") # collect kernels if "kernels" not in state: @@ -103,6 +105,8 @@ def struphy_compile( utils.save_state(state) # source files sources = " ".join(state["kernels"]) + + print(f"after {state = }") # actions if delete: @@ -222,6 +226,9 @@ def struphy_compile( flags += " --time-execution" # state + print(f"inside else {state = }") + print(f"{yes = }") + if state["last_used_language"] not in (language, None): if yes: yesno = "Y" @@ -230,6 +237,8 @@ def struphy_compile( f"Kernels compiled in language {state['last_used_language']} exist, will be deleted, continue (Y/n)?", ) + print(f"{yesno = }") + if yesno in ("", "Y", "y", "yes"): cmd = [ "struphy", @@ -246,6 +255,8 @@ def struphy_compile( state["last_used_omp_feec"] = flag_omp_feec utils.save_state(state) + + print(f"{state = }") # install psydac from wheel if not there source_install = False @@ -253,6 +264,8 @@ def struphy_compile( if "psydac" in req: source_install = True + print(f"{source_install = }") + struphy_ver = importlib.metadata.version("struphy") try: From 81db8ca5b678a8ad7d63a072515c5c7df8cf32c2 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 15:00:23 +0100 Subject: [PATCH 265/292] cat compile.py --- .github/workflows/test_our_image.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 9ad152ad1..96e81676b 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -45,7 +45,10 @@ jobs: which python3 source env/bin/activate which python3 + ls + cat src/struphy/concole/compile.py pip install .[test,dev,docs] + pip install .[dev] - name: Compile Struphy run: | which python3 From b91db0db3bec1d427bde47c47e912c6159a24280 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 15:02:06 +0100 Subject: [PATCH 266/292] fix typo --- .github/workflows/test_our_image.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 96e81676b..86fa5c5ea 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -46,7 +46,7 @@ jobs: source env/bin/activate which python3 ls - cat src/struphy/concole/compile.py + cat src/struphy/console/compile.py pip install .[test,dev,docs] pip install .[dev] - name: Compile Struphy From ffd261f481b76462164d9b3b35df4f130cdda10b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 15:09:27 +0100 Subject: [PATCH 267/292] check phys install --- .github/workflows/test_our_image.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 86fa5c5ea..119f75efb 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -47,8 +47,9 @@ jobs: which python3 ls cat src/struphy/console/compile.py + cat pyproject.toml + pip install .[phys] pip install .[test,dev,docs] - pip install .[dev] - name: Compile Struphy run: | which python3 @@ -86,5 +87,6 @@ jobs: which python3 source /struphy_c_/env_c_/bin/activate which python3 + struphy compile --status struphy compile From 991f2b2968c423170ebe1c40fad187a6b3c09c76 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 15:22:23 +0100 Subject: [PATCH 268/292] try to get some info from console.main --- .github/workflows/test_our_image.yml | 1 + src/struphy/console/main.py | 8 ++++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 119f75efb..b5f86ccd7 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -88,5 +88,6 @@ jobs: source /struphy_c_/env_c_/bin/activate which python3 struphy compile --status + struphy compile --status struphy compile diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index f6e03ff16..5a311fc70 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -197,6 +197,10 @@ def struphy(): else: raise ValueError(f"Unknown command: {args.command}") + print(f"{module_path = }") + print(f"{func_name = }") + print(f"{args = }") + # transform parser Namespace object to dictionary and remove "command" key kwargs = vars(args) for key in [ @@ -221,8 +225,8 @@ def struphy(): kwargs.pop(key, None) # start sub-command function with all parameters of that function - # for k, v in kwargs.items(): - # print(k, v) + for k, v in kwargs.items(): + print(k, v) func(**kwargs) From cf537053bbb5d9399f32af5aaf1ed2bbf0d336f3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 15:42:56 +0100 Subject: [PATCH 269/292] pull current branch --- .github/workflows/test_our_image.yml | 33 +++++++++++++++------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index b5f86ccd7..596d528ef 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -42,19 +42,17 @@ jobs: pip install -U pip - name: Install Struphy run: | - which python3 source env/bin/activate which python3 ls - cat src/struphy/console/compile.py - cat pyproject.toml - pip install .[phys] - pip install .[test,dev,docs] + pip install ".[dev,doc]" + struphy -p - name: Compile Struphy run: | which python3 source env/bin/activate which python3 + struphy compile --status struphy compile with-struphy: @@ -69,25 +67,30 @@ jobs: run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - name: Checkout repo uses: actions/checkout@v4 - - name: Activate pre-installed virtual environment - run: | - ls / -a - which python3 - source /struphy_c_/env_c_/bin/activate - which python3 - pip install -U pip + - name: Git branch name + id: git-branch-name + uses: EthanSK/git-branch-name-action@v1 + - name: Echo the branch name + run: echo "Branch name ${GIT_BRANCH_NAME}" - name: Install Struphy run: | + ls / -a which python3 - source /struphy_c_/env_c_/bin/activate + cd /struphy_c_ + git status + git fetch origin + echo ${GIT_BRANCH_NAME} + git checkout ${GIT_BRANCH_NAME} + git pull + source env_c_/bin/activate which python3 - pip install ".[test,dev,docs]" + struphy -p + pip install ".[dev,doc]" - name: Compile Struphy run: | which python3 source /struphy_c_/env_c_/bin/activate which python3 struphy compile --status - struphy compile --status struphy compile From b4cf3091bbd3a77b69097059085d82c74e598bf1 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 15:50:01 +0100 Subject: [PATCH 270/292] add -e flag --- .github/workflows/test_our_image.yml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 596d528ef..37e2fa487 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -45,14 +45,13 @@ jobs: source env/bin/activate which python3 ls - pip install ".[dev,doc]" + pip install ".[phys,mpi,doc]" struphy -p - name: Compile Struphy run: | which python3 source env/bin/activate which python3 - struphy compile --status struphy compile with-struphy: @@ -85,7 +84,7 @@ jobs: source env_c_/bin/activate which python3 struphy -p - pip install ".[dev,doc]" + pip install -e ".[phys,mpi,doc]" - name: Compile Struphy run: | which python3 From f7a6251cb12ec3e12ffd370776ed6501cfd2efa5 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 15:57:51 +0100 Subject: [PATCH 271/292] use spossann username --- .github/workflows/test_our_image.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml index 37e2fa487..c5667f1c1 100644 --- a/.github/workflows/test_our_image.yml +++ b/.github/workflows/test_our_image.yml @@ -25,7 +25,7 @@ jobs: container: image: ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest credentials: - username: ${{ github.actor }} + username: spossann password: ${{ secrets.GHCR_TOKEN }} steps: @@ -46,7 +46,7 @@ jobs: which python3 ls pip install ".[phys,mpi,doc]" - struphy -p + # struphy -p - name: Compile Struphy run: | which python3 From d1cff054a538c18f38284f5a1ffea2a3395a6a66 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 17:13:35 +0100 Subject: [PATCH 272/292] new action for Struphy install in container --- .../install/struphy_in_container/action.yml | 0 .../workflows/test_our_image_w_actions.yml | 78 +++++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 .github/actions/install/struphy_in_container/action.yml create mode 100644 .github/workflows/test_our_image_w_actions.yml diff --git a/.github/actions/install/struphy_in_container/action.yml b/.github/actions/install/struphy_in_container/action.yml new file mode 100644 index 000000000..e69de29bb diff --git a/.github/workflows/test_our_image_w_actions.yml b/.github/workflows/test_our_image_w_actions.yml new file mode 100644 index 000000000..02c11352d --- /dev/null +++ b/.github/workflows/test_our_image_w_actions.yml @@ -0,0 +1,78 @@ +# name: Deploy docs to GitHub Pages + +on: + pull_request: + branches: + - main + - devel + workflow_dispatch: + +defaults: + run: + shell: bash + +permissions: + contents: read +# id-token: write + +# concurrency: +# group: "pages" +# cancel-in-progress: false + +jobs: + with-reqs: + runs-on: ubuntu-latest + container: + image: ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest + credentials: + username: spossann + password: ${{ secrets.GHCR_TOKEN }} + + steps: + - name: Check for dockerenv file + run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) + - name: Checkout repo + uses: actions/checkout@v4 + - name: Create virtual environment + run: | + which python3 + python3 -m venv env + source env/bin/activate + which python3 + pip install -U pip + - name: Install Struphy + run: | + source env/bin/activate + which python3 + ls + pip install ".[phys,mpi,doc]" + # struphy -p + - name: Compile Struphy + run: | + which python3 + source env/bin/activate + which python3 + struphy compile + + with-struphy: + runs-on: ubuntu-latest + container: + image: ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest + credentials: + username: spossann + password: ${{ secrets.GHCR_TOKEN }} + steps: + - name: Check for dockerenv file + run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) + - name: Checkout repo + uses: actions/checkout@v4 + - name: Install Struphy in Container + uses: ./.github/actions/install/struphy_in_container + - name: Compile Struphy + run: | + which python3 + source /struphy_c_/env_c_/bin/activate + which python3 + struphy compile --status + struphy compile + From e63eaa3810e8be739d6da21300eee4f02b328077 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:00:58 +0100 Subject: [PATCH 273/292] add content to action --- .../install/struphy_in_container/action.yml | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/.github/actions/install/struphy_in_container/action.yml b/.github/actions/install/struphy_in_container/action.yml index e69de29bb..6fff4ad38 100644 --- a/.github/actions/install/struphy_in_container/action.yml +++ b/.github/actions/install/struphy_in_container/action.yml @@ -0,0 +1,25 @@ +name: "Install Struphy in Container" + +runs: + using: composite + steps: + - name: Git branch name + id: git-branch-name + uses: EthanSK/git-branch-name-action@v1 + - name: Echo the branch name + run: echo "Branch name ${GIT_BRANCH_NAME}" + - name: Install struphy + shell: bash + run: | + ls / -a + which python3 + cd /struphy_c_ + git status + git fetch origin + echo ${GIT_BRANCH_NAME} + git checkout ${GIT_BRANCH_NAME} + git pull + source env_c_/bin/activate + which python3 + struphy -p + pip install -e ".[phys,mpi,doc]" \ No newline at end of file From 589cf47f504e2c8df9e68ac5b2edf036cfbb29fb Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:03:59 +0100 Subject: [PATCH 274/292] add shell --- .github/actions/install/struphy_in_container/action.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/actions/install/struphy_in_container/action.yml b/.github/actions/install/struphy_in_container/action.yml index 6fff4ad38..0bfaf30c7 100644 --- a/.github/actions/install/struphy_in_container/action.yml +++ b/.github/actions/install/struphy_in_container/action.yml @@ -7,6 +7,7 @@ runs: id: git-branch-name uses: EthanSK/git-branch-name-action@v1 - name: Echo the branch name + shell: bash run: echo "Branch name ${GIT_BRANCH_NAME}" - name: Install struphy shell: bash From 103ea66a54a430a35b6000ac8563f960cf529130 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:10:28 +0100 Subject: [PATCH 275/292] revert console.main --- src/struphy/console/main.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/struphy/console/main.py b/src/struphy/console/main.py index 5a311fc70..f6e03ff16 100644 --- a/src/struphy/console/main.py +++ b/src/struphy/console/main.py @@ -197,10 +197,6 @@ def struphy(): else: raise ValueError(f"Unknown command: {args.command}") - print(f"{module_path = }") - print(f"{func_name = }") - print(f"{args = }") - # transform parser Namespace object to dictionary and remove "command" key kwargs = vars(args) for key in [ @@ -225,8 +221,8 @@ def struphy(): kwargs.pop(key, None) # start sub-command function with all parameters of that function - for k, v in kwargs.items(): - print(k, v) + # for k, v in kwargs.items(): + # print(k, v) func(**kwargs) From 0fe8b81f7e8edc89be9df2daddc1bdd08360ebb1 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:11:24 +0100 Subject: [PATCH 276/292] revert console.copile --- src/struphy/console/compile.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/src/struphy/console/compile.py b/src/struphy/console/compile.py index 98847e433..432e4fa1f 100644 --- a/src/struphy/console/compile.py +++ b/src/struphy/console/compile.py @@ -78,8 +78,6 @@ def struphy_compile( # Read struphy state file state = utils.read_state() - - print(f"before {state = }") # collect kernels if "kernels" not in state: @@ -105,8 +103,6 @@ def struphy_compile( utils.save_state(state) # source files sources = " ".join(state["kernels"]) - - print(f"after {state = }") # actions if delete: @@ -226,9 +222,6 @@ def struphy_compile( flags += " --time-execution" # state - print(f"inside else {state = }") - print(f"{yes = }") - if state["last_used_language"] not in (language, None): if yes: yesno = "Y" @@ -237,8 +230,6 @@ def struphy_compile( f"Kernels compiled in language {state['last_used_language']} exist, will be deleted, continue (Y/n)?", ) - print(f"{yesno = }") - if yesno in ("", "Y", "y", "yes"): cmd = [ "struphy", @@ -255,8 +246,6 @@ def struphy_compile( state["last_used_omp_feec"] = flag_omp_feec utils.save_state(state) - - print(f"{state = }") # install psydac from wheel if not there source_install = False @@ -264,8 +253,6 @@ def struphy_compile( if "psydac" in req: source_install = True - print(f"{source_install = }") - struphy_ver = importlib.metadata.version("struphy") try: From 64dca567a4ce54836015dd02ee424bd8d0bf9c01 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:16:22 +0100 Subject: [PATCH 277/292] update other images --- docker/almalinux-latest.dockerfile | 14 ++++++-------- docker/fedora-latest.dockerfile | 14 ++++++-------- docker/opensuse-latest.dockerfile | 7 ++----- 3 files changed, 14 insertions(+), 21 deletions(-) diff --git a/docker/almalinux-latest.dockerfile b/docker/almalinux-latest.dockerfile index acb1824fd..8d84d5e45 100644 --- a/docker/almalinux-latest.dockerfile +++ b/docker/almalinux-latest.dockerfile @@ -1,12 +1,13 @@ -# Here is how to build the image and upload it to the mpcdf gitlab registry: +# Here is how to build the image and upload it to the Github package registry: # # We suppose you are in the struphy repo directory. -# Start the docker engine and run "docker login" with the following token: +# Start the docker engine and login to the Github package registry using a github personal acces token (classic): # -# TOKEN=gldt-CgMRBMtePbSwdWTxKw4Q; echo "$TOKEN" | docker login gitlab-registry.mpcdf.mpg.de -u gitlab+deploy-token-162 --password-stdin +# export CR_PAT=YOUR_TOKEN +# echo $CR_PAT | docker login ghcr.io -u USERNAME --password-stdin # docker info -# docker build -t gitlab-registry.mpcdf.mpg.de/struphy/struphy/almalinux-latest --provenance=false -f docker/almalinux-latest.dockerfile . -# docker push gitlab-registry.mpcdf.mpg.de/struphy/struphy/almalinux-latest +# docker build -t ghcr.io/struphy-hub/struphy/almalinux-with-reqs:latest --provenance=false -f docker/almalinux-latest.dockerfile . +# docker push ghcr.io/struphy-hub/struphy/almalinux-with-reqs:latest FROM almalinux:latest @@ -42,9 +43,6 @@ RUN echo "Installing additional tools..." \ && export CC=`which gcc` \ && export CXX=`which g++` -# create new working dir -WORKDIR /install_struphy_here/ - # allow mpirun as root ENV OMPI_ALLOW_RUN_AS_ROOT=1 ENV OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 diff --git a/docker/fedora-latest.dockerfile b/docker/fedora-latest.dockerfile index 9cf384454..79c3ed6a7 100644 --- a/docker/fedora-latest.dockerfile +++ b/docker/fedora-latest.dockerfile @@ -1,12 +1,13 @@ -# Here is how to build the image and upload it to the mpcdf gitlab registry: +# Here is how to build the image and upload it to the Github package registry: # # We suppose you are in the struphy repo directory. -# Start the docker engine and run "docker login" with the following token: +# Start the docker engine and login to the Github package registry using a github personal acces token (classic): # -# TOKEN=gldt-CgMRBMtePbSwdWTxKw4Q; echo "$TOKEN" | docker login gitlab-registry.mpcdf.mpg.de -u gitlab+deploy-token-162 --password-stdin +# export CR_PAT=YOUR_TOKEN +# echo $CR_PAT | docker login ghcr.io -u USERNAME --password-stdin # docker info -# docker build -t gitlab-registry.mpcdf.mpg.de/struphy/struphy/fedora-latest --provenance=false -f docker/fedora-latest.dockerfile . -# docker push gitlab-registry.mpcdf.mpg.de/struphy/struphy/fedora-latest +# docker build -t ghcr.io/struphy-hub/struphy/fedora-with-reqs:latest --provenance=false -f docker/fedora-latest.dockerfile . +# docker push ghcr.io/struphy-hub/struphy/fedora-with-reqs:latest FROM fedora:latest @@ -34,9 +35,6 @@ RUN echo "Installing additional tools..." \ && export CC=`which gcc` \ && export CXX=`which g++` - # create new working dir -WORKDIR /install_struphy_here/ - # allow mpirun as root ENV OMPI_ALLOW_RUN_AS_ROOT=1 ENV OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 diff --git a/docker/opensuse-latest.dockerfile b/docker/opensuse-latest.dockerfile index 97d5a5aa0..ef7fc47f4 100644 --- a/docker/opensuse-latest.dockerfile +++ b/docker/opensuse-latest.dockerfile @@ -6,8 +6,8 @@ # export CR_PAT=YOUR_TOKEN # echo $CR_PAT | docker login ghcr.io -u USERNAME --password-stdin # docker info -# docker build -t gitlab-registry.mpcdf.mpg.de/struphy/struphy/opensuse-latest --provenance=false -f docker/opensuse-latest.dockerfile . -# docker push gitlab-registry.mpcdf.mpg.de/struphy/struphy/opensuse-latest +# docker build -t ghcr.io/struphy-hub/struphy/opensuse-with-reqs:latest --provenance=false -f docker/opensuse-latest.dockerfile . +# docker push ghcr.io/struphy-hub/struphy/opensuse-with-reqs:latest FROM opensuse/tumbleweed:latest @@ -43,9 +43,6 @@ RUN echo "Installing additional tools..." \ && export CXX=`which g++` \ && zypper clean --all -# Create a new working directory -WORKDIR /install_struphy_here/ - # Allow mpirun to run as root (for OpenMPI) ENV OMPI_ALLOW_RUN_AS_ROOT=1 ENV OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 From 64c9fa85cad64b5c90eebdba5c8532357891b37a Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:17:03 +0100 Subject: [PATCH 278/292] remove mpcdf image (too large anyways) --- .../mpcdf-gcc-openmpi-with-struphy.dockerfile | 56 ------------------- 1 file changed, 56 deletions(-) delete mode 100644 docker/mpcdf-gcc-openmpi-with-struphy.dockerfile diff --git a/docker/mpcdf-gcc-openmpi-with-struphy.dockerfile b/docker/mpcdf-gcc-openmpi-with-struphy.dockerfile deleted file mode 100644 index 1b64254f9..000000000 --- a/docker/mpcdf-gcc-openmpi-with-struphy.dockerfile +++ /dev/null @@ -1,56 +0,0 @@ -# Here is how to build the image and upload it to the mpcdf gitlab registry: -# -# We suppose you are in the struphy repo directory. -# Start the docker engine and run "docker login" with the following token: -# -# TOKEN=gldt-CgMRBMtePbSwdWTxKw4Q; echo "$TOKEN" | docker login gitlab-registry.mpcdf.mpg.de -u gitlab+deploy-token-162 --password-stdin -# docker info -# docker build -t gitlab-registry.mpcdf.mpg.de/struphy/struphy/mpcdf-gcc-openmpi-with-struphy --provenance=false -f docker/mpcdf-gcc-openmpi-with-struphy.dockerfile . -# docker push gitlab-registry.mpcdf.mpg.de/struphy/struphy/mpcdf-gcc-openmpi-with-struphy - -FROM gitlab-registry.mpcdf.mpg.de/mpcdf/ci-module-image/gcc_14-openmpi_5_0:latest - -RUN source ./mpcdf/soft/SLE_15/packages/x86_64/Modules/5.4.0/etc/profile.d/modules.sh \ - && module load gcc/14 openmpi/5.0 python-waterboa/2024.06 git graphviz/8 \ - && module load cmake netcdf-serial mkl hdf5-serial \ - && export FC=`which gfortran` \ - && export CC=`which gcc` \ - && export CXX=`which g++` \ - && git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git struphy_c_ \ - && cd struphy_c_ \ - && python3 -m venv env_c_ \ - && source env_c_/bin/activate \ - && pip install -U pip \ - && pip install -e .[phys] --no-cache-dir --no-binary mpi4py \ - && struphy compile \ - && deactivate - -RUN source ./mpcdf/soft/SLE_15/packages/x86_64/Modules/5.4.0/etc/profile.d/modules.sh \ - && module load gcc/14 openmpi/5.0 python-waterboa/2024.06 git graphviz/8 \ - && module load cmake netcdf-serial mkl hdf5-serial \ - && export FC=`which gfortran` \ - && export CC=`which gcc` \ - && export CXX=`which g++` \ - && git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git struphy_fortran_\ - && cd struphy_fortran_ \ - && python3 -m venv env_fortran_ \ - && source env_fortran_/bin/activate \ - && pip install -U pip \ - && pip install -e .[phys] --no-cache-dir --no-binary mpi4py \ - && struphy compile --language fortran -y \ - && deactivate - -RUN source ./mpcdf/soft/SLE_15/packages/x86_64/Modules/5.4.0/etc/profile.d/modules.sh \ - && module load gcc/14 openmpi/5.0 python-waterboa/2024.06 git graphviz/8 \ - && module load cmake netcdf-serial mkl hdf5-serial \ - && export FC=`which gfortran` \ - && export CC=`which gcc` \ - && export CXX=`which g++` \ - && git clone https://gitlab.mpcdf.mpg.de/struphy/struphy.git struphy_fortran_--omp-pic\ - && cd struphy_fortran_--omp-pic \ - && python3 -m venv env_fortran_--omp-pic \ - && source env_fortran_--omp-pic/bin/activate \ - && pip install -U pip \ - && pip install -e .[phys] --no-cache-dir --no-binary mpi4py \ - && struphy compile --language fortran --omp-pic -y \ - && deactivate \ No newline at end of file From 5d5c670ac0be86409dc1f6a472002494c877b0bd Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:19:45 +0100 Subject: [PATCH 279/292] revert two workflows --- .github/workflows/static_analysis.yml | 8 ++++---- .github/workflows/ubuntu-latest.yml | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/static_analysis.yml b/.github/workflows/static_analysis.yml index 1f18a0d09..2fd4c2051 100644 --- a/.github/workflows/static_analysis.yml +++ b/.github/workflows/static_analysis.yml @@ -5,10 +5,10 @@ on: branches: - main - devel - # pull_request: - # branches: - # - main - # - devel + pull_request: + branches: + - main + - devel # concurrency: # group: ${{ github.ref }} diff --git a/.github/workflows/ubuntu-latest.yml b/.github/workflows/ubuntu-latest.yml index b2ba00399..e66aeb3b3 100644 --- a/.github/workflows/ubuntu-latest.yml +++ b/.github/workflows/ubuntu-latest.yml @@ -4,10 +4,10 @@ on: branches: - main - devel - # pull_request: - # branches: - # - main - # - devel + pull_request: + branches: + - main + - devel # concurrency: # group: ${{ github.ref }} From e7b54d7105e1a240852d0b2aef44df461d09482b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:22:07 +0100 Subject: [PATCH 280/292] remove .gitlab folder --- .gitlab/issue_templates/struphy_issue.md | 13 ------------- .../struphy_merge_request.md | 15 --------------- 2 files changed, 28 deletions(-) delete mode 100644 .gitlab/issue_templates/struphy_issue.md delete mode 100644 .gitlab/merge_request_templates/struphy_merge_request.md diff --git a/.gitlab/issue_templates/struphy_issue.md b/.gitlab/issue_templates/struphy_issue.md deleted file mode 100644 index b2c4eb7bc..000000000 --- a/.gitlab/issue_templates/struphy_issue.md +++ /dev/null @@ -1,13 +0,0 @@ -**Bug description / feature request:** - -... - -**Expected behavior:** - -... - -**Proposed solution:** - -... - - diff --git a/.gitlab/merge_request_templates/struphy_merge_request.md b/.gitlab/merge_request_templates/struphy_merge_request.md deleted file mode 100644 index 87f9ff3cc..000000000 --- a/.gitlab/merge_request_templates/struphy_merge_request.md +++ /dev/null @@ -1,15 +0,0 @@ -**Solves the following issue(s):** - -Closes #... - -**Core changes:** - -None - -**Model-specific changes:** - -None - -**Documentation changes:** - -None \ No newline at end of file From fe8c7f3047c0238a40961dd1c715c3ca80e94149 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:38:43 +0100 Subject: [PATCH 281/292] set the Ubuntu testing as a cronjob at 1 am every Sunday --- .github/workflows/ubuntu-latest.yml | 17 ++++------------- 1 file changed, 4 insertions(+), 13 deletions(-) diff --git a/.github/workflows/ubuntu-latest.yml b/.github/workflows/ubuntu-latest.yml index e66aeb3b3..95f724d26 100644 --- a/.github/workflows/ubuntu-latest.yml +++ b/.github/workflows/ubuntu-latest.yml @@ -1,17 +1,8 @@ -name: Ubuntu +name: Ubuntu latest - cronjob on: - push: - branches: - - main - - devel - pull_request: - branches: - - main - - devel - -# concurrency: -# group: ${{ github.ref }} -# cancel-in-progress: true + schedule: + # run at 1 a.m. on Sunday + - cron: "0 1 * * 0" jobs: ubuntu-latest-build: From 9bda18da48095e7f5e725085c7f9db81e6320600 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:41:21 +0100 Subject: [PATCH 282/292] rename testing.yml -> reusable-testing.yml --- .github/workflows/{testing.yml => reusable-testing.yml} | 0 .github/workflows/ubuntu-latest.yml | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename .github/workflows/{testing.yml => reusable-testing.yml} (100%) diff --git a/.github/workflows/testing.yml b/.github/workflows/reusable-testing.yml similarity index 100% rename from .github/workflows/testing.yml rename to .github/workflows/reusable-testing.yml diff --git a/.github/workflows/ubuntu-latest.yml b/.github/workflows/ubuntu-latest.yml index 95f724d26..8fd73272e 100644 --- a/.github/workflows/ubuntu-latest.yml +++ b/.github/workflows/ubuntu-latest.yml @@ -6,6 +6,6 @@ on: jobs: ubuntu-latest-build: - uses: ./.github/workflows/testing.yml + uses: ./.github/workflows/reusable-testing.yml with: os: ubuntu-latest \ No newline at end of file From 074ff5a788bc2e10095d88d0dbd11ceca1e59661 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:51:10 +0100 Subject: [PATCH 283/292] new actions for PR unit and model tests --- .github/workflows/test-PR-unit.yml | 49 +++++++++++++++++++ .../workflows/test_our_image_w_actions.yml | 45 +++-------------- 2 files changed, 57 insertions(+), 37 deletions(-) create mode 100644 .github/workflows/test-PR-unit.yml diff --git a/.github/workflows/test-PR-unit.yml b/.github/workflows/test-PR-unit.yml new file mode 100644 index 000000000..cef5d13b1 --- /dev/null +++ b/.github/workflows/test-PR-unit.yml @@ -0,0 +1,49 @@ +name: PR - unit tests in Container + +on: + pull_request: + branches: + - main + - devel + workflow_dispatch: + +defaults: + run: + shell: bash + +permissions: + contents: read + +# concurrency: +# group: "pages" +# cancel-in-progress: false + +jobs: + container-with-struphy: + runs-on: ubuntu-latest + container: + image: ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest + credentials: + username: spossann + password: ${{ secrets.GHCR_TOKEN }} + steps: + + - name: Check for dockerenv file + run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) + + - name: Checkout repo + uses: actions/checkout@v4 + + - name: Install Struphy in Container + uses: ./.github/actions/install/struphy_in_container + + - name: Compile Struphy + run: | + which python3 + source /struphy_c_/env_c_/bin/activate + which python3 + struphy compile --status + struphy compile + + - name: Run unit tests + uses: ./.github/actions/tests/unit diff --git a/.github/workflows/test_our_image_w_actions.yml b/.github/workflows/test_our_image_w_actions.yml index 02c11352d..e9ddc592b 100644 --- a/.github/workflows/test_our_image_w_actions.yml +++ b/.github/workflows/test_our_image_w_actions.yml @@ -1,4 +1,4 @@ -# name: Deploy docs to GitHub Pages +name: PR - model tests in Container on: pull_request: @@ -13,48 +13,13 @@ defaults: permissions: contents: read -# id-token: write # concurrency: # group: "pages" # cancel-in-progress: false jobs: - with-reqs: - runs-on: ubuntu-latest - container: - image: ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest - credentials: - username: spossann - password: ${{ secrets.GHCR_TOKEN }} - - steps: - - name: Check for dockerenv file - run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - - name: Checkout repo - uses: actions/checkout@v4 - - name: Create virtual environment - run: | - which python3 - python3 -m venv env - source env/bin/activate - which python3 - pip install -U pip - - name: Install Struphy - run: | - source env/bin/activate - which python3 - ls - pip install ".[phys,mpi,doc]" - # struphy -p - - name: Compile Struphy - run: | - which python3 - source env/bin/activate - which python3 - struphy compile - - with-struphy: + container-with-struphy: runs-on: ubuntu-latest container: image: ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest @@ -62,12 +27,16 @@ jobs: username: spossann password: ${{ secrets.GHCR_TOKEN }} steps: + - name: Check for dockerenv file run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) + - name: Checkout repo uses: actions/checkout@v4 + - name: Install Struphy in Container uses: ./.github/actions/install/struphy_in_container + - name: Compile Struphy run: | which python3 @@ -76,3 +45,5 @@ jobs: struphy compile --status struphy compile + - name: Run model tests + uses: ./.github/actions/tests/models From 1ff17bf6c9024cd9fd449de4629ef40953cca8c3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 20:53:45 +0100 Subject: [PATCH 284/292] rename workflow for model testing --- ...image_w_actions.yml => test-PR-models.yml} | 0 .github/workflows/test_our_image.yml | 95 ------------------- 2 files changed, 95 deletions(-) rename .github/workflows/{test_our_image_w_actions.yml => test-PR-models.yml} (100%) delete mode 100644 .github/workflows/test_our_image.yml diff --git a/.github/workflows/test_our_image_w_actions.yml b/.github/workflows/test-PR-models.yml similarity index 100% rename from .github/workflows/test_our_image_w_actions.yml rename to .github/workflows/test-PR-models.yml diff --git a/.github/workflows/test_our_image.yml b/.github/workflows/test_our_image.yml deleted file mode 100644 index c5667f1c1..000000000 --- a/.github/workflows/test_our_image.yml +++ /dev/null @@ -1,95 +0,0 @@ -# name: Deploy docs to GitHub Pages - -on: - pull_request: - branches: - - main - - devel - workflow_dispatch: - -defaults: - run: - shell: bash - -permissions: - contents: read -# id-token: write - -# concurrency: -# group: "pages" -# cancel-in-progress: false - -jobs: - with-reqs: - runs-on: ubuntu-latest - container: - image: ghcr.io/struphy-hub/struphy/ubuntu-with-reqs:latest - credentials: - username: spossann - password: ${{ secrets.GHCR_TOKEN }} - - steps: - - name: Check for dockerenv file - run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - - name: Checkout repo - uses: actions/checkout@v4 - - name: Create virtual environment - run: | - which python3 - python3 -m venv env - source env/bin/activate - which python3 - pip install -U pip - - name: Install Struphy - run: | - source env/bin/activate - which python3 - ls - pip install ".[phys,mpi,doc]" - # struphy -p - - name: Compile Struphy - run: | - which python3 - source env/bin/activate - which python3 - struphy compile - - with-struphy: - runs-on: ubuntu-latest - container: - image: ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest - credentials: - username: ${{ github.actor }} - password: ${{ secrets.GHCR_TOKEN }} - steps: - - name: Check for dockerenv file - run: (ls /.dockerenv && echo Found dockerenv) || (echo No dockerenv) - - name: Checkout repo - uses: actions/checkout@v4 - - name: Git branch name - id: git-branch-name - uses: EthanSK/git-branch-name-action@v1 - - name: Echo the branch name - run: echo "Branch name ${GIT_BRANCH_NAME}" - - name: Install Struphy - run: | - ls / -a - which python3 - cd /struphy_c_ - git status - git fetch origin - echo ${GIT_BRANCH_NAME} - git checkout ${GIT_BRANCH_NAME} - git pull - source env_c_/bin/activate - which python3 - struphy -p - pip install -e ".[phys,mpi,doc]" - - name: Compile Struphy - run: | - which python3 - source /struphy_c_/env_c_/bin/activate - which python3 - struphy compile --status - struphy compile - From 52ff7402fc439c0d1d77ef8f4d91ad07f85b3444 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 28 Oct 2025 21:00:01 +0100 Subject: [PATCH 285/292] do not use actions in PR testing workflows; we must source a specific virtual env --- .github/workflows/test-PR-models.yml | 27 +++++++++++++++++++++++++-- .github/workflows/test-PR-unit.yml | 22 +++++++++++++++++++++- 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test-PR-models.yml b/.github/workflows/test-PR-models.yml index e9ddc592b..e9f9f03f7 100644 --- a/.github/workflows/test-PR-models.yml +++ b/.github/workflows/test-PR-models.yml @@ -45,5 +45,28 @@ jobs: struphy compile --status struphy compile - - name: Run model tests - uses: ./.github/actions/tests/models + - name: Model tests + shell: bash + run: | + which python3 + source /struphy_c_/env_c_/bin/activate + which python3 + struphy compile --status + struphy test LinearMHD + struphy test toy + struphy test models + struphy test verification + + - name: Model tests with MPI + shell: bash + run: | + which python3 + source /struphy_c_/env_c_/bin/activate + which python3 + struphy compile --status + struphy test models + struphy test models --mpi 2 + struphy test verification --mpi 1 + struphy test verification --mpi 4 + struphy test verification --mpi 4 --nclones 2 + struphy test VlasovAmpereOneSpecies --mpi 2 --nclones 2 diff --git a/.github/workflows/test-PR-unit.yml b/.github/workflows/test-PR-unit.yml index cef5d13b1..d79326322 100644 --- a/.github/workflows/test-PR-unit.yml +++ b/.github/workflows/test-PR-unit.yml @@ -45,5 +45,25 @@ jobs: struphy compile --status struphy compile + - name: Run unit tests with MPI + shell: bash + run: | + which python3 + source /struphy_c_/env_c_/bin/activate + which python3 + struphy compile --status + struphy --refresh-models + struphy test unit --mpi 2 + - name: Run unit tests - uses: ./.github/actions/tests/unit + shell: bash + run: | + which python3 + source /struphy_c_/env_c_/bin/activate + which python3 + struphy compile --status + struphy --refresh-models + pip show mpi4py + pip uninstall -y mpi4py + pip list + struphy test unit From 388619359fc8c91020ff79a237c1ee71e75a85bb Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 29 Oct 2025 08:03:44 +0100 Subject: [PATCH 286/292] add gvec env vars --- .github/workflows/test-PR-unit.yml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/.github/workflows/test-PR-unit.yml b/.github/workflows/test-PR-unit.yml index d79326322..cf3c26a15 100644 --- a/.github/workflows/test-PR-unit.yml +++ b/.github/workflows/test-PR-unit.yml @@ -51,6 +51,15 @@ jobs: which python3 source /struphy_c_/env_c_/bin/activate which python3 + echo $FC + echo $CC + echo $CXX + export FC=`which gfortran` + export CC=`which gcc` + export CXX=`which g++` + echo $FC + echo $CC + echo $CXX struphy compile --status struphy --refresh-models struphy test unit --mpi 2 @@ -61,6 +70,15 @@ jobs: which python3 source /struphy_c_/env_c_/bin/activate which python3 + echo $FC + echo $CC + echo $CXX + export FC=`which gfortran` + export CC=`which gcc` + export CXX=`which g++` + echo $FC + echo $CC + echo $CXX struphy compile --status struphy --refresh-models pip show mpi4py From ae4c71a24851347fdac657b134948942493b1167 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 29 Oct 2025 08:08:57 +0100 Subject: [PATCH 287/292] try serial unit tests first --- .github/workflows/test-PR-models.yml | 2 +- .github/workflows/test-PR-unit.yml | 19 +++++++++++-------- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/.github/workflows/test-PR-models.yml b/.github/workflows/test-PR-models.yml index e9f9f03f7..9dfeb3ebc 100644 --- a/.github/workflows/test-PR-models.yml +++ b/.github/workflows/test-PR-models.yml @@ -19,7 +19,7 @@ permissions: # cancel-in-progress: false jobs: - container-with-struphy: + model-tests-in-container-with-struphy: runs-on: ubuntu-latest container: image: ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest diff --git a/.github/workflows/test-PR-unit.yml b/.github/workflows/test-PR-unit.yml index cf3c26a15..82833a16d 100644 --- a/.github/workflows/test-PR-unit.yml +++ b/.github/workflows/test-PR-unit.yml @@ -19,7 +19,7 @@ permissions: # cancel-in-progress: false jobs: - container-with-struphy: + unit-tests-in-container-with-struphy: runs-on: ubuntu-latest container: image: ghcr.io/struphy-hub/struphy/ubuntu-with-struphy:latest @@ -45,7 +45,7 @@ jobs: struphy compile --status struphy compile - - name: Run unit tests with MPI + - name: Run unit tests shell: bash run: | which python3 @@ -62,9 +62,12 @@ jobs: echo $CXX struphy compile --status struphy --refresh-models - struphy test unit --mpi 2 + pip show mpi4py + pip uninstall -y mpi4py + pip list + struphy test unit - - name: Run unit tests + - name: Run unit tests with MPI shell: bash run: | which python3 @@ -79,9 +82,9 @@ jobs: echo $FC echo $CC echo $CXX + pip install -U mpi4py struphy compile --status struphy --refresh-models - pip show mpi4py - pip uninstall -y mpi4py - pip list - struphy test unit + struphy test unit --mpi 2 + + From a1ebe58b20491e7462b13b695e9fb1365a4a1737 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 29 Oct 2025 08:19:10 +0100 Subject: [PATCH 288/292] run test_l2_projectors in separate step --- .github/workflows/test-PR-unit.yml | 32 +++++++++++++++----- src/struphy/feec/tests/test_l2_projectors.py | 2 +- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/.github/workflows/test-PR-unit.yml b/.github/workflows/test-PR-unit.yml index 82833a16d..2af61e45b 100644 --- a/.github/workflows/test-PR-unit.yml +++ b/.github/workflows/test-PR-unit.yml @@ -45,7 +45,7 @@ jobs: struphy compile --status struphy compile - - name: Run unit tests + - name: run test_l2_projector shell: bash run: | which python3 @@ -62,10 +62,9 @@ jobs: echo $CXX struphy compile --status struphy --refresh-models - pip show mpi4py - pip uninstall -y mpi4py - pip list - struphy test unit + cd /struphy_c_ + ls + python3 src/struphy/feec/tests/test_l2_projectors.py - name: Run unit tests with MPI shell: bash @@ -82,9 +81,28 @@ jobs: echo $FC echo $CC echo $CXX - pip install -U mpi4py struphy compile --status struphy --refresh-models struphy test unit --mpi 2 - + - name: Run unit tests + shell: bash + run: | + which python3 + source /struphy_c_/env_c_/bin/activate + which python3 + echo $FC + echo $CC + echo $CXX + export FC=`which gfortran` + export CC=`which gcc` + export CXX=`which g++` + echo $FC + echo $CC + echo $CXX + struphy compile --status + struphy --refresh-models + pip show mpi4py + pip uninstall -y mpi4py + pip list + struphy test unit diff --git a/src/struphy/feec/tests/test_l2_projectors.py b/src/struphy/feec/tests/test_l2_projectors.py index 7da42eff4..89caa8941 100644 --- a/src/struphy/feec/tests/test_l2_projectors.py +++ b/src/struphy/feec/tests/test_l2_projectors.py @@ -260,5 +260,5 @@ def f(x, y, z): spl_kind = [False, True, True] array_input = True test_l2_projectors_mappings(Nel, p, spl_kind, array_input, do_plot=False, with_desc=False) - # test_l2_projectors_convergence(0, 1, True, do_plot=True) + test_l2_projectors_convergence(0, 1, True, do_plot=False) # test_l2_projectors_convergence(1, 1, False, do_plot=True) From e00bbc58bc47ad4cccadad76fdf2479cf1a63931 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 29 Oct 2025 08:45:50 +0100 Subject: [PATCH 289/292] exclude gvec from l2 projector test --- src/struphy/feec/tests/test_l2_projectors.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/tests/test_l2_projectors.py b/src/struphy/feec/tests/test_l2_projectors.py index 89caa8941..2e9f611eb 100644 --- a/src/struphy/feec/tests/test_l2_projectors.py +++ b/src/struphy/feec/tests/test_l2_projectors.py @@ -15,7 +15,7 @@ @pytest.mark.parametrize("p", [[2, 1, 1], [3, 2, 1]]) @pytest.mark.parametrize("spl_kind", [[False, True, True]]) @pytest.mark.parametrize("array_input", [False, True]) -def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_desc, do_plot=False): +def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_gvec=False, with_desc=False, do_plot=False): """Tests the L2-projectors for all available mappings. Both callable and array inputs to the projectors are tested. @@ -50,6 +50,10 @@ def test_l2_projectors_mappings(Nel, p, spl_kind, array_input, with_desc, do_plo print(f"Testing {dom_class =}") print("#" * 80) + if "GVEC" in dom_type and not with_gvec: + print(f"Attention: {with_gvec =}, GVEC not tested here !!") + continue + if "DESC" in dom_type and not with_desc: print(f"Attention: {with_desc =}, DESC not tested here !!") continue From a2071de161288f765e2f60e9e2523f40c3eea43c Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 29 Oct 2025 14:04:20 +0100 Subject: [PATCH 290/292] revert deletion of gitlb templates --- .gitlab/issue_templates/struphy_issue.md | 13 +++++++++++++ .../struphy_merge_request.md | 15 +++++++++++++++ 2 files changed, 28 insertions(+) create mode 100644 .gitlab/issue_templates/struphy_issue.md create mode 100644 .gitlab/merge_request_templates/struphy_merge_request.md diff --git a/.gitlab/issue_templates/struphy_issue.md b/.gitlab/issue_templates/struphy_issue.md new file mode 100644 index 000000000..b2c4eb7bc --- /dev/null +++ b/.gitlab/issue_templates/struphy_issue.md @@ -0,0 +1,13 @@ +**Bug description / feature request:** + +... + +**Expected behavior:** + +... + +**Proposed solution:** + +... + + diff --git a/.gitlab/merge_request_templates/struphy_merge_request.md b/.gitlab/merge_request_templates/struphy_merge_request.md new file mode 100644 index 000000000..87f9ff3cc --- /dev/null +++ b/.gitlab/merge_request_templates/struphy_merge_request.md @@ -0,0 +1,15 @@ +**Solves the following issue(s):** + +Closes #... + +**Core changes:** + +None + +**Model-specific changes:** + +None + +**Documentation changes:** + +None \ No newline at end of file From c7165bede18a91609a44d0c2122b83c7ab8eee59 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 29 Oct 2025 14:07:30 +0100 Subject: [PATCH 291/292] remove separate l2-projector test --- .github/workflows/test-PR-unit.yml | 40 ------------------------------ 1 file changed, 40 deletions(-) diff --git a/.github/workflows/test-PR-unit.yml b/.github/workflows/test-PR-unit.yml index 2af61e45b..1f6df8cf4 100644 --- a/.github/workflows/test-PR-unit.yml +++ b/.github/workflows/test-PR-unit.yml @@ -45,42 +45,12 @@ jobs: struphy compile --status struphy compile - - name: run test_l2_projector - shell: bash - run: | - which python3 - source /struphy_c_/env_c_/bin/activate - which python3 - echo $FC - echo $CC - echo $CXX - export FC=`which gfortran` - export CC=`which gcc` - export CXX=`which g++` - echo $FC - echo $CC - echo $CXX - struphy compile --status - struphy --refresh-models - cd /struphy_c_ - ls - python3 src/struphy/feec/tests/test_l2_projectors.py - - name: Run unit tests with MPI shell: bash run: | which python3 source /struphy_c_/env_c_/bin/activate which python3 - echo $FC - echo $CC - echo $CXX - export FC=`which gfortran` - export CC=`which gcc` - export CXX=`which g++` - echo $FC - echo $CC - echo $CXX struphy compile --status struphy --refresh-models struphy test unit --mpi 2 @@ -91,18 +61,8 @@ jobs: which python3 source /struphy_c_/env_c_/bin/activate which python3 - echo $FC - echo $CC - echo $CXX - export FC=`which gfortran` - export CC=`which gcc` - export CXX=`which g++` - echo $FC - echo $CC - echo $CXX struphy compile --status struphy --refresh-models pip show mpi4py pip uninstall -y mpi4py - pip list struphy test unit From c7806185ae4ae86562987573a2bd2054e76e6284 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Wed, 29 Oct 2025 14:24:15 +0100 Subject: [PATCH 292/292] do not run verification model test in unit tests --- src/struphy/console/test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index ebcff34d4..ab64707e3 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -42,14 +42,14 @@ def struphy_test( str(mpi), "pytest", "-k", - "not _models and not _tutorial and not pproc", + "not _models and not _tutorial and not pproc and not _verif_", "--with-mpi", ] else: cmd = [ "pytest", "-k", - "not _models and not _tutorial and not pproc", + "not _models and not _tutorial and not pproc and not _verif_", ] if with_desc: