From 451d8c9fc1c79baec93845db66de15526a7646b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mike=20Mu=CC=88ller?= <mmueller@python-academy.de>
Date: Sat, 6 Apr 2024 15:45:17 +0200
Subject: [PATCH] Move yield of `States.stress_period_end` to right place

---
 src/pymf6/api.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/src/pymf6/api.py b/src/pymf6/api.py
index 8ec769d..54c1dda 100644
--- a/src/pymf6/api.py
+++ b/src/pymf6/api.py
@@ -75,8 +75,6 @@ def loop(self):
         verbose = self.verbose
         sim = self.api
 
-        yield sim, States.initialize
-
         current_time = mf6.get_current_time()
         end_time = mf6.get_end_time()
         kperold = [0 for _ in range(sim.subcomponent_count)]
@@ -97,11 +95,7 @@ def loop(self):
                 kperold=kperold)
             mf6.finalize_time_step()
             current_time = mf6.get_current_time()
-            sim_grp = self._sim_grp
-            if sim_grp.nstp == sim_grp.kstp + 1:
-                yield sim_grp, States.stress_period_end
         try:
-            yield sim, States.finalize
             mf6.finalize()
         except Exception as err:
             msg = 'MF6 simulation failed, check listing file'
@@ -155,6 +149,8 @@ def _solutions_loop(self, sim, mf6, current_time, kperold):
                         break
             yield sim_grp, States.timestep_end
             mf6.finalize_solve(sol_id)
+            if sim_grp.nstp == sim_grp.kstp + 1:
+                yield sim_grp, States.stress_period_end
         if not has_converged:
             print(f"Simulation group: {sim_grp} DID NOT CONVERGE")
         self._sim_grp = sim_grp