flatten ipi driver file structure

Author: PicoCentauri

README.md

@@ -12,7 +12,11 @@ Requires Python >= 3.11.
 pip install .
 ```
 
-Note: At the moment you need `marathon-dev` from `https://github.com/sirmarcel/marathon-dev` on the `hermes-reorg` branch. This will be fixed soon. 
+Note: At the moment you need `marathon-dev` from `https://github.com/sirmarcel/marathon-dev` on the `hermes-reorg` branch. This will be fixed soon.
+
+```bash
+pip install git+https://github.com/sirmarcel/marathon-dev.git@hermes-reorg
+```
 
 ## Usage
 

pyproject.toml

@@ -20,7 +20,7 @@ name = "lorem-jax"
 requires-python = ">=3.11"
 
 [project.scripts]
-lorem-install-ipi-driver = "lorem.dynamics._install_ipi_driver:install_ipi_driver"
+lorem-install-ipi-driver = "lorem._install_ipi_driver:install_ipi_driver"
 lorem-train = "lorem.train:main"
 
 [tool.pytest.ini_options]

src/lorem/__init__.py

@@ -7,8 +7,6 @@
 except ModuleNotFoundError:
     comms = None
 
-from .dynamics import LOREMCalculator
-
 # These modules require marathon and jaxpme at import time.
 # When those are not installed (e.g. CI), we skip them.
 try:
@@ -21,7 +19,6 @@
 
 __all__ = [
     "__version__",
-    "LOREMCalculator",
     "Lorem",
     "LoremBEC",
     "to_sample",

src/lorem/dynamics/_install_ipi_driver.py src/lorem/_install_ipi_driver.pysrc/lorem/dynamics/_install_ipi_driver.py renamed to src/lorem/_install_ipi_driver.py

@@ -18,9 +18,9 @@ def install_ipi_driver():
     pes_dir = Path(ipi_spec.submodule_search_locations[0]) / "pes"
 
     # Locate source driver module
-    source_spec = find_spec("lorem.dynamics.ipi")
+    source_spec = find_spec("lorem.ipi")
     if source_spec is None or source_spec.origin is None:
-        raise FileNotFoundError("Could not locate lorem.dynamics.ipi module")
+        raise FileNotFoundError("Could not locate lorem.ipi module")
     source_path = Path(source_spec.origin)
 
     target_path = pes_dir / "lorem.py"

src/lorem/dynamics/__init__.py

@@ -4,20 +4,18 @@
 
 from ipi.pes.ase import ASEDriver
 
-from lorem import LOREMCalculator
+from lorem.calculator import Calculator
 
 __DRIVER_NAME__ = "lorem"
 __DRIVER_CLASS__ = "LOREM_driver"
 
 
 class LOREM_driver(ASEDriver):
-    def __init__(self, template, model_path, device="cpu", *args, **kwargs):
-        super().__init__(template, *args, **kwargs)
-
+    def __init__(self, template, model_path, *args, **kwargs):
         self.model_path = model_path
-        self.device = device
+        super().__init__(template, *args, **kwargs)
         self.capabilities.append("BEC")
 
     def check_parameters(self):
         super().check_parameters()
-        self.ase_calculator = LOREMCalculator.from_checkpoint(self.model_path)
+        self.ase_calculator = Calculator.from_checkpoint(self.model_path)