Merge branch 'main' into 94-Making-Outlines-Prefilter

mdagent/tools/base_tools/preprocess_tools/packing.py

@@ -67,7 +67,7 @@ def __init__(
     ):
         self.path_registry = path_registry
         self.molecules = []
-        self.file_number = 1
+        self.file_number = file_number
         self.file_description = file_description
         self.final_name = None
 
@@ -205,7 +205,6 @@ def packmol_wrapper(
 
 
 """Args schema for packmol_wrapper tool. Useful for OpenAI functions"""
-##TODO
 
 
 class PackmolInput(BaseModel):
@@ -369,15 +368,20 @@ def validate_input(cls, values: Union[str, Dict[str, Any]]) -> Dict:
                         "must be equal to the number of species in the system. "
                         f"You have {number_of_species} "
                         f"from {len(pdbfiles)} pdbfiles and {len(small_molecules)} "
-                        "small molecules"
+                        "small molecules. You have included "
+                        f"{len(number_of_molecules)} values for "
+                        f"number_of_molecules and {len(instructions)}"
+                        "instructions."
                     )
                 }
             return {
                 "error": (
                     "The length of number_of_molecules must be equal to the "
                     f"number of species in the system. You have {number_of_species} "
                     f"from {len(pdbfiles)} pdbfiles and {len(small_molecules)} "
-                    "small molecules"
+                    f"small molecules. You have included "
+                    f"{len(number_of_molecules)} values "
+                    "for number_of_molecules"
                 )
             }
         elif not number_of_species == len(instructions):
@@ -386,7 +390,8 @@ def validate_input(cls, values: Union[str, Dict[str, Any]]) -> Dict:
                     "The length of instructions must be equal to the "
                     f"number of species in the system. You have {number_of_species} "
                     f"from {len(pdbfiles)} pdbfiles and {len(small_molecules)} "
-                    "small molecules"
+                    "small molecules. You have included "
+                    f"{len(instructions)} instructions."
                 )
             }
         registry = PathRegistry.get_instance()
@@ -405,8 +410,9 @@ def validate_input(cls, values: Union[str, Dict[str, Any]]) -> Dict:
                 if len(instruction[0].split(" ")) == 1:
                     return {
                         "error": (
-                            "The instruction 'center' must be accompanied by more "
-                            "instructions. Example 'fixed 0. 0. 0. 0. 0. 0.' "
+                            "The instruction 'center' must be "
+                            "accompanied by more instructions. "
+                            "Example 'fixed 0. 0. 0. 0. 0. 0.' "
                             "The complete instruction would be: 'center \n fixed 0. 0. "
                             "0. 0. 0. 0.' with a newline separating the two "
                             "instructions."

mdagent/tools/base_tools/simulation_tools/setup_and_run.py

@@ -1311,6 +1311,8 @@ def _parse_cutoff(self, cutoff):
 
         # Convert to string in case it's not (e.g., int or float)
         cutoff = str(cutoff)
+        if cutoff[-1] == "s":
+            cutoff = cutoff[:-1]
 
         # Remove spaces and convert to lowercase for easier parsing
         cutoff = cutoff.replace(" ", "").lower()
@@ -1389,12 +1391,20 @@ def _parse_parameter(self, parameter, default_unit, possible_units):
         if "*" in parameter_str:
             num_part, unit_part = parameter_str.split("*")
             num_value = float(num_part)
+        elif "poundforce/inch^2" in parameter_str:
+            num_value = float(parameter_str.replace("poundforce/inch^2", ""))
+            unit_part = "poundforce/inch^2"
         # Check for division symbol and split if necessary
         # e.g. "1/ps" or "1/ps^-1"
         elif "/" in parameter_str:
             num_part, unit_part = parameter_str.split("/")
             num_value = float(num_part)
             unit_part = "/" + unit_part
+        elif "^-1" in parameter_str:
+            parameter_str = parameter_str.replace("^-1", "")
+            match = re.match(r"^(\d+(?:\.\d+)?)([a-zA-Z]+)$", parameter_str)
+            num_value = float(match.group(1))
+            unit_part = "/" + match.group(2)
         else:
             # Attempt to convert directly to float; if it fails,
             # it must have a unit like "K", "ps", etc.
@@ -1410,16 +1420,15 @@ def _parse_parameter(self, parameter, default_unit, possible_units):
                     error_msg += f"Invalid format for parameter: '{parameter_str}'."
 
         # Convert the unit part to an OpenMM unit
-        if unit_part in possible_units:
-            return num_value * possible_units[unit_part], error_msg
+        if unit_part.lower() in possible_units:
+            return num_value * possible_units[unit_part.lower()], error_msg
         else:
             # If the unit is not recognized, raise an error
             error_msg += f"""Unknown unit '{unit_part}' for parameter.
             Valid units include: {list(possible_units.keys())}."""
 
             return parameter, error_msg
 
-    # Example method to use _parse_parameter for specific parameter
     def parse_temperature(self, temperature):
         possible_units = {
             "k": unit.kelvin,
@@ -1461,7 +1470,7 @@ def parse_pressure(self, pressure):
             "atmosphere": unit.atmospheres,
             "pascal": unit.pascals,
             "pascals": unit.pascals,
-            "Pa": unit.pascals,
+            "pa": unit.pascals,
             "poundforce/inch^2": unit.psi,
             "psi": unit.psi,
         }

tests/test_packing.py

@@ -0,0 +1,166 @@
+import os
+
+import pytest
+
+from mdagent.tools.base_tools.preprocess_tools.packing import (
+    Molecule,
+    PackmolBox,
+    PackMolTool,
+)
+from mdagent.utils import PathRegistry
+
+
+@pytest.fixture
+def get_registry():
+    return PathRegistry()
+
+
+@pytest.fixture
+def packmolbox(get_registry):
+    return PackmolBox(get_registry)
+
+
+@pytest.fixture
+def packmoltool(get_registry):
+    return PackMolTool(get_registry)
+
+
+@pytest.fixture
+def dummy_molecule():
+    return Molecule(
+        filename="water", file_id=100, number_of_molecules=2, instructions=None
+    )
+
+
+@pytest.fixture
+def packmol_valid_input():
+    return {
+        "pdbfiles_id": ["3pqr_test"],
+        "small_molecules": ["water"],
+        "number_of_molecules": [1, 2],
+        "instructions": [
+            ["fixed 0. 0. 0. 0. 0. 0. centerofmass"],
+            ["inside box 0. 0. 0. 90. 90. 90."],
+        ],
+    }
+
+
+def test_packmol_add_molecule(packmolbox, dummy_molecule):
+    initial_length = len(packmolbox.molecules)
+    packmolbox.add_molecule(dummy_molecule)
+    assert len(packmolbox.molecules) == initial_length + 1
+
+
+def test_packmol_generate_input_header(packmolbox):
+    packmolbox.generate_input_header()
+    # assert file packmol.inp exists
+    assert os.path.isfile("packmol.inp")
+    os.remove("packmol.inp")
+
+
+def test_packmol_generate_input(packmolbox, dummy_molecule):
+    packmolbox.add_molecule(dummy_molecule)
+    output = packmolbox.generate_input()
+    assert "structure water" in output
+    assert "number 2" in output
+    assert "end structure" in output
+
+
+def test_packmol_validate_input_missing_info(
+    packmoltool, packmol_valid_input, get_registry
+):
+    example_input = packmol_valid_input
+    example_input["pdbfiles_id"] = []
+    input_valid = packmoltool.validate_input(example_input)
+    assert "error" in input_valid.keys()
+    assert (
+        "The length of number_of_molecules AND "
+        "instructions must be equal to the number "
+        "of species in the system. You have 1 from 0 "
+        "pdbfiles and 1 small molecules" in input_valid["error"]
+    )
+
+    example_input["pdbfiles_id"] = ["nonsense"]
+    input_valid = packmoltool.validate_input(example_input)
+    assert "error" in input_valid.keys()
+    assert (
+        input_valid["error"]
+        == "nonsense is not a valid pdbfile_id in the path_registry"
+    )
+
+    example_input["pdbfiles_id"] = ["3pqr_test"]
+    example_input["small_molecules"] = ["water", "urea"]
+    input_valid = packmoltool.validate_input(example_input)
+    assert "error" in input_valid.keys()
+    assert (
+        "The length of number_of_molecules AND "
+        "instructions must be equal to the number "
+        "of species in the system. You have 3 from 1 "
+        "pdbfiles and 2 small molecules" in input_valid["error"]
+    )
+
+    get_registry.map_path("3pqr_test", "3pqr.cif", "cif_test_file")
+    packmoltool = PackMolTool(get_registry)
+    example_input["pdbfiles_id"] = ["3pqr_test"]
+    example_input["small_molecules"] = ["nonsense"]
+    input_valid = packmoltool.validate_input(example_input)
+    assert "error" in input_valid.keys()
+    assert "nonsense could not be converted to a pdb file" in input_valid["error"]
+
+
+def test_pacmol_validate_input_instruction_fail(packmoltool, packmol_valid_input):
+    example_input = packmol_valid_input
+    example_input["small_molecules"] = ["water"]
+    example_input["instructions"] = [
+        ["fail 0. 0. 0. 0. 0. 0. centerofmass"],
+        ["inside box 0. 0. 0. 90. 90. 90."],
+    ]
+    input_valid = packmoltool.validate_input(example_input)
+    assert "error" in input_valid.keys()
+    assert "The first word of each instruction must be one of" in input_valid["error"]
+
+    example_input["instructions"] = [["center"], [["inside box 0. 0. 0. 90. 90. 90."]]]
+    input_valid = packmoltool.validate_input(example_input)
+    assert "error" in input_valid.keys()
+    assert (
+        "The instruction 'center' must be accompanied by more instructions"
+        in input_valid["error"]
+    )
+
+    example_input = packmol_valid_input
+    example_input["instructions"] = [
+        example_input["instructions"],
+        example_input["instructions"],
+    ]
+    assert len(example_input["instructions"]) == 2
+    input_valid = packmoltool.validate_input(example_input)
+    assert "error" in input_valid.keys()
+    assert "Each instruction must be a single string" in input_valid["error"]
+
+
+def test_packmol_validate_input_valid(get_registry):
+    get_registry.map_path("3pqr_test", "3pqr.cif", "cif_test_file")
+    get_registry.map_path("water", "water.pdb", "fake_water_test_file")
+    packmoltool = PackMolTool(get_registry)
+    example_input = {
+        "pdbfiles_id": ["3pqr_test"],
+        "small_molecules": ["water"],
+        "number_of_molecules": [1, 2],
+        "instructions": [
+            ["fixed 0. 0. 0. 0. 0. 0. centerofmass"],
+            ["inside box 0. 0. 0. 90. 90. 90."],
+        ],
+    }
+    input_valid = packmoltool.validate_input(example_input)
+    assert input_valid == example_input
+
+    example_input["small_molecules"] = ["water", "urea"]
+    get_registry.map_path("urea", "urea.pdb", "fake_urea_test_file")
+    example_input["instructions"] = [
+        example_input["instructions"][0],
+        example_input["instructions"][0],
+        example_input["instructions"][0],
+    ]
+    example_input["number_of_molecules"] = [1, 2, 2]
+    input_valid = packmoltool.validate_input(example_input)
+    assert input_valid == example_input