Merge pull request #77 from Roy-Haolin-Du/main

Fix: Add timestep unit conversion for stat inefficiencies (fs to ns)

Author: Roy-Haolin-Du

a3fe/__init__.py

@@ -13,26 +13,25 @@
 import warnings as _warnings
 
 from ._version import __version__
-from .run import (
-    CalcSet,
-    Calculation,
-    LamWindow,
-    Leg,
-    Simulation,
-    Stage,
-)
-
 from .configuration import (
-    SlurmConfig,
-    _EngineConfig,
-    SomdConfig,
     EngineType,
     JobStatus,
     LegType,
     PreparationStage,
+    SlurmConfig,
+    SomdConfig,
+    SomdSystemPreparationConfig,
     StageType,
+    _EngineConfig,
     enums,
-    SomdSystemPreparationConfig,
+)
+from .run import (
+    CalcSet,
+    Calculation,
+    LamWindow,
+    Leg,
+    Simulation,
+    Stage,
 )
 
 # A3FE can open many files due to the use of multiprocessing and

a3fe/_version.py

@@ -1 +1 @@
-__version__ = "0.4.0"
+__version__ = "0.4.1"

a3fe/analyse/mbar.py

@@ -15,6 +15,7 @@
 
 import numpy as _np
 
+from ..configuration import SlurmConfig as _SlurmConfig
 from ..read._process_somd_files import read_mbar_gradients as _read_mbar_gradients
 from ..read._process_somd_files import read_mbar_result as _read_mbar_result
 from ..read._process_somd_files import (
@@ -23,8 +24,6 @@
 from ..run._virtual_queue import Job as _Job
 from ..run._virtual_queue import VirtualQueue as _VirtualQueue
 
-from ..configuration import SlurmConfig as _SlurmConfig
-
 
 def run_mbar(
     output_dir: str,

a3fe/analyse/process_grads.py

@@ -136,7 +136,8 @@ def __init__(
         stat_ineffs_all_winds = (
             _np.array(stat_ineffs_all_winds)
             * lam_winds[0].sims[0].engine_config.timestep
-        )  # Timestep should be same for all sims
+            / 1_000_000
+        )  # Timestep should be same for all sims, fs to ns
 
         # Get the SEMs of the free energy changes from the inter-run SEMs of the gradients
         lam_weights = _np.array([lam.lam_val_weight for lam in lam_winds])

a3fe/configuration/__init__.py

@@ -1,9 +1,9 @@
 """Pydantic configuration classes for the a3fe package."""
 
+from .engine_config import SomdConfig, _EngineConfig
+from .enums import EngineType, JobStatus, LegType, PreparationStage, StageType
+from .slurm_config import SlurmConfig
 from .system_prep_config import (
-    _BaseSystemPreparationConfig,
     SomdSystemPreparationConfig,
+    _BaseSystemPreparationConfig,
 )
-from .slurm_config import SlurmConfig
-from .engine_config import _EngineConfig, SomdConfig
-from .enums import EngineType, JobStatus, LegType, PreparationStage, StageType

a3fe/configuration/engine_config.py

@@ -5,23 +5,36 @@
 ]
 
 import os as _os
+from abc import ABC as _ABC
+from abc import abstractmethod as _abstractmethod
 from decimal import Decimal as _Decimal
 from typing import (
     Dict as _Dict,
-    Literal as _Literal,
+)
+from typing import (
     List as _List,
-    Union as _Union,
+)
+from typing import (
+    Literal as _Literal,
+)
+from typing import (
     Optional as _Optional,
 )
+from typing import (
+    Union as _Union,
+)
+
+import yaml as _yaml
 from pydantic import (
     BaseModel as _BaseModel,
+)
+from pydantic import (
     Field as _Field,
+)
+from pydantic import (
     model_validator as _model_validator,
 )
 
-import yaml as _yaml
-from abc import ABC as _ABC, abstractmethod as _abstractmethod
-
 
 class _EngineConfig(_BaseModel, _ABC):
     """Base class for engine runner configurations."""

a3fe/configuration/enums.py

@@ -1,12 +1,13 @@
 """Enums required for Classes in the run package."""
 
 from enum import Enum as _Enum
+from typing import Any as _Any
 from typing import List as _List
-import yaml as _yaml
 
-from .engine_config import _EngineConfig, SomdConfig as _SomdConfig
+import yaml as _yaml
 
-from typing import Any as _Any
+from .engine_config import SomdConfig as _SomdConfig
+from .engine_config import _EngineConfig
 
 __all__ = [
     "JobStatus",
@@ -141,11 +142,17 @@ def file_suffix(self) -> str:
     @property
     def prep_fn(self):
         """The function to use to prepare the input files for this stage."""
+        from ..run.system_prep import (
+            heat_and_preequil_input as _heat_and_preequil_input,
+        )
+        from ..run.system_prep import (
+            minimise_input as _minimise_input,
+        )
         from ..run.system_prep import (
             parameterise_input as _parameterise_input,
+        )
+        from ..run.system_prep import (
             solvate_input as _solvate_input,
-            minimise_input as _minimise_input,
-            heat_and_preequil_input as _heat_and_preequil_input,
         )
 
         prep_fns = {

a3fe/configuration/slurm_config.py

@@ -4,17 +4,16 @@
     "SlurmConfig",
 ]
 
-import yaml as _yaml
-import subprocess as _subprocess
+import os as _os
 import re as _re
+import subprocess as _subprocess
+from typing import Dict as _Dict
+from typing import List as _List
 
+import yaml as _yaml
 from pydantic import BaseModel as _BaseModel
-from pydantic import Field as _Field
 from pydantic import ConfigDict as _ConfigDict
-
-import os as _os
-
-from typing import List as _List, Dict as _Dict
+from pydantic import Field as _Field
 
 
 class SlurmConfig(_BaseModel):
@@ -33,8 +32,12 @@ class SlurmConfig(_BaseModel):
     extra_options: _Dict[str, str] = _Field(
         {}, description="Extra options to pass to SLURM. For example, {'account': 'qt'}"
     )
-    queue_check_interval: int = _Field(30, ge=1, description="Interval in seconds between SLURM queue status checks.")
-    job_submission_wait: int = _Field(300, ge=1, description="Wait time in seconds for job submission to SLURM queue.")
+    queue_check_interval: int = _Field(
+        30, ge=1, description="Interval in seconds between SLURM queue status checks."
+    )
+    job_submission_wait: int = _Field(
+        300, ge=1, description="Wait time in seconds for job submission to SLURM queue."
+    )
 
     model_config = _ConfigDict(validate_assignment=True)
 

a3fe/configuration/system_prep_config.py

@@ -6,18 +6,17 @@
     "SomdSystemPreparationConfig",
 ]
 
-import yaml as _yaml
-
 from abc import ABC as _ABC
+from typing import Dict as _Dict
+from typing import List as _List
 
+import yaml as _yaml
 from pydantic import BaseModel as _BaseModel
-from pydantic import Field as _Field
 from pydantic import ConfigDict as _ConfigDict
+from pydantic import Field as _Field
 
-from .enums import StageType as _StageType
 from .enums import LegType as _LegType
-
-from typing import List as _List, Dict as _Dict
+from .enums import StageType as _StageType
 
 
 class _BaseSystemPreparationConfig(_ABC, _BaseModel):

a3fe/read/_read_exp_dgs.py

@@ -2,11 +2,10 @@
 This must have the columns: calc_base_dir, name, exp_dg, exp_er"""
 
 import os as _os
-
 from typing import Optional as _Optional
 
-import pandas as _pd
 import numpy as _np
+import pandas as _pd
 
 
 def read_exp_dgs(