Cleanup

gem · Aug 26, 2024 · 01c1169 · 01c1169
1 parent c4c4699
commit 01c1169
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/openquake/calculators/classical.py b/openquake/calculators/classical.py
@@ -128,8 +128,10 @@ def classical(sources, sitecol, cmaker, dstore, monitor):
             yield result
     else:
         result = hazclassical(sources, sitecol, cmaker)
+        # print(f"{monitor.task_no=} {result['pnemap'].size_mb=}")
         result['allsources'] = allsources
         rmap = result.pop('pnemap').remove_zeros().to_rates()
+        # print(f"{monitor.task_no=} {rmap.size_mb=}")
         if cmaker.tiling and cmaker.save_on_tmp:
             # tested in case_22
             scratch = parallel.scratch_dir(monitor.calc_id)

diff --git a/openquake/hazardlib/map_array.py b/openquake/hazardlib/map_array.py
@@ -30,6 +30,7 @@
 F32 = numpy.float32
 F64 = numpy.float64
 BYTES_PER_FLOAT = 8
+TWO20 = 2 ** 20  # 1 MB
 TWO24 = 2 ** 24
 rates_dt = numpy.dtype([('sid', U32), ('lid', U16), ('gid', U16),
                         ('rate', F32)])
@@ -248,6 +249,12 @@ def sidx(self):
             idxs[sid] = idx
         return idxs
 
+    @property
+    def size_mb(self):
+        if hasattr(self, 'array'):
+            return self.array.nbytes / TWO20
+        return sum(arr.nbytes / TWO20 for arr in self.acc.values())
+
     def new(self, acc):
         new = copy.copy(self)
         if isinstance(acc, numpy.ndarray):