Update petab.py

AMICI-dev · Nov 16, 2024 · 94aa679 · 94aa679
1 parent d79cfc1
commit 94aa679
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/python/sdist/amici/jax/petab.py b/python/sdist/amici/jax/petab.py
@@ -23,7 +23,7 @@
     ParameterMappingForCondition,
     create_parameter_mapping,
 )
-from amici.jax.model import JAXModel, simulate_condition
+from amici.jax.model import JAXModel
 
 
 def jax_unscale(
@@ -239,10 +239,9 @@ def run_simulation(
             if len(simulation_condition) > 1
             else jnp.array([])
         )
-        return simulate_condition(
+        return self.model.simulate_condition(
             p,
             p_preeq,
-            self.model,
             jax.lax.stop_gradient(jnp.array(ts_preeq)),
             jax.lax.stop_gradient(jnp.array(ts_dyn)),
             jax.lax.stop_gradient(jnp.array(ts_posteq)),