Merge pull request #10216 from gem/RuptureGetter

Simplified the RuptureGetter

openquake/calculators/event_based.py

@@ -619,6 +619,7 @@ def agg_dicts(self, acc, result):
     def _read_scenario_ruptures(self):
         oq = self.oqparam
         gsim_lt = read_gsim_lt(oq)
+        trts = list(gsim_lt.values)
         if (str(gsim_lt.branches[0].gsim) == '[FromFile]'
                 and 'gmfs' not in oq.inputs):
             raise InvalidFile('%s: missing gsim or gsim_logic_tree_file' %
@@ -633,7 +634,8 @@ def _read_scenario_ruptures(self):
                 raise InvalidFile(
                     '%s for a scenario calculation must contain a single '
                     'branchset, found %d!' % (oq.inputs['job_ini'], bsets))
-            [(trt, rlzs_by_gsim)] = gsim_lt.get_rlzs_by_gsim_trt().items()
+            [(trt_smr, rlzs_by_gsim)] = gsim_lt.get_rlzs_by_gsim_dic().items()
+            trt = trts[trt_smr // TWO24]
             rup = readinput.get_rupture(oq)
             oq.mags_by_trt = {trt: ['%.2f' % rup.mag]}
             self.cmaker = ContextMaker(trt, rlzs_by_gsim, oq)

openquake/calculators/extract.py

@@ -1406,13 +1406,14 @@ def extract_rupture_info(dstore, what):
               ('strike', F32), ('dip', F32), ('rake', F32)]
     rows = []
     boundaries = []
+    rlzs_by_gsim = dstore['full_lt'].get_rlzs_by_gsim_dic()
     for rgetter in getters.get_rupture_getters(dstore):
         proxies = rgetter.get_proxies(min_mag)
         if 'source_mags' not in dstore:  # ruptures import from CSV
             mags = numpy.unique(dstore['ruptures']['mag'])
         else:
             mags = dstore[f'source_mags/{rgetter.trt}'][:]
-        rdata = RuptureData(rgetter.trt, rgetter.rlzs_by_gsim, mags)
+        rdata = RuptureData(rgetter.trt, rlzs_by_gsim[rgetter.trt_smr], mags)
         arr = rdata.to_array(proxies)
         for r in arr:
             if source_id is None:

openquake/calculators/getters.py

@@ -368,9 +368,8 @@ def get_rupture_getters(dstore, ct=0, srcfilter=None, rupids=None):
             proxies, maxweight, operator.itemgetter('n_occ'),
             key=operator.itemgetter('trt_smr')):
         trt_smr = block[0]['trt_smr']
-        rbg = full_lt.get_rlzs_by_gsim(trt_smr)
         rg = RuptureGetter(block, dstore.filename, trt_smr,
-                           full_lt.trt_by(trt_smr), rbg)
+                           full_lt.trt_by(trt_smr))
         rgetters.append(rg)
     return rgetters
 
@@ -435,16 +434,13 @@ class RuptureGetter(object):
         source group index
     :param trt:
         tectonic region type string
-    :param rlzs_by_gsim:
-        dictionary gsim -> rlzs for the group
     """
-    def __init__(self, proxies, filename, trt_smr, trt, rlzs_by_gsim):
+    def __init__(self, proxies, filename, trt_smr, trt):
         self.proxies = proxies
         self.weight = sum(proxy['n_occ'] for proxy in proxies)
         self.filename = filename
         self.trt_smr = trt_smr
         self.trt = trt
-        self.rlzs_by_gsim = rlzs_by_gsim
         self.num_events = sum(int(proxy['n_occ']) for proxy in proxies)
 
     @property
@@ -482,8 +478,7 @@ def split(self, srcfilter, maxw):
                 proxies.append(proxy)
         rgetters = []
         for block in general.block_splitter(proxies, maxw, weight):
-            rg = RuptureGetter(block, self.filename, self.trt_smr, self.trt,
-                               self.rlzs_by_gsim)
+            rg = RuptureGetter(block, self.filename, self.trt_smr, self.trt)
             rgetters.append(rg)
         return rgetters
 

openquake/commonlib/calc.py

@@ -175,6 +175,7 @@ class RuptureImporter(object):
     def __init__(self, dstore):
         self.datastore = dstore
         self.oqparam = dstore['oqparam']
+        self.full_lt = dstore['full_lt']
         self.scenario = 'scenario' in self.oqparam.calculation_mode
         try:
             self.N = len(dstore['sitecol'])
@@ -232,8 +233,9 @@ def _save_events(self, rup_array, rgetters):
         # build the associations eid -> rlz sequentially or in parallel
         # this is very fast: I saw 30 million events associated in 1 minute!
         iterargs = []
+        rlzs_by_gsim = self.full_lt.get_rlzs_by_gsim_dic()
         for i, rg in enumerate(rgetters):
-            iterargs.append((rg.proxies, rg.rlzs_by_gsim, i))
+            iterargs.append((rg.proxies, rlzs_by_gsim[rg.trt_smr], i))
         if len(events) < 1E5:
             acc = general.AccumDict()  # ordinal -> eid_rlz
             for args in iterargs:

openquake/hazardlib/gsim_lt.py

@@ -38,6 +38,8 @@
 from openquake.hazardlib.gsim.base import GMPE, CoeffsTable
 from openquake.hazardlib.imt import from_string
 
+U32 = numpy.uint32
+TWO24 = 2**24
 
 @dataclass
 class GsimBranch:
@@ -570,7 +572,7 @@ def sample(self, n, seed, sampling_method='early_weights'):
             rlzs.append(rlz)
         return rlzs
 
-    def get_rlzs_by_gsim_trt(self, samples=0, seed=42,
+    def get_rlzs_by_gsim_dic(self, samples=0, seed=42,
                              sampling_method='early_weights'):
         """
         :param samples:
@@ -580,17 +582,17 @@ def get_rlzs_by_gsim_trt(self, samples=0, seed=42,
         :param sampling_method:
             sampling method, by default 'early_weights'
         :returns:
-            dictionary trt -> gsim -> all_rlz_ordinals for each gsim in the trt
+            dictionary trt_smr -> gsim -> rlz_ordinals
         """
         if samples:
             rlzs = self.sample(samples, seed, sampling_method)
         else:
             rlzs = list(self)
         ddic = {}
         for i, trt in enumerate(self.values):
-            ddic[trt] = {gsim: [rlz.ordinal for rlz in rlzs
-                                if rlz.value[i] == gsim]
-                         for gsim in self.values[trt]}
+            ddic[i*TWO24] = {gsim: U32([rlz.ordinal for rlz in rlzs
+                                        if rlz.value[i] == gsim])
+                             for gsim in self.values[trt]}
         return ddic
 
     def __iter__(self):

openquake/hazardlib/logictree.py

@@ -1340,29 +1340,33 @@ def get_realizations(self):
 
     def _rlzs_by_gsim(self, trt_smr):
         # return dictionary gsim->rlzs
-        if not hasattr(self, '_rlzs_by'):
-            rlzs = self.get_realizations()
-            trtis = range(len(self.gsim_lt.values))
-            smrs = numpy.array([sm.ordinal for sm in self.sm_rlzs])
-            if self.source_model_lt.filename == 'fake.xml':  # scenario
-                smr_by_ltp = {'~'.join(sm_rlz.lt_path): i
-                              for i, sm_rlz in enumerate(self.sm_rlzs)}
-                smidx = numpy.zeros(self.get_num_paths(), int)
-                for rlz in rlzs:
-                    smidx[rlz.ordinal] = smr_by_ltp['~'.join(rlz.sm_lt_path)]
-                self._rlzs_by = _ddic(trtis, smrs,
-                                      lambda smr: rlzs[smidx == smr])
-            else:  # classical and event based
-                start = 0
-                slices = []
-                for sm in self.sm_rlzs:
-                    slices.append(slice(start, start + sm.samples))
-                    start += sm.samples
-                self._rlzs_by = _ddic(trtis, smrs,
-                                      lambda smr: rlzs[slices[smr]])
-        if not self._rlzs_by:
-            return {}
-        return self._rlzs_by[trt_smr]
+        return self.get_rlzs_by_gsim_dic()[trt_smr]
+
+    def get_rlzs_by_gsim_dic(self):
+        """
+        :returns: a dictionary trt_smr -> gsim -> rlz ordinals
+        """
+        if hasattr(self, '_rlzs_by'):
+            return self._rlzs_by
+        rlzs = self.get_realizations()
+        trtis = range(len(self.gsim_lt.values))
+        smrs = numpy.array([sm.ordinal for sm in self.sm_rlzs])
+        if self.source_model_lt.filename == 'fake.xml':  # scenario
+            smr_by_ltp = {'~'.join(sm_rlz.lt_path): i
+                          for i, sm_rlz in enumerate(self.sm_rlzs)}
+            smidx = numpy.zeros(self.get_num_paths(), int)
+            for rlz in rlzs:
+                smidx[rlz.ordinal] = smr_by_ltp['~'.join(rlz.sm_lt_path)]
+            self._rlzs_by = _ddic(trtis, smrs,
+                                  lambda smr: rlzs[smidx == smr])
+        else:  # classical and event based
+            start = 0
+            slices = []
+            for sm in self.sm_rlzs:
+                slices.append(slice(start, start + sm.samples))
+                start += sm.samples
+            self._rlzs_by = _ddic(trtis, smrs, lambda smr: rlzs[slices[smr]])
+        return self._rlzs_by
 
     def get_rlzs_by_gsim(self, trt_smr):
         """