Refactor for twofluid 1D (#79)

ergodicio · Sep 27, 2024 · dc82c25 · dc82c25
1 parent 6a18eff
commit dc82c25
Show file tree

Hide file tree

Showing 13 changed files with 11 additions and 10 deletions.
diff --git a/adept/_base_.py b/adept/_base_.py
@@ -252,8 +252,7 @@ def _get_adept_module_(self, cfg: Dict) -> ADEPTModule:
         """
 
         if cfg["solver"] == "tf-1d":
-            from adept.tf1d.modules import BaseTwoFluid1D as this_module
-            from adept.tf1d.datamodel import ConfigModel
+            from adept.tf1d import BaseTwoFluid1D as this_module, ConfigModel
 
             # config = ConfigModel(**cfg)
 

diff --git a/adept/tf1d/__init__.py → adept/_tf1d/__init__.py b/adept/tf1d/__init__.py → adept/_tf1d/__init__.py
diff --git a/adept/tf1d/datamodel.py → adept/_tf1d/datamodel.py b/adept/tf1d/datamodel.py → adept/_tf1d/datamodel.py
diff --git a/adept/tf1d/modules.py → adept/_tf1d/modules.py b/adept/tf1d/modules.py → adept/_tf1d/modules.py
@@ -6,8 +6,8 @@
 from jax import numpy as jnp, tree_util as jtu
 
 from adept import ADEPTModule
-from adept.tf1d.solvers.vector_field import VF
-from adept.tf1d.storage import save_arrays, plot_xrs
+from adept._tf1d.solvers.vector_field import VF
+from adept._tf1d.storage import save_arrays, plot_xrs
 
 
 class BaseTwoFluid1D(ADEPTModule):

diff --git a/adept/tf1d/solvers/pushers.py → adept/_tf1d/solvers/pushers.py b/adept/tf1d/solvers/pushers.py → adept/_tf1d/solvers/pushers.py
diff --git a/adept/tf1d/solvers/vector_field.py → adept/_tf1d/solvers/vector_field.py b/adept/tf1d/solvers/vector_field.py → adept/_tf1d/solvers/vector_field.py
@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 import equinox as eqx
 
-from adept.tf1d.solvers import pushers
+from adept._tf1d.solvers import pushers
 
 
 class VF(eqx.Module):

diff --git a/adept/tf1d/storage.py → adept/_tf1d/storage.py b/adept/tf1d/storage.py → adept/_tf1d/storage.py
diff --git a/adept/tf1d/train_damping.py → adept/_tf1d/train_damping.py b/adept/tf1d/train_damping.py → adept/_tf1d/train_damping.py
@@ -18,7 +18,7 @@
 import equinox as eqx
 from tqdm import tqdm
 
-from adept.tf1d import helpers
+from adept._tf1d import helpers
 from diffrax import diffeqsolve, ODETerm, SaveAt, Tsit5
 from utils import misc, plotters
 

diff --git a/adept/tf1d.py b/adept/tf1d.py
@@ -0,0 +1,2 @@
+from ._tf1d.modules import BaseTwoFluid1D
+from ._tf1d.datamodel import ConfigModel
diff --git a/adept/vlasov1d2v/helpers.py b/adept/vlasov1d2v/helpers.py
@@ -15,7 +15,7 @@
 
 from adept.vlasov1d2v.integrator import VlasovMaxwell, Stepper
 from adept.vlasov1d2v.storage import store_f, store_fields, get_save_quantities
-from adept.tf1d.pushers import get_envelope
+from adept._base_ import get_envelope
 
 gamma_da = xarray.open_dataarray(os.path.join(os.path.dirname(__file__), "..", "vlasov1d", "gamma_func_for_sg.nc"))
 m_ax = gamma_da.coords["m"].data

diff --git a/adept/vlasov1d2v/integrator.py b/adept/vlasov1d2v/integrator.py
@@ -6,7 +6,7 @@
 import diffrax
 
 from adept.vlasov1d2v.pushers import field, fokker_planck, vlasov
-from adept.tf1d.pushers import get_envelope
+from adept._base_ import get_envelope
 
 
 class Stepper(diffrax.Euler):

diff --git a/adept/vlasov1d2v/pushers/field.py b/adept/vlasov1d2v/pushers/field.py
@@ -6,7 +6,7 @@
 
 from jax import numpy as jnp
 
-from adept.tf1d.pushers import get_envelope
+from adept._base_ import get_envelope
 
 
 class Driver:

diff --git a/tests/test_tf1d/test_resonance_search.py b/tests/test_tf1d/test_resonance_search.py
@@ -20,7 +20,7 @@
 from tqdm import tqdm
 
 from adept import ergoExo
-from adept.tf1d.modules import BaseTwoFluid1D
+from adept.tf1d import BaseTwoFluid1D
 from adept.electrostatic import get_roots_to_electrostatic_dispersion