fixups

AMICI-dev · Dec 6, 2024 · 4eb18d9 · 4eb18d9
1 parent 7b5383c
commit 4eb18d9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/python/sdist/amici/jax/petab.py b/python/sdist/amici/jax/petab.py
@@ -83,7 +83,7 @@ class JAXProblem(eqx.Module):
             np.ndarray,
         ],
     ]
-    _petab_measurement_indices: dict[tuple[str, ...], pd.Index]
+    _petab_measurement_indices: dict[tuple[str, ...], tuple[int, ...]]
     _petab_problem: petab.Problem
 
     def __init__(self, model: JAXModel, petab_problem: petab.Problem):
@@ -247,7 +247,7 @@ def _get_measurements(
                 iys,
                 iy_trafos,
             )
-            indices[tuple(simulation_condition)] = index
+            indices[tuple(simulation_condition)] = tuple(index.tolist())
         return measurements, indices
 
     def get_all_simulation_conditions(self) -> tuple[tuple[str, ...], ...]:

diff --git a/python/tests/test_jax.py b/python/tests/test_jax.py
@@ -177,6 +177,7 @@ def check_fields_jax(
     my = my.flatten()
     ts = ts.flatten()
     iys = iys.flatten()
+    iy_trafos = np.zeros_like(iys)
 
     ts_init = ts[ts == 0]
     ts_dyn = ts[ts > 0]
@@ -194,6 +195,7 @@ def check_fields_jax(
         "ts_posteq": jnp.array(ts_posteq),
         "my": jnp.array(my),
         "iys": jnp.array(iys),
+        "iy_trafos": jnp.array(iy_trafos),
         "x_preeq": jnp.array([]),
         "solver": diffrax.Kvaerno5(),
         "controller": diffrax.PIDController(atol=ATOL_SIM, rtol=RTOL_SIM),