fix doc typehints

AMICI-dev · Dec 5, 2024 · 7a780ad · 7a780ad
1 parent 08afd8b
commit 7a780ad
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/python/sdist/amici/jax/petab.py b/python/sdist/amici/jax/petab.py
@@ -154,7 +154,7 @@ def _get_parameter_mappings(
     def _get_measurements(
         self, simulation_conditions: pd.DataFrame
     ) -> dict[
-        str,
+        tuple[str, ...],
         tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray],
     ]:
         """
@@ -320,7 +320,7 @@ def run_simulation(
             Step size controller to use for simulation
         :param max_steps:
             Maximum number of steps to take during simulation
-        :param preeq:
+        :param x_preeq:
             Pre-equilibration state if available
         :return:
             Tuple of log-likelihood and simulation statistics
@@ -375,7 +375,7 @@ def run_preequilibration(
 
 def run_simulations(
     problem: JAXProblem,
-    simulation_conditions: Iterable[tuple] | None = None,
+    simulation_conditions: Iterable[tuple[str, ...]] | None = None,
     solver: diffrax.AbstractSolver = diffrax.Kvaerno5(),
     controller: diffrax.AbstractStepSizeController = diffrax.PIDController(
         rtol=1e-8,