Move petab simulation code, private petab helpers

python/sdist/amici/adapters/fiddy.py

@@ -20,7 +20,6 @@
 from fiddy import CachedFunction, Type, fiddy_array
 from petab.v1.C import LIN, LOG, LOG10
 
-from amici.importers.petab.v1.parameter_mapping import create_parameter_mapping
 from amici.sim.sundials import (
     AmiciExpData,
     AmiciModel,
@@ -31,6 +30,9 @@
     run_simulation,
 )
 from amici.sim.sundials.petab.v1 import LLH, SLLH, create_edatas
+from amici.sim.sundials.petab.v1._parameter_mapping import (
+    create_parameter_mapping,
+)
 
 if TYPE_CHECKING:
     from amici.sim.sundials.petab import PetabSimulationResult, PetabSimulator

python/sdist/amici/importers/petab/_cli/import_petab.py

@@ -1,7 +1,7 @@
 import argparse
 
 import petab.v1 as petab
-from amici.importers.petab.v1.petab_import import import_model_sbml
+from amici.importers.petab.v1._petab_import import import_model_sbml
 
 
 def _parse_cli_args():

python/sdist/amici/importers/petab/_petab_importer.py

@@ -26,7 +26,7 @@
 from amici.logging import get_logger
 from amici.sim.jax.petab import JAXProblem
 
-from .v1.sbml_import import _add_global_parameter
+from .v1._sbml_import import _add_global_parameter
 
 if TYPE_CHECKING:
     import pysb
@@ -329,7 +329,7 @@ def _do_import_sbml(self):
             for change in self.petab_problem[condition_id].changes
         }
 
-        from .v1.sbml_import import show_model_info
+        from .v1._sbml_import import show_model_info
 
         show_model_info(self.petab_problem.model.sbml_model)
         sbml_importer = amici.SbmlImporter(

python/sdist/amici/importers/petab/v1/__init__.py

@@ -4,7 +4,7 @@
 #  preequilibration
 PREEQ_INDICATOR_ID = "preequilibration_indicator"
 
-from .petab_import import import_petab_problem
+from ._petab_import import import_petab_problem
 
 __all__ = [
     "import_petab_problem",

python/sdist/amici/importers/petab/v1/petab_import.py → python/sdist/amici/importers/petab/v1/_petab_import.py

@@ -18,16 +18,16 @@
 import amici
 from amici.logging import get_logger
 
-from .import_helpers import (
+from ._import_helpers import (
     _can_import_model,
     _create_model_name,
     _get_package_name_and_path,
     check_model,
 )
-from .sbml_import import import_model_sbml
+from ._sbml_import import import_model_sbml
 
 try:
-    from .pysb_import import import_model_pysb
+    from ._pysb_import import import_model_pysb
 except ModuleNotFoundError:
     # pysb not available
     import_model_pysb = None

python/sdist/amici/importers/petab/v1/pysb_import.py → python/sdist/amici/importers/petab/v1/_pysb_import.py

@@ -24,11 +24,11 @@
 from amici.logging import get_logger, log_execution_time, set_log_level
 
 from . import PREEQ_INDICATOR_ID
-from .import_helpers import (
+from ._import_helpers import (
     get_fixed_parameters,
     petab_noise_distribution_to_amici,
 )
-from .util import get_states_in_condition_table
+from ._util import get_states_in_condition_table
 
 logger = get_logger(__name__, logging.WARNING)
 

python/sdist/amici/importers/petab/v1/sbml_import.py → python/sdist/amici/importers/petab/v1/_sbml_import.py

@@ -16,12 +16,12 @@
 from amici.logging import log_execution_time, set_log_level
 
 from . import PREEQ_INDICATOR_ID
-from .import_helpers import (
+from ._import_helpers import (
     check_model,
     get_fixed_parameters,
     get_observation_model,
 )
-from .util import get_states_in_condition_table
+from ._util import get_states_in_condition_table
 
 logger = logging.getLogger(__name__)
 

python/sdist/amici/sim/__init__.py

@@ -1 +1,3 @@
 """Functionality for simulating AMICI models."""
+
+from __future__ import annotations

python/sdist/amici/sim/_parameter_mapping.py

@@ -0,0 +1,148 @@
+"""Parameter mapping for conditions.
+
+Only code that is independent of JAX or SUNDIALS objects.
+"""
+
+from __future__ import annotations
+
+import numbers
+from collections.abc import Sequence
+from itertools import chain
+
+from petab.v1 import LIN
+
+SingleParameterMapping = dict[str, numbers.Number | str]
+SingleScaleMapping = dict[str, str]
+
+
+class ParameterMappingForCondition:
+    """Parameter mapping for condition.
+
+    Contains mappings for free parameters, fixed parameters, and fixed
+    pre-equilibration parameters, both for parameters and scales.
+
+    In the scale mappings, for each simulation parameter, the scale
+    on which the value is passed (and potentially gradients are to be
+    returned) is given. In the parameter mappings, for each simulation
+    parameter a corresponding optimization parameter (or a numeric value)
+    is given.
+
+    If a mapping is not passed, the parameter mappings are assumed to be empty,
+    and if a scale mapping is not passed, all scales are set to linear.
+
+    :param map_sim_var:
+        Mapping for free simulation parameters.
+    :param scale_map_sim_var:
+        Scales for free simulation parameters.
+    :param map_preeq_fix:
+        Mapping for fixed pre-equilibration parameters.
+    :param scale_map_preeq_fix:
+        Scales for fixed pre-equilibration parameters.
+    :param map_sim_fix:
+        Mapping for fixed simulation parameters.
+    :param scale_map_sim_fix:
+        Scales for fixed simulation parameters.
+    """
+
+    def __init__(
+        self,
+        map_sim_var: SingleParameterMapping = None,
+        scale_map_sim_var: SingleScaleMapping = None,
+        map_preeq_fix: SingleParameterMapping = None,
+        scale_map_preeq_fix: SingleScaleMapping = None,
+        map_sim_fix: SingleParameterMapping = None,
+        scale_map_sim_fix: SingleScaleMapping = None,
+    ):
+        if map_sim_var is None:
+            map_sim_var = {}
+        self.map_sim_var = map_sim_var
+
+        if scale_map_sim_var is None:
+            scale_map_sim_var = {key: LIN for key in map_sim_var}
+        self.scale_map_sim_var = scale_map_sim_var
+
+        if map_preeq_fix is None:
+            map_preeq_fix = {}
+        self.map_preeq_fix = map_preeq_fix
+
+        if scale_map_preeq_fix is None:
+            scale_map_preeq_fix = {key: LIN for key in map_preeq_fix}
+        self.scale_map_preeq_fix = scale_map_preeq_fix
+
+        if map_sim_fix is None:
+            map_sim_fix = {}
+        self.map_sim_fix = map_sim_fix
+
+        if scale_map_sim_fix is None:
+            scale_map_sim_fix = {key: LIN for key in map_sim_fix}
+        self.scale_map_sim_fix = scale_map_sim_fix
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}("
+            f"map_sim_var={repr(self.map_sim_var)},"
+            f"scale_map_sim_var={repr(self.scale_map_sim_var)},"
+            f"map_preeq_fix={repr(self.map_preeq_fix)},"
+            f"scale_map_preeq_fix={repr(self.scale_map_preeq_fix)},"
+            f"map_sim_fix={repr(self.map_sim_fix)},"
+            f"scale_map_sim_fix={repr(self.scale_map_sim_fix)})"
+        )
+
+    @property
+    def free_symbols(self) -> set[str]:
+        """Get IDs of all (symbolic) parameters present in this mapping"""
+        return {
+            p
+            for p in chain(
+                self.map_sim_var.values(),
+                self.map_preeq_fix.values(),
+                self.map_sim_fix.values(),
+            )
+            if isinstance(p, str)
+        }
+
+
+class ParameterMapping(Sequence):
+    r"""Parameter mapping for multiple conditions.
+
+    This can be used like a list of :class:`ParameterMappingForCondition`\ s.
+
+    :param parameter_mappings:
+        List of parameter mappings for specific conditions.
+    """
+
+    def __init__(
+        self, parameter_mappings: list[ParameterMappingForCondition] = None
+    ):
+        super().__init__()
+        if parameter_mappings is None:
+            parameter_mappings = []
+        self.parameter_mappings = parameter_mappings
+
+    def __iter__(self):
+        yield from self.parameter_mappings
+
+    def __getitem__(
+        self, item
+    ) -> ParameterMapping | ParameterMappingForCondition:
+        result = self.parameter_mappings[item]
+        if isinstance(result, ParameterMappingForCondition):
+            return result
+        return ParameterMapping(result)
+
+    def __len__(self):
+        return len(self.parameter_mappings)
+
+    def append(
+        self, parameter_mapping_for_condition: ParameterMappingForCondition
+    ):
+        """Append a condition-specific parameter mapping."""
+        self.parameter_mappings.append(parameter_mapping_for_condition)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({repr(self.parameter_mappings)})"
+
+    @property
+    def free_symbols(self) -> set[str]:
+        """Get IDs of all (symbolic) parameters present in this mapping"""
+        return set.union(*(mapping.free_symbols for mapping in self))

python/sdist/amici/sim/jax/petab.py

@@ -22,10 +22,8 @@
 from optimistix import AbstractRootFinder
 
 from amici import _module_from_path
-from amici.importers.petab.v1.parameter_mapping import (
-    ParameterMappingForCondition,
-)
 from amici.logging import get_logger
+from amici.sim._parameter_mapping import ParameterMappingForCondition
 from amici.sim.jax.model import JAXModel, ReturnValue
 
 DEFAULT_CONTROLLER_SETTINGS = {

python/sdist/amici/sim/sundials/petab/v1/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from ._conditions import (
     create_edatas,
     create_parameterized_edatas,

python/sdist/amici/sim/sundials/petab/v1/_conditions.py

@@ -19,19 +19,21 @@
     TIME,
 )
 
-from amici.importers.petab.v1.parameter_mapping import (
+from amici.importers.petab.v1._util import get_states_in_condition_table
+from amici.sim._parameter_mapping import (
     ParameterMapping,
     ParameterMappingForCondition,
-    petab_to_amici_scale,
-    scale_parameters_dict,
-    unscale_parameters_dict,
 )
-from amici.importers.petab.v1.util import get_states_in_condition_table
 from amici.sim.sundials import (
     AmiciModel,
     ExpData,
     parameter_scaling_from_int_vector,
 )
+from amici.sim.sundials.petab.v1._parameter_scaling import (
+    petab_to_amici_scale,
+    scale_parameters_dict,
+    unscale_parameters_dict,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -283,7 +285,7 @@ def create_parameterized_edatas(
 
     # Get parameter mapping
     if parameter_mapping is None:
-        from amici.importers.petab.v1.parameter_mapping import (
+        from amici.sim.sundials.petab.v1._parameter_mapping import (
             create_parameter_mapping,
         )