Use Updates type information in Solver loop

gprMax/solvers.py

@@ -17,18 +17,19 @@
 # along with gprMax.  If not, see <http://www.gnu.org/licenses/>.
 
 import gprMax.config as config
-from gprMax.grid.cuda_grid import CUDAGrid
-from gprMax.grid.fdtd_grid import FDTDGrid
-from gprMax.grid.opencl_grid import OpenCLGrid
 
-from .grid import CUDAGrid, FDTDGrid, OpenCLGrid
+from .grid.cuda_grid import CUDAGrid
+from .grid.fdtd_grid import FDTDGrid
+from .grid.opencl_grid import OpenCLGrid
+from .subgrids.updates import SubgridUpdates
 from .subgrids.updates import create_updates as create_subgrid_updates
 from .updates.cpu_updates import CPUUpdates
 from .updates.cuda_updates import CUDAUpdates
 from .updates.opencl_updates import OpenCLUpdates
+from .updates.updates import Updates
 
 
-def create_G():
+def create_G() -> FDTDGrid:
     """Create grid object according to solver.
 
     Returns:
@@ -45,7 +46,53 @@ def create_G():
     return G
 
 
-def create_solver(G):
+class Solver:
+    """Generic solver for Update objects"""
+
+    def __init__(self, updates: Updates, hsg=False):
+        """
+        Args:
+            updates: Updates contains methods to run FDTD algorithm.
+            hsg: boolean to use sub-gridding.
+        """
+
+        self.updates = updates
+        self.hsg = hsg
+        self.solvetime = 0
+        self.memused = 0
+
+    def solve(self, iterator):
+        """Time step the FDTD model.
+
+        Args:
+            iterator: can be range() or tqdm()
+        """
+
+        self.updates.time_start()
+
+        for iteration in iterator:
+            self.updates.store_outputs()
+            self.updates.store_snapshots(iteration)
+            self.updates.update_magnetic()
+            self.updates.update_magnetic_pml()
+            self.updates.update_magnetic_sources()
+            if isinstance(self.updates, SubgridUpdates):
+                self.updates.hsg_2()
+            self.updates.update_electric_a()
+            self.updates.update_electric_pml()
+            self.updates.update_electric_sources()
+            if isinstance(self.updates, SubgridUpdates):
+                self.updates.hsg_1()
+            self.updates.update_electric_b()
+            if isinstance(self.updates, CUDAUpdates):
+                self.memused = self.updates.calculate_memory_used(iteration)
+
+        self.updates.finalise()
+        self.solvetime = self.updates.calculate_solve_time()
+        self.updates.cleanup()
+
+
+def create_solver(G: FDTDGrid) -> Solver:
     """Create configured solver object.
 
     N.B. A large range of different functions exist to advance the time step for
@@ -84,50 +131,4 @@ def create_solver(G):
         updates = OpenCLUpdates(G)
         solver = Solver(updates)
 
-    return solver
-
-
-class Solver:
-    """Generic solver for Update objects"""
-
-    def __init__(self, updates, hsg=False):
-        """
-        Args:
-            updates: Updates contains methods to run FDTD algorithm.
-            hsg: boolean to use sub-gridding.
-        """
-
-        self.updates = updates
-        self.hsg = hsg
-        self.solvetime = 0
-        self.memused = 0
-
-    def solve(self, iterator):
-        """Time step the FDTD model.
-
-        Args:
-            iterator: can be range() or tqdm()
-        """
-
-        self.updates.time_start()
-
-        for iteration in iterator:
-            self.updates.store_outputs()
-            self.updates.store_snapshots(iteration)
-            self.updates.update_magnetic()
-            self.updates.update_magnetic_pml()
-            self.updates.update_magnetic_sources()
-            if self.hsg:
-                self.updates.hsg_2()
-            self.updates.update_electric_a()
-            self.updates.update_electric_pml()
-            self.updates.update_electric_sources()
-            if self.hsg:
-                self.updates.hsg_1()
-            self.updates.update_electric_b()
-            if config.sim_config.general["solver"] == "cuda":
-                self.memused = self.updates.calculate_memory_used(iteration)
-
-        self.updates.finalise()
-        self.solvetime = self.updates.calculate_solve_time()
-        self.updates.cleanup()
+    return solver