Made the GmfComputer a bit faster

openquake/hazardlib/calc/conditioned_gmfs.py

@@ -123,7 +123,10 @@
 from openquake.hazardlib.geo.geodetic import geodetic_distance
 from openquake.hazardlib.gsim.base import ContextMaker
 
+U32 = numpy.uint32
 F32 = numpy.float32
+
+
 class NoInterIntraStdDevs(Exception):
     def __init__(self, gsim):
         self.gsim = gsim
@@ -245,6 +248,11 @@ def compute_all(self, dstore, sig_eps=None, max_iml=None,
             self.update(data, array, sig, eps, eid_, rlz_, rlzs,
                         [mea, tau+phi, tau, phi], sig_eps, max_iml)
 
+        for key, val in sorted(data.items()):
+            if key in 'eid sid rlz':
+                data[key] = numpy.concatenate(data[key], dtype=U32)
+            else:
+                data[key] = numpy.concatenate(data[key], dtype=F32)
         return pandas.DataFrame(data)
 
     def compute(self, gsim, num_events, mea, tau, phi, rng):

openquake/hazardlib/calc/gmf.py

@@ -58,6 +58,22 @@ def exp(vals, imt):
     return numpy.exp(vals)
 
 
+def set_max_min(array, mean, max_iml, min_iml, imts):
+    M, N, E = array.shape
+
+    # manage max_iml
+    for m, im in enumerate(imts):
+        if (array[m] > max_iml[m]).any():
+            for n in range(N):
+                bad = array[m, n] > max_iml[m]  # shape E
+                array[m, n, bad] = exp(mean[m, n], im)
+
+    # manage min_iml
+    for n in range(N):
+        for e in range(E):
+            if (array[:, n, e] < min_iml).all():
+                array[:, n, e] = 0
+
 
 class GmfComputer(object):
     """
@@ -132,58 +148,40 @@ def __init__(self, rupture, sitecol, cmaker, correlation_model=None,
             self.sites = sitecol.complete.filtered(self.ctx.sids)
         self.cross_correl = cross_correl or NoCrossCorrelation(
             cmaker.truncation_level)
+        self.gmv_fields = [f'gmv_{m}' for m in range(len(cmaker.imts))]
 
     def update(self, data, array, sig, eps, eid_, rlz_, rlzs,
                mean_stds, sig_eps, max_iml):
         sids = self.ctx.sids
         min_iml = self.cmaker.min_iml
         mag = self.ebrupture.rupture.mag
-        M, N, E = array.shape  # sig and eps have shapes (M, E) instead
-
-        # manage max_iml
-        if max_iml is not None:
-            for m, im in enumerate(self.cmaker.imtls):
-                if (array[m] > max_iml[m]).any():
-                    for n in range(N):
-                        bad = array[m, n] > max_iml[m]  # shape E
-                        mean = mean_stds[0][m, n]
-                        if len(mean.shape) == 2:  # conditioned GMFs
-                            mean = mean[:, 0]  # shape (N, 1)
-                        array[m, n, bad] = exp(mean, im)
-
-        # manage min_iml
-        for n in range(N):
-            for e in range(E):
-                if (array[:, n, e] < min_iml).all():
-                    array[:, n, e] = 0
-
+        mean = mean_stds[0]
+        if len(mean.shape) == 3:  # shape (M, N, 1) for conditioned gmfs
+            mean = mean[:, :, 0]
+        set_max_min(array, mean, max_iml, min_iml, self.cmaker.imts)
         array = array.transpose(1, 0, 2)  # from M, N, E to N, M, E
+        N = len(array)
         n = 0
         for rlz in rlzs:
             eids = eid_[rlz_ == rlz]
-            for ei, eid in enumerate(eids):
-                gmfa = array[:, :, n + ei]  # shape (N, M)
-                if sig_eps is not None:
+            if sig_eps is not None:
+                for ei, eid in enumerate(eids):
                     tup = tuple([eid, rlz] + list(sig[:, n + ei]) +
                                 list(eps[:, n + ei]))
                     sig_eps.append(tup)
-                items = []
+            for ei, eid in enumerate(eids):
+                gmfa = array[:, :, n + ei]  # shape (N, M)
+                # gmv can be zero due to the minimum_intensity, coming
+                # from the job.ini or from the vulnerability functions
+                data['sid'].append(sids)
+                data['eid'].append(numpy.full(N, eid, U32))
+                data['rlz'].append(numpy.full(N, rlz, U32))
                 for sp in self.sec_perils:
                     o = sp.compute(mag, zip(self.imts, gmfa.T), self.ctx)
                     for outkey, outarr in zip(sp.outputs, o):
-                        items.append((outkey, outarr))
-                for i, gmv in enumerate(gmfa):
-                    if gmv.sum() == 0:
-                        continue
-                    data['sid'].append(sids[i])
-                    data['eid'].append(eid)
-                    data['rlz'].append(rlz)  # used in compute_gmfs_curves
-                    for m in range(M):
-                        data[f'gmv_{m}'].append(gmv[m])
-                    for outkey, outarr in items:
-                        data[outkey].append(outarr[i])
-                    # gmv can be zero due to the minimum_intensity, coming
-                    # from the job.ini or from the vulnerability functions
+                        data[outkey].append(outarr)
+                for m, gmv_field in enumerate(self.gmv_fields):
+                    data[gmv_field].append(gmfa[:, m])
             n += len(eids)
 
     def compute_all(self, scenario, sig_eps=None, max_iml=None, mon=Monitor()):
@@ -196,6 +194,9 @@ def compute_all(self, scenario, sig_eps=None, max_iml=None, mon=Monitor()):
         data = AccumDict(accum=[])
         mean_stds = self.cmaker.get_mean_stds([self.ctx])  # (4, G, M, N)
         rng = numpy.random.default_rng(self.seed)
+        if max_iml is None:
+            M = len(self.cmaker.imts)
+            max_iml = numpy.full(M, numpy.inf, float)
         for g, (gs, rlzs) in enumerate(rlzs_by_gsim.items()):
             num_events = numpy.isin(rlz_, rlzs).sum()
             if num_events == 0:  # it may happen
@@ -211,9 +212,9 @@ def compute_all(self, scenario, sig_eps=None, max_iml=None, mon=Monitor()):
 
         for key, val in sorted(data.items()):
             if key in 'eid sid rlz':
-                data[key] = U32(data[key])
+                data[key] = numpy.concatenate(data[key], dtype=U32)
             else:
-                data[key] = F32(data[key])
+                data[key] = numpy.concatenate(data[key], dtype=F32)
         return pandas.DataFrame(data)
 
     def compute(self, gsim, num_events, mean_stds, rng):