Update optimization observable selection and add integration test

Refactored observable state selection in optimize() to use dataset observables and model state order for consistency. Added an integration test for the optimization workflow using a Michaelis-Menten model to verify correct execution and output types.

Author: JR-1991

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))