Add imports for polygons at top level of yaml.

Author: Eli

schimpy/prepare_schism.py

@@ -84,7 +84,7 @@ def create_hgrid(s, inputs, logger):
             )
         small_area_param = section.get("small_areas")
         if small_area_param is not None:
-            small_area_param["prepro_output_dir"] = inputs["prepro_output_dir"]
+            small_area_param["prepro_output_dir"] = inputs["diagnostics_dir"]
             # This just emits warnings unless the fail threshold is met
             small_areas(s.mesh, logger=logger, **small_area_param)
 
@@ -151,7 +151,7 @@ def create_hgrid(s, inputs, logger):
                 s.mesh.nodes[:, 2] = stacked_dem_fill(
                     dem_list,
                     s.mesh.nodes[:, :2],
-                    inputs["prepro_output_dir"],
+                    inputs["diagnostics_dir"],
                     require_all=False,
                     na_fill=default_depth_for_missing_dem,
                     negate=True,
@@ -173,7 +173,8 @@ def create_hgrid(s, inputs, logger):
                         smooth = href_spec.get("smooth", None)
                         # reuse the exact machinery used by the 'gr3' section:
                         s.create_node_partitioning(
-                            href_path, polygons, default, smooth
+                            href_path, polygons, default, smooth,
+                            global_imports=inputs.get("imports"),
                         )  # writes GR3
                         href_arg = href_path
                     else:
@@ -244,7 +245,7 @@ def create_hgrid(s, inputs, logger):
                     stacked_dem_fill(
                         dem_list,
                         s.mesh.nodes[:, :2],
-                        inputs["prepro_output_dir"],
+                        inputs["diagnostics_dir"],
                         require_all=False,
                         na_fill=default_depth_for_missing_dem,
                     )
@@ -259,13 +260,16 @@ def create_hgrid(s, inputs, logger):
 
         depth_enforce_params = section.get("depth_enforcement")
         if depth_enforce_params is not None:
+            global_imports = inputs.get("imports")
             if "polygons" in depth_enforce_params:
                 s.mesh.nodes[:, 2] = s.apply_polygons(
-                    default=None, polygons=depth_enforce_params["polygons"]
+                    default=None, polygons=depth_enforce_params["polygons"],
+                    global_imports=global_imports,
                 )
             if "linestrings" in depth_enforce_params:
                 s.mesh.nodes[:, 2] = s.apply_linestring_ops(
-                    default=None, linestrings=depth_enforce_params["linestrings"]
+                    default=None, linestrings=depth_enforce_params["linestrings"],
+                    global_imports=global_imports,
                 )
 
         # Write hgrid.gr3
@@ -295,6 +299,7 @@ def create_vgrid(s, inputs, logger):
         return
 
     output_dir = inputs["prepro_output_dir"]
+    diagnostics_dir = inputs["diagnostics_dir"]
 
     # --- Resolve hgrid (shared by v1 and v2) ---
     if "hgrid" in section:
@@ -333,6 +338,7 @@ def create_vgrid(s, inputs, logger):
             )
             for k in stray:
                 section.pop(k)
+        section["diagnostics_dir"] = diagnostics_dir
         vgrid_gen(hgrid, **section)
 
     elif gen_version == "v2":
@@ -343,6 +349,10 @@ def create_vgrid(s, inputs, logger):
         )
         # v2 passes named arguments
         section.pop("hgrid", None)
+        # Route debug_prefix under diagnostics_dir if set
+        debug_prefix = section.get("debug_prefix")
+        if debug_prefix:
+            debug_prefix = os.path.join(diagnostics_dir, debug_prefix)
         vgrid_gen_v2(
             hgrid=hgrid,
             vgrid_out=section["vgrid_out"],
@@ -353,7 +363,7 @@ def create_vgrid(s, inputs, logger):
             region_constraints=section.get("region_constraints"),
             constraint_taper_rings=int(section.get("constraint_taper_rings", 3)),
             dz_scale_gr3=section.get("dz_scale_gr3"),
-            debug_prefix=section.get("debug_prefix"),
+            debug_prefix=debug_prefix,
             pileup_log=section.get("pileup_log"),
         )
     else:
@@ -394,6 +404,7 @@ def create_gr3_with_polygons(s, inputs, logger):
     if dict_gr3 is None:
         return
     logger.info("Processing gr3 outputs...")
+    global_imports = inputs.get("imports")
     expected_items = ("polygons", "default", "smooth")
     output_dir = inputs["prepro_output_dir"]
     for fname, item in dict_gr3.items():
@@ -419,7 +430,7 @@ def create_gr3_with_polygons(s, inputs, logger):
         default = item.get("default")
         logger.info("Creating %s..." % fname)
         smooth = item.get("smooth")
-        s.create_node_partitioning(fname, polygons, default, smooth)
+        s.create_node_partitioning(fname, polygons, default, smooth, global_imports=global_imports)
 
 
 def create_prop_with_polygons(s, inputs, logger):
@@ -428,6 +439,7 @@ def create_prop_with_polygons(s, inputs, logger):
     if dict_prop is None:
         return
     logger.info("Processing prop outputs...")
+    global_imports = inputs.get("imports")
     expected_items = ("default", "polygons")
     output_dir = inputs["prepro_output_dir"]
     for fname, item in dict_prop.items():
@@ -445,7 +457,7 @@ def create_prop_with_polygons(s, inputs, logger):
         polygons = item.get("polygons", [])
         default = item.get("default")
         logger.info("Creating %s..." % fname)
-        s.create_prop_partitioning(fname, polygons, default)
+        s.create_prop_partitioning(fname, polygons, default, global_imports=global_imports)
 
 
 def create_structures(s, inputs, logger):
@@ -761,6 +773,7 @@ def process_output_dir(inputs):
             os.mkdir(outdir)
         else:
             raise ValueError("Output directory (output_dir) does not exist")
+
     return outdir
 
 
@@ -811,6 +824,7 @@ def process_prepare_yaml(in_fname, use_logging=True, write_echo=True, envvar=Non
     keys_top_level = [
         "config",
         "env",
+        "imports",
         "min_schimpy_version",
         "prepro_output_dir",
         "mesh",
@@ -827,6 +841,11 @@ def process_prepare_yaml(in_fname, use_logging=True, write_echo=True, envvar=Non
     logger.info("Processing the top level...")
     check_and_suggest(list(inputs.keys()), keys_top_level, logger)
 
+    # Create diagnostics subdirectory for intermediate/debug outputs
+    diagnostics_dir = os.path.join(outdir, "diagnostics")
+    os.makedirs(diagnostics_dir, exist_ok=True)
+    inputs["diagnostics_dir"] = diagnostics_dir
+
     if write_echo:
         out_fname = (
             os.path.splitext(in_fname)[0] + "_echo" + os.path.splitext(in_fname)[1]

schimpy/schism_setup.py

@@ -507,7 +507,7 @@ def _parse_attribute(self, expr):
         expr = re.sub(r"(\b)max(\b)", r"\g<1>numpy.maximum\g<2>", expr)
         return expr
 
-    def apply_linestring_ops(self, default, linestrings):
+    def apply_linestring_ops(self, default, linestrings, global_imports=None):
         mesh = self.mesh
         attr = np.copy(mesh.nodes[:, 2])
 
@@ -569,8 +569,19 @@ def apply_linestring_ops(self, default, linestrings):
                             node_local_nds[inode] = set(elnodes)
 
                 newglobals = {}
-                for item in newglobals:
-                    print(newglobals[item])
+                _ls_imports = ["numpy"] + (global_imports if global_imports else [])
+                for imp in _ls_imports:
+                    iname = imp.split(".")[-1]
+                    try:
+                        newglobals[iname] = importlib.import_module(imp)
+                    except (ModuleNotFoundError, ImportError):
+                        parent = ".".join(imp.split(".")[:-1])
+                        if parent:
+                            try:
+                                parent_mod = importlib.import_module(parent)
+                                newglobals[iname] = getattr(parent_mod, iname)
+                            except Exception:
+                                pass
                 for inode in node_local_nds:
                     # inode is the node being calculated
                     support_nodes = list(node_local_nds[inode])
@@ -594,7 +605,7 @@ def apply_linestring_ops(self, default, linestrings):
                     attr[inode] = valreplace
         return attr
 
-    def apply_polygons(self, polygons, default):
+    def apply_polygons(self, polygons, default, global_imports=None):
         """Partition the grid with the given polygons.
         Each node (not element) will be assigned with an integer ID
         which is the index of the polygons.
@@ -607,6 +618,9 @@ def apply_polygons(self, polygons, default):
         polygons: list
             a list of polygon dict (from YAML most of time)
 
+        global_imports: list, optional
+            list of module paths available to all polygon expressions
+
         Returns
         -------
         numpy.ndarray
@@ -625,7 +639,7 @@ def apply_polygons(self, polygons, default):
         for polygon in polygons:
             name = polygon.get("name")
             imports_str = polygon.get("imports")
-            imports = ["numpy"]
+            imports = ["numpy"] + (global_imports if global_imports else [])
             if isinstance(imports_str, str):
                 imports.extend(imports_str.split())
             vertices = np.array(polygon.get("vertices"))
@@ -697,16 +711,23 @@ def _evaluate_in_polygon(self, expr, polygon, globals):
         mesh = globals["mesh"]
         imports = globals["imports"]
         mod_names = {}
-        # if imports is not None:
         for imp in imports:
             iname = imp.split(".")[-1]
-            # ipack = '.'.join(imp.split('.')[0:-1])
             try:
-                # if len(ipack) > 0 else importlib.import_module(iname)
                 mod = importlib.import_module(imp)
                 mod_names[iname] = mod
-            except:
-                print("Failed to import {}".format(imp))
+            except (ModuleNotFoundError, ImportError):
+                # Support "from parent import attr" style, e.g.
+                # "schimpy.sav_density_raster.sav_density" -> sav_density function
+                parent = ".".join(imp.split(".")[:-1])
+                if parent:
+                    try:
+                        parent_mod = importlib.import_module(parent)
+                        mod_names[iname] = getattr(parent_mod, iname)
+                    except Exception:
+                        print("Failed to import {}".format(imp))
+                else:
+                    print("Failed to import {}".format(imp))
         globals.update(mod_names)
 
         box = np.array(polygon.bounds)[[0, 2, 1, 3]]
@@ -720,15 +741,15 @@ def _evaluate_in_polygon(self, expr, polygon, globals):
         vals = eval(expr, globals)
         return nodes_in_polygon, vals
 
-    def create_node_partitioning(self, gr3_fname, polygons, default, smooth):
+    def create_node_partitioning(self, gr3_fname, polygons, default, smooth, global_imports=None):
         """Create a gr3 file with node partitioning using
         polygons in the polygon file.
         gr3_fname = output gr3 file name
         polygons = polygons
         """
         from schimpy.laplace_smooth_data import laplace_smooth_data2
 
-        attr = self.apply_polygons(polygons, default)
+        attr = self.apply_polygons(polygons, default, global_imports=global_imports)
         if smooth is not None:
             self._logger.info("Smoothing data. This process can take several minutes")
             attr = laplace_smooth_data2(self.mesh, data=attr, **smooth)
@@ -737,7 +758,7 @@ def create_node_partitioning(self, gr3_fname, polygons, default, smooth):
         else:
             self.write_hgrid(gr3_fname, attr)
 
-    def create_prop_partitioning(self, prop_fname, polygons, default):
+    def create_prop_partitioning(self, prop_fname, polygons, default, global_imports=None):
         """Create a prop file with element partitioning using
         polygons in the polygon file.
         prop_fname = output prop file name
@@ -747,7 +768,7 @@ def create_prop_partitioning(self, prop_fname, polygons, default):
         # # default = polygon_data[option_name]
         # polygons = polygon_data['polygons']
         # attr = self._partition_nodes_with_polygons(polygons, default)
-        attr = self.apply_polygons(polygons, default)
+        attr = self.apply_polygons(polygons, default, global_imports=global_imports)
         mesh = self.mesh
         elementflags = np.empty((mesh.n_elems(), 1))
         elementflags.fill(0.0)