Merge pull request #9097 from gem/faster_gmfs

Refactored the GMF generation

Author: micheles

openquake/commonlib/calc.py

@@ -24,7 +24,7 @@
 from openquake.baselib import performance, parallel, hdf5, general
 from openquake.hazardlib.source import rupture
 from openquake.hazardlib import probability_map
-from openquake.hazardlib.source.rupture import events_dt, get_eid_rlz
+from openquake.hazardlib.source.rupture import get_events
 from openquake.commonlib import util
 
 TWO16 = 2 ** 16
@@ -185,8 +185,7 @@ def get_eid_rlz(self, proxies, rlzs_by_gsim, ordinal):
         :returns: a composite array with the associations eid->rlz
         """
         rlzs = numpy.concatenate(list(rlzs_by_gsim.values()))
-        eid_rlz = get_eid_rlz(proxies, rlzs, self.scenario)
-        return {ordinal: numpy.array(eid_rlz, events_dt)}
+        return {ordinal: get_events(proxies, rlzs, self.scenario)}
 
     def import_rups_events(self, rup_array, get_rupture_getters):
         """

openquake/hazardlib/calc/conditioned_gmfs.py

@@ -116,6 +116,7 @@
 from openquake.baselib.general import AccumDict
 from openquake.baselib.performance import Monitor
 from openquake.hazardlib import correlation, cross_correlation
+from openquake.hazardlib.source.rupture import get_eid_rlz
 from openquake.hazardlib.imt import from_string
 from openquake.hazardlib.calc.gmf import GmfComputer, exp
 from openquake.hazardlib.const import StdDev
@@ -229,12 +230,11 @@ def compute_all(self, dstore, sig_eps=None, max_iml=None,
         """
         min_iml = self.cmaker.min_iml
         rlzs_by_gsim = self.cmaker.gsims
-        eid_rlz = self.ebrupture.get_eid_rlz(rlzs_by_gsim, scenario=True)
+        rlzs = numpy.concatenate(list(rlzs_by_gsim.values()))
+        eid_, rlz_ = get_eid_rlz(vars(self.ebrupture), rlzs, scenario=True)
         mag = self.ebrupture.rupture.mag
         data = AccumDict(accum=[])
         rng = numpy.random.default_rng(self.seed)
-        num_events = self.num_events
-        assert num_events
         # NB: ms is a dictionary gsim -> [imt -> array]
         sids = dstore['conditioned/sids'][:]
         for g, (gsim, rlzs) in enumerate(rlzs_by_gsim.items()):
@@ -244,7 +244,7 @@ def compute_all(self, dstore, sig_eps=None, max_iml=None,
                 phi = dstore['conditioned/gsim_%d/phi' % g][:]
             with mon2:
                 array, sig, eps = self.compute(
-                    gsim, num_events, mea, tau, phi, rng)
+                    gsim, self.num_events, mea, tau, phi, rng)
             M, N, E = array.shape  # sig and eps have shapes (M, E) instead
             assert len(sids) == N, (len(sids), N)
 
@@ -265,7 +265,7 @@ def compute_all(self, dstore, sig_eps=None, max_iml=None,
             array = array.transpose(1, 0, 2)  # from M, N, E to N, M, E
             n = 0
             for rlz in rlzs:
-                eids = eid_rlz[eid_rlz['rlz'] == rlz]['eid']
+                eids = eid_[rlz_ == rlz]
                 for ei, eid in enumerate(eids):
                     gmfa = array[:, :, n + ei]  # shape (N, M)
                     if sig_eps is not None:

openquake/hazardlib/calc/gmf.py

@@ -26,6 +26,7 @@
 from openquake.baselib.general import AccumDict
 from openquake.baselib.performance import Monitor
 from openquake.hazardlib.const import StdDev
+from openquake.hazardlib.source.rupture import get_eid_rlz
 from openquake.hazardlib.cross_correlation import NoCrossCorrelation
 from openquake.hazardlib.gsim.base import ContextMaker, FarAwayRupture
 from openquake.hazardlib.imt import from_string
@@ -138,12 +139,13 @@ def compute_all(self, scenario, sig_eps=None, max_iml=None, mon=Monitor()):
         min_iml = self.cmaker.min_iml
         rlzs_by_gsim = self.cmaker.gsims
         sids = self.ctx.sids
-        eid_rlz = self.ebrupture.get_eid_rlz(rlzs_by_gsim, scenario)
+        rlzs = numpy.concatenate(list(rlzs_by_gsim.values()))
+        eid_, rlz_ = get_eid_rlz(vars(self.ebrupture), rlzs, scenario)
         mag = self.ebrupture.rupture.mag
         data = AccumDict(accum=[])
         mean_stds = self.cmaker.get_mean_stds([self.ctx])  # (4, G, M, N)
         for g, (gs, rlzs) in enumerate(rlzs_by_gsim.items()):
-            num_events = numpy.isin(eid_rlz['rlz'], rlzs).sum()
+            num_events = numpy.isin(rlz_, rlzs).sum()
             if num_events == 0:  # it may happen
                 continue
             # NB: the trick for performance is to keep the call to
@@ -171,7 +173,7 @@ def compute_all(self, scenario, sig_eps=None, max_iml=None, mon=Monitor()):
             array = array.transpose(1, 0, 2)  # from M, N, E to N, M, E
             n = 0
             for rlz in rlzs:
-                eids = eid_rlz[eid_rlz['rlz'] == rlz]['eid']
+                eids = eid_[rlz_ == rlz]
                 for ei, eid in enumerate(eids):
                     gmfa = array[:, :, n + ei]  # shape (N, M)
                     if sig_eps is not None:

openquake/hazardlib/source/rupture.py

@@ -755,25 +755,6 @@ def __init__(self, rupture, source_id, trt_smr, n_occ=1, id=None, e0=0):
     def tectonic_region_type(self):
         return self.rupture.tectonic_region_type
 
-    # TODO: replace with the function get_eid_rlz
-    def get_eid_rlz(self, rlzs_by_gsim, scenario):
-        """
-        :param rlzs_by_gsim: a dictionary gsims -> rlzs array
-        :param scenario: if true distribute the rlzs evenly else randomly
-        :returns: an array with fields (eid, rlz)
-        """
-        out = numpy.zeros(self.n_occ, [('eid', U32), ('rlz', U32)])
-        rlzs = numpy.concatenate(list(rlzs_by_gsim.values()))
-        out['eid'] = numpy.arange(self.e0, self.e0 + self.n_occ, dtype=U32)
-        if scenario:
-            # the rlzs are distributed evenly
-            div = self.n_occ // len(rlzs)
-            out['rlz'] = rlzs[numpy.arange(self.n_occ) // div]
-        else:
-            # event_based: the rlzs are distributed randomly
-            out['rlz'] = general.random_choice(rlzs, self.n_occ, 0, self.seed)
-        return out
-
     def get_eids(self):
         """
         :returns: an array of event IDs
@@ -785,7 +766,25 @@ def __repr__(self):
             self.__class__.__name__, self.id, self.n_occ)
 
 
-def get_eid_rlz(recs, rlzs, scenario):
+def get_eid_rlz(rec, rlzs, scenario):
+    """
+    :param rlzs: an array of realization indices
+    :param scenario: if true distribute the rlzs evenly else randomly
+    :returns: two arrays (eid, rlz)
+    """
+    e0 = rec['e0']
+    n = rec['n_occ']
+    eid = numpy.arange(e0, e0 + n, dtype=U32)
+    if scenario:
+        # the rlzs are distributed evenly
+        rlz = rlzs[numpy.arange(rec['n_occ']) // (n // len(rlzs))]
+    else:
+        # event_based: the rlzs are distributed randomly
+        rlz = general.random_choice(rlzs, n, 0, rec['seed'])
+    return eid, rlz
+
+
+def get_events(recs, rlzs, scenario):
     """
     Build the associations event_id -> rlz_id for each rup_id.
 
@@ -798,15 +797,10 @@ def get_eid_rlz(recs, rlzs, scenario):
         n = rec['n_occ']
         stop = start + n
         slc = out[start:stop]
-        slc['id'] = numpy.arange(rec['e0'], rec['e0'] + n, dtype=U32)
+        eid, rlz = get_eid_rlz(rec, rlzs, scenario)
+        slc['id'] = eid
+        slc['rlz_id'] = rlz
         slc['rup_id'] = rec['id']
-        if scenario:
-            # the rlzs are distributed evenly
-            div = n // len(rlzs)
-            slc['rlz_id'] = rlzs[numpy.arange(n) // div]
-        else:
-            # event_based: the rlzs are distributed randomly
-            slc['rlz_id'] = general.random_choice(rlzs, n, 0, rec['seed'])
         start = stop
     return out