Merge branch 'master' of https://github.com/rustychris/stompy

Author: gei-rusty

stompy/grid/unstructured_grid.py

@@ -7,6 +7,7 @@
 
 import sys,os,types
 import logging
+import pathlib # python 3 only, but python 2 is not worth supporting anymore.
 
 try:
     from osgeo import ogr,osr
@@ -76,6 +77,14 @@ def __init__(self,*a,meshes=[],**k):
         super().__init__(*a,**k)
         self.meshes=meshes
 
+class DuplicateMeshElement(GridException):
+    def __init__(self, dest_mesh, *a, dest_cell=-1, src_mesh=None, src_cell=-1, **k):
+        super().__init__(*a,**k)
+        self.dest_mesh = dest_mesh
+        self.src_mesh = src_mesh
+        self.dest_cell = dest_cell
+        self.src_cell=src_cell
+
 def request_square(ax,max_bounds=None):
     """
     Attempt to set a square aspect ratio on matplotlib axes ax
@@ -521,6 +530,9 @@ def modify_max_sides(self,max_sides):
             if not np.all( self.cells['nodes'][:,max_sides:] == self.UNDEFINED ):
                 raise GridException("Some cells cannot fit in requested max_sides")
 
+        # DFM hack, sometimes we have other cell:edge variables.
+        cell_side_fields=['edges','nodes','NetElemLink']
+            
         old_max_sides=self.max_sides
         self.max_sides=max_sides
 
@@ -537,7 +549,7 @@ def modify_max_sides(self,max_sides):
             else:
                 shape=None
 
-            if name in ['edges','nodes']:
+            if name in cell_side_fields: # 'edges','nodes','NetElemLink'
                 new_cell_dtype.append( (name,vtype,self.max_sides) )
             else:
                 new_cell_dtype.append( typeinfo ) # just copy
@@ -549,7 +561,7 @@ def modify_max_sides(self,max_sides):
 
         for typeinfo in old_dtype:
             name=typeinfo[0]
-            if name in ['edges','nodes']:
+            if name in cell_side_fields: 
                 if old_max_sides > self.max_sides:
                     new_cells[name][:,:] = self.cells[name][:,:self.max_sides]
                 else:
@@ -721,6 +733,9 @@ def from_ugrid(nc,mesh_name=None,skip_edges=False,fields='auto',
          face_node_connectivity, edge_node_connectivity, and optionally face_edge
          connectivity, edge_face_connectivity, etc.)
         """
+        if isinstance(nc,pathlib.Path):
+            nc=str(nc)
+            
         if isinstance(nc,six.string_types):
             filename=nc
             if dialect=='fvcom':
@@ -2906,7 +2921,13 @@ def bad_fields(Adata,Bdata): # field froms B which get dropped
             for i,edge in enumerate(kwargs['edges']):
                 kwargs['edges'][i]=edge_map[edge]
 
-            cell_map[n]=self.add_cell(**kwargs)
+            try:
+                cell_map[n]=self.add_cell(**kwargs)
+            except DuplicateMeshElement as exc:
+                exc.src_mesh = ugB
+                exc.src_cell = n
+                exc.dest_mesh = self
+                raise
 
         return node_map,edge_map,cell_map
 
@@ -4911,13 +4932,17 @@ def add_cell(self,**kwargs):
             if ( (n1==self.edges['nodes'][j,0]) and
                  (n2==self.edges['nodes'][j,1]) ):
                 # this cell is on the 'left' side of the edge
-                assert self.edges['cells'][j,0]<0
+                # assert self.edges['cells'][j,0]<0
+                if self.edges['cells'][j,0]>=0:
+                    raise DuplicateMeshElement(dest_mesh=self, dest_cell=self.edges['cells'][j,0])
                 # TODO: probably this ought to be using modify_edge
                 self.edges['cells'][j,0]=i
             elif ( (n1==self.edges['nodes'][j,1]) and
                    (n2==self.edges['nodes'][j,0]) ):
                 # this cell is on the 'right' side of the edge
-                assert self.edges['cells'][j,1]<0
+                # assert self.edges['cells'][j,1]<0
+                if self.edges['cells'][j,1]>=0:
+                    raise DuplicateMeshElement(dest_mesh=self, dest_cell=self.edges['cells'][j,1])
                 # TODO: probably this ought to be using modify_edge
                 self.edges['cells'][j,1]=i
             else:
@@ -5925,6 +5950,32 @@ def smooth_matrix(self,f=0.5,K='scaled',dx='grid',V='grid',A='grid',dt=None):
 
         return D.tocsr()
 
+    def fill_by_smoothing_func(self,cells,count=7):
+        """
+        Return a function that takes a subset of cell values and returns a dense, interpolated
+        field across all the cells of a grid.
+        cells: cell indexes where data is known
+        count: iterations of filling
+        """
+        M = self.smooth_matrix()
+        def fill_func(subset,cells=cells,M=M,count=count):
+            v=np.full(self.Ncells(), 0.0) # values
+            w=np.full(self.Ncells(), 0.0) # weights
+            valid=np.isfinite(subset)
+            v[cells[valid]] = subset[valid]
+            w[cells[valid]] = 1.0
+
+            for _ in range(count):
+                v=M.dot(v)
+                w=M.dot(w)
+
+            thresh=0.01
+            result = v/w.clip(thresh)
+            result[ w<thresh ]=np.nan
+            return result
+        return fill_func
+
+    
     def edge_clip_mask(self,xxyy,ends=False):
         """
         return a bitmask over edges falling in the boundiny box.
@@ -6014,7 +6065,7 @@ def plot_cells(self,ax=None,mask=None,values=None,clip=None,centers=False,labele
                 values = np.asanyarray(values)
 
         if ragged_edges is None:
-            ragged_edges='edgecolor' in kwargs
+            ragged_edges='edgecolor' in kwargs or 'ec' in kwargs
 
         if mask is None:
             mask=~self.cells['deleted']

stompy/model/delft/dflow_model.py

@@ -185,7 +185,7 @@ def copy_files_for_restart(self):
 
         for fn in os.listdir(self.restart_from.run_dir):
             _,suffix = os.path.splitext(fn)
-            do_copy = ( (suffix in ['.tim','.pli','.pliz','.ext','.xyz','.ini','.xyn'])
+            do_copy = ( (suffix in ['.tim','.pli','.pliz','.ext','.xyz','.ini','.xyn','.pol'])
                         or (fn in flowfm_ext)
                         # sources.sub is explicitly handled in write_config
                         or (fn in flowfm_mdu and fn !='sources.sub')
@@ -687,6 +687,11 @@ def partition(self,partition_grid=None):
             # slashes in the path to the mdu (ver 1.4.4)
             # since run_dflowfm uses run_dir as the working directory
             # here we strip to the basename
+            # Note that it may be tempting to add genpolygon=1 here, but
+            # in tests with 2023.01 that led to corrupt partitioning
+            # with ghost cells modified on some partitions but not 
+            # other partitions, including creating egregious concave
+            # cells.
             cmd=["--partition:ndomains=%d:icgsolver=6"%self.num_procs,
                  os.path.basename(self.mdu.filename)]
             print(f"About to call {cmd}")
@@ -707,7 +712,13 @@ def partition(self,partition_grid=None):
 
             # not a cross platform solution!
             gen_parallel=os.path.join(self.dfm_bin_dir,"generate_parallel_mdu.sh")
-            cmd=[gen_parallel,os.path.basename(self.mdu.filename),"%d"%self.num_procs,'6']
+            
+            cmd=[gen_parallel,os.path.basename(self.mdu.filename),"%d"%self.num_procs]
+            part_file =self.mdu['geometry','PartitioneFile']
+            if part_file not in [None,""]:
+                self.log.info("Passing partition file to generate_parallel_mdu.sh")
+                cmd.append(part_file)
+            cmd.append('6') # icgsolver
             print(f"About to call {cmd}")
             return utils.call_with_path(cmd,self.run_dir)
 
@@ -1508,7 +1519,7 @@ def map_outputs(self):
 
     _mu=None
     def map_dataset(self,force_multi=False,grid=None,chain=False,xr_kwargs={},
-                    clone_from=None):
+                    clone_from=None,refresh=False):
         """
         Return map dataset. For MPI runs, this will emulate a single, merged
         global dataset via multi_ugrid. For serial runs it directly opens
@@ -1527,10 +1538,17 @@ def map_dataset(self,force_multi=False,grid=None,chain=False,xr_kwargs={},
           which should speedup the loading process. Currently only handled when
           chain is False.
         """
+        default_xr_kwargs = dict(decode_timedelta=True)
+        xr_kwargs = default_xr_kwargs | xr_kwargs
+        
         if not chain:
             if self.num_procs<=1 and not force_multi:
                 # xarray caches this.
                 ds=xr.open_dataset(self.map_outputs()[0],**xr_kwargs)
+                if refresh:
+                    ds.close()
+                    ds=xr.open_dataset(self.map_outputs()[0],**xr_kwargs)
+                    
                 # put the grid in as an attribute so that the non-multiugrid
                 # and multiugrid return values can be used the same way.
                 ds.attrs['grid'] = ugrid.UnstructuredGrid.read_ugrid(ds)
@@ -1541,6 +1559,8 @@ def map_dataset(self,force_multi=False,grid=None,chain=False,xr_kwargs={},
                 if self._mu is None:
                     self._mu=multi_ugrid.MultiUgrid(self.map_outputs(),grid=grid,xr_kwargs=xr_kwargs,
                                                     clone_from=clone_from)
+                if refresh:
+                    self._mu.reload() # not sufficient if the grid changes, though.
                 return self._mu
         else:
             # as with his_dataset(), caching is not aware of options like chain.
@@ -1796,6 +1816,13 @@ def set_restart_from(self,model,deep=True,mdu_suffix=""):
         else:
             assert self.run_dir == model.run_dir
 
+        part_files = glob.glob(os.path.join(model.run_dir,"*_part.pol"))
+        if len(part_files)>0:
+            fn=part_files[0]
+            # Should get copied in copy_for_restart
+            self.log.info("Automatically adding {os.path.basename(fn)} as a partition file")
+            self.mdu['geometry','PartitionFile']=os.path.basename(fn)
+
         self.set_restart_file()
     def set_restart_file(self):
         """ 

stompy/model/delft/io.py

@@ -403,7 +403,7 @@ def read_pli(fn,one_per_line=True):
     """
     features=[]
 
-    with open(fn,'rt') as fp:
+    with open(fn,'rt',encoding='latin-1') as fp:
         if not one_per_line:
             toker=inp_tok(fp)
             token=lambda: six.next(toker)
@@ -1372,7 +1372,7 @@ def format_section(self,s):
     def get_value(self,sec_key):
         """
         return the string-valued settings for a given key.
-        if they key is not found, returns None.
+        if the key is not found, returns None.
         If the key is present but with no value, returns the empty string
         """
         section=self.format_section(sec_key[0])

stompy/model/delft/waq_scenario.py

@@ -9274,9 +9274,19 @@ def bloom_path(self):
     def original_bloominp_path(self):
         # this gets copied into the model run directory
         return os.path.join(self.share_path,'bloominp.d09')
+    
+    _proc_path=None
     @property
     def proc_path(self):
+        if self._proc_path is not None:
+            return self._proc_path
         return os.path.join(self.share_path,'proc_def')
+    
+    @proc_path.setter
+    def proc_path(self,value):
+        if value.endswith('.def') or value.endswith('.dat'):
+            raise ValueError("proc_path should not include the extension")
+        self._proc_path=value
 
     # plot process diagrams
     def cmd_plot_process(self,run_name='dwaq'):

stompy/plot/plot_utils.py

@@ -1503,6 +1503,14 @@ def plot_with_isolated(x,y,*a,**k):
     k.pop('label','')
     ax.plot(x[isolated],y[isolated],marker=marker, *a,**k)
 
+def zoom(fac, ax=None):
+    # fac:0 no zoom
+    # fac:1 double width, zooming out
+    # fac:-1 collapse to empty, zoom in infinitely.
+    ax = ax or plt.gca()
+    xxyy = utils.expand_xxyy(ax.axis(),fac)
+    ax.axis(xxyy)
+
 def add_scroller(ax):
     def on_scroll(event):
         if event.inaxes != ax:
@@ -1520,3 +1528,4 @@ def on_scroll(event):
     for cb in list(fig.canvas.callbacks.callbacks['scroll_event'])[2:]:
         fig.canvas.mpl_disconnect(cb)
     fig.canvas.mpl_connect('scroll_event', on_scroll)
+

stompy/spatial/join_features.py

@@ -510,8 +510,10 @@ def lines_to_polygons(new_features,close_arc=False,single_feature=True,force_ori
 
     log.info("Building index")
     # Because the index only hands back the poly, not an index.
+    # Shapely no longer likes user attributes, use a dict instead.
+    poly_to_join_id = {}
     for i,p in enumerate(simple_polys):
-        p.join_id=i
+        poly_to_join_id[p]=i  # p.join_id=i
 
     index=STRtree(simple_polys)
     log.info("done building index")
@@ -535,7 +537,8 @@ def lines_to_polygons(new_features,close_arc=False,single_feature=True,force_ori
         prep_ext_poly = prepare_geometry(ext_poly)
 
         hits=index.query(ext_poly)
-        hit_indexes=[p.join_id for p in hits]
+        hit_indexes=[poly_to_join_id(p) # p.join_id
+                     for p in hits]
         # this keeps us comparing large->small, needed to avoid
         # confusing islands in lake with lakes
         hit_indexes.sort()

stompy/utils.py

@@ -307,16 +307,16 @@ def within_2d(vecs,xxyy):
 
 
 def expand_xyxy(xyxy,factor):
-    dx=xyxy[2] - xyxy[0]
-    dy=xyxy[3] - xyxy[1]
+    dx=(xyxy[2] - xyxy[0])/2.0
+    dy=(xyxy[3] - xyxy[1])/2.0
     return [ xyxy[0] - dx*factor,
              xyxy[1] - dy*factor,
              xyxy[2] + dx*factor,
              xyxy[3] + dy*factor]
 
 def expand_xxyy(xxyy,factor):
-    dx=xxyy[1] - xxyy[0]
-    dy=xxyy[3] - xxyy[2]
+    dx=(xxyy[1] - xxyy[0])/2.0
+    dy=(xxyy[3] - xxyy[2])/2.0
     return [ xxyy[0] - dx*factor,
              xxyy[1] + dx*factor,
              xxyy[2] - dy*factor,