test for EH. doesnt work with ee yet though. hmm

adept/vfp1d/fokker_planck.py

@@ -104,7 +104,7 @@ def calc_ros_j(self, flm, power):
             * self.dv
         )
 
-    def get_offdiagonal_contrib(self, t, y, args):
+    def get_ee_offdiagonal_contrib(self, t, y, args):
         ddv = args["ddvf0"]
         d2dv2 = args["d2dv2f0"]
         il = args["il"]
@@ -115,9 +115,9 @@ def get_offdiagonal_contrib(self, t, y, args):
         term3 = (self.a2[il] * d2dv2 + self.b3[il] * ddv) * self.calc_ros_i(flm, power=il)
         term4 = (self.a2[il] * d2dv2 + self.b4[il] * ddv) * self.calc_ros_j(flm, power=1.0 - il)
 
-        return self.nuee_coeff * (term1 + term2 + term3 + term4)
+        return term1 + term2 + term3 + term4
 
-    def get_ee_contrib(self, f0):
+    def get_ee_diagonal_contrib(self, f0):
         i0 = 4 * jnp.pi * jnp.cumsum(self.v[None, :] ** 2.0 * f0, axis=1) * self.dv
         jm1 = 4 * jnp.pi / self.v[None, :] * jnp.cumsum((self.v[None, :] * f0)[:, ::-1], axis=1)[:, ::-1] * self.dv
         i2 = 4 * jnp.pi * self.v[None, :] ** 2.0 * jnp.cumsum(self.v[None, :] ** 4.0 * f0, axis=1) * self.dv
@@ -191,7 +191,7 @@ def __implicit_ee__(this_flm):
         return lx.linear_solve(op, f10, solver=lx.SVD()).value
 
     def __call__(self, Z, ni, f0, f10, dt):
-        ee_diag, ee_lower, ee_upper = self.get_ee_contrib(f0)
+        ee_diag, ee_lower, ee_upper = self.get_ee_diagonal_contrib(f0)
 
         for il in range(1, self.nl + 1):
             ei_diag = -il * (il + 1) / 2.0 * (Z[:, None] ** 2.0) * ni[:, None] / self.v[None, :] ** 3.0
@@ -209,12 +209,12 @@ def __call__(self, Z, ni, f0, f10, dt):
 
                 new_f10 = vmap(self._solve_one_x_tridiag_, in_axes=(0, 0, 0, 0))(diag, upper, lower, f10)
 
-                new_f10 = new_f10 + dt * self.get_offdiagonal_contrib(
+                new_f10 = new_f10 + dt * self.nuee_coeff * self.get_ee_offdiagonal_contrib(
                     None, f10, {"ddvf0": ddv, "d2dv2f0": d2dv2, "il": il}
                 )
 
                 # new_f10 = diffrax.diffeqsolve(
-                #     diffrax.ODETerm(self.get_offdiagonal_contrib),
+                #     diffrax.ODETerm(self.get_ee_offdiagonal_contrib),
                 #     solver=diffrax.Tsit5(),
                 #     t0=0.0,
                 #     t1=dt,

tests/test_vfp1d/epp-short.yaml

@@ -0,0 +1,73 @@
+units:
+  laser wavelength: 351nm
+  reference electron temperature: 300eV
+  reference ion temperature: 300eV
+  reference electron density: 1.5e21/cm^3
+  Z: 30
+  Ion: Au+
+  logLambda: nrl
+
+
+density:
+  quasineutrality: true
+  species-background:
+    noise_seed: 420
+    noise_type: gaussian
+    noise_val: 0.0
+    v0: 0.0
+    T0: 1.0
+    m: 2.0
+    n: 
+      basis: uniform
+      baseline: 1.0
+      bump_or_trough: bump
+      center: 0.0
+      rise: 25.0
+      bump_height: 0.0
+      width: 100000.0
+    T:
+      basis: sine
+      baseline: 1.0
+      amplitude: 1.0e-3
+      wavelength: 50000.0
+
+grid:
+  dt: 25.0
+  nv: 256
+  nx: 32
+  tmin: 0.
+  tmax: 50000.0
+  vmax: 8.0
+  xmax: 100000.0
+  xmin: 0.0
+  nl: 1
+
+save:
+  fields:
+    t:
+      tmin: 0.0
+      tmax: 50000.0
+      nt: 11
+  electron:
+    t:
+      tmin: 0.0
+      tmax: 50000.0
+      nt: 6
+
+mode: vfp-2d
+
+mlflow:
+  experiment: vfp2d
+  run: epperlein-short
+
+drivers:
+  ex: {}
+  ey: {}
+
+terms:
+  fokker_planck:
+    flm:
+      ee: false
+    f00:
+      type: Dougherty
+  e_solver: oshun

tests/test_vfp1d/test_kappa_eh.py

@@ -0,0 +1,30 @@
+import os, yaml, mlflow, numpy as np
+
+from utils.runner import run
+
+
+def _run_(Z):
+    # with open("configs/tf-1d/damping.yaml", "r") as fi:
+    with open(f"{os.path.join(os.getcwd(), 'tests/test_vfp1d/epp-short')}.yaml", "r") as fi:
+        cfg = yaml.safe_load(fi)
+
+    cfg["units"]["Z"] = Z
+    mlflow.set_experiment(cfg["mlflow"]["experiment"])
+    # modify config
+    with mlflow.start_run(run_name=cfg["mlflow"]["run"]) as mlflow_run:
+        result, datasets = run(cfg)
+
+    kappa_eh = mlflow.get_run(mlflow_run.info.run_id).data.metrics["kappa_eh"]
+    kappa = mlflow.get_run(mlflow_run.info.run_id).data.metrics["kappa"]
+    return kappa, kappa_eh
+
+
+def test_kappa_eh():
+    actual = []
+    desired = []
+    for Z in list(range(1, 21, 4)) + [40, 60, 80]:
+        kappa, kappa_eh = _run_(Z)
+        actual.append(kappa)
+        desired.append(kappa_eh)
+
+    np.testing.assert_allclose(actual, desired, rtol=0.2)