From 67c4d79189a34689df9a7048c831aa8248fe4b01 Mon Sep 17 00:00:00 2001
From: Dilan Pathirana <59329744+dilpath@users.noreply.github.com>
Date: Mon, 6 Jan 2025 17:22:37 +0100
Subject: [PATCH] Store all calibrated models in `Problem.state`; update all
 notebooks (#134)

* fix type error list[Model] -> Models

* support provision of subspace petab problem to `Model.from_yaml`. backwards compatibility: add `Model.get_hash`, support `predecessor_model_hash=None` input

* backwards compatibility: support `Models` construction from single model YAML

* add all iteration calibrated models to problem state in `end_iteration`

* add `Problem._state`, currently just stores all calibrated models from all iterations

* workflow_python notebook: demo `Models.df`; demo `Problem.state.models`

* start tracking iteration in `State`

* Update test_cases/0009/README.md

---------

Co-authored-by: Daniel Weindl <dweindl@users.noreply.github.com>
---
 doc/examples/README.md                        |  1 +
 .../calibrated_models/calibrated_models.yaml  | 16 ++---
 doc/examples/example_cli_famos.ipynb          |  3 +-
 .../example_cli_famos_calibration_tool.py     |  2 +-
 .../model_selection/calibrated_M1_4.yaml      | 16 ++---
 .../model_selection/calibrated_M1_7.yaml      | 16 ++---
 .../model_selection/calibrated_models_1.yaml  | 55 +++++++-------
 doc/examples/visualization.ipynb              |  8 +--
 doc/examples/workflow_python.ipynb            | 53 +++++++++++---
 doc/standard/model.yaml                       |  5 +-
 doc/standard/models.yaml                      | 20 +-----
 doc/standard/problem.yaml                     | 12 +---
 petab_select/candidate_space.py               |  2 +-
 petab_select/cli.py                           |  1 +
 petab_select/model.py                         | 72 ++++++++++++-------
 petab_select/models.py                        | 50 ++++++-------
 petab_select/plot.py                          | 20 +++---
 petab_select/problem.py                       | 69 ++++++++++++------
 petab_select/ui.py                            |  9 ++-
 pyproject.toml                                |  2 +
 test/candidate_space/test_famos.py            | 10 +--
 test/cli/input/models.yaml                    | 10 ---
 test/ui/test_ui.py                            |  2 +
 23 files changed, 253 insertions(+), 201 deletions(-)
 create mode 100644 doc/examples/README.md

diff --git a/doc/examples/README.md b/doc/examples/README.md
new file mode 100644
index 0000000..85a3b2d
--- /dev/null
+++ b/doc/examples/README.md
@@ -0,0 +1 @@
+These notebooks need to be run to see the output. Pre-computed output can be viewed in the documentation, at https://petab-select.readthedocs.io/en/stable/examples.html
diff --git a/doc/examples/calibrated_models/calibrated_models.yaml b/doc/examples/calibrated_models/calibrated_models.yaml
index 8b96baf..53abe83 100644
--- a/doc/examples/calibrated_models/calibrated_models.yaml
+++ b/doc/examples/calibrated_models/calibrated_models.yaml
@@ -15,8 +15,8 @@
     k1: 0
     k2: 0
     k3: 0
-  petab_yaml: petab_problem.yaml
-  predecessor_model_hash: virtual_initial_model
+  model_subspace_petab_yaml: petab_problem.yaml
+  predecessor_model_hash: virtual_initial_model-
 - criteria:
     AICc: -0.17540608110890332
     NLLH: -4.087703040554452
@@ -35,7 +35,7 @@
     k1: 0.2
     k2: 0.1
     k3: estimate
-  petab_yaml: petab_problem.yaml
+  model_subspace_petab_yaml: petab_problem.yaml
   predecessor_model_hash: M_0-000
 - criteria:
     AICc: -0.27451438069575573
@@ -55,7 +55,7 @@
     k1: 0.2
     k2: estimate
     k3: 0
-  petab_yaml: petab_problem.yaml
+  model_subspace_petab_yaml: petab_problem.yaml
   predecessor_model_hash: M_0-000
 - criteria:
     AICc: -0.7053270766271886
@@ -75,7 +75,7 @@
     k1: estimate
     k2: 0.1
     k3: 0
-  petab_yaml: petab_problem.yaml
+  model_subspace_petab_yaml: petab_problem.yaml
   predecessor_model_hash: M_0-000
 - criteria:
     AICc: 9.294672923372811
@@ -96,7 +96,7 @@
     k1: estimate
     k2: 0.1
     k3: estimate
-  petab_yaml: petab_problem.yaml
+  model_subspace_petab_yaml: petab_problem.yaml
   predecessor_model_hash: M_3-000
 - criteria:
     AICc: 7.8521704398854
@@ -117,7 +117,7 @@
     k1: estimate
     k2: estimate
     k3: 0
-  petab_yaml: petab_problem.yaml
+  model_subspace_petab_yaml: petab_problem.yaml
   predecessor_model_hash: M_3-000
 - criteria:
     AICc: 35.94352968170024
@@ -139,5 +139,5 @@
     k1: estimate
     k2: estimate
     k3: estimate
-  petab_yaml: petab_problem.yaml
+  model_subspace_petab_yaml: petab_problem.yaml
   predecessor_model_hash: M_3-000
diff --git a/doc/examples/example_cli_famos.ipynb b/doc/examples/example_cli_famos.ipynb
index b5895ac..7a2e147 100644
--- a/doc/examples/example_cli_famos.ipynb
+++ b/doc/examples/example_cli_famos.ipynb
@@ -38,7 +38,8 @@
     "\n",
     "output_path = Path().resolve() / \"output_famos\"\n",
     "output_path_str = str(output_path)\n",
-    "shutil.rmtree(output_path_str)\n",
+    "if output_path.exists():\n",
+    "    shutil.rmtree(output_path_str)\n",
     "output_path.mkdir(exist_ok=False, parents=True)"
    ]
   },
diff --git a/doc/examples/example_cli_famos_calibration_tool.py b/doc/examples/example_cli_famos_calibration_tool.py
index f5b58c2..5254c5b 100644
--- a/doc/examples/example_cli_famos_calibration_tool.py
+++ b/doc/examples/example_cli_famos_calibration_tool.py
@@ -12,7 +12,7 @@
 for model in models:
     calibrate(model=model)
     predecessor_model_hashes |= {model.predecessor_model_hash}
-models.to_yaml(output_yaml=calibrated_models_yaml)
+models.to_yaml(filename=calibrated_models_yaml)
 
 if len(predecessor_model_hashes) == 0:
     pass
diff --git a/doc/examples/model_selection/calibrated_M1_4.yaml b/doc/examples/model_selection/calibrated_M1_4.yaml
index c1e94f1..a5df0b7 100644
--- a/doc/examples/model_selection/calibrated_M1_4.yaml
+++ b/doc/examples/model_selection/calibrated_M1_4.yaml
@@ -1,19 +1,19 @@
+model_subspace_id: M1_4
+model_subspace_indices:
+- 0
+- 0
+- 0
 criteria:
   AIC: 15
+model_hash: M1_4-000
+model_subspace_petab_yaml: ../model_selection/petab_problem.yaml
 estimated_parameters:
   k2: 0.15
   k3: 0.0
 iteration: 1
-model_hash: M1_4-000
 model_id: M1_4-000
-model_subspace_id: M1_4
-model_subspace_indices:
-- 0
-- 0
-- 0
 parameters:
   k1: 0
   k2: estimate
   k3: estimate
-petab_yaml: ../model_selection/petab_problem.yaml
-predecessor_model_hash: null
+predecessor_model_hash: M1_2-000
diff --git a/doc/examples/model_selection/calibrated_M1_7.yaml b/doc/examples/model_selection/calibrated_M1_7.yaml
index 48c64c6..f955db3 100644
--- a/doc/examples/model_selection/calibrated_M1_7.yaml
+++ b/doc/examples/model_selection/calibrated_M1_7.yaml
@@ -1,20 +1,20 @@
+model_subspace_id: M1_7
+model_subspace_indices:
+- 0
+- 0
+- 0
 criteria:
   AIC: 20
+model_hash: M1_7-000
+model_subspace_petab_yaml: ../model_selection/petab_problem.yaml
 estimated_parameters:
   k1: 0.25
   k2: 0.1
   k3: 0.0
 iteration: 2
-model_hash: M1_7-000
 model_id: M1_7-000
-model_subspace_id: M1_7
-model_subspace_indices:
-- 0
-- 0
-- 0
 parameters:
   k1: estimate
   k2: estimate
   k3: estimate
-petab_yaml: ../model_selection/petab_problem.yaml
-predecessor_model_hash: null
+predecessor_model_hash: M1_4-000
diff --git a/doc/examples/model_selection/calibrated_models_1.yaml b/doc/examples/model_selection/calibrated_models_1.yaml
index 9e3a39f..343e774 100644
--- a/doc/examples/model_selection/calibrated_models_1.yaml
+++ b/doc/examples/model_selection/calibrated_models_1.yaml
@@ -1,51 +1,52 @@
-- criteria:
-    AIC: 180
-  estimated_parameters: {}
-  iteration: 1
-  model_hash: M1_0-000
-  model_id: M1_0-000
-  model_subspace_id: M1_0
+- model_subspace_id: M1_0
   model_subspace_indices:
   - 0
   - 0
   - 0
+  criteria:
+    AIC: 180
+  model_hash: M1_0-000
+  model_subspace_petab_yaml: ../model_selection/petab_problem.yaml
+  estimated_parameters: {}
+  iteration: 1
+  model_id: M1_0-000
   parameters:
     k1: 0
     k2: 0
     k3: 0
-  petab_yaml: ../model_selection/petab_problem.yaml
-  predecessor_model_hash: null
-- criteria:
-    AIC: 100
-  estimated_parameters: {}
-  iteration: 1
-  model_hash: M1_1-000
-  model_id: M1_1-000
-  model_subspace_id: M1_1
+  predecessor_model_hash: virtual_initial_model-
+- model_subspace_id: M1_1
   model_subspace_indices:
   - 0
   - 0
   - 0
+  criteria:
+    AIC: 100
+  model_hash: M1_1-000
+  model_subspace_petab_yaml: ../model_selection/petab_problem.yaml
+  estimated_parameters: {}
+  iteration: 1
+  model_id: M1_1-000
   parameters:
     k1: 0.2
     k2: 0.1
     k3: estimate
-  petab_yaml: ../model_selection/petab_problem.yaml
-  predecessor_model_hash: null
-- criteria:
-    AIC: 50
-  estimated_parameters: {}
-  iteration: 1
-  model_hash: M1_2-000
-  model_id: M1_2-000
-  model_subspace_id: M1_2
+  predecessor_model_hash: virtual_initial_model-
+
+- model_subspace_id: M1_2
   model_subspace_indices:
   - 0
   - 0
   - 0
+  criteria:
+    AIC: 50
+  model_hash: M1_2-000
+  model_subspace_petab_yaml: ../model_selection/petab_problem.yaml
+  estimated_parameters: {}
+  iteration: 1
+  model_id: M1_2-000
   parameters:
     k1: 0.2
     k2: estimate
     k3: 0
-  petab_yaml: ../model_selection/petab_problem.yaml
-  predecessor_model_hash: null
+  predecessor_model_hash: virtual_initial_model-
diff --git a/doc/examples/visualization.ipynb b/doc/examples/visualization.ipynb
index 13f36b4..16bd485 100644
--- a/doc/examples/visualization.ipynb
+++ b/doc/examples/visualization.ipynb
@@ -31,7 +31,7 @@
     "\n",
     "import petab_select\n",
     "import petab_select.plot\n",
-    "from petab_select import VIRTUAL_INITIAL_MODEL_HASH\n",
+    "from petab_select import VIRTUAL_INITIAL_MODEL\n",
     "\n",
     "models = petab_select.Models.from_yaml(\n",
     "    \"calibrated_models/calibrated_models.yaml\"\n",
@@ -69,12 +69,12 @@
     "# Custom labels\n",
     "labels = {}\n",
     "for model in models:\n",
-    "    labels[model.get_hash()] = \"M_\" + \"\".join(\n",
+    "    labels[model.hash] = \"M_\" + \"\".join(\n",
     "        \"1\" if value == petab_select.ESTIMATE else \"0\"\n",
     "        for value in model.parameters.values()\n",
     "    )\n",
-    "labels[VIRTUAL_INITIAL_MODEL_HASH] = \"\\n\".join(\n",
-    "    petab_select.VIRTUAL_INITIAL_MODEL.split(\"_\")\n",
+    "labels[VIRTUAL_INITIAL_MODEL.hash] = \"\\n\".join(\n",
+    "    petab_select.VIRTUAL_INITIAL_MODEL.model_subspace_id.split(\"_\")\n",
     ").title()\n",
     "\n",
     "# Custom colors for some models\n",
diff --git a/doc/examples/workflow_python.ipynb b/doc/examples/workflow_python.ipynb
index ba5d7e7..c943565 100644
--- a/doc/examples/workflow_python.ipynb
+++ b/doc/examples/workflow_python.ipynb
@@ -65,12 +65,12 @@
     "    print(\n",
     "        f\"\"\"\\\n",
     "Model subspace ID: {model.model_subspace_id}\n",
-    "PEtab YAML location: {model.petab_yaml}\n",
+    "PEtab YAML location: {model.model_subspace_petab_yaml}\n",
     "Custom model parameters: {model.parameters}\n",
-    "Model hash: {model.get_hash()}\n",
+    "Model hash: {model.hash}\n",
     "Model ID: {model.model_id}\n",
     "{select_problem.criterion}: {model.get_criterion(select_problem.criterion, compute=False)}\n",
-    "Model calibrated in iteration: {model.iteration}\n",
+    "Model was calibrated in iteration: {model.iteration}\n",
     "\"\"\"\n",
     "    )\n",
     "\n",
@@ -177,6 +177,7 @@
     "    calibrate(candidate_model)\n",
     "\n",
     "iteration_results = petab_select.ui.end_iteration(\n",
+    "    problem=select_problem,\n",
     "    candidate_space=iteration[CANDIDATE_SPACE],\n",
     "    calibrated_models=iteration[UNCALIBRATED_MODELS],\n",
     ")"
@@ -235,6 +236,7 @@
     "\n",
     "    # Finalize iteration\n",
     "    iteration_results = petab_select.ui.end_iteration(\n",
+    "        problem=select_problem,\n",
     "        candidate_space=iteration[CANDIDATE_SPACE],\n",
     "        calibrated_models=iteration[UNCALIBRATED_MODELS],\n",
     "    )\n",
@@ -257,7 +259,7 @@
     ")\n",
     "\n",
     "for candidate_model in iteration_results[MODELS]:\n",
-    "    if candidate_model.get_hash() == local_best_model.get_hash():\n",
+    "    if candidate_model.hash == local_best_model.hash:\n",
     "        print(BOLD_TEXT + \"BEST MODEL OF CURRENT ITERATION\" + NORMAL_TEXT)\n",
     "    print_model(candidate_model)"
    ]
@@ -285,7 +287,7 @@
     ")\n",
     "\n",
     "for candidate_model in iteration_results[MODELS]:\n",
-    "    if candidate_model.get_hash() == local_best_model.get_hash():\n",
+    "    if candidate_model.hash == local_best_model.hash:\n",
     "        print(BOLD_TEXT + \"BEST MODEL OF CURRENT ITERATION\" + NORMAL_TEXT)\n",
     "    print_model(candidate_model)"
    ]
@@ -313,7 +315,7 @@
     ")\n",
     "\n",
     "for candidate_model in iteration_results[MODELS]:\n",
-    "    if candidate_model.get_hash() == local_best_model.get_hash():\n",
+    "    if candidate_model.hash == local_best_model.hash:\n",
     "        print(BOLD_TEXT + \"BEST MODEL OF CURRENT ITERATION\" + NORMAL_TEXT)\n",
     "    print_model(candidate_model)"
    ]
@@ -398,7 +400,7 @@
     "    criterion=select_problem.criterion\n",
     ")\n",
     "# candidate_space.calibrated_models = iteration_results[CANDIDATE_SPACE].calibrated_models\n",
-    "petab_select.ui.start_iteration(\n",
+    "iteration = petab_select.ui.start_iteration(\n",
     "    problem=select_problem,\n",
     "    candidate_space=candidate_space,\n",
     ");"
@@ -411,9 +413,44 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "for candidate_model in candidate_space.models:\n",
+    "for candidate_model in iteration[UNCALIBRATED_MODELS]:\n",
+    "    calibrate(candidate_model)\n",
     "    print_model(candidate_model)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "924ee6af-f72a-455d-8985-13a6f00a14a1",
+   "metadata": {},
+   "source": [
+    "All calibrated models across all iterations can be retrieved via `Problem.state.models`, which is set by `ui.end_iteration`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10b62dd7-a3c3-420e-a88d-9ac44e367145",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Add latest model to all models\n",
+    "petab_select.ui.end_iteration(\n",
+    "    problem=select_problem,\n",
+    "    candidate_space=iteration[CANDIDATE_SPACE],\n",
+    "    calibrated_models=iteration[UNCALIBRATED_MODELS],\n",
+    ");"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "68d1f89a-31a2-4f96-8fe0-82d1ad130832",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Print all models\n",
+    "select_problem.state.models.df"
+   ]
   }
  ],
  "metadata": {
diff --git a/doc/standard/model.yaml b/doc/standard/model.yaml
index a49a042..71b5901 100644
--- a/doc/standard/model.yaml
+++ b/doc/standard/model.yaml
@@ -1,10 +1,7 @@
 $defs:
   ModelHash:
     type: string
-description: "A model.\n\nSee :class:`ModelBase` for the standardized attributes.\
-  \ Additional\nattributes are available in ``Model`` to improve usability.\n\nAttributes:\n\
-  \    _model_subspace_petab_problem:\n        The PEtab problem of the model subspace\
-  \ of this model.\n        If not provided, this is reconstructed from\n        :attr:`model_subspace_petab_yaml`."
+description: A model.
 properties:
   model_subspace_id:
     title: Model Subspace Id
diff --git a/doc/standard/models.yaml b/doc/standard/models.yaml
index a90f3d6..b66a9e1 100644
--- a/doc/standard/models.yaml
+++ b/doc/standard/models.yaml
@@ -1,10 +1,6 @@
 $defs:
   Model:
-    description: "A model.\n\nSee :class:`ModelBase` for the standardized attributes.\
-      \ Additional\nattributes are available in ``Model`` to improve usability.\n\n\
-      Attributes:\n    _model_subspace_petab_problem:\n        The PEtab problem of\
-      \ the model subspace of this model.\n        If not provided, this is reconstructed\
-      \ from\n        :attr:`model_subspace_petab_yaml`."
+    description: A model.
     properties:
       model_subspace_id:
         title: Model Subspace Id
@@ -67,19 +63,7 @@ $defs:
     type: object
   ModelHash:
     type: string
-description: 'A collection of models.
-
-
-  Provide a PEtab Select ``problem`` to the constructor or via
-
-  ``set_problem``, to use add models by hashes. This means that all models
-
-  must belong to the same PEtab Select problem.
-
-
-  This permits both ``list`` and ``dict`` operations -- see
-
-  :class:``ListDict`` for further details.'
+description: A collection of models.
 items:
   $ref: '#/$defs/Model'
 title: Models
diff --git a/doc/standard/problem.yaml b/doc/standard/problem.yaml
index b552986..ceba7a9 100644
--- a/doc/standard/problem.yaml
+++ b/doc/standard/problem.yaml
@@ -1,14 +1,4 @@
-description: "Handle everything related to the model selection problem.\n\nAttributes:\n\
-  \    model_space:\n        The model space.\n    calibrated_models:\n        Calibrated\
-  \ models. Will be used to augment the model selection\n        problem (e.g. by\
-  \ excluding them from the model space).\n    candidate_space_arguments:\n      \
-  \  Arguments are forwarded to the candidate space constructor.\n    compare:\n \
-  \       A method that compares models by selection criterion. See\n        :func:`petab_select.model.default_compare`\
-  \ for an example.\n    criterion:\n        The criterion used to compare models.\n\
-  \    method:\n        The method used to search the model space.\n    version:\n\
-  \        The version of the PEtab Select format.\n    yaml_path:\n        The location\
-  \ of the selection problem YAML file. Used for relative\n        paths that exist\
-  \ in e.g. the model space files."
+description: The model selection problem.
 properties:
   format_version:
     default: 1.0.0
diff --git a/petab_select/candidate_space.py b/petab_select/candidate_space.py
index a9a3be3..f3ed08a 100644
--- a/petab_select/candidate_space.py
+++ b/petab_select/candidate_space.py
@@ -1301,7 +1301,7 @@ def jump_to_most_distant(
         # self.predecessor_model = None
         self.set_predecessor_model(None)
         self.best_model_of_current_run = None
-        self.models = [predecessor_model]
+        self.models = Models([predecessor_model])
 
         self.write_summary_tsv("Jumped to the most distant model.")
         self.update_method(self.method_scheme[(Method.MOST_DISTANT,)])
diff --git a/petab_select/cli.py b/petab_select/cli.py
index e5bbdbf..c630fd7 100644
--- a/petab_select/cli.py
+++ b/petab_select/cli.py
@@ -277,6 +277,7 @@ def end_iteration(
 
     # Finalize iteration results
     iteration_results = ui.end_iteration(
+        problem=problem,
         candidate_space=candidate_space,
         calibrated_models=calibrated_models,
     )
diff --git a/petab_select/model.py b/petab_select/model.py
index 149a68c..cd1bdb9 100644
--- a/petab_select/model.py
+++ b/petab_select/model.py
@@ -25,9 +25,6 @@
 
 from .constants import (
     ESTIMATE,
-    CRITERIA,
-    ESTIMATED_PARAMETERS,
-    ITERATION,
     MODEL_HASH,
     MODEL_HASH_DELIMITER,
     MODEL_ID,
@@ -68,18 +65,18 @@ class ModelHash(BaseModel):
     back into the corresponding model. Currently, if two models from two
     different model subspaces are actually the same PEtab problem, they will
     still have different model hashes.
-
-    Attributes:
-        model_subspace_id:
-            The ID of the model subspace of the model. Unique up to a single
-            PEtab Select problem model space.
-        model_subspace_indices_hash:
-            A hash of the location of the model in its model
-            subspace. Unique up to a single model subspace.
     """
 
     model_subspace_id: str
+    """The ID of the model subspace of the model.
+
+    Unique up to a single model space.
+    """
     model_subspace_indices_hash: str
+    """A hash of the location of the model in its model subspace.
+
+    Unique up to a single model subspace.
+    """
 
     @model_validator(mode="wrap")
     def _check_kwargs(
@@ -91,18 +88,20 @@ def _check_kwargs(
 
         See documentation of Pydantic wrap validators.
         """
+        # Deprecate? Or introduce some `UNKNOWN_MODEL`?
+        if kwargs is None:
+            kwargs = VIRTUAL_INITIAL_MODEL.hash
+
         if isinstance(kwargs, ModelHash):
             return kwargs
-
-        if isinstance(kwargs, dict):
+        elif isinstance(kwargs, dict):
             kwargs[MODEL_SUBSPACE_INDICES_HASH] = (
                 ModelHash.hash_model_subspace_indices(
                     kwargs[MODEL_SUBSPACE_INDICES]
                 )
             )
             del kwargs[MODEL_SUBSPACE_INDICES]
-
-        if isinstance(kwargs, str):
+        elif isinstance(kwargs, str):
             kwargs = ModelHash.kwargs_from_str(hash_str=kwargs)
 
         expected_model_hash = None
@@ -397,19 +396,17 @@ def resolve_paths(self, root_path: str | Path) -> None:
 
 
 class Model(ModelBase):
-    """A model.
+    """A model."""
 
-    See :class:`ModelBase` for the standardized attributes. Additional
-    attributes are available in ``Model`` to improve usability.
-
-    Attributes:
-        _model_subspace_petab_problem:
-            The PEtab problem of the model subspace of this model.
-            If not provided, this is reconstructed from
-            :attr:`model_subspace_petab_yaml`.
-    """
+    # See :class:`ModelBase` for the standardized attributes. Additional
+    # attributes are available in ``Model`` to improve usability.
 
     _model_subspace_petab_problem: petab.Problem = PrivateAttr(default=None)
+    """The PEtab problem of the model subspace of this model.
+
+    If not provided, this is reconstructed from
+    :attr:`model_subspace_petab_yaml`.
+    """
 
     @model_validator(mode="after")
     def _fix_petab_problem(self: Model) -> Model:
@@ -679,14 +676,35 @@ def get_parameter_values(
         ]
 
     @staticmethod
-    def from_yaml(filename: str | Path) -> Model:
-        """Load a model from a YAML file."""
+    def from_yaml(
+        filename: str | Path,
+        model_subspace_petab_problem: petab.Problem | None = None,
+    ) -> Model:
+        """Load a model from a YAML file.
+
+        Args:
+            filename:
+                The filename.
+            model_subspace_petab_problem:
+                A preloaded copy of the PEtab problem of the model subspace
+                that this model belongs to.
+        """
         model = ModelStandard.load_data(
             filename=filename,
             root_path=Path(filename).parent,
+            _model_subspace_petab_problem=model_subspace_petab_problem,
         )
         return model
 
+    def get_hash(self) -> ModelHash:
+        """Deprecated. Use `Model.hash` instead."""
+        warnings.warn(
+            "Use `Model.hash` instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.hash
+
 
 def default_compare(
     model0: Model,
diff --git a/petab_select/models.py b/petab_select/models.py
index 74c3828..477d8cb 100644
--- a/petab_select/models.py
+++ b/petab_select/models.py
@@ -75,20 +75,17 @@ class _ListDict(RootModel, MutableSequence):
 
     The typing is currently based on PEtab Select objects. Hence, objects are
     in ``_models``, and metadata (model hashes) are in ``_hashes``.
-
-    Attributes:
-        _models:
-            The list of objects (list items/dictionary values)
-            (PEtab Select models).
-        _hashes:
-            The list of metadata (dictionary keys) (model hashes).
-        _problem:
-            The PEtab Select problem.
     """
 
     root: list[Model] = Field(default_factory=list)
+    """The list of models."""
     _hashes: list[ModelHash] = PrivateAttr(default_factory=list)
+    """The list of model hashes."""
     _problem: Problem | None = PrivateAttr(default=None)
+    """The PEtab Select problem that all models belong to.
+
+    If this is provided, then you can add models by hashes.
+    """
 
     @model_validator(mode="wrap")
     def _check_kwargs(
@@ -102,12 +99,12 @@ def _check_kwargs(
         if isinstance(kwargs, list):
             _models = kwargs
         elif isinstance(kwargs, dict):
-            # Identify the argument with the models
+            # Identify the models
             if "models" in kwargs and "root" in kwargs:
                 raise ValueError("Provide only one of `root` and `models`.")
             _models = kwargs.get("models") or kwargs.get("root") or []
 
-            # Identify the argument with the PEtab Select problem
+            # Identify the PEtab Select problem
             if "problem" in kwargs and "_problem" in kwargs:
                 raise ValueError(
                     "Provide only one of `problem` and `_problem`."
@@ -393,15 +390,7 @@ def values(self) -> Models:
 
 
 class Models(_ListDict):
-    """A collection of models.
-
-    Provide a PEtab Select ``problem`` to the constructor or via
-    ``set_problem``, to use add models by hashes. This means that all models
-    must belong to the same PEtab Select problem.
-
-    This permits both ``list`` and ``dict`` operations -- see
-    :class:``ListDict`` for further details.
-    """
+    """A collection of models."""
 
     def set_problem(self, problem: Problem) -> None:
         """Set the PEtab Select problem for this set of models."""
@@ -426,7 +415,7 @@ def lint(self):
     @staticmethod
     def from_yaml(
         filename: TYPE_PATH,
-        petab_problem: petab.Problem = None,
+        model_subspace_petab_problem: petab.Problem = None,
         problem: Problem = None,
     ) -> Models:
         """Load models from a YAML file.
@@ -434,10 +423,12 @@ def from_yaml(
         Args:
             filename:
                 Location of the YAML file.
-            petab_problem:
-                Provide a preloaded copy of the PEtab problem. N.B.:
+            model_subspace_petab_problem:
+                A preloaded copy of the PEtab problem. N.B.:
                 all models should share the same PEtab problem if this is
-                provided.
+                provided (e.g. all models belong to the same model subspace,
+                or all model subspaces have the same
+                ``model_subspace_petab_yaml`` in the model space file(s)).
             problem:
                 The PEtab Select problem. N.B.: all models should belong to the
                 same PEtab Select problem if this is provided.
@@ -445,12 +436,21 @@ def from_yaml(
         Returns:
             The models.
         """
+        # Handle single-model files, for backwards compatibility.
+        try:
+            model = Model.from_yaml(
+                filename=filename,
+                model_subspace_petab_problem=model_subspace_petab_problem,
+            )
+            return Models([model])
+        except:  # noqa: S110
+            pass
         return ModelsStandard.load_data(
             filename=filename,
             _problem=problem,
             model_kwargs={
                 ROOT_PATH: Path(filename).parent,
-                MODEL_SUBSPACE_PETAB_PROBLEM: petab_problem,
+                MODEL_SUBSPACE_PETAB_PROBLEM: model_subspace_petab_problem,
             },
         )
 
diff --git a/petab_select/plot.py b/petab_select/plot.py
index e485ba0..1f2f539 100644
--- a/petab_select/plot.py
+++ b/petab_select/plot.py
@@ -18,7 +18,7 @@
     group_by_iteration,
 )
 from .constants import Criterion
-from .model import VIRTUAL_INITIAL_MODEL_HASH, Model, ModelHash
+from .model import VIRTUAL_INITIAL_MODEL, Model
 from .models import Models
 
 RELATIVE_LABEL_FONTSIZE = -2
@@ -144,6 +144,7 @@ def line_best_by_iteration(
 
     ax.get_xticks()
     ax.set_xticks(list(range(len(criterion_values))))
+    ax.set_xlabel("Iteration and model", fontsize=fz)
     ax.set_ylabel((r"$\Delta$" if relative else "") + criterion, fontsize=fz)
     # could change to compared_model_ids, if all models are plotted
     ax.set_xticklabels(
@@ -217,7 +218,7 @@ def graph_history(
             for model in models
         }
     labels = labels.copy()
-    labels[VIRTUAL_INITIAL_MODEL_HASH] = "Virtual\nInitial\nModel"
+    labels[VIRTUAL_INITIAL_MODEL.hash] = "Virtual\nInitial\nModel"
 
     G = nx.DiGraph()
     edges = []
@@ -502,8 +503,8 @@ def graph_iteration_layers(
         [model.hash for model in iteration_models]
         for iteration_models in group_by_iteration(models).values()
     ]
-    if VIRTUAL_INITIAL_MODEL_HASH in ancestry.values():
-        ordering.insert(0, [VIRTUAL_INITIAL_MODEL_HASH])
+    if VIRTUAL_INITIAL_MODEL.hash in ancestry.values():
+        ordering.insert(0, [VIRTUAL_INITIAL_MODEL.hash])
 
     model_estimated_parameters = {
         model.hash: set(model.estimated_parameters) for model in models
@@ -529,6 +530,9 @@ def graph_iteration_layers(
     }
 
     labels = labels or {}
+    labels[VIRTUAL_INITIAL_MODEL.hash] = labels.get(
+        VIRTUAL_INITIAL_MODEL.hash, "Virtual\nInitial\nModel"
+    )
     labels = (
         labels
         | {
@@ -537,11 +541,9 @@ def graph_iteration_layers(
             if model.hash not in labels
         }
         | {
-            ModelHash.from_hash(
-                model.predecessor_model_hash
-            ): model.predecessor_model_hash
+            model.predecessor_model_hash: model.predecessor_model_hash
             for model in models
-            if ModelHash.from_hash(model.predecessor_model_hash) not in labels
+            if model.predecessor_model_hash not in labels
         }
     )
     if augment_labels:
@@ -582,7 +584,7 @@ def __getitem__(self, key):
         labels = {
             model_hash: (
                 label0
-                if model_hash == VIRTUAL_INITIAL_MODEL_HASH
+                if model_hash == VIRTUAL_INITIAL_MODEL.hash
                 else "\n".join(
                     [
                         label0,
diff --git a/petab_select/problem.py b/petab_select/problem.py
index c874132..e30511c 100644
--- a/petab_select/problem.py
+++ b/petab_select/problem.py
@@ -41,42 +41,59 @@
 ]
 
 
+class State(BaseModel):
+    """Carry the state of applying model selection methods to the problem."""
+
+    models: Models = Field(default_factory=Models)
+    """All calibrated models."""
+    iteration: int = Field(default=0)
+    """The latest iteration of model selection."""
+
+    def increment_iteration(self) -> None:
+        """Start the next iteration."""
+        self.iteration += 1
+
+    def reset(self) -> None:
+        """Reset the state.
+
+        N.B.: does not reset all state information, which currently also exists
+        in other classes. Open a GitHub issue if you see unusual behavior. A
+        quick fix is to simply recreate the PEtab Select problem, and any other
+        objects that you use, e.g. the candidate space, whenever you need a
+        full reset.
+        https://github.com/PEtab-dev/petab_select/issues
+        """
+        # FIXME state information is currently distributed across multiple
+        # classes, e.g. exclusions in model subspaces and candidate spaces.
+        # move all state information here.
+        self.models = Models()
+        self.iteration = 0
+
+
 class Problem(BaseModel):
-    """Handle everything related to the model selection problem.
-
-    Attributes:
-        model_space:
-            The model space.
-        calibrated_models:
-            Calibrated models. Will be used to augment the model selection
-            problem (e.g. by excluding them from the model space).
-        candidate_space_arguments:
-            Arguments are forwarded to the candidate space constructor.
-        compare:
-            A method that compares models by selection criterion. See
-            :func:`petab_select.model.default_compare` for an example.
-        criterion:
-            The criterion used to compare models.
-        method:
-            The method used to search the model space.
-        version:
-            The version of the PEtab Select format.
-        yaml_path:
-            The location of the selection problem YAML file. Used for relative
-            paths that exist in e.g. the model space files.
-    """
+    """The model selection problem."""
 
     format_version: str = Field(default="1.0.0")
+    """The file format version."""
     criterion: Annotated[
         Criterion, PlainSerializer(lambda x: x.value, return_type=str)
     ]
+    """The criterion used to compare models."""
     method: Annotated[
         Method, PlainSerializer(lambda x: x.value, return_type=str)
     ]
+    """The method used to search the model space."""
     model_space_files: list[Path]
+    """The files that define the model space."""
     candidate_space_arguments: dict[str, Any] = Field(default_factory=dict)
+    """Method-specific arguments.
+
+    These are forwarded to the candidate space constructor.
+    """
 
     _compare: Callable[[Model, Model], bool] = PrivateAttr(default=None)
+    """The method by which models are compared."""
+    _state: State = PrivateAttr(default_factory=State)
 
     @model_validator(mode="wrap")
     def _check_input(
@@ -88,6 +105,8 @@ def _check_input(
             return data
 
         compare = data.pop("compare", None) or data.pop("_compare", None)
+        if "state" in data:
+            data["_state"] = data["state"]
         root_path = Path(data.pop(ROOT_PATH, ""))
 
         problem = handler(data)
@@ -111,6 +130,10 @@ def _check_input(
 
         return problem
 
+    @property
+    def state(self) -> State:
+        return self._state
+
     @staticmethod
     def from_yaml(filename: TYPE_PATH) -> Problem:
         """Load a problem from a YAML file."""
diff --git a/petab_select/ui.py b/petab_select/ui.py
index 34abc14..f77bbbd 100644
--- a/petab_select/ui.py
+++ b/petab_select/ui.py
@@ -47,6 +47,7 @@ def start_iteration_result(candidate_space: CandidateSpace) -> dict[str, Any]:
     # will see. All models (user-supplied and newly-calibrated) will
     # have their iteration set (again) in `end_iteration`, via
     # `CandidateSpace.get_iteration_calibrated_models`
+    # TODO use problem.state.iteration instead
     for model in candidate_space.models:
         model.iteration = candidate_space.iteration
     return {
@@ -138,7 +139,8 @@ def start_iteration(
     candidate_space.criterion = criterion
 
     # Start a new iteration
-    candidate_space.iteration += 1
+    problem.state.increment_iteration()
+    candidate_space.iteration = problem.state.iteration
 
     # Set the predecessor model to the previous predecessor model.
     predecessor_model = candidate_space.previous_predecessor_model
@@ -241,6 +243,7 @@ def start_iteration(
 
 
 def end_iteration(
+    problem: Problem,
     candidate_space: CandidateSpace,
     calibrated_models: Models,
 ) -> dict[str, Models | bool | CandidateSpace]:
@@ -253,6 +256,8 @@ def end_iteration(
     ends.
 
     Args:
+        problem:
+            The PEtab Select problem.
         candidate_space:
             The candidate space.
         calibrated_models:
@@ -298,6 +303,8 @@ def end_iteration(
 
     iteration_results[CANDIDATE_SPACE] = candidate_space
 
+    problem.state.models.extend(iteration_results[MODELS])
+
     return iteration_results
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 65e8a77..934171d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -53,9 +53,11 @@ doc = [
     "nbsphinx>=0.9.5",
     "pandoc>=2.4",
     "nbconvert>=7.16.4",
+    "ipykernel>= 6.23.1",
     "ipython>=7.21.0",
     "readthedocs-sphinx-ext>=2.2.5",
     "sphinx-autodoc-typehints",
+    "petab_select[plot]",
 ]
 
 [project.scripts]
diff --git a/test/candidate_space/test_famos.py b/test/candidate_space/test_famos.py
index 5036dc6..e82cfa8 100644
--- a/test/candidate_space/test_famos.py
+++ b/test/candidate_space/test_famos.py
@@ -5,11 +5,10 @@
 from more_itertools import one
 
 import petab_select
-from petab_select import Method, ModelHash, Models
+from petab_select import Method, ModelHash
 from petab_select.constants import (
     CANDIDATE_SPACE,
     MODEL_HASH,
-    MODELS,
     TERMINATE,
     UNCALIBRATED_MODELS,
     Criterion,
@@ -126,7 +125,6 @@ def parse_summary_to_progress_list(summary_tsv: str) -> tuple[Method, set]:
         return progress_list
 
     progress_list = []
-    all_calibrated_models = Models()
 
     candidate_space = petab_select_problem.new_candidate_space()
     expected_repeated_model_hash0 = candidate_space.predecessor_model.hash
@@ -145,17 +143,15 @@ def parse_summary_to_progress_list(summary_tsv: str) -> tuple[Method, set]:
             )
 
             # Calibrate candidate models
-            calibrated_models = Models()
             for candidate_model in iteration[UNCALIBRATED_MODELS]:
                 calibrate(candidate_model)
-                calibrated_models[candidate_model.hash] = candidate_model
 
             # Finalize iteration
             iteration_results = petab_select.ui.end_iteration(
+                problem=petab_select_problem,
                 candidate_space=iteration[CANDIDATE_SPACE],
-                calibrated_models=calibrated_models,
+                calibrated_models=iteration[UNCALIBRATED_MODELS],
             )
-            all_calibrated_models += iteration_results[MODELS]
             candidate_space = iteration_results[CANDIDATE_SPACE]
 
             # Stop iteration if there are no candidate models
diff --git a/test/cli/input/models.yaml b/test/cli/input/models.yaml
index b9d12b8..b6daa42 100644
--- a/test/cli/input/models.yaml
+++ b/test/cli/input/models.yaml
@@ -10,11 +10,6 @@
     k2: 0.15
     k3: 0.0
   model_id: model_1
-  model_subspace_id: M
-  model_subspace_indices:
-  - 0
-  - 1
-  - 1
   parameters:
     k1: 0.2
     k2: estimate
@@ -28,11 +23,6 @@
   model_hash: M-110
   model_subspace_petab_yaml: ../../../doc/examples/model_selection/petab_problem.yaml
   model_id: model_2
-  model_subspace_id: M
-  model_subspace_indices:
-  - 1
-  - 1
-  - 0
   parameters:
     k1: estimate
     k2: estimate
diff --git a/test/ui/test_ui.py b/test/ui/test_ui.py
index 7efb442..f153cff 100644
--- a/test/ui/test_ui.py
+++ b/test/ui/test_ui.py
@@ -43,6 +43,7 @@ def test_user_calibrated_models(petab_select_problem):
     assert model_M1_0.model_subspace_id == "M1_0"
     model_M1_0.set_criterion(petab_select_problem.criterion, 100)
     iteration_results = petab_select.ui.end_iteration(
+        problem=petab_select_problem,
         candidate_space=iteration[CANDIDATE_SPACE],
         calibrated_models=[model_M1_0],
     )
@@ -61,6 +62,7 @@ def test_user_calibrated_models(petab_select_problem):
     for uncalibrated_model in iteration[UNCALIBRATED_MODELS]:
         uncalibrated_model.set_criterion(petab_select_problem.criterion, 50)
     iteration_results = petab_select.ui.end_iteration(
+        problem=petab_select_problem,
         candidate_space=iteration[CANDIDATE_SPACE],
         calibrated_models=iteration[UNCALIBRATED_MODELS],
     )