From bf06f2bef3ad2f3892042c9e734de91e02e87893 Mon Sep 17 00:00:00 2001 From: Archis Joglekar Date: Fri, 22 Dec 2023 11:11:06 -0800 Subject: [PATCH] Vlasov1d - moment scalars (#25) * moment scalars for vlasov 1d * ln * entropy moments spatio-temporal collision frequencies * tested collisions * working wavepackets and absorbing boundaries --- adept/vlasov1d/helpers.py | 32 ++++++++++- adept/vlasov1d/integrator.py | 43 ++++++++++++-- adept/vlasov1d/pushers/fokker_planck.py | 35 ++++++------ adept/vlasov1d/storage.py | 44 ++++++++++++--- configs/vlasov-1d/epw.yaml | 45 +++++++++++++-- configs/vlasov-1d/wavepacket.yaml | 65 +++++++++++++++++----- env.yaml | 4 +- env_gpu.yaml | 4 +- tests/test_vlasov1d/configs/resonance.yaml | 42 +++++++++++++- utils/runner.py | 2 + 10 files changed, 262 insertions(+), 54 deletions(-) diff --git a/adept/vlasov1d/helpers.py b/adept/vlasov1d/helpers.py index 075a131..471651a 100644 --- a/adept/vlasov1d/helpers.py +++ b/adept/vlasov1d/helpers.py @@ -264,12 +264,18 @@ def post_process(result, cfg: Dict, td: str): os.makedirs(os.path.join(td, "plots"), exist_ok=True) os.makedirs(os.path.join(td, "plots", "fields"), exist_ok=True) os.makedirs(os.path.join(td, "plots", "fields", "lineouts"), exist_ok=True) + os.makedirs(os.path.join(td, "plots", "fields", "logplots"), exist_ok=True) + + os.makedirs(os.path.join(td, "plots", "scalars"), exist_ok=True) + + binary_dir = os.path.join(td, "binary") + os.makedirs(binary_dir) # merge # flds_paths = [os.path.join(flds_path, tf) for tf in flds_list] # arr = xarray.open_mfdataset(flds_paths, combine="by_coords", parallel=True) for k in result.ys.keys(): if k.startswith("field"): - fields_xr = store_fields(cfg, td, result.ys[k], result.ts[k], k) + fields_xr = store_fields(cfg, binary_dir, result.ys[k], result.ts[k], k) t_skip = int(fields_xr.coords["t"].data.size // 8) t_skip = t_skip if t_skip > 1 else 1 tslice = slice(0, -1, t_skip) @@ -279,12 +285,34 @@ def post_process(result, cfg: Dict, td: str): plt.savefig(os.path.join(td, "plots", "fields", f"spacetime-{nm[7:]}.png"), bbox_inches="tight") plt.close() + np.log10(np.abs(fld)).plot() + plt.savefig( + os.path.join(td, "plots", "fields", "logplots", f"spacetime-log-{nm[7:]}.png"), bbox_inches="tight" + ) + plt.close() + fld[tslice].T.plot(col="t", col_wrap=4) plt.savefig(os.path.join(td, "plots", "fields", "lineouts", f"{nm[7:]}.png"), bbox_inches="tight") plt.close() + elif k.startswith("default"): + scalars_xr = xarray.Dataset( + {k: xarray.DataArray(v, coords=(("t", result.ts["default"]),)) for k, v in result.ys["default"].items()} + ) + scalars_xr.to_netcdf(os.path.join(binary_dir, f"scalars-t={round(scalars_xr.coords['t'].data[-1], 4)}.nc")) + + for nm, srs in scalars_xr.items(): + fig, ax = plt.subplots(1, 2, figsize=(10, 4), tight_layout=True) + srs.plot(ax=ax[0]) + ax[0].grid() + np.log10(np.abs(srs)).plot(ax=ax[1]) + ax[1].grid() + ax[1].set_ylabel("$log_{10}$(|" + nm + "|)") + fig.savefig(os.path.join(td, "plots", "scalars", f"{nm}.png"), bbox_inches="tight") + plt.close() + f_xr = store_f(cfg, result.ts, td, result.ys) mlflow.log_metrics({"postprocess_time_min": round((time() - t0) / 60, 3)}) - return {"fields": fields_xr, "dists": f_xr} + return {"fields": fields_xr, "dists": f_xr, "scalars": scalars_xr} diff --git a/adept/vlasov1d/integrator.py b/adept/vlasov1d/integrator.py index c9be17a..d894c8d 100644 --- a/adept/vlasov1d/integrator.py +++ b/adept/vlasov1d/integrator.py @@ -6,6 +6,7 @@ import diffrax from adept.vlasov1d.pushers import field, fokker_planck, vlasov +from adept.tf1d.pushers import get_envelope class Stepper(diffrax.Euler): @@ -124,15 +125,37 @@ def __init__(self, cfg): self.compute_charges = partial(jnp.trapz, dx=cfg["grid"]["dv"], axis=1) self.dt = self.cfg["grid"]["dt"] self.driver = field.Driver(cfg["grid"]["x"]) - self.nu_prof = cfg["nuee"] - self.kr_prof = 0.0 + + def nu_prof(self, t, nu_args): + t_L = nu_args["time"]["center"] - nu_args["time"]["width"] * 0.5 + t_R = nu_args["time"]["center"] + nu_args["time"]["width"] * 0.5 + t_wL = nu_args["time"]["rise"] + t_wR = nu_args["time"]["rise"] + x_L = nu_args["space"]["center"] - nu_args["space"]["width"] * 0.5 + x_R = nu_args["space"]["center"] + nu_args["space"]["width"] * 0.5 + x_wL = nu_args["space"]["rise"] + x_wR = nu_args["space"]["rise"] + + nu_time = get_envelope(t_wL, t_wR, t_L, t_R, t) + if nu_args["time"]["bump_or_trough"] == "trough": + nu_time = 1 - nu_time + nu_time = nu_args["time"]["baseline"] + nu_args["time"]["bump_height"] * nu_time + + nu_prof = get_envelope(x_wL, x_wR, x_L, x_R, self.cfg["grid"]["x"]) + if nu_args["space"]["bump_or_trough"] == "trough": + nu_prof = 1 - nu_prof + nu_prof = nu_args["space"]["baseline"] + nu_args["space"]["bump_height"] * nu_prof + + return nu_time * nu_prof def __call__(self, t, y, args): """ This is just a wrapper around a Vlasov-Poisson + Fokker-Planck timestep - :param loop_carry: - :param current_params: + :param t: + :param y: + :param args: + :return: """ @@ -140,8 +163,18 @@ def __call__(self, t, y, args): dex = [val[0] for val in de] djy = de[0][1] + if self.cfg["terms"]["fokker_planck"]["is_on"]: + nu_fp_prof = self.nu_prof(t=t, nu_args=args["terms"]["fokker_planck"]) + else: + nu_fp_prof = None + + if self.cfg["terms"]["krook"]["is_on"]: + nu_K_prof = self.nu_prof(t=t, nu_args=args["terms"]["krook"]) + else: + nu_K_prof = None + electron_density_n = self.compute_charges(y["electron"]) - e, f, force, pond = self.vpfp(y["electron"], y["a"], y["e"], dex, self.nu_prof, self.kr_prof) + e, f, force, pond = self.vpfp(y["electron"], y["a"], y["e"], dex, nu_fp_prof, nu_K_prof) electron_density_np1 = self.compute_charges(f) a = self.wave_solver( diff --git a/adept/vlasov1d/pushers/fokker_planck.py b/adept/vlasov1d/pushers/fokker_planck.py index 09e20eb..fd89127 100644 --- a/adept/vlasov1d/pushers/fokker_planck.py +++ b/adept/vlasov1d/pushers/fokker_planck.py @@ -17,22 +17,23 @@ def __init__(self, cfg): self.td_solver = TridiagonalSolver(self.cfg) def __init_fp_operator__(self): - # if self.cfg["solver"]["fp_operator"] == "lenard_bernstein": - # return LenardBernstein(self.cfg) - # elif self.cfg["solver"]["fp_operator"] == "dougherty": - return Dougherty(self.cfg) - # else: - # raise NotImplementedError - - def __call__(self, nu_fp: jnp.float64, nu_K: jnp.float64, f: jnp.ndarray, dt: jnp.float64) -> jnp.ndarray: - if np.any(self.cfg["nuee"] > 0.0): + if self.cfg["terms"]["fokker_planck"]["type"].casefold() == "lenard_bernstein": + return LenardBernstein(self.cfg) + elif self.cfg["terms"]["fokker_planck"]["type"].casefold() == "dougherty": + return Dougherty(self.cfg) + else: + raise NotImplementedError + + def __call__(self, nu_fp: jnp.ndarray, nu_K: jnp.ndarray, f: jnp.ndarray, dt: jnp.float64) -> jnp.ndarray: + if self.cfg["terms"]["fokker_planck"]["is_on"]: # The three diagonals representing collision operator for all x cee_a, cee_b, cee_c = self.fp(nu=nu_fp, f_xv=f, dt=dt) # Solve over all x f = self.td_solver(cee_a, cee_b, cee_c, f) - # if (np.any(self.cfg["grid"]["kr_prof"] > 0.0)) and (np.any(self.cfg["grid"]["kt_prof"] > 0.0)): - # f = self.krook(nu_K, f, dt) + if self.cfg["terms"]["krook"]["is_on"]: + f = self.krook(nu_K, f, dt) + return f @@ -72,9 +73,9 @@ def __call__( """ v0t_sq = self.vx_moment(f_xv * self.v[None, :] ** 2.0) - a = nu * dt * (-v0t_sq[:, None] / self.dv**2.0 + jnp.roll(self.v, 1)[None, :] / 2 / self.dv) - b = 1.0 + nu * dt * self.ones * (2.0 * v0t_sq[:, None] / self.dv**2.0) - c = nu * dt * (-v0t_sq[:, None] / self.dv**2.0 - jnp.roll(self.v, -1)[None, :] / 2 / self.dv) + a = nu[:, None] * dt * (-v0t_sq[:, None] / self.dv**2.0 + jnp.roll(self.v, 1)[None, :] / 2 / self.dv) + b = 1.0 + nu[:, None] * dt * self.ones * (2.0 * v0t_sq[:, None] / self.dv**2.0) + c = nu[:, None] * dt * (-v0t_sq[:, None] / self.dv**2.0 - jnp.roll(self.v, -1)[None, :] / 2 / self.dv) return a, b, c @@ -101,13 +102,13 @@ def __call__( v0t_sq = self.vx_moment(f_xv * (self.v[None, :] - vbar[:, None]) ** 2.0) a = ( - nu + nu[:, None] * dt * (-v0t_sq[:, None] / self.dv**2.0 + (jnp.roll(self.v, 1)[None, :] - vbar[:, None]) / 2.0 / self.dv) ) - b = 1.0 + nu * dt * self.ones * (2.0 * v0t_sq[:, None] / self.dv**2.0) + b = 1.0 + nu[:, None] * dt * self.ones * (2.0 * v0t_sq[:, None] / self.dv**2.0) c = ( - nu + nu[:, None] * dt * (-v0t_sq[:, None] / self.dv**2.0 - (jnp.roll(self.v, -1)[None, :] - vbar[:, None]) / 2.0 / self.dv) ) diff --git a/adept/vlasov1d/storage.py b/adept/vlasov1d/storage.py index 46435ae..7524bf6 100644 --- a/adept/vlasov1d/storage.py +++ b/adept/vlasov1d/storage.py @@ -8,7 +8,7 @@ import xarray as xr -def store_fields(cfg: Dict, td: str, fields: Dict, this_t: np.ndarray, prefix: str) -> xr.Dataset: +def store_fields(cfg: Dict, binary_dir: str, fields: Dict, this_t: np.ndarray, prefix: str) -> xr.Dataset: """ Stores fields to netcdf @@ -19,8 +19,6 @@ def store_fields(cfg: Dict, td: str, fields: Dict, this_t: np.ndarray, prefix: s :param this_t: :return: """ - binary_dir = os.path.join(td, "binary") - os.makedirs(binary_dir, exist_ok=True) if any(x in ["x", "kx"] for x in cfg["save"][prefix].keys()): crds = set(cfg["save"][prefix].keys()) - {"t", "func"} @@ -103,14 +101,16 @@ def store_f(cfg: Dict, this_t: Dict, td: str, ys: Dict) -> xr.Dataset: def get_field_save_func(cfg, k): if {"t"} == set(cfg["save"][k].keys()): + def _calc_moment_(inp): + return jnp.trapz(inp, dx=cfg["grid"]["dv"], axis=1) + def fields_save_func(t, y, args): - temp = { - "n": jnp.trapz(y["electron"], dx=cfg["grid"]["dv"], axis=1), - "v": jnp.trapz(y["electron"] * cfg["grid"]["v"][None, :], dx=cfg["grid"]["dv"], axis=1), - } + temp = {"n": _calc_moment_(y["electron"]), "v": _calc_moment_(y["electron"] * cfg["grid"]["v"][None, :])} v_m_vbar = cfg["grid"]["v"][None, :] - temp["v"][:, None] - temp["p"] = jnp.trapz(y["electron"] * v_m_vbar**2.0, dx=cfg["grid"]["dv"], axis=1) - temp["q"] = jnp.trapz(y["electron"] * v_m_vbar**3.0, dx=cfg["grid"]["dv"], axis=1) + temp["p"] = _calc_moment_(y["electron"] * v_m_vbar**2.0) + temp["q"] = _calc_moment_(y["electron"] * v_m_vbar**3.0) + temp["-flogf"] = _calc_moment_(y["electron"] * jnp.log(jnp.abs(y["electron"]))) + temp["f^2"] = _calc_moment_(y["electron"] * y["electron"]) temp["e"] = y["e"] temp["de"] = y["de"] @@ -172,4 +172,30 @@ def get_save_quantities(cfg: Dict) -> Dict: elif k.startswith("electron"): cfg["save"][k]["func"] = get_dist_save_func(cfg, k) + cfg["save"]["default"] = {"t": {"ax": cfg["grid"]["t"]}, "func": get_default_save_func(cfg)} + return cfg + + +def get_default_save_func(cfg): + v = cfg["grid"]["v"][None, :] + dv = cfg["grid"]["dv"] + + def _calc_mean_moment_(inp): + return jnp.mean(jnp.trapz(inp, dx=dv, axis=1)) + + def save(t, y, args): + scalars = { + "mean_P": _calc_mean_moment_(y["electron"] * v**2.0), + "mean_j": _calc_mean_moment_(y["electron"] * v), + "mean_n": _calc_mean_moment_(y["electron"]), + "mean_q": _calc_mean_moment_(y["electron"] * v**3.0), + "mean_-flogf": _calc_mean_moment_(-jnp.log(jnp.abs(y["electron"])) * jnp.abs(y["electron"])), + "mean_f2": _calc_mean_moment_(y["electron"] * y["electron"]), + "mean_de2": jnp.mean(y["de"] ** 2.0), + "mean_e2": jnp.mean(y["e"] ** 2.0), + } + + return scalars + + return save diff --git a/configs/vlasov-1d/epw.yaml b/configs/vlasov-1d/epw.yaml index c372760..cf25c81 100644 --- a/configs/vlasov-1d/epw.yaml +++ b/configs/vlasov-1d/epw.yaml @@ -27,8 +27,8 @@ density: grid: c_light: 10 dt: 0.5 - nv: 512 - nx: 32 + nv: 4096 + nx: 128 tmin: 0. tmax: 480.0 vmax: 6.4 @@ -56,7 +56,7 @@ mlflow: drivers: ex: '0': - a0: 1.e-3 + a0: 1.e-4 k0: 0.3 t_center: 40.0 t_rise: 5.0 @@ -68,4 +68,41 @@ drivers: x_width: 4000000.0 ey: {} -nuee: 1.0e-5 \ No newline at end of file +terms: + fokker_planck: + is_on: True + type: Dougherty + time: + baseline: 1.0e-5 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 + space: + baseline: 1.0 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 + krook: + is_on: False + time: + baseline: 1.0 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 + space: + baseline: 1.0 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 \ No newline at end of file diff --git a/configs/vlasov-1d/wavepacket.yaml b/configs/vlasov-1d/wavepacket.yaml index 8039396..9e25f4a 100644 --- a/configs/vlasov-1d/wavepacket.yaml +++ b/configs/vlasov-1d/wavepacket.yaml @@ -22,13 +22,14 @@ density: center: 0.0 rise: 25.0 slope: 0.0 - wall_height: 0.0 + bump_height: 0.0 width: 100000.0 + grid: c_light: 10 dt: 0.5 - nv: 1024 - nx: 1024 + nv: 512 + nx: 4096 tmin: 0. tmax: 1000.0 vmax: 6.4 @@ -50,22 +51,60 @@ save: mode: vlasov-1d mlflow: - experiment: epw-finite-length-compare - run: vlasov + experiment: vlasov1d-epw-finite-length-compare + run: vlasov-absorbing drivers: ex: '0': - a0: 1.e-3 - k0: 0.3 - t_center: 70.0 + a0: 1.e-2 + k0: 0.4 + t_center: 60.0 t_rise: 5.0 - t_width: 100.0 - w0: 1.13 + t_width: 80.0 + w0: 1.28 dw0: 0. - x_center: 1400.0 + x_center: 600.0 x_rise: 20.0 - x_width: 600.0 + x_width: 900.0 ey: {} -nuee: 0.0 \ No newline at end of file +terms: + fokker_planck: + is_on: False + type: Dougherty + time: + baseline: 1.0e-5 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 + space: + baseline: 1.0 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 + + krook: + is_on: True + time: + baseline: 1.0 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 + space: + baseline: 0.0 + bump_or_trough: trough + center: 2000.0 + rise: 20.0 + slope: 0.0 + bump_height: 1.0 + width: 3800.0 \ No newline at end of file diff --git a/env.yaml b/env.yaml index c6495ae..7ec51bb 100644 --- a/env.yaml +++ b/env.yaml @@ -25,4 +25,6 @@ dependencies: - mlflow_export_import - pint - diffrax - - plasmapy \ No newline at end of file + - plasmapy + - parsl + - interpax \ No newline at end of file diff --git a/env_gpu.yaml b/env_gpu.yaml index f794516..c86d445 100644 --- a/env_gpu.yaml +++ b/env_gpu.yaml @@ -28,4 +28,6 @@ dependencies: - pint - diffrax - flatdict - - plasmapy \ No newline at end of file + - plasmapy + - parsl + - interpax \ No newline at end of file diff --git a/tests/test_vlasov1d/configs/resonance.yaml b/tests/test_vlasov1d/configs/resonance.yaml index 45bb224..9b1ab13 100644 --- a/tests/test_vlasov1d/configs/resonance.yaml +++ b/tests/test_vlasov1d/configs/resonance.yaml @@ -22,8 +22,9 @@ density: center: 0.0 rise: 25.0 slope: 0.0 - wall_height: 0.0 + bump_height: 0.0 width: 100000.0 + grid: c_light: 10 dt: 0.5 @@ -68,4 +69,41 @@ drivers: x_width: 4000000.0 ey: {} -nuee: 1.0e-5 \ No newline at end of file +terms: + fokker_planck: + is_on: True + type: Dougherty + time: + baseline: 1.0 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 + space: + baseline: 1.0 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 + krook: + is_on: True + time: + baseline: 1.0 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 + space: + baseline: 1.0 + bump_or_trough: bump + center: 0.0 + rise: 25.0 + slope: 0.0 + bump_height: 0.0 + width: 100000.0 \ No newline at end of file diff --git a/utils/runner.py b/utils/runner.py index cd35377..e6d423e 100644 --- a/utils/runner.py +++ b/utils/runner.py @@ -122,6 +122,8 @@ def _run_(these_models, time_quantities: Dict): args = {"drivers": cfg["drivers"]} if these_models is not None: args["models"] = these_models + if "terms" in cfg.keys(): + args["terms"] = cfg["terms"] return diffeqsolve( terms=diffeqsolve_quants["terms"],