Merge pull request #28 from JR-1991/fix-optimization

Fix `lmfit` optimization function

Author: JR-1991

catalax/dataset/measurement.py

@@ -295,6 +295,11 @@ def to_jax_arrays(
         Raises:
             ValueError: If any state in state_order is missing from the measurement data.
         """
+        if len(self.data) == 0:
+            raise ValueError(
+                f"The measurement data is empty. Please add data to the measurement {self.id}"
+            )
+
         unused_states = [
             state for state in self.data.keys() if state not in state_order
         ]
@@ -307,7 +312,7 @@ def to_jax_arrays(
 
         if missing_states:
             raise ValueError(
-                f"The measurement state are inconsistent with the dataset state. "
+                f"The measurement states are inconsistent with the dataset states."
                 f"Missing {missing_states} in measurement {self.id}"
             )
 

catalax/tools/optimization.py

@@ -50,16 +50,15 @@ def optimize(
     """
 
     params = _initialize_params(model, global_upper_bound, global_lower_bound)
-    observables = jnp.array(
-        [
-            index
-            for index, ode in enumerate(model.odes.values())
-            if ode.observable is True
-        ]
-    )
+    dataset_observables = dataset.get_observable_states_order()
+    observables = [
+        index
+        for index, state in enumerate(model.get_state_order())
+        if state in dataset_observables
+    ]
 
     # Extract data arrays for the residual computation
-    data, times, _ = dataset.to_jax_arrays(model.get_observable_state_order())
+    data, times, _ = dataset.to_jax_arrays(model.get_state_order())
 
     # Create simulation config from dataset
     config = dataset.to_config()

tests/integration/test_optimization.py

@@ -0,0 +1,34 @@
+import catalax as ctx
+
+
+class TestOptimization:
+    def test_optimize_runs(self):
+        # Create a Michaelis-Menten enzyme kinetics model
+        model = ctx.Model(name="Enzyme Kinetics")
+
+        # Define states
+        model.add_states(S="Substrate", P="Product")
+
+        # Add reaction kinetics via schemes
+        model.add_reaction(
+            "S -> P",
+            symbol="r1",
+            equation="v_max * S / (K_m + S)",
+        )
+
+        # Set parameter values
+        model.parameters["v_max"].value = 0.2
+        model.parameters["v_max"].initial_value = 0.1
+        model.parameters["K_m"].value = 0.1
+        model.parameters["K_m"].initial_value = 0.05
+
+        # Create dataset and simulate
+        dataset = ctx.Dataset.from_model(model)
+        dataset.add_initial(S=1.0, P=0.0)
+        config = ctx.SimulationConfig(t1=10, nsteps=10)
+        dataset = model.simulate(dataset=dataset, config=config)
+
+        # Run optimization and check that it completes and returns expected types
+        result, new_model = ctx.optimize(model, dataset)
+        assert hasattr(result, "params")
+        assert isinstance(new_model, type(model))