FormatNXMX: Avoid storing duplicate static masks (#789)

When trying to dials.import a large number of experiment files, the RAM
usage on Linux was unnecessarily high due to separate static masks being
created for each file. Since masks tend to be limited in number for each
experiment, it makes more sense to store masks with unique values only,
which is what this PR aims to do.

Closes dials/dials#2227.

Author: yash4karan

AUTHORS

@@ -29,4 +29,4 @@ Robert Oeffner
 Takanori Nakane
 Tara Michels-Clark
 Viktor Bengtsson
-
+Yash Karan

newsfragments/789.bugfix

@@ -0,0 +1 @@
+``dials.import``: Reduce excessive memory usage when importing many (>100s) FormatNXMX files.

src/dxtbx/format/FormatNXmx.py

@@ -1,12 +1,41 @@
 from __future__ import annotations
 
+import weakref
+
 import h5py
 import nxmx
 
+import scitbx.array_family.flex as flex
+
 import dxtbx.nexus
 from dxtbx.format.FormatNexus import FormatNexus
 
 
+class _MaskCache:
+    """A singleton to hold unique static_mask objects to avoid duplications"""
+
+    def __init__(self):
+        self.local_mask_cache = weakref.WeakValueDictionary()
+
+    def _mask_hasher(self, mask: flex.bool) -> int:
+        return hash(mask.as_numpy_array().tobytes())
+
+    def store_unique_and_get(
+        self, mask_tuple: tuple[flex.bool, ...] | None
+    ) -> tuple[flex.bool, ...] | None:
+        if mask_tuple is None:
+            return None
+        output = []
+        for mask in mask_tuple:
+            mask_hash = self._mask_hasher(mask)
+            mask = self.local_mask_cache.setdefault(mask_hash, mask)
+            output.append(mask)
+        return tuple(output)
+
+
+mask_cache = _MaskCache()
+
+
 def detector_between_sample_and_source(detector, beam):
     """Check if the detector is perpendicular to beam and
     upstream of the sample."""
@@ -83,7 +112,9 @@ def _start(self):
             self._detector_model = inverted_distance_detector(self._detector_model)
 
         self._scan_model = dxtbx.nexus.get_dxtbx_scan(nxsample, nxdetector)
-        self._static_mask = dxtbx.nexus.get_static_mask(nxdetector)
+        self._static_mask = mask_cache.store_unique_and_get(
+            dxtbx.nexus.get_static_mask(nxdetector)
+        )
         self._bit_depth_readout = nxdetector.bit_depth_readout
 
         if self._scan_model: