Implement multispecies support in vlasov1d

adept/_vlasov1d/datamodel.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class SpaceProfileModel(BaseModel):
@@ -23,8 +23,12 @@ class SpeciesBackgroundModel(BaseModel):
 
 
 class DensityModel(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
     quasineutrality: bool
-    species_background: SpeciesBackgroundModel
+    # Note: Allow arbitrary species-* fields to be defined dynamically
+    # e.g., species_background, species_beam, etc.
+    # These are validated as SpeciesBackgroundModel when accessed
 
 
 class UnitsModel(BaseModel):
@@ -37,11 +41,11 @@ class UnitsModel(BaseModel):
 
 class GridModel(BaseModel):
     dt: float
-    nv: int
+    nv: int | None = None  # Optional: for backward compatibility with single-species config files
     nx: int
     tmin: float
     tmax: float
-    vmax: float
+    vmax: float | None = None  # Optional: for backward compatibility with single-species config files
     xmax: float
     xmin: float
 
@@ -53,8 +57,11 @@ class TimeSaveModel(BaseModel):
 
 
 class SaveModel(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
     fields: dict[str, TimeSaveModel]
-    electron: dict[str, TimeSaveModel]
+    # Note: Allow arbitrary species save sections (electron, ion, etc.)
+    # These are validated as dict[str, TimeSaveModel] when accessed
 
 
 class ExDriverModel(BaseModel):
@@ -112,10 +119,35 @@ class KrookModel(BaseModel):
     space: SpaceTermModel
 
 
+class SpeciesConfig(BaseModel):
+    """Configuration for a physical species in multi-species simulations.
+
+    Each species has its own charge-to-mass ratio and velocity grid, allowing
+    for electron-ion physics and other multi-species interactions.
+
+    Attributes:
+        name: Species name (e.g., 'electron', 'ion')
+        charge: Charge in units of fundamental charge (e.g., -1.0 for electrons, 10.0 for Z=10 ions)
+        mass: Mass in units of electron mass (e.g., 1.0 for electrons, 1836.0 for protons)
+        vmax: Velocity grid maximum for this species
+        nv: Number of velocity grid points for this species
+        density_components: List of density component names from the 'density:' section
+                          that contribute to this species' distribution function
+    """
+
+    name: str
+    charge: float
+    mass: float
+    vmax: float
+    nv: int
+    density_components: list[str]
+
+
 class TermsModel(BaseModel):
     field: str
     edfdv: str
     time: str
+    species: list[SpeciesConfig] | None = None
     fokker_planck: FokkerPlanckModel
     krook: KrookModel
 

adept/_vlasov1d/helpers.py

@@ -1,6 +1,7 @@
 #  Copyright (c) Ergodic LLC 2023
 #  research@ergodic.io
 import os
+import warnings
 from time import time
 
 import numpy as np
@@ -92,101 +93,147 @@ def _initialize_distribution_(
 
 
 def _initialize_total_distribution_(cfg, cfg_grid):
-    params = cfg["density"]
-    n_prof_total = np.zeros([cfg_grid["nx"]])
-    f = np.zeros([cfg_grid["nx"], cfg_grid["nv"]])
-    species_found = False
-    for name, species_params in cfg["density"].items():
-        if name.startswith("species-"):
-            v0 = species_params["v0"]
-            T0 = species_params["T0"]
-            m = species_params["m"]
-            if name in params:
-                if "v0" in params[name]:
-                    v0 = params[name]["v0"]
+    """
+    Initialize distribution functions for all species.
 
-                if "T0" in params[name]:
-                    T0 = params[name]["T0"]
+    The species config is normalized in modules.py:get_derived_quantities() so that
+    a species config always exists (for backward compatibility with single-species
+    config files, a default electron species is generated).
 
-                if "m" in params[name]:
-                    m = params[name]["m"]
+    Returns:
+        dict mapping species_name -> (n_prof, f_s, v_ax)
+    """
+    species_configs = cfg["terms"]["species"]
+    species_distributions = {}
+    species_found = False
 
-            if species_params["basis"] == "uniform":
-                nprof = np.ones_like(n_prof_total)
+    for species_cfg in species_configs:
+        species_name = species_cfg["name"]
+        density_components = species_cfg["density_components"]
+        vmax = species_cfg["vmax"]
+        nv = species_cfg["nv"]
 
-            elif species_params["basis"] == "linear":
-                left = species_params["center"] - species_params["width"] * 0.5
-                right = species_params["center"] + species_params["width"] * 0.5
-                rise = species_params["rise"]
-                mask = get_envelope(rise, rise, left, right, cfg_grid["x"])
+        # Initialize arrays for this species
+        n_prof_species = np.zeros([cfg_grid["nx"]])
+        dv = 2.0 * vmax / nv
+        vax = np.linspace(-vmax + dv / 2.0, vmax - dv / 2.0, nv)
+        f_species = np.zeros([cfg_grid["nx"], nv])
 
-                ureg = pint.UnitRegistry()
-                _Q = ureg.Quantity
+        # Sum contributions from all density components
+        for component_name in density_components:
+            if component_name not in cfg["density"]:
+                raise ValueError(f"Density component '{component_name}' not found in config")
 
-                L = (
-                    _Q(species_params["gradient scale length"]).to("nm").magnitude
-                    / cfg["units"]["derived"]["x0"].to("nm").magnitude
-                )
-                nprof = species_params["val at center"] + (cfg_grid["x"] - species_params["center"]) / L
-                nprof = mask * nprof
-            elif species_params["basis"] == "exponential":
-                left = species_params["center"] - species_params["width"] * 0.5
-                right = species_params["center"] + species_params["width"] * 0.5
-                rise = species_params["rise"]
-                mask = get_envelope(rise, rise, left, right, cfg_grid["x"])
-
-                ureg = pint.UnitRegistry()
-                _Q = ureg.Quantity
-
-                L = (
-                    _Q(species_params["gradient scale length"]).to("nm").magnitude
-                    / cfg["units"]["derived"]["x0"].to("nm").magnitude
+            species_params = cfg["density"][component_name]
+            v0 = species_params["v0"]
+            T0 = species_params["T0"]
+
+            # Handle mass parameter with deprecation
+            if "m" in species_params:
+                warnings.warn(
+                    f"Density component '{component_name}': The 'm' parameter in density config is deprecated. "
+                    f"Please use the mass from the species config instead.",
+                    DeprecationWarning,
+                    stacklevel=2,
                 )
-                nprof = species_params["val at center"] * np.exp((cfg_grid["x"] - species_params["center"]) / L)
-                nprof = mask * nprof
-
-            elif species_params["basis"] == "tanh":
-                left = species_params["center"] - species_params["width"] * 0.5
-                right = species_params["center"] + species_params["width"] * 0.5
-                rise = species_params["rise"]
-                nprof = get_envelope(rise, rise, left, right, cfg_grid["x"])
-
-                if species_params["bump_or_trough"] == "trough":
-                    nprof = 1 - nprof
-                nprof = species_params["baseline"] + species_params["bump_height"] * nprof
-
-            elif species_params["basis"] == "sine":
-                baseline = species_params["baseline"]
-                amp = species_params["amplitude"]
-                kk = species_params["wavenumber"]
-                nprof = baseline * (1.0 + amp * jnp.sin(kk * cfg_grid["x"]))
+                m = species_params["m"]
+                # Check if it matches the species config mass
+                if abs(m - species_cfg["mass"]) > 1e-10:
+                    warnings.warn(
+                        f"Density component '{component_name}': Mass mismatch! "
+                        f"Using m={m} from density config, but species '{species_name}' "
+                        f"has mass={species_cfg['mass']}. This may lead to inconsistent physics.",
+                        UserWarning,
+                        stacklevel=2,
+                    )
             else:
-                raise NotImplementedError
+                # Use mass from species config
+                m = species_cfg["mass"]
 
-            n_prof_total += nprof
+            # Get density profile
+            nprof = _get_density_profile(species_params, cfg, cfg_grid)
+            n_prof_species += nprof
 
-            # Distribution function
+            # Distribution function for this component
             temp_f, _ = _initialize_distribution_(
                 nx=int(cfg_grid["nx"]),
-                nv=int(cfg_grid["nv"]),
+                nv=nv,
                 v0=v0,
                 m=m,
                 T0=T0,
-                vmax=cfg_grid["vmax"],
+                vmax=vmax,
                 n_prof=nprof,
                 noise_val=species_params["noise_val"],
                 noise_seed=int(species_params["noise_seed"]),
                 noise_type=species_params["noise_type"],
             )
-            f += temp_f
+            f_species += temp_f
             species_found = True
-        else:
-            pass
+
+        species_distributions[species_name] = (n_prof_species, f_species, vax)
 
     if not species_found:
         raise ValueError("No species found! Check the config")
 
-    return n_prof_total, f
+    return species_distributions
+
+
+def _get_density_profile(species_params, cfg, cfg_grid):
+    """Extract density profile generation logic into a helper function."""
+    if species_params["basis"] == "uniform":
+        nprof = np.ones([cfg_grid["nx"]])
+
+    elif species_params["basis"] == "linear":
+        left = species_params["center"] - species_params["width"] * 0.5
+        right = species_params["center"] + species_params["width"] * 0.5
+        rise = species_params["rise"]
+        mask = get_envelope(rise, rise, left, right, cfg_grid["x"])
+
+        ureg = pint.UnitRegistry()
+        _Q = ureg.Quantity
+
+        L = (
+            _Q(species_params["gradient scale length"]).to("nm").magnitude
+            / cfg["units"]["derived"]["x0"].to("nm").magnitude
+        )
+        nprof = species_params["val at center"] + (cfg_grid["x"] - species_params["center"]) / L
+        nprof = mask * nprof
+
+    elif species_params["basis"] == "exponential":
+        left = species_params["center"] - species_params["width"] * 0.5
+        right = species_params["center"] + species_params["width"] * 0.5
+        rise = species_params["rise"]
+        mask = get_envelope(rise, rise, left, right, cfg_grid["x"])
+
+        ureg = pint.UnitRegistry()
+        _Q = ureg.Quantity
+
+        L = (
+            _Q(species_params["gradient scale length"]).to("nm").magnitude
+            / cfg["units"]["derived"]["x0"].to("nm").magnitude
+        )
+        nprof = species_params["val at center"] * np.exp((cfg_grid["x"] - species_params["center"]) / L)
+        nprof = mask * nprof
+
+    elif species_params["basis"] == "tanh":
+        left = species_params["center"] - species_params["width"] * 0.5
+        right = species_params["center"] + species_params["width"] * 0.5
+        rise = species_params["rise"]
+        nprof = get_envelope(rise, rise, left, right, cfg_grid["x"])
+
+        if species_params["bump_or_trough"] == "trough":
+            nprof = 1 - nprof
+        nprof = species_params["baseline"] + species_params["bump_height"] * nprof
+
+    elif species_params["basis"] == "sine":
+        baseline = species_params["baseline"]
+        amp = species_params["amplitude"]
+        kk = species_params["wavenumber"]
+        nprof = baseline * (1.0 + amp * jnp.sin(kk * cfg_grid["x"]))
+    else:
+        raise NotImplementedError
+
+    return nprof
 
 
 def post_process(result: Solution, cfg: dict, td: str, args: dict):