refactor: extract FMU definition builder and simplify get_definition

modelon/impact/client/entities/_initialize_from.py

@@ -10,11 +10,11 @@
     from modelon.impact.client.sal.service import Service
 
 
-def _resolve_initialize_from(
+def _resolve_extension_initialize_from(
     workspace_id: str,
     sal: Service,
     modifiers: Dict[str, Any],
-) -> Optional[Union[Case, Experiment, ExternalResult]]:
+) -> Optional[Union[Case, Experiment]]:
     if "initializeFrom" in modifiers:
         from modelon.impact.client.entities.experiment import Experiment
 
@@ -27,8 +27,16 @@ def _resolve_initialize_from(
         case_id = modifiers["initializeFromCase"]["caseId"]
         case_data = sal.experiment.case_get(workspace_id, exp_id, case_id)
         return Case(case_data["id"], workspace_id, exp_id, sal, case_data)
-    elif "initializeFromExternalResult" in modifiers:
+    return None
+
+
+def _resolve_initialize_from(
+    workspace_id: str,
+    sal: Service,
+    modifiers: Dict[str, Any],
+) -> Optional[Union[Case, Experiment, ExternalResult]]:
+    if "initializeFromExternalResult" in modifiers:
         return ExternalResult(
             result_id=modifiers["initializeFromExternalResult"], service=sal
         )
-    return None
+    return _resolve_extension_initialize_from(workspace_id, sal, modifiers)

modelon/impact/client/entities/experiment.py

@@ -6,14 +6,16 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
-from modelon.impact.client.entities._initialize_from import _resolve_initialize_from
+from modelon.impact.client.entities._initialize_from import (
+    _resolve_extension_initialize_from,
+    _resolve_initialize_from,
+)
 from modelon.impact.client.entities.asserts import assert_variable_in_result
 from modelon.impact.client.entities.case import Case
 from modelon.impact.client.entities.custom_function import (
     CustomFunction,
     _build_custom_function,
 )
-from modelon.impact.client.entities.external_result import ExternalResult
 from modelon.impact.client.entities.interfaces.experiment import ExperimentReference
 from modelon.impact.client.entities.model import (
     Model,
@@ -194,6 +196,62 @@ def label(self) -> Optional[str]:
         return self._meta_data.get("label")
 
 
+def _build_extensions_from_dict(
+    extensions: List[Dict[str, Any]],
+    workspace_id: str,
+    sal: "Service",
+) -> List[SimpleExperimentExtension]:
+    sim_exts = []
+    for extension in extensions:
+        analysis = extension.get("analysis", {})
+        ext_custom_function_params = {
+            param["name"]: param["value"] for param in analysis.get("parameters", [])
+        }
+        sim_ext = SimpleExperimentExtension(
+            parameter_modifiers=ext_custom_function_params,
+            solver_options=analysis.get("solverOptions"),
+            simulation_options=analysis.get("simulationOptions"),
+            simulation_log_level=analysis.get("simulationLogLevel"),
+            initialize_from=_resolve_extension_initialize_from(
+                workspace_id, sal, extension.get("modifiers", {})
+            ),
+        )
+        ext_modifiers = {
+            mod["name"]: to_domain_parameter_value(mod)
+            for mod in extension.get("modifiers", {}).get("variables", [])
+        }
+        sim_ext = sim_ext.with_modifiers(modifiers=ext_modifiers)
+        case_data = extension.get("caseData", [])
+        case_labels = [data.get("label") for data in case_data]
+        if case_labels:
+            sim_ext = sim_ext.with_case_label(case_labels[0])
+        sim_exts.append(sim_ext)
+    return sim_exts
+
+
+def _build_simple_fmu_experiment_definition(
+    base: Dict[str, Any],
+    custom_function: "CustomFunction",
+    workspace_id: str,
+    sal: "Service",
+) -> SimpleFMUExperimentDefinition:
+    analysis = base["analysis"]
+    modifiers = base.get("modifiers", {})
+    fmu_id = base["model"]["fmu"]["id"]
+    variable_modifiers = {
+        mod["name"]: get_operator_from_dict(mod)
+        for mod in modifiers.get("variables", [])
+    }
+    return SimpleFMUExperimentDefinition(
+        fmu=ModelExecutable(workspace_id, fmu_id, sal),
+        custom_function=custom_function,
+        solver_options=analysis.get("solverOptions", {}),
+        simulation_options=analysis.get("simulationOptions", {}),
+        simulation_log_level=analysis.get("simulationLogLevel", "WARNING"),
+        initialize_from=_resolve_initialize_from(workspace_id, sal, modifiers),
+    ).with_modifiers(modifiers=variable_modifiers)
+
+
 class Experiment(ExperimentReference):
     """Class containing Experiment functionalities."""
 
@@ -613,35 +671,6 @@ def get_solver_options(self) -> SolverOptions:
         analysis = self._get_info()["experiment"]["base"]["analysis"]
         return SolverOptions(analysis.get("solverOptions", {}), self.custom_function)
 
-    def _get_initialize_from_case(self, experiment_id: str, case_id: str) -> Case:
-        case_data = self._sal.experiment.case_get(
-            self._workspace_id, experiment_id, case_id
-        )
-        return Case(
-            case_data["id"], self._workspace_id, experiment_id, self._sal, case_data
-        )
-
-    def _get_initialize_from_experiment(self, experiment_id: str) -> Experiment:
-        resp = self._sal.workspace.experiment_get(self._workspace_id, experiment_id)
-        return Experiment(self._workspace_id, resp["id"], self._sal, resp)
-
-    def _get_initialize_from(
-        self, modifiers: Dict[str, Any]
-    ) -> Optional[Union[Case, Experiment, ExternalResult]]:
-        return _resolve_initialize_from(self._workspace_id, self._sal, modifiers)
-
-    def _get_extension_initialize_from(
-        self, modifiers: Dict[str, Any]
-    ) -> Optional[Union[Case, Experiment]]:
-        if "initializeFrom" in modifiers:
-            exp_id = modifiers["initializeFrom"]
-            return self._get_initialize_from_experiment(exp_id)
-        elif "initializeFromCase" in modifiers:
-            case_id = modifiers["initializeFromCase"]["caseId"]
-            exp_id = modifiers["initializeFromCase"]["experimentId"]
-            return self._get_initialize_from_case(exp_id, case_id)
-        return None
-
     def get_definition(self) -> ValidExperimentDefinitions:
         """Get an experiment definition that can be used to reproduce this experiment
         result.
@@ -655,9 +684,12 @@ def get_definition(self) -> ValidExperimentDefinitions:
             definition = experiment.get_definition()
 
         """
-        base = self._get_info(cached=False)["experiment"]["base"]
+        info = self._get_info(cached=False)["experiment"]
+        base = info["base"]
+        extensions_data = info.get("extensions", [])
         analysis = base["analysis"]
         custom_function = self._get_custom_function(analysis)
+        definition: ValidExperimentDefinitions
         if self._get_workflow() == _Workflow.CLASS_BASED:
             model = Model(
                 self.get_model_name(),
@@ -669,52 +701,15 @@ def get_definition(self) -> ValidExperimentDefinitions:
                 model, base, custom_function, self._workspace_id, self._sal
             )
         else:
-            fmu_id = base["model"]["fmu"]["id"]
-            definition = SimpleFMUExperimentDefinition(
-                fmu=ModelExecutable(self._workspace_id, fmu_id, self._sal),
-                custom_function=custom_function,
-                solver_options=self.get_solver_options(),
-                simulation_options=self.get_simulation_options(),
-                simulation_log_level=analysis["simulationLogLevel"],
-                initialize_from=self._get_initialize_from(base["modifiers"]),
-            )  # type: ignore
-            modifiers = {
-                mod["name"]: get_operator_from_dict(mod)
-                for mod in base["modifiers"]["variables"]
-            }
-            definition = definition.with_modifiers(modifiers=modifiers)
-        extensions = self._get_info()["experiment"].get("extensions")
-        if extensions:
-            sim_exts = []
-            for extension in extensions:
-                analysis = extension.get("analysis", {})
-                ext_custom_function_params = {
-                    param["name"]: param["value"]
-                    for param in analysis.get("parameters", [])
-                }
-                sim_ext = SimpleExperimentExtension(
-                    parameter_modifiers=ext_custom_function_params,
-                    solver_options=analysis.get("solverOptions"),
-                    simulation_options=analysis.get("simulationOptions"),
-                    simulation_log_level=analysis.get("simulationLogLevel"),
-                    initialize_from=self._get_extension_initialize_from(
-                        extension["modifiers"]
-                    )
-                    if extension.get("modifiers")
-                    else None,
+            definition = _build_simple_fmu_experiment_definition(
+                base, custom_function, self._workspace_id, self._sal
+            )
+        if extensions_data:
+            definition = definition.with_extensions(
+                _build_extensions_from_dict(
+                    extensions_data, self._workspace_id, self._sal
                 )
-                ext_modifiers = {
-                    mod["name"]: to_domain_parameter_value(mod)
-                    for mod in extension.get("modifiers", {}).get("variables", [])
-                }
-                sim_ext = sim_ext.with_modifiers(modifiers=ext_modifiers)
-                case_data = extension.get("caseData", [])
-                case_labels = [data.get("label") for data in case_data]
-                if case_labels:
-                    case_label = case_labels[0]
-                    sim_ext = sim_ext.with_case_label(case_label)
-                sim_exts.append(sim_ext)
-            definition = definition.with_extensions(sim_exts)
+            )
         return definition
 
     def _get_custom_function(self, analysis: Dict[str, Any]) -> CustomFunction:

modelon/impact/client/experiment_definition/expansion.py

@@ -113,9 +113,9 @@ def get_parameters_as_dict(self) -> Optional[Dict[str, Any]]:
         return asdict(self)
 
 
-def expansion_from_dict(
-    algorithm: str, parameters: Dict[str, Any]
-) -> ExpansionAlgorithm:
+def expansion_from_dict(expansion_dict: Dict[str, Any]) -> ExpansionAlgorithm:
+    algorithm = expansion_dict.get("algorithm", "")
+    parameters = expansion_dict.get("parameters", {})
     if algorithm == "SOBOL":
         return Sobol(samples=parameters["samples"])
     elif algorithm == "LATINHYPERCUBE":

modelon/impact/client/experiment_definition/model_based.py

@@ -78,10 +78,7 @@ def _build_simple_modelica_experiment_definition(
         mod["name"]: get_operator_from_dict(mod)
         for mod in base.get("modifiers", {}).get("variables", [])
     }
-    expansion_dict = base.get("expansion", {})
-    expansion = expansion_from_dict(
-        expansion_dict.get("algorithm", ""), expansion_dict.get("parameters", {})
-    )
+    expansion = expansion_from_dict(base.get("expansion", {}))
     return (
         SimpleModelicaExperimentDefinition(
             model=model,