Allow ensemble type to be passed as a string from the command line

fjclark · May 1, 2024 · 3cd1a11 · 3cd1a11
1 parent fb59df8
commit 3cd1a11
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 10 deletions.
diff --git a/maize/steps/exs/biosimspace/_base.py b/maize/steps/exs/biosimspace/_base.py
@@ -62,8 +62,9 @@ def required_callables(cls) -> list[str]:
     save_name: FileParameter[Path] = FileParameter(optional=True)
     """
     If supplied, the output files will be saved with
-    this name. E.g., if set to Path("output"), the output
-    files will be output.rst7 and output.prm7.
+    this name. E.g., if set to Path("output") (or "output"
+    through the CLI), the output files will be output.rst7
+    and output.prm7.
     """
 
     # Output

diff --git a/maize/steps/exs/biosimspace/afe.py b/maize/steps/exs/biosimspace/afe.py
@@ -235,7 +235,7 @@ def _get_protocol(self) -> "BSS.Protocol._protocol.Protocol":
 
         pressure = (
             None
-            if self.ensemble.value == Ensemble.NVT
+            if Ensemble(self.ensemble.value) == Ensemble.NVT
             else self.pressure.value * BSS.Units.Pressure.atm
         )
 
@@ -441,7 +441,7 @@ class _AFEBase(_BioSimSpaceBase, ABC):
     temperature: Parameter[float] = Parameter(default=300.0)
     """The temperature, in K. Default = 300.0 K."""
 
-    ensemble: Parameter[Ensemble] = Parameter(default=Ensemble.NPT)
+    ensemble: Parameter[Literal["NVT", "NPT"]] = Parameter(default="NVT")
     """The ensemble to use. Default = NPT."""
 
     pressure: Parameter[float] = Parameter(default=1)

diff --git a/maize/steps/exs/biosimspace/equilibrate.py b/maize/steps/exs/biosimspace/equilibrate.py
@@ -4,7 +4,7 @@
 
 from abc import ABC
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import pytest
 
@@ -54,7 +54,7 @@ class _EquilibrateBase(_BioSimSpaceBase, ABC):
     temperature_end: Parameter[float] = Parameter(default=300.0)
     """The ending temperature, in K. Default = 300.0 K."""
 
-    ensemble: Parameter[Ensemble] = Parameter(default=Ensemble.NVT)
+    ensemble: Parameter[Literal["NVT", "NPT"]] = Parameter(default="NVT")
     """The ensemble to use for the equilibration. Default = NVT."""
 
     pressure: Parameter[float] = Parameter(default=1)
@@ -113,7 +113,7 @@ def _get_protocol(self) -> "BSS.Protocol._protocol.Protocol":
 
         pressure = (
             None
-            if self.ensemble.value == Ensemble.NVT
+            if Ensemble(self.ensemble.value) == Ensemble.NVT
             else self.pressure.value * BSS.Units.Pressure.atm
         )
 

diff --git a/maize/steps/exs/biosimspace/production.py b/maize/steps/exs/biosimspace/production.py
@@ -4,7 +4,7 @@
 
 from abc import ABC
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import pytest
 
@@ -64,7 +64,7 @@ class _ProductionBase(_BioSimSpaceBase, ABC):
     temperature: Parameter[float] = Parameter(default=300.0)
     """The temperature, in K. Default = 300.0 K."""
 
-    ensemble: Parameter[Ensemble] = Parameter(default=Ensemble.NPT)
+    ensemble: Parameter[Literal["NVT", "NPT"]] = Parameter(default="NVT")
     """The ensemble to use. Default = NPT."""
 
     pressure: Parameter[float] = Parameter(default=1)
@@ -127,7 +127,7 @@ def _get_protocol(self) -> "BSS.Protocol._protocol.Protocol":
 
         pressure = (
             None
-            if self.ensemble.value == Ensemble.NVT
+            if Ensemble(self.ensemble.value) == Ensemble.NVT
             else self.pressure.value * BSS.Units.Pressure.atm
         )