STAMP v0.11.2 release

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stamp"
-version = "0.11.1"
+version = "0.11.2"
 description = "A workflow for Sub-Tomogram Averaging Membrane Proteins"
 readme = "README.md"
 authors = [

src/stamp/schemas/subvolumes.py

@@ -4,28 +4,35 @@
 
 # Import external dependencies
 import numpy as np
+from collections import OrderedDict
 from collections.abc import Iterator, Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
 
+# _MAX_OPEN_MMAPS: cap on concurrently open per-tomogram mmaps, well under default OS fd limits (macOS ulimit -n defaults to 256)
+_MAX_OPEN_MMAPS = 64
+
 # Subvolumes: particle -> (tomogram_id, row) index over mmap'd per-tomogram .npy caches
 @dataclass(frozen=True)
 class Subvolumes:
     cache_dir: Path
     index: list[tuple[str, int]]  # index[i] = (tomogram_id, row) for particle i, in extraction order
     box_voxels: int
-    _mmaps: dict[str, np.ndarray] = field(default_factory=dict, repr=False, compare=False)
+    _mmaps: OrderedDict[str, np.ndarray] = field(default_factory=OrderedDict, repr=False, compare=False)
 
     def __len__(self) -> int:
         return len(self.index)
 
-    # _mmap_for: open (or reuse) the mmap for one tomogram's cache file
+    # _mmap_for: open (or reuse) the mmap for one tomogram's cache file, evicting the least-recently-used entry once the cache is full so open file handles stay bounded
     def _mmap_for(self, tomogram_id: str) -> np.ndarray:
         cached = self._mmaps.get(tomogram_id)
         if cached is not None:
+            self._mmaps.move_to_end(tomogram_id)
             return cached
         array = np.load(self.cache_dir / f'{tomogram_id}.npy', mmap_mode='r')
         self._mmaps[tomogram_id] = array  # mutating the dict's contents, not the frozen field itself
+        if len(self._mmaps) > _MAX_OPEN_MMAPS:
+            self._mmaps.popitem(last=False)  # drop oldest; np.memmap closes its file on garbage collection
         return array
 
     def __getitem__(self, i: int) -> np.ndarray:

tests/unit/test_schemas.py

@@ -178,6 +178,19 @@ def test_iter_batches_covers_every_index_grouped_by_tomogram(self, tmp_path):
         for i in seen:
             np.testing.assert_array_equal(seen[i], subvols[i])
 
+    def test_mmap_cache_stays_bounded_across_many_tomograms(self, tmp_path):
+        box_voxels = 2
+        n_tomograms = 200  # well past _MAX_OPEN_MMAPS
+        index: list[tuple[str, int]] = []
+        for t in range(n_tomograms):
+            tomogram_id = f'tomo{t}'
+            _write_cache(tmp_path, tomogram_id, 1, box_voxels, offset=t * 100)
+            index.append((tomogram_id, 0))
+        subvols = Subvolumes(tmp_path, index, box_voxels)
+        for i in range(len(subvols)):
+            _ = subvols[i]  # touch every tomogram once, sequentially
+        assert len(subvols._mmaps) <= 64
+
     def test_write_rolled_produces_a_readable_subvols(self, tmp_path):
         box_voxels = 2
         _write_cache(tmp_path, 'tomoA', 3, box_voxels)