diff --git a/adept/__init__.py b/adept/__init__.py index cbabaa9..2ca51a7 100644 --- a/adept/__init__.py +++ b/adept/__init__.py @@ -9,9 +9,6 @@ from jax import numpy as jnp -from utils import misc - - def get_envelope(p_wL, p_wR, p_L, p_R, ax): return 0.5 * (jnp.tanh((ax - p_L) / p_wL) - jnp.tanh((ax - p_R) / p_wR)) @@ -224,9 +221,11 @@ def setup(self, cfg: Dict, adept_module: ADEPTModule = None) -> Dict[str, Module self.mlflow_run_id = mlflow_run.info.run_id else: + from adept.utils.misc import get_cfg + with mlflow.start_run(run_id=self.mlflow_run_id, nested=self.mlflow_nested) as mlflow_run: with tempfile.TemporaryDirectory(dir=self.base_tempdir) as temp_path: - cfg = misc.get_cfg(artifact_uri=mlflow_run.info.artifact_uri, temp_path=temp_path) + cfg = get_cfg(artifact_uri=mlflow_run.info.artifact_uri, temp_path=temp_path) modules = self._setup_(cfg, td, adept_module) mlflow.log_artifacts(td) # logs the temporary directory to mlflow @@ -261,6 +260,8 @@ def _get_adept_module_(self, cfg: Dict) -> ADEPTModule: return this_module(cfg) def _setup_(self, cfg: Dict, td: str, adept_module: ADEPTModule = None, log: bool = True) -> Dict[str, Module]: + from adept.utils.misc import log_params + if adept_module is None: self.adept_module = self._get_adept_module_(cfg) else: @@ -281,7 +282,7 @@ def _setup_(self, cfg: Dict, td: str, adept_module: ADEPTModule = None, log: boo self.adept_module.get_derived_quantities() # gets the derived quantities if log: - misc.log_params(self.adept_module.cfg) # logs the parameters to mlflow + log_params(self.adept_module.cfg) # logs the parameters to mlflow with open(os.path.join(td, "derived_config.yaml"), "w") as fi: yaml.dump(self.adept_module.cfg, fi) diff --git a/utils/__init__.py b/adept/utils/__init__.py similarity index 100% rename from utils/__init__.py rename to adept/utils/__init__.py diff --git a/utils/misc.py b/adept/utils/misc.py similarity index 100% rename from utils/misc.py rename to adept/utils/misc.py diff --git a/utils/plotters.py b/adept/utils/plotters.py similarity index 100% rename from utils/plotters.py rename to adept/utils/plotters.py diff --git a/utils/runner.py b/adept/utils/runner.py similarity index 100% rename from utils/runner.py rename to adept/utils/runner.py diff --git a/run.py b/run.py index 0dbcb1a..0fea6c7 100644 --- a/run.py +++ b/run.py @@ -11,7 +11,7 @@ import yaml from adept import ergoExo -from utils.misc import export_run +from adept.utils.misc import export_run if __name__ == "__main__": diff --git a/tests/test_lpse2d/test_epw_frequency.py b/tests/test_lpse2d/test_epw_frequency.py index 1eaf88f..231bfce 100644 --- a/tests/test_lpse2d/test_epw_frequency.py +++ b/tests/test_lpse2d/test_epw_frequency.py @@ -9,7 +9,6 @@ from numpy import testing -# from utils.runner import run from adept.theory.electrostatic import get_roots_to_electrostatic_dispersion, get_nlfs from adept import ergoExo diff --git a/tests/test_lpse2d/test_tpd_threshold.py b/tests/test_lpse2d/test_tpd_threshold.py index 7b60c60..50eed5f 100644 --- a/tests/test_lpse2d/test_tpd_threshold.py +++ b/tests/test_lpse2d/test_tpd_threshold.py @@ -2,7 +2,7 @@ from adept.lpse2d.helpers import calc_threshold_intensity import numpy as np -from utils.misc import setup_parsl +from adept.utils.misc import setup_parsl # from parsl.app.app import python_app @@ -11,7 +11,7 @@ def run_once(L, Te, I0): import yaml from adept import ergoExo - from utils.misc import export_run + from adept.utils.misc import export_run with open("tests/test_lpse2d/configs/tpd.yaml", "r") as fi: cfg = yaml.safe_load(fi) diff --git a/tests/test_tf1d/test_against_vlasov.py b/tests/test_tf1d/test_against_vlasov.py index b248ed4..cc0b794 100644 --- a/tests/test_tf1d/test_against_vlasov.py +++ b/tests/test_tf1d/test_against_vlasov.py @@ -13,8 +13,6 @@ from adept.theory import electrostatic from adept import ergoExo -# from utils.runner import run - def _modify_defaults_(defaults): rand_k0 = 0.358 diff --git a/tests/test_vfp1d/test_kappa_eh.py b/tests/test_vfp1d/test_kappa_eh.py index 7e641e4..5ee6b9e 100644 --- a/tests/test_vfp1d/test_kappa_eh.py +++ b/tests/test_vfp1d/test_kappa_eh.py @@ -2,8 +2,6 @@ from adept import ergoExo -# from utils.runner import run - def _run_(Z, ee): # with open("configs/tf-1d/damping.yaml", "r") as fi: