Improved slow tasks for MEX

openquake/baselib/general.py

@@ -352,8 +352,7 @@ def split_in_blocks(sequence, hint, weight=lambda item: 1, key=nokey):
 
      >>> items = 'ABCDE'
      >>> list(split_in_blocks(items, 3))
-     [<WeightedSequence ['A', 'B'], weight=2>, <WeightedSequence ['C', 'D'], weight=2>, <WeightedSequence ['E'], weight=1>]
-
+     [<WeightedSequence ['A'], weight=1>, <WeightedSequence ['B'], weight=1>, <WeightedSequence ['C'], weight=1>, <WeightedSequence ['D'], weight=1>, <WeightedSequence ['E'], weight=1>]
     """
     if isinstance(sequence, pandas.DataFrame):
         num_elements = len(sequence)
@@ -373,7 +372,7 @@ def split_in_blocks(sequence, hint, weight=lambda item: 1, key=nokey):
     assert hint > 0, hint
     assert len(items) > 0, len(items)
     total_weight = float(sum(weight(item) for item in items))
-    return block_splitter(items, math.ceil(total_weight / hint), weight, key)
+    return block_splitter(items, total_weight / hint, weight, key)
 
 
 def assert_close(a, b, rtol=1e-07, atol=0, context=None):

openquake/calculators/tests/classical_test.py

@@ -81,7 +81,7 @@ def test_case_01(self):
 
             slow = view('task:classical:-1', self.calc.datastore)
             self.assertIn('taskno', slow)
-            self.assertIn('duration', slow)
+            self.assertIn('time', slow)
 
         # there is a single source
         self.assertEqual(len(self.calc.datastore['source_info']), 1)
@@ -803,4 +803,4 @@ def test_case_87(self):
             'hazard_curve-mean-SA(0.2).csv',
             'hazard_curve-mean-SA(1.0).csv',
             'hazard_curve-mean-SA(2.0).csv'],
-            case_87.__file__)
+            case_87.__file__)

openquake/calculators/tests/logictree_test.py

@@ -155,12 +155,6 @@ def test_case_07(self):
              'hazard_curve-smltp_b3-gsimltp_b1.csv'],
             case_07.__file__)
 
-        # check the weights of the sources
-        info = self.calc.datastore.read_df('source_info', 'source_id')
-        self.assertEqual(info.loc[b'1'].weight, 276)
-        self.assertEqual(info.loc[b'2'].weight, 177)
-        self.assertEqual(info.loc[b'3'].weight, 5871)
-
         # testing view_relevant_sources
         arr = view('relevant_sources:PGA', self.calc.datastore)
         self.assertEqual(decode(arr['src_id']), ['1', '2'])

openquake/calculators/views.py

@@ -37,6 +37,7 @@
 from openquake.baselib.performance import performance_view, Monitor
 from openquake.baselib.python3compat import encode, decode
 from openquake.hazardlib import logictree, calc, source, geo
+from openquake.hazardlib.valid import basename
 from openquake.hazardlib.contexts import ContextMaker
 from openquake.commonlib import util
 from openquake.risklib import riskmodels
@@ -804,13 +805,20 @@ def view_task_hazard(token, dstore):
     rec = data[int(index)]
     taskno = rec['task_no']
     if len(dstore['source_data/src_id']):
-        sdata = dstore.read_df('source_data', 'taskno').loc[taskno]
-        num_ruptures = sdata.nrupts.sum()
-        eff_sites = sdata.nsites.sum()
-        msg = ('taskno={:_d}, fragments={:_d}, num_ruptures={:_d}, '
-               'eff_sites={:_d}, weight={:.1f}, duration={:.1f}s').format(
-                     taskno, len(sdata), num_ruptures, eff_sites,
-                     rec['weight'], rec['duration'])
+        sdata = dstore.read_df('source_data')
+        sd = sdata[sdata.taskno == taskno]
+        acc = AccumDict(accum=numpy.zeros(5))
+        for src_id, nsites, esites, nrupts, weight, ctimes in zip(
+                sd.src_id, sd.nsites, sd.esites, sd.nrupts, sd.weight, sd.ctimes):
+            acc[basename(src_id, ';:.')] += numpy.array(
+                [nsites, esites, nrupts, weight, ctimes])
+        df = pandas.DataFrame(dict(src_id=list(acc)))
+        for i, name in enumerate(['nsites', 'esites', 'nrupts', 'weight', 'ctimes']):
+            df[name] = [arr[i] for arr in acc.values()]
+        time = df.ctimes.sum()
+        weight = df.weight.sum()
+        msg = f'{taskno=}, {weight=}, {time=}s\n%s' % df.set_index('src_id')
+        return msg
     else:
         msg = ''
     return msg

openquake/hazardlib/contexts.py

@@ -65,6 +65,7 @@
 DIST_BINS = sqrscale(80, 1000, NUM_BINS)
 MEA = 0
 STD = 1
+EPS = 1E-3
 bymag = operator.attrgetter('mag')
 # These coordinates were provided by M Gerstenberger (personal
 # communication, 10 August 2018)
@@ -980,8 +981,13 @@ def genctxs(self, same_mag_rups, sites, src_id):
             '''
             rparams = self.get_rparams(rup)
             dd = self.defaultdict.copy()
-            np = len(rparams.get('probs_occur', []))
-            dd['probs_occur'] = numpy.zeros(np)
+            try:
+                po = rparams['probs_occur']
+            except KeyError:
+                dd['probs_occur'] = numpy.zeros(0)
+            else:
+                L = len(po) if len(po.shape) == 1 else po.shape[1]
+                dd['probs_occur'] = numpy.zeros(L)
             ctx = RecordBuilder(**dd).zeros(len(r_sites))
             for par, val in rparams.items():
                 ctx[par] = val
@@ -1061,9 +1067,9 @@ def get_ctx_iter(self, src, sitecol, src_id=0, step=1):
                 allrups = sorted([rup for rup in allrups
                                   if minmag < rup.mag < maxmag],
                                  key=bymag)
+                self.num_rups = len(allrups) or 1
                 if not allrups:
                     return iter([])
-                self.num_rups = len(allrups)
                 # sorted by mag by construction
                 u32mags = U32([rup.mag * 100 for rup in allrups])
                 rups_sites = [(rups, sitecol) for rups in split_array(
@@ -1279,79 +1285,40 @@ def get_att_curves(self, site, msr, mag, aratio=1., strike=0.,
                 interp1d(ctx.rrup, tau),
                 interp1d(ctx.rrup, phi))
 
-    def estimate_sites(self, src, sites):
-        """
-        :param src: a (Collapsed)PointSource
-        :param sites: a filtered SiteCollection
-        :returns: how many sites are impacted overall
-        """
-        magdist = {mag: self.maximum_distance(mag)
-                   for mag, rate in src.get_annual_occurrence_rates()}
-        nphc = src.count_nphc()
-        dists = sites.get_cdist(src.location)
-        planardict = src.get_planar(iruptures=True)
-        esites = 0
-        for m, (mag, [planar]) in enumerate(planardict.items()):
-            rrup = dists[dists < magdist[mag]]
-            nclose = (rrup < src.get_psdist(m, mag, self.pointsource_distance,
-                                            magdist)).sum()
-            nfar = len(rrup) - nclose
-            esites += nclose * nphc + nfar
-        return esites
-
     # tested in test_collapse_small
     def estimate_weight(self, src, srcfilter, multiplier=1):
         """
         :param src: a source object
         :param srcfilter: a SourceFilter instance
         :returns: (weight, estimate_sites)
         """
+        if src.nsites == 0:  # was discarded by the prefiltering
+            return EPS, 0
         sites = srcfilter.get_close_sites(src)
         if sites is None:
             # may happen for CollapsedPointSources
-            return 0, 0
+            return EPS, 0
         src.nsites = len(sites)
-        N = len(srcfilter.sitecol.complete)  # total sites
-        if (hasattr(src, 'location') and src.count_nphc() > 1 and
-                self.pointsource_distance < 1000):
-            # cps or pointsource with nontrivial nphc
-            esites = self.estimate_sites(src, sites) * multiplier
-        else:
-            step = 100 if src.code == b'F' else 10
-            ctxs = list(self.get_ctx_iter(src, sites, step=step))  # reduced
-            if not ctxs:
-                return src.num_ruptures if N == 1 else 0, 0
-            esites = (sum(len(ctx) for ctx in ctxs) * src.num_ruptures /
-                      self.num_rups * multiplier)  # num_rups from get_ctx_iter
-        weight = esites / N  # the weight is the effective number of ruptures
-        return weight, int(esites)
+        t0 = time.time()
+        ctxs = list(self.get_ctx_iter(src, sites, step=8))  # reduced
+        src.dt = time.time() - t0
+        if not ctxs:
+            return EPS, 0
+        esites = (sum(len(ctx) for ctx in ctxs) * src.num_ruptures /
+                  self.num_rups * multiplier)  # num_rups from get_ctx_iter
+        weight = src.dt * src.num_ruptures / self.num_rups
+        return weight or EPS, int(esites)
 
     def set_weight(self, sources, srcfilter, multiplier=1, mon=Monitor()):
         """
         Set the weight attribute on each prefiltered source
         """
         if hasattr(srcfilter, 'array'):  # a SiteCollection was passed
             srcfilter = SourceFilter(srcfilter, self.maximum_distance)
-        G = len(self.gsims)
-        for src in sources:
-            if src.nsites == 0:  # was discarded by the prefiltering
-                src.esites = 0
-                src.weight = .01
-            else:
-                with mon:
-                    src.weight, src.esites = self.estimate_weight(
-                        src, srcfilter, multiplier)
-                    if src.weight == 0:
-                        src.weight = 0.001
-                    src.weight *= G
-                    if src.code == b'P':
-                        src.weight += .1
-                    elif src.code == b'C':
-                        src.weight += 10.
-                    elif src.code == b'F':
-                        src.weight += .25  * src.num_ruptures
-                    else:
-                        src.weight += 1.
+        with mon:
+            for src in sources:
+                src.weight, src.esites = self.estimate_weight(
+                    src, srcfilter, multiplier)
 
 
 def by_dists(gsim):

openquake/hazardlib/source/point.py

@@ -301,8 +301,8 @@ def _gen_ruptures(self, shift_hypo=False, step=1, iruptures=False):
             npd_ = list(enumerate(npd))
             hdd_ = list(enumerate(hdd))
             for m, (mrate, mag) in magd_[::step]:
-                for n, (nrate, np) in npd_[::step]:
-                    for d, (drate, cdep) in hdd_[::step]:
+                for n, (nrate, np) in npd_:
+                    for d, (drate, cdep) in hdd_:
                         rate = mrate * nrate * drate
                         yield PointRupture(
                             mag, self.tectonic_region_type,