From 601c67ee4703991081c3a3b0652c03695152af0a Mon Sep 17 00:00:00 2001 From: nmannall Date: Thu, 18 Apr 2024 17:41:35 +0100 Subject: [PATCH] Add type information --- gprMax/scene.py | 41 ++++++++++++++++++--------------- gprMax/subgrids/user_objects.py | 30 +++++++++++++++--------- 2 files changed, 42 insertions(+), 29 deletions(-) diff --git a/gprMax/scene.py b/gprMax/scene.py index 28b93cde..a2691e7a 100644 --- a/gprMax/scene.py +++ b/gprMax/scene.py @@ -17,6 +17,7 @@ # along with gprMax. If not, see . 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) diff --git a/gprMax/subgrids/user_objects.py b/gprMax/subgrids/user_objects.py index 3baa1028..a62b98b4 100644 --- a/gprMax/subgrids/user_objects.py +++ b/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