Add type information

gprMax/scene.py

@@ -17,6 +17,7 @@
 # along with gprMax.  If not, see <http://www.gnu.org/licenses/>.
 
 import logging
+from typing import List, Optional, Union
 
 from gprMax.cmds_geometry.add_grass import AddGrass
 from gprMax.cmds_geometry.add_surface_roughness import AddSurfaceRoughness
@@ -26,6 +27,8 @@ from gprMax.cmds_geometry.fractal_box import FractalBox
 from gprMax.cmds_multiuse import UserObjectMulti
 from gprMax.cmds_singleuse import Discretisation, Domain, TimeWindow, UserObjectSingle
 from gprMax.materials import create_built_in_materials
+from gprMax.model import Model
+from gprMax.subgrids.grid import SubGridBaseGrid
 from gprMax.subgrids.user_objects import SubGridBase as SubGridUserBase
 from gprMax.user_inputs import create_user_input_points
 
@@ -36,9 +39,9 @@ class Scene:
     """Scene stores all of the user created objects."""
 
     def __init__(self):
-        self.multiple_cmds = []
-        self.single_cmds = []
-        self.geometry_cmds = []
+        self.multiple_cmds: List[UserObjectMulti] = []
+        self.single_cmds: List[UserObjectSingle] = []
+        self.geometry_cmds: List[UserObjectGeometry] = []
         self.essential_cmds = [Domain, TimeWindow, Discretisation]
 
     def add(self, user_object):
@@ -72,17 +75,13 @@ class Scene:
             logger.exception("Error creating user input object")
             raise
 
-    def process_subgrid_cmds(self):
+    def process_subgrid_cmds(self, model: Model):
         """Process all commands in any sub-grids."""
 
-        def func(obj):
-            if isinstance(obj, SubGridUserBase):
-                return True
-            else:
-                return False
-
         # Subgrid user objects
-        subgrid_cmds = list(filter(func, self.multiple_cmds))
+        subgrid_cmds = [
+            sg_cmd for sg_cmd in self.multiple_cmds if isinstance(sg_cmd, SubGridUserBase)
+        ]
 
         # Iterate through the user command objects under the subgrid user object
         for sg_cmd in subgrid_cmds:
@@ -90,11 +89,17 @@ class Scene:
             # object. This reference allows the multi and geo user objects
             # to build in the correct subgrid.
             sg = sg_cmd.subgrid
-            self.process_cmds(sg_cmd.children_multiple, sg)
+            self.process_cmds(sg_cmd.children_multiple, model, sg)
             self.process_geocmds(sg_cmd.children_geometry, sg)
 
-    def process_cmds(self, commands, grid):
+    def process_cmds(
+        self,
+        commands: Union[List[UserObjectMulti], List[UserObjectSingle]],
+        model: Model,
+        subgrid: Optional[SubGridBaseGrid] = None,
+    ):
         """Process list of commands."""
+        grid = model.G if subgrid is None else subgrid
         cmds_sorted = sorted(commands, key=lambda cmd: cmd.order)
         for obj in cmds_sorted:
             self.build_obj(obj, grid)
@@ -118,7 +123,7 @@ class Scene:
 
         return self
 
-    def process_singlecmds(self, G):
+    def process_singlecmds(self, model: Model):
         # Check for duplicate commands and warn user if they exist
         cmds_unique = list(set(self.single_cmds))
         if len(cmds_unique) != len(self.single_cmds):
@@ -136,7 +141,7 @@ class Scene:
                 )
                 raise ValueError
 
-        self.process_cmds(cmds_unique, G)
+        self.process_cmds(cmds_unique, model)
 
     def create_internal_objects(self, G):
         """Calls the UserObject.build() function in the correct way - API
@@ -148,10 +153,10 @@ class Scene:
         create_built_in_materials(G)
 
         # Process commands that can only have a single instance
-        self.process_singlecmds(G)
+        self.process_singlecmds(model)
 
         # Process main grid multiple commands
-        self.process_cmds(self.multiple_cmds, G)
+        self.process_cmds(self.multiple_cmds, model)
 
         # Initialise geometry arrays for main and subgrids
         for grid in [G] + G.subgrids:
@@ -161,4 +166,4 @@ class Scene:
         self.process_geocmds(self.geometry_cmds, G)
 
         # Process all the commands for subgrids
-        self.process_subgrid_cmds()
+        self.process_subgrid_cmds(model)

gprMax/subgrids/user_objects.py

@@ -18,9 +18,15 @@
 
 import logging
 from copy import copy
+from typing import List, Tuple, Union
 
 import numpy as np
 
+from gprMax.grid.fdtd_grid import FDTDGrid
+from gprMax.model import Model
+from gprMax.subgrids.grid import SubGridBaseGrid
+from gprMax.user_inputs import MainGridUserInput
+
 from ..cmds_geometry.cmds_geometry import UserObjectGeometry
 from ..cmds_multiuse import UserObjectMulti
 from .subgrid_hsg import SubGridHSG as SubGridHSGUser
@@ -35,10 +41,10 @@ class SubGridBase(UserObjectMulti):
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.children_multiple = []
-        self.children_geometry = []
+        self.children_multiple: List[UserObjectMulti] = []
+        self.children_geometry: List[UserObjectGeometry] = []
 
-    def add(self, node):
+    def add(self, node: Union[UserObjectMulti, UserObjectGeometry]):
         """Adds other user objects. Geometry and multi only."""
         if isinstance(node, UserObjectMulti):
             self.children_multiple.append(node)
@@ -48,13 +54,15 @@ class SubGridBase(UserObjectMulti):
             logger.exception(f"{str(node)} this Object can not be added to a sub grid")
             raise ValueError
 
-    def set_discretisation(self, sg, grid):
+    def set_discretisation(self, sg: SubGridBaseGrid, grid: FDTDGrid):
         sg.dx = grid.dx / sg.ratio
         sg.dy = grid.dy / sg.ratio
         sg.dz = grid.dz / sg.ratio
         sg.dl = np.array([sg.dx, sg.dy, sg.dz])
 
-    def set_main_grid_indices(self, sg, uip, p1, p2):
+    def set_main_grid_indices(
+        self, sg: SubGridBaseGrid, uip: MainGridUserInput, p1: Tuple[int], p2: Tuple[int]
+    ):
         """Sets subgrid indices related to main grid placement."""
         # Location of the IS
         sg.i0, sg.j0, sg.k0 = p1
@@ -63,26 +71,26 @@ class SubGridBase(UserObjectMulti):
         sg.x1, sg.y1, sg.z1 = uip.round_to_grid(p1)
         sg.x2, sg.y2, sg.z2 = uip.round_to_grid(p2)
 
-    def set_name(self, sg):
+    def set_name(self, sg: SubGridBaseGrid):
         sg.name = self.kwargs["id"]
 
-    def set_working_region_cells(self, sg):
+    def set_working_region_cells(self, sg: SubGridBaseGrid):
         """Number of cells in each dimension for the working region."""
         sg.nwx = (sg.i1 - sg.i0) * sg.ratio
         sg.nwy = (sg.j1 - sg.j0) * sg.ratio
         sg.nwz = (sg.k1 - sg.k0) * sg.ratio
 
-    def set_total_cells(self, sg):
+    def set_total_cells(self, sg: SubGridBaseGrid):
         """Number of cells in each dimension for the whole region."""
         sg.nx = 2 * sg.n_boundary_cells_x + sg.nwx
         sg.ny = 2 * sg.n_boundary_cells_y + sg.nwy
         sg.nz = 2 * sg.n_boundary_cells_z + sg.nwz
 
-    def set_iterations(self, sg, main):
+    def set_iterations(self, sg: SubGridBaseGrid, main: FDTDGrid):
         """Sets number of iterations that will take place in the subgrid."""
         sg.iterations = main.iterations * sg.ratio
 
-    def setup(self, sg, grid, uip):
+    def setup(self, sg: SubGridBaseGrid, grid: FDTDGrid, uip: MainGridUserInput):
         """ "Common setup to both all subgrid types."""
         p1 = self.kwargs["p1"]
         p2 = self.kwargs["p2"]
@@ -192,7 +200,7 @@ class SubGridHSG(SubGridBase):
         self.order = 18
         self.hash = "#subgrid_hsg"
 
-    def build(self, grid, uip):
+    def build(self, grid: FDTDGrid, uip: MainGridUserInput) -> SubGridHSGUser:
         sg = SubGridHSGUser(**self.kwargs)
         self.setup(sg, grid, uip)
         return sg