Fix progress bar not updating realization count for new iterations

This commit fixes the bug introduced in 31e607b, where the status reporting in GUI was done the same way when rerunning failed realizations, and running new iterations. This is an issue because when rerunning failed realizations, we want to show all realizations and add the finished/failed count from the previous run, while new iterations should drop the failed realizations altogether.

src/ert/run_models/base_run_model.py

@@ -188,6 +188,7 @@ def __init__(
         self.minimum_required_realizations = minimum_required_realizations
         self.active_realizations = copy.copy(active_realizations)
         self.start_iteration = start_iteration
+        self.restart = False
         self.validate_active_realizations_count()
 
     def log_at_startup(self) -> None:
@@ -404,7 +405,10 @@ def get_current_status(self) -> dict[str, int]:
                 for real in all_realizations.values():
                     status[str(real["status"])] += 1
 
-        status["Finished"] += self._get_number_of_finished_realizations_from_reruns()
+        if self.restart:
+            status["Finished"] += (
+                self._get_number_of_finished_realizations_from_reruns()
+            )
         return status
 
     def _get_number_of_finished_realizations_from_reruns(self) -> int:
@@ -644,7 +648,11 @@ def get_number_of_existing_runpaths(self) -> int:
         return [real_path.exists() for real_path in realization_set].count(True)
 
     def get_number_of_active_realizations(self) -> int:
-        return self._initial_realizations_mask.count(True)
+        return (
+            self._initial_realizations_mask.count(True)
+            if self.restart
+            else self.active_realizations.count(True)
+        )
 
     def get_number_of_successful_realizations(self) -> int:
         return self.active_realizations.count(True)

src/ert/run_models/ensemble_experiment.py

@@ -63,6 +63,7 @@ def run_experiment(
         restart: bool = False,
     ) -> None:
         self.log_at_startup()
+        self.restart = restart
         if not restart:
             self.experiment = self._storage.create_experiment(
                 name=self.experiment_name,

src/ert/run_models/ensemble_smoother.py

@@ -62,6 +62,7 @@ def run_experiment(
         self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False
     ) -> None:
         self.log_at_startup()
+        self.restart = restart
         ensemble_format = self.target_ensemble_format
         experiment = self._storage.create_experiment(
             parameters=self.ert_config.ensemble_config.parameter_configuration,

src/ert/run_models/evaluate_ensemble.py

@@ -64,6 +64,7 @@ def run_experiment(
         self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False
     ) -> None:
         self.log_at_startup()
+        self.restart = restart
         ensemble = self.ensemble
         experiment = ensemble.experiment
         self.set_env_key("_ERT_EXPERIMENT_ID", str(experiment.id))

src/ert/run_models/everest_run_model.py

@@ -391,6 +391,7 @@ def run_experiment(
         self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False
     ) -> None:
         self.log_at_startup()
+        self.restart = restart
         simulator = Simulator(
             self.everest_config,
             self.ert_config,

src/ert/run_models/iterated_ensemble_smoother.py

@@ -122,7 +122,7 @@ def run_experiment(
         self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False
     ) -> None:
         self.log_at_startup()
-
+        self.restart = restart
         target_ensemble_format = self.target_ensemble_format
         experiment = self._storage.create_experiment(
             parameters=self.ert_config.ensemble_config.parameter_configuration,

src/ert/run_models/multiple_data_assimilation.py

@@ -85,6 +85,7 @@ def run_experiment(
         self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False
     ) -> None:
         self.log_at_startup()
+        self.restart = restart
         if self.restart_run:
             id = self.prior_ensemble_id
             try:

tests/ert/unit_tests/run_models/test_base_run_model.py

@@ -202,6 +202,49 @@ def test_num_cpu_is_propagated_from_config_to_ensemble(run_args):
     assert ensemble.reals[1].num_cpu == 42
 
 
+@pytest.mark.parametrize(
+    "real_status_dict, expected_result",
+    [
+        pytest.param(
+            {"0": "Finished", "1": "Finished", "2": "Finished"},
+            {"Finished": 3},
+            id="ran_all_realizations_and_all_succeeded",
+        ),
+        pytest.param(
+            {"0": "Finished", "1": "Finished", "2": "Failed"},
+            {"Finished": 2, "Failed": 1},
+            id="ran_all_realizations_and_some_failed",
+        ),
+        pytest.param(
+            {"0": "Finished", "1": "Running", "2": "Failed"},
+            {"Finished": 1, "Failed": 1, "Running": 1},
+            id="ran_all_realizations_and_result_was_mixed",
+        ),
+    ],
+)
+def test_get_current_status(
+    real_status_dict,
+    expected_result,
+):
+    config = ErtConfig.from_file_contents("NUM_REALIZATIONS 3")
+    initial_active_realizations = [True] * 3
+    new_active_realizations = [True] * 3
+    brm = BaseRunModel(
+        config=config,
+        storage=MagicMock(spec=Storage),
+        queue_config=config.queue_config,
+        status_queue=MagicMock(spec=SimpleQueue),
+        active_realizations=initial_active_realizations,
+    )
+    snapshot_dict_reals = {}
+    for index, realization_status in real_status_dict.items():
+        snapshot_dict_reals[index] = {"status": realization_status}
+    iter_snapshot = EnsembleSnapshot.from_nested_dict({"reals": snapshot_dict_reals})
+    brm._iter_snapshot[0] = iter_snapshot
+    brm.active_realizations = new_active_realizations
+    assert dict(brm.get_current_status()) == expected_result
+
+
 @pytest.mark.parametrize(
     "initial_active_realizations, new_active_realizations, real_status_dict, expected_result",
     [
@@ -240,27 +283,6 @@ def test_num_cpu_is_propagated_from_config_to_ensemble(run_args):
             {"Finished": 0},
             id="did_not_run_all_realizations_and_none_succeeded",
         ),
-        pytest.param(
-            [True, True, True],
-            [True, True, True],
-            {"0": "Finished", "1": "Finished", "2": "Finished"},
-            {"Finished": 3},
-            id="ran_all_realizations_and_all_succeeded",
-        ),
-        pytest.param(
-            [True, True, True],
-            [True, True, True],
-            {"0": "Finished", "1": "Finished", "2": "Failed"},
-            {"Finished": 2, "Failed": 1},
-            id="ran_all_realizations_and_some_failed",
-        ),
-        pytest.param(
-            [True, True, True],
-            [True, True, True],
-            {"0": "Finished", "1": "Running", "2": "Failed"},
-            {"Finished": 1, "Failed": 1, "Running": 1},
-            id="ran_all_realizations_and_result_was_mixed",
-        ),
         pytest.param(
             [True, True, True],
             [True, True, False],
@@ -277,7 +299,7 @@ def test_num_cpu_is_propagated_from_config_to_ensemble(run_args):
         ),
     ],
 )
-def test_get_current_status(
+def test_get_current_status_when_rerun(
     initial_active_realizations,
     new_active_realizations,
     real_status_dict,
@@ -292,10 +314,73 @@ def test_get_current_status(
         status_queue=MagicMock(spec=SimpleQueue),
         active_realizations=initial_active_realizations,
     )
+    brm.restart = True
     snapshot_dict_reals = {}
     for index, realization_status in real_status_dict.items():
         snapshot_dict_reals[index] = {"status": realization_status}
     iter_snapshot = EnsembleSnapshot.from_nested_dict({"reals": snapshot_dict_reals})
     brm._iter_snapshot[0] = iter_snapshot
     brm.active_realizations = new_active_realizations
     assert dict(brm.get_current_status()) == expected_result
+
+
+def test_get_current_status_for_new_iteration_when_realization_failed_in_previous_run():
+    """Active realizations gets changed when we run next iteration, and the failed realizations from
+    the previous run should not be present in the current_status."""
+    initial_active_realizations = [True] * 5
+    # Realization 0,1, and 3 failed in the previous iteration
+    new_active_realizations = [False, False, True, False, True]
+    config = ErtConfig.from_file_contents("NUM_REALIZATIONS 5")
+    brm = BaseRunModel(
+        config=config,
+        storage=MagicMock(spec=Storage),
+        queue_config=config.queue_config,
+        status_queue=MagicMock(spec=SimpleQueue),
+        active_realizations=initial_active_realizations,
+    )
+    snapshot_dict_reals = {
+        "2": {"status": "Running"},
+        "4": {"status": "Finished"},
+    }
+    iter_snapshot = EnsembleSnapshot.from_nested_dict({"reals": snapshot_dict_reals})
+    brm._iter_snapshot[0] = iter_snapshot
+    brm.active_realizations = new_active_realizations
+
+    assert brm.restart is False
+    assert dict(brm.get_current_status()) == {"Running": 1, "Finished": 1}
+
+
+@pytest.mark.parametrize(
+    "new_active_realizations, was_rerun, expected_result",
+    [
+        pytest.param(
+            [False, False, False, True, False],
+            True,
+            5,
+            id="rerun_so_total_realization_count_is_not_affected_by_previous_failed_realizations",
+        ),
+        pytest.param(
+            [True, True, False, False, False],
+            False,
+            2,
+            id="new_iteration_so_total_realization_count_is_only_previously_successful_realizations",
+        ),
+    ],
+)
+def test_get_number_of_active_realizations_varies_when_rerun_or_new_iteration(
+    new_active_realizations, was_rerun, expected_result
+):
+    """When rerunning, we include all realizations in the total amount of active realization.
+    When running a new iteration based on the result of the previous iteration, we only include the successful realizations."""
+    initial_active_realizations = [True] * 5
+    config = ErtConfig.from_file_contents("NUM_REALIZATIONS 5")
+    brm = BaseRunModel(
+        config=config,
+        storage=MagicMock(spec=Storage),
+        queue_config=config.queue_config,
+        status_queue=MagicMock(spec=SimpleQueue),
+        active_realizations=initial_active_realizations,
+    )
+    brm.active_realizations = new_active_realizations
+    brm.restart = was_rerun
+    assert brm.get_number_of_active_realizations() == expected_result