From 543b6ea00948c58a5c567b90350a2535656c2740 Mon Sep 17 00:00:00 2001
From: nmannall <n.mannall@epcc.ed.ac.uk>
Date: Mon, 20 May 2024 16:16:53 +0100
Subject: [PATCH] Pass iteration to update functions that need it

---
 gprMax/fields_outputs.py      | 5 +++--
 gprMax/solvers.py             | 5 +++--
 gprMax/updates/cpu_updates.py | 9 ++++-----
 gprMax/updates/updates.py     | 4 ++--
 4 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/gprMax/fields_outputs.py b/gprMax/fields_outputs.py
index 648491e7..c34d6efd 100644
--- a/gprMax/fields_outputs.py
+++ b/gprMax/fields_outputs.py
@@ -21,12 +21,14 @@ from pathlib import Path
 
 import h5py
 
+from gprMax.grid.fdtd_grid import FDTDGrid
+
 from ._version import __version__
 
 logger = logging.getLogger(__name__)
 
 
-def store_outputs(G):
+def store_outputs(G: FDTDGrid, iteration: int):
     """Stores field component values for every receiver and transmission line.
 
     Args:
@@ -34,7 +36,6 @@ def store_outputs(G):
     """
 
     # Assign iteration and fields to local variables
-    iteration = G.iteration
     Ex, Ey, Ez, Hx, Hy, Hz = G.Ex, G.Ey, G.Ez, G.Hx, G.Hy, G.Hz
 
     for rx in G.rxs:
diff --git a/gprMax/solvers.py b/gprMax/solvers.py
index 91dab762..09d39dcc 100644
--- a/gprMax/solvers.py
+++ b/gprMax/solvers.py
@@ -53,7 +53,7 @@ class Solver:
         self.updates.time_start()
 
         for iteration in iterator:
-            self.updates.store_outputs()
+            self.updates.store_outputs(iteration)
             self.updates.store_snapshots(iteration)
             self.updates.update_magnetic()
             self.updates.update_magnetic_pml()
@@ -62,7 +62,8 @@ class Solver:
                 self.updates.hsg_2()
             self.updates.update_electric_a()
             self.updates.update_electric_pml()
-            self.updates.update_electric_sources()
+            self.updates.update_electric_sources(iteration)
+            # TODO: Increment iteration here if add Model to Solver
             if isinstance(self.updates, SubgridUpdates):
                 self.updates.hsg_1()
             self.updates.update_electric_b()
diff --git a/gprMax/updates/cpu_updates.py b/gprMax/updates/cpu_updates.py
index 043e70a9..17660f93 100644
--- a/gprMax/updates/cpu_updates.py
+++ b/gprMax/updates/cpu_updates.py
@@ -38,9 +38,9 @@ class CPUUpdates(Updates):
 
         self.grid = G
 
-    def store_outputs(self):
+    def store_outputs(self, iteration):
         """Stores field component values for every receiver and transmission line."""
-        store_outputs_cpu(self.grid)
+        store_outputs_cpu(self.grid, iteration)
 
     def store_snapshots(self, iteration):
         """Stores any snapshots.
@@ -134,7 +134,7 @@ class CPUUpdates(Updates):
         for pml in self.grid.pmls["slabs"]:
             pml.update_electric()
 
-    def update_electric_sources(self):
+    def update_electric_sources(self, iteration):
         """Updates electric field components from sources -
         update any Hertzian dipole sources last.
         """
@@ -142,7 +142,7 @@ class CPUUpdates(Updates):
             self.grid.voltagesources + self.grid.transmissionlines + self.grid.hertziandipoles
         ):
             source.update_electric(
-                self.grid.iteration,
+                iteration,
                 self.grid.updatecoeffsE,
                 self.grid.ID,
                 self.grid.Ex,
@@ -150,7 +150,6 @@ class CPUUpdates(Updates):
                 self.grid.Ez,
                 self.grid,
             )
-        self.grid.iteration += 1
 
     def update_electric_b(self):
         """If there are any dispersive materials do 2nd part of dispersive
diff --git a/gprMax/updates/updates.py b/gprMax/updates/updates.py
index 182abb7f..065318a0 100644
--- a/gprMax/updates/updates.py
+++ b/gprMax/updates/updates.py
@@ -24,7 +24,7 @@ class Updates(ABC):
     """Defines update functions for a solver."""
 
     @abstractmethod
-    def store_outputs(self) -> None:
+    def store_outputs(self, iteration: int) -> None:
         """Stores field component values for every receiver and transmission line."""
         pass
 
@@ -63,7 +63,7 @@ class Updates(ABC):
         pass
 
     @abstractmethod
-    def update_electric_sources(self) -> None:
+    def update_electric_sources(self, iteration: int) -> None:
         """Updates electric field components from sources -
         update any Hertzian dipole sources last.
         """