Fix for #602

autotest/t009_test.py

@@ -399,7 +399,7 @@ def test_ds_6d_6e_disordered():
     sfr2 = m2.get_package("SFR")
 
 
-    if len(sfr.all_segments) != len(sfr2.all_segments):
+    if len(sfr.graph) != len(sfr2.graph):
         raise AssertionError
 
     if len(sfr.segment_data[0]) != len(sfr2.segment_data[0]):

autotest/t010_test.py

@@ -79,9 +79,9 @@ def test_sfrcheck():
     m.sfr.reach_data['ireach'][3] += 1
 
     # create circular routing instance
-    m.sfr.segment_data[0]['outseg'][1] = 1
-    m.sfr.segment_data[0]['outseg']
-    
+    m.sfr.segment_data[0]['outseg'][0] = 1
+    m.sfr._graph = None  # weak, but the above shouldn't happen
+
     chk = check(m.sfr)
     chk.numbering()
     assert 'continuity in segment and reach numbering' in chk.errors

autotest/t064_test_performance.py

@@ -0,0 +1,94 @@
+"""
+Tests to prevent performance regressions
+"""
+import os
+import shutil
+import time
+import numpy as np
+import flopy.modflow as fm
+
+
+class TestModflowPerformance():
+    """Test flopy.modflow performance with realistic model/package sizes,
+    in a reasonable timeframe.
+    """
+    @classmethod
+    def setup_class(cls):
+        """Make a modflow model."""
+        print('setting up model...')
+        t0 = time.time()
+        size = 100
+        nlay = 10
+        nper = 10
+        nsfr = int((size ** 2)/5)
+
+        cls.modelname = 'junk'
+        cls.model_ws = 'temp/t064'
+        external_path = 'external/'
+
+        if not os.path.isdir(cls.model_ws):
+            os.makedirs(cls.model_ws)
+        if not os.path.isdir(os.path.join(cls.model_ws, external_path)):
+            os.makedirs(os.path.join(cls.model_ws, external_path))
+
+        m = fm.Modflow(cls.modelname, model_ws=cls.model_ws, external_path=external_path)
+
+        dis = fm.ModflowDis(m, nper=nper, nlay=nlay, nrow=size, ncol=size,
+                            top=nlay, botm=list(range(nlay)))
+
+        rch = fm.ModflowRch(m, rech={k: .001 - np.cos(k) * .001 for k in range(nper)})
+
+        ra = fm.ModflowWel.get_empty(size ** 2)
+        well_spd = {}
+        for kper in range(nper):
+            ra_per = ra.copy()
+            ra_per['k'] = 1
+            ra_per['i'] = (np.ones((size, size)) * np.arange(size)).transpose().ravel().astype(int)
+            ra_per['j'] = list(range(size)) * size
+            well_spd[kper] = ra
+        wel = fm.ModflowWel(m, stress_period_data=well_spd)
+
+        # SFR package
+        rd = fm.ModflowSfr2.get_empty_reach_data(nsfr)
+        rd['iseg'] = range(len(rd))
+        rd['ireach'] = 1
+        sd = fm.ModflowSfr2.get_empty_segment_data(nsfr)
+        sd['nseg'] = range(len(sd))
+        sfr = fm.ModflowSfr2(reach_data=rd, segment_data=sd, model=m)
+        cls.init_time = time.time() - t0
+        cls.m = m
+
+    def test_init_time(self):
+        """test model and package init time(s)."""
+        mfp = TestModflowPerformance()
+        target = 0.3 # seconds
+        assert mfp.init_time < target, "model init took {:.2fs}, should take {:.1f}s".format(mfp.init_time, target)
+        print('setting up model took {:.2f}s'.format(mfp.init_time))
+
+    def test_0_write_time(self):
+        """test write time"""
+        print('writing files...')
+        mfp = TestModflowPerformance()
+        target = 5
+        t0 = time.time()
+        mfp.m.write_input()
+        t1 = time.time() - t0
+        assert t1 < target, "model write took {:.2fs}, should take {:.1f}s".format(t1, target)
+        print('writing input took {:.2f}s'.format(t1))
+
+    def test_9_load_time(self):
+        """test model load time"""
+        print('loading model...')
+        mfp = TestModflowPerformance()
+        target = 3
+        t0 = time.time()
+        m = fm.Modflow.load('{}.nam'.format(mfp.modelname),
+                            model_ws=mfp.model_ws, check=False)
+        t1 = time.time() - t0
+        assert t1 < target, "model load took {:.2fs}, should take {:.1f}s".format(t1, target)
+        print('loading the model took {:.2f}s'.format(t1))
+
+    @classmethod
+    def teardown_class(cls):
+        # cleanup
+        shutil.rmtree(cls.model_ws)

flopy/modflow/mfsfr2.py

@@ -370,6 +370,7 @@ def __init__(self, model, nstrm=-2, nss=1, nsfrpar=0, nparseg=0,
                          unit_number=units, extra=extra, filenames=fname)
 
         self.url = 'sfr2.htm'
+        self._graph = None # dict of routing connections
 
         # Dataset 0
         self.heading = '# {} package for '.format(self.name[0]) + \
@@ -483,14 +484,14 @@ def __init__(self, model, nstrm=-2, nss=1, nsfrpar=0, nparseg=0,
                     self.segment_data[i][n] = segment_data[i][n]
         # compute outreaches if nseg and outseg columns have non-default values
         if np.diff(self.reach_data.iseg).max() != 0 and \
-                np.diff(self.all_segments.nseg).max() != 0 \
-                and np.diff(self.all_segments.outseg).max() != 0:
-            if len(self.all_segments) == 1:
+                np.max(list(set(self.graph.keys()))) != 0 \
+                and np.max(list(set(self.graph.values()))) != 0:
+            if len(self.graph) == 1:
                 self.segment_data[0]['nseg'] = 1
                 self.reach_data['iseg'] = 1
 
             consistent_seg_numbers = len(set(self.reach_data.iseg).difference(
-                set(self.all_segments.nseg))) == 0
+                set(self.graph.keys()))) == 0
             if not consistent_seg_numbers:
                 warnings.warn(
                     "Inconsistent segment numbers of reach_data and segment_data")
@@ -582,44 +583,12 @@ def dataset_5(self):
                     ds5[per] = [len(sd), irdflag[per], iptflag[per]]
         return ds5
 
-    @property
-    def all_segments(self):
-        """
-        Method to get a list of all segments in the simulation,
-        since all segments do not have to be active in a given stress
-        period
-
-        returns:
-        -------
-        ra : (np.recarray)
-            recarray contains a single entry of segment data for each stream segment
-        """
-
-        ra = self.get_empty_segment_data(self.nss)
-        i = 0
-        for _, recarray in sorted(self.segment_data.items()):
-            if i == self.nss:
-                break
-            else:
-                for rec in recarray:
-                    if rec.nseg in ra.nseg:
-                        pass
-                    else:
-                        ra[i] = rec
-                        i += 1
-
-                    if i == self.nss:
-                        break
-        return ra
-
     @property
     def graph(self):
-        graph = dict(
-            zip(self.all_segments.nseg, self.all_segments.outseg))
-        outlets = set(graph.values()).difference(
-            set(graph.keys()))  # including lakes
-        graph.update({o: 0 for o in outlets})
-        return graph
+        """Dictionary of routing connections between segments."""
+        if self._graph is None:
+            self._graph = self._make_graph()
+        return self._graph
 
     @property
     def paths(self):
@@ -630,9 +599,10 @@ def paths(self):
         nseg = np.array(sorted(self._paths.keys()), dtype=int)
         nseg = nseg[nseg > 0].copy()
         outseg = np.array([self._paths[k][1] for k in nseg])
-        sd = self.all_segments
-        if not np.array_equal(nseg, sd.nseg) or not np.array_equal(outseg,
-                                                                   sd.outseg):
+        existing_nseg = sorted(list(self.graph.keys()))
+        existing_outseg = [self.graph[k] for k in existing_nseg]
+        if not np.array_equal(nseg, existing_nseg) or \
+                not np.array_equal(outseg, existing_outseg):
             self._set_paths()
         return self._paths
 
@@ -644,6 +614,17 @@ def df(self):
             msg = 'ModflowSfr2.df: pandas not available'
             raise ImportError(msg)
 
+    def _make_graph(self):
+        # get all segments and their outseg
+        graph = {}
+        for recarray in self.segment_data.values():
+            graph.update(dict(zip(recarray['nseg'], recarray['outseg'])))
+
+        outlets = set(graph.values()).difference(
+            set(graph.keys()))  # including lakes
+        graph.update({o: 0 for o in outlets if o != 0})
+        return graph
+
     def _set_paths(self):
         graph = self.graph
         self._paths = {seg: find_path(graph, seg) for seg in graph.keys()}
@@ -973,6 +954,7 @@ def check(self, f=None, verbose=True, level=1):
         >>> m = flopy.modflow.Modflow.load('model.nam')
         >>> m.sfr2.check()
         """
+        self._graph = None  # remake routing graph from segment data
         chk = check(self, verbose=verbose, level=level)
         chk.for_nans()
         chk.numbering()
@@ -1176,10 +1158,6 @@ def reset_reaches(self):
         self.reach_data.sort(order=['iseg', 'ireach'])
         reach_data = self.reach_data
         segment_data = list(set(self.reach_data.iseg))# self.segment_data[0]
-        # ireach = []
-        # for iseg in segment_data.nseg:
-        #    nreaches = np.sum(reach_data.iseg == iseg)
-        #    ireach += list(range(1, nreaches + 1))
         reach_counts = np.bincount(reach_data.iseg)[1:]
         reach_counts = dict(zip(range(1, len(reach_counts) + 1),
                                 reach_counts))
@@ -1336,6 +1314,7 @@ def repair_outsegs(self):
                              self.segment_data[0].nseg)
         isasegment = isasegment | (self.segment_data[0].outseg < 0)
         self.segment_data[0]['outseg'][~isasegment] = 0.
+        self._graph = None
 
     def renumber_segments(self):
         """
@@ -1348,11 +1327,8 @@ def renumber_segments(self):
         r : dictionary mapping old segment numbers to new
         """
 
-        segments = self.all_segments
-        segments.sort(order="nseg")
-        # get renumbering info from per=0
-        nseg = segments.nseg
-        outseg = segments.outseg
+        nseg = sorted(list(self.graph.keys()))
+        outseg = [self.graph[k] for k in nseg]
 
         # explicitly fix any gaps in the numbering
         # (i.e. from removing segments)
@@ -1400,6 +1376,7 @@ def reassign_upsegs(r, nexts, upsegs):
             assert not np.any(inds)
             assert len(self.segment_data[per]['nseg']) == \
                    self.segment_data[per]['nseg'].max()
+        self._graph = None # reset routing dict
 
         # renumber segments in reach_data
         self.reach_data['iseg'] = [r.get(s, s) for s in self.reach_data.iseg]
@@ -1981,7 +1958,6 @@ def __init__(self, sfrpackage, verbose=True, level=1):
 
         self.reach_data = sfrpackage.reach_data
         self.segment_data = sfrpackage.segment_data
-        self.all_segments = sfrpackage.all_segments
         self.verbose = verbose
         self.level = level
         self.passed = []
@@ -2202,7 +2178,7 @@ def routing(self):
 
         # check reach connections for proximity
         if self.mg is not None or self.mg is not None:
-            rd = self.sfr.reach_data
+            rd = self.sfr.reach_data.copy()
             rd.sort(order=['reachID'])
             try:
                 xcentergrid, ycentergrid, zc = self.mg.get_cellcenters()
@@ -3096,6 +3072,7 @@ def _parse_6bc(line, icalc, nstrm, isfropt, reachinput, per=0):
 
 def find_path(graph, start, end=0, path=()):
 
+    graph = graph.copy()
     path = list(path) + [start]
     if start == end:
         return path