Feature decouple somd 3 #29

Roy-Haolin-Du · 2025-01-08T20:26:17Z

Description

Replace template_config.cfg with Pydantic config

Todos

Remove the template_config.cfg and replace with a Pydantic SomdConfig class (which derived from a general EngineConfig class which can be configured in the code, or else read from a yaml/json file if not supplied.

…class _engine_runner_config. Write out the A3FE version, and make the output logs more concise and suitable for debugging

fjclark

Looks great, thanks! Just left a few comments.

fjclark · 2025-01-09T11:11:34Z

a3fe/configuration/_engine_runner_config.py

+        self._stream_log_level = stream_log_level
+        self._set_up_logging()
+
+    def _set_up_logging(self) -> None:


Can remove logging as it's not used. Thanks

fjclark · 2025-01-09T11:12:46Z

a3fe/configuration/_engine_runner_config.py

+        config = self.get_config()
+        with open(file_path, "w") as f:
+            _yaml.safe_dump(config, f, default_flow_style=False)
+        self._logger.info(f"Configuration dumped to {file_path}")


Can probably remove, thanks

fjclark · 2025-01-09T11:15:04Z

a3fe/configuration/_engine_runner_config.py

+        """
+        pass
+
+    def __eq__(self, other: object) -> bool:


No need to define if this derived from pdyantic base model.

fjclark · 2025-01-09T11:15:21Z

a3fe/configuration/_engine_runner_config.py

+from ..run._logging_formatters import _A3feStreamFormatter
+
+
+class EngineRunnerConfig(ABC):


Inherit from base model please.

fjclark · 2025-01-09T11:16:12Z

a3fe/configuration/engine_config.py

+
+from ._engine_runner_config import EngineRunnerConfig as _EngineRunnerConfig
+
+


Probably shift the base class into the same file - have a think about it - up to you!

fjclark · 2025-01-09T15:08:59Z

docs/guides.rst

@@ -191,8 +192,8 @@ three replicates. Note that this is expected to produce an erroneously favourabl
  cfg.runtime_npt = 50 # ps
  cfg.ensemble_equilibration_time = 100 # ps
  calc = a3.Calculation(ensemble_size = 3)
-  calc_set.setup(bound_leg_sysprep_config = cfg, free_leg_sysprep_config = cfg)
-  calc_set.run(adaptive = False, runtime=0.1) # ns
+  calc.setup(bound_leg_sysprep_config = cfg, free_leg_sysprep_config = cfg)


fjclark · 2025-01-09T15:11:24Z

docs/guides.rst

@@ -140,7 +141,7 @@ Individual Simulation settings
 -------------------------------

 To customise the specifics of how each lambda window is run (e.g. timestep), you can use the ``set_simfile_option`` method. For example, to set the timestep to 2 fs, run
-``calc.set_simfile_option("timestep", "2 * femtosecond")``. This will change parameters from the defaults given in ``template_config.cfg`` in the ``input`` directory, and warn
+``calc.set_simfile_option("timestep", "2 * femtosecond")``. This will change parameters from the defaults generated by ``engine_config``, and warn


We need to update this entire function to update the configs for all sub-simulation runners. Please rename it to something more appropriate e.g. set_config_option and test that it updates all configs.

fjclark · 2025-01-09T15:12:00Z

docs/guides.rst

@@ -304,12 +305,12 @@ You can run sets of calculations using the :class:`a3fe.run.CalcSet` class. To d
 ABFE with Charged Ligands
 *************************

-Since A3FE 0.2.0, ABFE calculations with charged ligands are supported using a co-alchemical ion approach. The charge of the ligand will be automatically detected, assuming that this is correctly specified in the input sdf. The only change in the input required is that the use of PME, rather than reaction field electrostatics, should be specified in ``template_config.cfg`` as e.g.:
+Since A3FE 0.2.0, ABFE calculations with charged ligands are supported using a co-alchemical ion approach. The charge of the ligand will be automatically detected, assuming that this is correctly specified in the input sdf. The only change in the input required is that the use of PME, rather than reaction field electrostatics, should be specified in ``somd_config.cfg`` as e.g.:


Please change to SomdEngineConfig (or whatever it was). Thanks

fjclark · 2025-01-09T15:12:25Z

docs/guides.rst

Please also change to point to the the config class rather than the template. Thanks!

fjclark · 2025-01-09T15:13:57Z

a3fe/run/simulation.py

@@ -338,19 +338,6 @@ def run(self, runtime: float = 2.5) -> None:
        -------
        None
        """
-        # Need to make sure that runtime is a multiple of the time per cycle


Please ensure that the somd.cfg file is written out before here, to ensure that the most up-to-date options are written out before the simulation is run.

Roy-Haolin-Du · 2025-01-16T15:42:43Z

Hi Finlay,

I hope you're doing well!

So far, the template_config.cfg has been replaced with a Pydantic SomdConfig class, which can now be generated automatically when the simulation requires it.

Additionally, we aim to avoid writing the somd.cfg file until it’s ready to be submitted to Slurm. The only place where the file should be written is in the run method of the simulation.py file. The workflow should only interact with the somd_config object, avoiding read from file. Information should flow one way: from the Pydantic classes to the files, without reading data back from files.

To support this, I have removed all read_simfile_option calls and definitions, as well as the write_simfile_option function. However, the generated somd.cfg file is still missing some extended parameters, such as morph, top, crdFiles, and restraints.
Here is the related commit for reference: [8d7b957]
8d7b957.

Perhaps in our next meeting, we can discuss this part in detail and explore potential solutions. I’d appreciate any advice you have.

Thanks a lot for your support!

Best regards

fjclark · 2025-01-19T11:45:27Z

Hi Roy, I'm doing well thanks, hope you are too! Thanks for all your work - this is looking great! Also thanks for the clear description of your changes and issues - let's discuss this tomorrow.

…tions

a3fe/configuration/_engine_runner_config.py

a3fe/configuration/engine_config.py

fjclark · 2025-01-20T17:12:48Z

a3fe/configuration/engine_config.py

+    ### Constraints ###
+    constraint: str = _Field("hbonds", description="Constraint type, must be hbonds or allbonds")
+    hydrogen_mass_factor: float = _Field(3.0, alias="hydrogen mass repartitioning factor", description="Hydrogen mass repartitioning factor")
+    integrator: _Literal["langevinmiddle", "leapfrogverlet"] = _Field("langevinmiddle", description="Integration algorithm")


a3fe/configuration/engine_config.py

fjclark · 2025-01-20T17:15:59Z

a3fe/configuration/engine_config.py

-        Set the number of cycles in the SOMD config file.
+    ### Non-Bonded Interactions ###
+    cutoff_type: _Literal["cutoffperiodic", "PME"] = _Field("cutoffperiodic", description="Type of cutoff to use. Options: PME, cutoffperiodic")
+    cutoff_distance: float = _Field(12.0, alias="cutoff distance", ge=6.0, le=18.0, description="Cutoff distance in angstroms (6-18). Default 12.0 for cutoffperiodic and 10.0 for PME")


We probably just want one default, otherwise we're going to need some more complex logic and may end up overwriting user input. Let's stick with the RF defaults of RF and 12 A.

fjclark · 2025-01-20T17:35:17Z

a3fe/configuration/engine_config.py

+    ligand_charge: int = _Field(0, description="Net charge of the ligand. If non-zero, must use PME for electrostatics.")
+    charge_difference: int = _Field(0, description="Charge difference between the ligand and the system. If non-zero, must use co-alchemical ion approach.")


These are the same for ABFE - let's just keep the ligand charge option. We can write it out to the config file as charge_difference = - ligand charge. Thanks

fjclark · 2025-01-20T17:36:14Z

a3fe/configuration/engine_config.py

+        default_factory=_get_default_lambda_array,
+        description="Lambda array for alchemical perturbation, varies from 0.0 to 1.0 across stages"
+    )
+    lambda_val: _Optional[float] = _Field(None, description="Lambda value for current stage")


Another example of something that we'd rather not expose to the user, since it varies based on the simulation. We need a good strategy for this in general (haven't got a good solution yet - let me know if you think of one!)

a3fe/configuration/engine_config.py

…d only in one place

fjclark · 2025-02-19T10:20:52Z

a3fe/run/stage.py

@@ -175,6 +175,10 @@ def __init__(
                self.lam_vals = lambda_values
            else:
                self.lam_vals = self._get_lam_vals()


Please modify this to completely remove any reference to "somd.cfg" which should only be written by the engine config before it is used.

Please check why somd.cfg is appearing in the stage input directories - it shouldn't.

fjclark · 2025-02-19T10:30:17Z

a3fe/run/stage.py

+
+            if self.engine_config is not None:
+                self.engine_config.lambda_values = self.lam_vals
+


Please remove all of this logic. Do not allow the user to pass a set of lambda values - this violates the single source of information principle because then the values are also in the config. Please just create a property which gets the lambda values from the engine config.

Please add a similar property for systemprepconfig, then delete the mapping dictionaries and replace with this functionality. Thanks!

…in-Du/a3fe into feature-decouple-somd-3

…s are only stored in one place.

…da stored in right place

Roy-Haolin-Du · 2025-02-24T15:14:48Z

a3fe/run/leg.py

            for file in _glob.glob(f"{stage_input_dir}/lambda_0.0000/*"):
-                _shutil.copy(file, stage_input_dir)
+                if _os.path.basename(file) != "somd.cfg": 
+                    _shutil.copy(file, stage_input_dir)
            for file in _glob.glob(f"{stage_input_dir}/lambda_*"):


Beacuse somd.cfg would be generated by BSS automately, copy exclude somd.cfg if we dont want.

Roy-Haolin-Du · 2025-02-24T15:16:45Z

a3fe/run/simulation.py

@@ -405,7 +403,8 @@ def failed(self) -> bool:

    @property
    def slurm_output_files(self) -> _List[str]:
-        return _glob.glob(f"{self.slurm_file_base}*")
+        """Get a list of all slurm output files for this simulation."""
+        return _glob.glob(_os.path.join(self.output_dir, "slurm-*.out"))



We have analyses slurm and simulation slurm result files, just clearly use the correct to do analyse.

Roy-Haolin-Du · 2025-02-24T15:18:32Z

a3fe/run/stage.py

-            if lambda_values is not None:
-                self.lam_vals = lambda_values
+            if self.engine_config.lambda_values is None:
+                self.engine_config.lambda_values = self.get_lambda_values
            else:


Keep the single source of information principle, the values are also in the config files.

Single source of information - always read the lambda values from the engine config, so make self.lam_vals a property which reads self.engine_config e.g.

@Property
def lam_vals(self):
return self.engine_config.lambda_values

Roy-Haolin-Du · 2025-02-24T15:19:04Z

a3fe/run/stage.py

+    def get_lambda_values(self) -> _List[float]:
+        """Return list of lambda values from the engine configuration."""
+        return _SomdConfig._from_config_file(_os.path.join(self.input_dir, "somd.cfg")).lambda_values
+


A property which gets the lambda values from the engine config.

Please remove this.

Roy-Haolin-Du · 2025-02-24T15:20:42Z

a3fe/tests/test_run.py

@@ -443,6 +440,9 @@ def test_setup_calc_stages(self, setup_calc):
                    cfg.lambda_values[leg.leg_type][stage.stage_type]
                )
                assert lam_vals == expected_lam_vals
+
+                if leg.leg_type == a3.LegType.BOUND:
+                    assert stage.engine_config.boresch_restraints_dictionary is not None



test restraint files would be generated from engine_config if bound leg.

fjclark · 2025-02-24T16:45:32Z

a3fe/run/system_prep.py

@@ -53,7 +51,7 @@ def parameterise_input(
    parameterised_system : _BSS._SireWrappers._system.System
        Parameterised system.
    """
-    cfg = _ENGINE_TYPE_TO_SYSPREP_CONFIG[engine_type].load(input_dir, leg_type)
+    cfg = _SomdSystemPreparationConfig.load(input_dir, leg_type)


So this is no longer generic, but specific to SOMD - please make it generic again by reading the config from the passed engine type enum.

fjclark · 2025-02-24T16:47:02Z

a3fe/run/leg.py

@@ -231,15 +231,15 @@ def setup(
        cfg = (
            sysprep_config
            if sysprep_config is not None
-            else _ENGINE_TYPE_TO_SYSPREP_CONFIG[self.engine_type]()
+            else _SomdSystemPreparationConfig()


Please make generic.

else self.engine_config.sysprep_config

fjclark · 2025-02-24T16:50:21Z

a3fe/run/leg.py

@@ -673,9 +672,10 @@ def setup_stages(
                property_map={"velocity": "foo"},
            )  # We will run outside of BSS

-            # Copy input written by BSS to the stage input directory
+            # Copy input written by BSS to the stage input directory, excluding only somd.cfg
            for file in _glob.glob(f"{stage_input_dir}/lambda_0.0000/*"):


I think this would be clearer (please modify as required)

files = [
file for file in glob.glob(f"{stage_input_dir}/lambda_0.0000/*")
if not file.endswith(".cfg")
]

fjclark · 2025-02-25T16:03:39Z

a3fe/run/simulation.py

@@ -405,7 +403,8 @@ def failed(self) -> bool:

    @property
    def slurm_output_files(self) -> _List[str]:
-        return _glob.glob(f"{self.slurm_file_base}*")
+        """Get a list of all slurm output files for this simulation."""
+        return _glob.glob(_os.path.join(self.output_dir, "slurm-*.out"))


Please revert - we've fixed this now!

fjclark · 2025-02-25T16:11:18Z

a3fe/run/stage.py

@@ -171,13 +167,10 @@ def __init__(
        )

        if not self.loaded_from_pickle:
-            if lambda_values is not None:
-                self.lam_vals = lambda_values
+            if self.engine_config.lambda_values is None:


Please remove this

fjclark · 2025-02-25T16:13:37Z

a3fe/run/stage.py

@@ -229,24 +222,29 @@ def lam_windows(self) -> _List[_LamWindow]:
        return self._sub_sim_runners

    @lam_windows.setter
-    def legs(self, value) -> None:
+    def lam_windows(self, value) -> None:


fjclark · 2025-02-25T16:14:41Z

a3fe/run/stage.py

    @property
    def lam_val_weights(self) -> _List[float]:
        """Return the weights for each lambda window. These are calculated
        according to how each windows contributes to the overall free energy
        estimate, as given by TI and the trapezoidal rule."""
        lam_val_weights = []
-        for i, lam_val in enumerate(self.lam_vals):
+        for i, lam_val in enumerate(self.engine_config.lambda_values):


If lam_vals is a property which reads self.engine_config, no need to change this. This may be clearer though - so up to you.

fjclark · 2025-02-25T16:15:08Z

a3fe/run/stage.py

@@ -1228,7 +1204,7 @@ def update(self, save_name: str = "output_saved") -> None:
        self._logger.info("Deleting old lambda windows and creating new ones...")
        self._sub_sim_runners = []
        lam_val_weights = self.lam_val_weights
-        for i, lam_val in enumerate(self.lam_vals):
+        for i, lam_val in enumerate(self.engine_config.lambda_values):


Can also change back (if you want to)

Roy-Haolin-Du added 5 commits January 6, 2025 12:09

Replace template_config.cfg with Pydantic config

8dc92c7

test and format for removing run_somd.sh and template_config.cfg

b654a04

Update engine_config validation, ncycle, runtime and test passed

fb7346c

keep partition as default

ec4858f

Flexibly handle SOMD configuration files and create an abstract base …

abc0b9e

…class _engine_runner_config. Write out the A3FE version, and make the output logs more concise and suitable for debugging

fjclark reviewed Jan 9, 2025

View reviewed changes

Roy-Haolin-Du added 3 commits January 12, 2025 19:45

engine_confi and test

624576e

engine_confi and test

f0730aa

only consult somd_config object, not read from files

8d7b957

consult engine_config to update somd.cfg with the necessary configura…

5c1c09f

…tions