Merge pull request #119 from slaclab/fix-gp-input

Adjust inputs order for prob models/gps, add mlflow to main requs

Author: pluflou

README.md

@@ -8,6 +8,7 @@ LUME-model holds data structures used in the LUME modeling toolset. Variables an
 * pydantic
 * numpy
 * pyyaml
+* mlflow
 
 ## Install
 

dev-environment.yml

@@ -8,6 +8,7 @@ dependencies:
   - numpy
   - pyyaml
   - botorch
+  - mlflow
 
   # dev requirements
   - pytest

environment.yml

@@ -7,3 +7,4 @@ dependencies:
   - pydantic>2.3
   - numpy
   - pyyaml
+  - mlflow

lume_model/models/gp_model.py

@@ -142,6 +142,8 @@ def _get_predictions(
         Returns:
             Dictionary of output variable names to distributions.
         """
+        # Reorder the input dictionary to match the model's input order
+        input_dict = super()._arrange_inputs(input_dict)
         # Create tensor from input_dict
         x = super()._create_tensor_from_dict(input_dict)
         # Transform the input

lume_model/models/prob_model_base.py

@@ -58,6 +58,19 @@ def dtype(
                 f"expected one of ['double', 'single']."
             )
 
+    def _arrange_inputs(
+        self, d: dict[str, Union[float, torch.Tensor]]
+    ) -> dict[str, Union[float, torch.Tensor]]:
+        """Enforces order of input variables before creating a tensor.
+
+        Args:
+            d: Dictionary of input variable names to tensors.
+
+        Returns:
+            Ordered input tensor.
+        """
+        return {k: d[k] for k in self.input_names}
+
     @staticmethod
     def _create_tensor_from_dict(
         d: dict[str, Union[float, torch.Tensor]],

pyproject.toml

@@ -23,7 +23,8 @@ requires-python = ">=3.10"
 dependencies = [
     "pydantic",
     "numpy",
-    "pyyaml"
+    "pyyaml",
+    "mlflow"
 ]
 dynamic = ["version"]
 [tool.setuptools_scm]
@@ -34,8 +35,7 @@ dev = [
     "botorch",
     "torch",
     "pre-commit",
-    "pytest",
-    "mlflow"
+    "pytest"
 ]
 docs = [
     "mkdocs",

requirements.txt

@@ -1,3 +1,4 @@
 pydantic
 numpy
 pyyaml
+mlflow