diff --git a/adept/_base_.py b/adept/_base_.py index 3330dcb..b810a13 100644 --- 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 b/adept/_tf1d/__init__.py similarity index 100% rename from adept/tf1d/__init__.py rename to adept/_tf1d/__init__.py diff --git a/adept/tf1d/datamodel.py b/adept/_tf1d/datamodel.py similarity index 100% rename from adept/tf1d/datamodel.py rename to adept/_tf1d/datamodel.py diff --git a/adept/tf1d/modules.py b/adept/_tf1d/modules.py similarity index 98% rename from adept/tf1d/modules.py rename to adept/_tf1d/modules.py index 49ad9bd..1b5ec83 100644 --- a/adept/tf1d/modules.py +++ b/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 b/adept/_tf1d/solvers/pushers.py similarity index 100% rename from adept/tf1d/solvers/pushers.py rename to adept/_tf1d/solvers/pushers.py diff --git a/adept/tf1d/solvers/vector_field.py b/adept/_tf1d/solvers/vector_field.py similarity index 99% rename from adept/tf1d/solvers/vector_field.py rename to adept/_tf1d/solvers/vector_field.py index 6e49c05..d73b53c 100644 --- a/adept/tf1d/solvers/vector_field.py +++ b/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 b/adept/_tf1d/storage.py similarity index 100% rename from adept/tf1d/storage.py rename to adept/_tf1d/storage.py diff --git a/adept/tf1d/train_damping.py b/adept/_tf1d/train_damping.py similarity index 99% rename from adept/tf1d/train_damping.py rename to adept/_tf1d/train_damping.py index bbd45f3..9eedf7e 100644 --- a/adept/tf1d/train_damping.py +++ b/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 new file mode 100644 index 0000000..40d258a --- /dev/null +++ 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 index 5e2d230..6631c66 100644 --- 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 index 94905bc..c18f636 100644 --- 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 index a366ec9..fdf1372 100644 --- 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 index d3df613..0b6a6e5 100644 --- 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