Restructure helpers directory

stfc · Oct 31, 2024 · 5aad9e7 · 5aad9e7
1 parent 8d53a1f
commit 5aad9e7
Show file tree

Hide file tree

Showing 18 changed files with 410 additions and 394 deletions.
diff --git a/janus_core/calculations/base.py b/janus_core/calculations/base.py
@@ -12,7 +12,8 @@
     PathLike,
 )
 from janus_core.helpers.log import config_logger, config_tracker
-from janus_core.helpers.utils import FileNameMixin, input_structs, none_to_dict
+from janus_core.helpers.struct_io import input_structs
+from janus_core.helpers.utils import FileNameMixin, none_to_dict
 
 
 class BaseCalculation(FileNameMixin):

diff --git a/janus_core/calculations/descriptors.py b/janus_core/calculations/descriptors.py
@@ -15,7 +15,9 @@
     MaybeSequence,
     PathLike,
 )
-from janus_core.helpers.utils import check_calculator, none_to_dict, output_structs
+from janus_core.helpers.mlip_calculators import check_calculator
+from janus_core.helpers.struct_io import output_structs
+from janus_core.helpers.utils import none_to_dict
 
 
 class Descriptors(BaseCalculation):

diff --git a/janus_core/calculations/eos.py b/janus_core/calculations/eos.py
@@ -19,7 +19,8 @@
     OutputKwargs,
     PathLike,
 )
-from janus_core.helpers.utils import none_to_dict, output_structs
+from janus_core.helpers.struct_io import output_structs
+from janus_core.helpers.utils import none_to_dict
 
 
 class EoS(BaseCalculation):

diff --git a/janus_core/calculations/geom_opt.py b/janus_core/calculations/geom_opt.py
@@ -11,16 +11,17 @@
 from numpy import linalg
 
 from janus_core.calculations.base import BaseCalculation
-from janus_core.helpers.janus_types import ASEOptArgs, OutputKwargs, PathLike
-from janus_core.helpers.utils import (
+from janus_core.helpers.janus_types import (
     Architectures,
+    ASEOptArgs,
     ASEReadArgs,
     Devices,
-    none_to_dict,
-    output_structs,
-    snap_symmetry,
-    spacegroup,
+    OutputKwargs,
+    PathLike,
 )
+from janus_core.helpers.struct_io import output_structs
+from janus_core.helpers.utils import none_to_dict
+from janus_core.processing.symmetry import snap_symmetry, spacegroup
 
 
 class GeomOpt(BaseCalculation):

diff --git a/janus_core/calculations/md.py b/janus_core/calculations/md.py
@@ -26,24 +26,20 @@
 
 from janus_core.calculations.base import BaseCalculation
 from janus_core.calculations.geom_opt import GeomOpt
-from janus_core.helpers.correlator import Correlation
 from janus_core.helpers.janus_types import (
+    Architectures,
+    ASEReadArgs,
     CorrelationKwargs,
+    Devices,
     Ensembles,
     OutputKwargs,
     PathLike,
     PostProcessKwargs,
 )
-from janus_core.helpers.post_process import compute_rdf, compute_vaf
-from janus_core.helpers.utils import (
-    Architectures,
-    ASEReadArgs,
-    Devices,
-    input_structs,
-    none_to_dict,
-    output_structs,
-    write_table,
-)
+from janus_core.helpers.struct_io import input_structs, output_structs
+from janus_core.helpers.utils import none_to_dict, write_table
+from janus_core.processing.correlator import Correlation
+from janus_core.processing.post_process import compute_rdf, compute_vaf
 
 DENS_FACT = (units.m / 1.0e2) ** 3 / units.mol
 

diff --git a/janus_core/calculations/single_point.py b/janus_core/calculations/single_point.py
@@ -18,7 +18,9 @@
     PathLike,
     Properties,
 )
-from janus_core.helpers.utils import check_calculator, none_to_dict, output_structs
+from janus_core.helpers.mlip_calculators import check_calculator
+from janus_core.helpers.struct_io import output_structs
+from janus_core.helpers.utils import none_to_dict
 
 
 class SinglePoint(BaseCalculation):

diff --git a/janus_core/cli/train.py b/janus_core/cli/train.py
@@ -49,7 +49,7 @@ def train(
         is Path("train-summary.yml").
     """
     from janus_core.cli.utils import carbon_summary, end_summary, start_summary
-    from janus_core.helpers.train import train as run_train
+    from janus_core.training.train import train as run_train
 
     with open(mlip_config, encoding="utf8") as config_file:
         config = yaml.safe_load(config_file)

diff --git a/janus_core/helpers/mlip_calculators.py b/janus_core/helpers/mlip_calculators.py
@@ -11,6 +11,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, get_args
 
+from ase.calculators.mixing import SumCalculator
+
 from janus_core.helpers.janus_types import Architectures, Devices, PathLike
 
 if TYPE_CHECKING:
@@ -231,3 +233,33 @@ def choose_calculator(
     calculator.parameters["arch"] = arch
 
     return calculator
+
+
+def check_calculator(calc: Calculator, attribute: str) -> None:
+    """
+    Ensure calculator has ability to calculate properties.
+
+    If the calculator is a SumCalculator that inlcudes the TorchDFTD3Calculator, this
+    also sets the relevant function so that the MLIP component of the calculator is
+    used for properties unrelated to dispersion.
+
+    Parameters
+    ----------
+    calc : Calculator
+        ASE Calculator to check.
+    attribute : str
+        Attribute to check calculator for.
+    """
+    # If dispersion added to MLIP calculator, use only MLIP calculator for calculation
+    if isinstance(calc, SumCalculator):
+        if (
+            len(calc.mixer.calcs) == 2
+            and calc.mixer.calcs[1].name == "TorchDFTD3Calculator"
+            and hasattr(calc.mixer.calcs[0], attribute)
+        ):
+            setattr(calc, attribute, getattr(calc.mixer.calcs[0], attribute))
+
+    if not hasattr(calc, attribute) or not callable(getattr(calc, attribute)):
+        raise NotImplementedError(
+            f"The attached calculator does not currently support {attribute}"
+        )