wim/ndbioimage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

- as_array_dyn lru caching

Browse Source
This commit is contained in:
w.pomp
2026-08-04 18:29:56 +02:00
parent 9f5f9debc5
commit e4de76f26c
2 changed files with 296 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -75,6 +75,7 @@ docs/_build/
/tests/files/*
AGENTS.md
.agentbridge
.agent-work
py/ndbioimage/jassets
py/ndbioimage/deps
src/view.rs
+295 -38
View File
@@ -13,7 +13,7 @@ use num::traits::ToBytes;
use num::{Bounded, FromPrimitive, ToPrimitive, Zero};
use serde::{Deserialize, Serialize};
use serde_with::serde_as;
use std::any::type_name;
use std::any::{Any, type_name};
use std::collections::{HashMap, HashSet};
use std::fmt::{Debug, Display, Formatter};
use std::hash::{Hash, Hasher};
@@ -48,11 +48,22 @@ fn slc_bnd(idx: isize, bnd: isize) -> Result<isize, Error> {
}
pub trait Number:
'static + AddAssign + Bounded + Clone + Div<Self, Output = Self> + FromPrimitive + PartialOrd + Zero
'static
+ Send
+ Sync
+ AddAssign
+ Bounded
+ Clone
+ Div<Self, Output = Self>
+ FromPrimitive
+ PartialOrd
+ Zero
{
}
impl<T> Number for T where
T: 'static
+ Send
+ Sync
+ AddAssign
+ Bounded
+ Clone
@@ -66,6 +77,9 @@ impl<T> Number for T where
/// maximum number of frames held in the cache
const DEFAULT_FRAME_CACHE_SIZE: usize = 128;
/// maximum number of materialized arrays held in the cache
const DEFAULT_ARRAY_CACHE_SIZE: usize = 2;
/// identity of the reader a frame was read from
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct ReaderKey {
@@ -75,17 +89,52 @@ struct ReaderKey {
position: usize,
}
type FrameKey = (ReaderKey, usize, usize, usize);
/// borrowed view of [`FrameKey`] for cache lookups without allocation.
/// hashes byte-identically to [`FrameKey`] (str/String and Path/PathBuf hash
/// borrowed view of [`ReaderKey`] for cache lookups without allocation.
/// hashes byte-identically to [`ReaderKey`] (str/String and Path/PathBuf hash
/// the same), so it only matches the same reader identity.
#[derive(Debug, PartialEq, Eq)]
struct FrameKeyRef<'a> {
struct ReaderKeyRef<'a> {
name: &'a str,
path: &'a Path,
series: usize,
position: usize,
}
impl Hash for ReaderKeyRef<'_> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.name.hash(state);
self.path.hash(state);
self.series.hash(state);
self.position.hash(state);
}
}
impl Equivalent<ReaderKey> for ReaderKeyRef<'_> {
fn equivalent(&self, key: &ReaderKey) -> bool {
self.name == key.name
&& self.path == key.path.as_path()
&& self.series == key.series
&& self.position == key.position
}
}
impl ReaderKeyRef<'_> {
fn to_owned(&self) -> ReaderKey {
ReaderKey {
name: self.name.to_string(),
path: self.path.to_path_buf(),
series: self.series,
position: self.position,
}
}
}
type FrameKey = (ReaderKey, usize, usize, usize);
/// borrowed view of [`FrameKey`] for cache lookups without allocation.
#[derive(Debug, PartialEq, Eq)]
struct FrameKeyRef<'a> {
reader: ReaderKeyRef<'a>,
c: usize,
z: usize,
t: usize,
@@ -93,10 +142,7 @@ struct FrameKeyRef<'a> {
impl Hash for FrameKeyRef<'_> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.name.hash(state);
self.path.hash(state);
self.series.hash(state);
self.position.hash(state);
self.reader.hash(state);
self.c.hash(state);
self.z.hash(state);
self.t.hash(state);
@@ -106,29 +152,78 @@ impl Hash for FrameKeyRef<'_> {
impl Equivalent<FrameKey> for FrameKeyRef<'_> {
fn equivalent(&self, key: &FrameKey) -> bool {
let (rk, c, z, t) = key;
self.name == rk.name
&& self.path == rk.path.as_path()
&& self.series == rk.series
&& self.position == rk.position
&& self.c == *c
&& self.z == *z
&& self.t == *t
self.reader.equivalent(rk) && self.c == *c && self.z == *z && self.t == *t
}
}
impl FrameKeyRef<'_> {
fn to_owned(&self) -> FrameKey {
(
ReaderKey {
name: self.name.to_string(),
path: self.path.to_path_buf(),
series: self.series,
position: self.position,
},
self.c,
self.z,
self.t,
)
(self.reader.to_owned(), self.c, self.z, self.t)
}
}
/// identity of a materialized array in the process-wide cache
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct ArrayKey {
reader: ReaderKey,
dtype: &'static str,
slice: Vec<SliceInfoElem>,
axes: Vec<Axis>,
operations: Vec<(Axis, Operation)>,
}
/// borrowed view of [`ArrayKey`] for cache lookups without allocation.
#[derive(Debug, PartialEq, Eq)]
struct ArrayKeyRef<'a> {
reader: ReaderKeyRef<'a>,
dtype: &'static str,
slice: &'a [SliceInfoElem],
axes: &'a [Axis],
operations: &'a IndexMap<Axis, Operation>,
}
impl Hash for ArrayKeyRef<'_> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.reader.hash(state);
self.dtype.hash(state);
self.slice.hash(state);
self.axes.hash(state);
self.operations.len().hash(state);
for (ax, op) in self.operations.iter() {
ax.hash(state);
op.hash(state);
}
}
}
impl Equivalent<ArrayKey> for ArrayKeyRef<'_> {
fn equivalent(&self, key: &ArrayKey) -> bool {
self.reader.equivalent(&key.reader)
&& self.dtype == key.dtype
&& self.slice == key.slice
&& self.axes == key.axes
&& self.operations.len() == key.operations.len()
&& self
.operations
.iter()
.zip(&key.operations)
.all(|((ax, op), (key_ax, key_op))| ax == key_ax && op == key_op)
}
}
impl ArrayKeyRef<'_> {
fn to_owned(&self) -> ArrayKey {
ArrayKey {
reader: self.reader.to_owned(),
dtype: self.dtype,
slice: self.slice.to_vec(),
axes: self.axes.to_vec(),
operations: self
.operations
.iter()
.map(|(ax, op)| (*ax, op.clone()))
.collect(),
}
}
}
@@ -211,6 +306,86 @@ impl Debug for FrameCache {
}
}
/// process-wide LRU cache of materialized arrays, shared between all views and threads
static GLOBAL_ARRAY_CACHE: OnceLock<ArrayCache> = OnceLock::new();
/// thread-safe LRU cache of materialized arrays produced by `as_array_dyn`
struct ArrayCache {
inner: Mutex<ArrayCacheInner>,
}
struct ArrayCacheInner {
map: IndexMap<ArrayKey, Arc<dyn Any + Send + Sync>>,
capacity: usize,
}
impl Default for ArrayCache {
fn default() -> Self {
Self::new(DEFAULT_ARRAY_CACHE_SIZE)
}
}
impl ArrayCache {
fn new(capacity: usize) -> Self {
Self {
inner: Mutex::new(ArrayCacheInner {
map: IndexMap::with_capacity(capacity),
capacity,
}),
}
}
fn global() -> &'static ArrayCache {
GLOBAL_ARRAY_CACHE.get_or_init(ArrayCache::default)
}
fn capacity(&self) -> usize {
self.inner.lock().unwrap().capacity
}
fn set_capacity(&self, capacity: usize) {
let mut inner = self.inner.lock().unwrap();
inner.capacity = capacity;
while inner.map.len() > inner.capacity {
inner.map.shift_remove_index(0);
}
}
fn get<T, Q>(&self, key: &Q) -> Option<Arc<ArrayD<T>>>
where
Q: ?Sized + Hash + Equivalent<ArrayKey>,
T: Any + Send + Sync,
{
let mut inner = self.inner.lock().unwrap();
if let Some(idx) = inner.map.get_index_of(key) {
let (key, array) = inner.map.shift_remove_index(idx).unwrap();
inner.map.insert(key, array.clone());
array.downcast::<ArrayD<T>>().ok()
} else {
None
}
}
fn insert<T: Any + Send + Sync>(&self, key: ArrayKey, array: ArrayD<T>) {
let mut inner = self.inner.lock().unwrap();
inner.map.insert(key, Arc::new(array));
while inner.map.len() > inner.capacity {
inner.map.shift_remove_index(0);
}
}
#[cfg(test)]
fn len(&self) -> usize {
self.inner.lock().unwrap().map.len()
}
}
impl Debug for ArrayCache {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ArrayCache").finish_non_exhaustive()
}
}
/// sliceable view on an image file
#[serde_as]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -337,6 +512,18 @@ impl<D: Dimension, R: Reader> View<D, R> {
FrameCache::global().capacity()
}
/// set the maximum number of materialized arrays the process-wide cache
/// may hold, evicting the least recently used arrays
pub fn with_array_cache_capacity(self, capacity: usize) -> Self {
ArrayCache::global().set_capacity(capacity);
self
}
/// the maximum number of materialized arrays the process-wide cache may hold
pub fn array_cache_capacity(&self) -> usize {
ArrayCache::global().capacity()
}
fn with_operations(mut self, operations: IndexMap<Axis, Operation>) -> Self {
self.operations = operations;
self
@@ -799,6 +986,10 @@ impl<D: Dimension, R: Reader> View<D, R> {
Array1<T>: MinMax<Output = Array0<T>>,
Array2<T>: MinMax<Output = Array1<T>>,
{
let key = self.array_key::<T>();
if let Some(arr) = ArrayCache::global().get(&key) {
return Ok(arr.as_ref().clone());
}
let mut op_xy = IndexMap::new();
if let Some((&ax, op)) = self.operations.first()
&& ((ax == Axis::X) || (ax == Axis::Y))
@@ -954,7 +1145,7 @@ impl<D: Dimension, R: Reader> View<D, R> {
}
}
_ => {
panic!("xy cannot be 3d or more");
unreachable!("xy cannot be 3d or more");
}
};
if let Some((_, op)) = op_czt.first() {
@@ -1020,6 +1211,7 @@ impl<D: Dimension, R: Reader> View<D, R> {
let m = T::from_usize(n).unwrap_or_else(|| T::zero());
out.take().unwrap().mapv(|x| x / m.clone())
};
ArrayCache::global().insert(key.to_owned(), array.clone());
Ok(array)
}
@@ -1051,16 +1243,33 @@ impl<D: Dimension, R: Reader> View<D, R> {
fn frame_key(&self, c: usize, z: usize, t: usize) -> FrameKeyRef<'_> {
FrameKeyRef {
name: self.reader.reader_name(),
path: self.reader.path(),
series: self.reader.series(),
position: self.reader.position(),
reader: ReaderKeyRef {
name: self.reader.reader_name(),
path: self.reader.path(),
series: self.reader.series(),
position: self.reader.position(),
},
c,
z,
t,
}
}
fn array_key<T: 'static>(&self) -> ArrayKeyRef<'_> {
ArrayKeyRef {
reader: ReaderKeyRef {
name: self.reader.reader_name(),
path: self.reader.path(),
series: self.reader.series(),
position: self.reader.position(),
},
dtype: type_name::<T>(),
slice: &self.slice,
axes: &self.axes,
operations: &self.operations,
}
}
fn get_cached_frame(&self, c: usize, z: usize, t: usize) -> Result<Arc<Frame>, Error> {
let key = self.frame_key(c, z, t);
if let Some(frame) = FrameCache::global().get(&key) {
@@ -1297,14 +1506,17 @@ to_bytes_vec_impl!(
#[cfg(test)]
mod tests {
use crate::axes::Axis;
use crate::axes::{Axis, Operation};
use crate::error::Error;
use crate::readers::{DynReader, Frame, Reader};
use crate::stats::MinMax;
use crate::view::{FrameCache, Item, ReaderKey};
use crate::view::{
ArrayCache, ArrayKey, ArrayKeyRef, FrameCache, Item, ReaderKey, ReaderKeyRef,
};
use indexmap::IndexMap;
use ndarray::{Array, Array4, Array5, NewAxis};
use ndarray::{Array2, s};
use std::path::PathBuf;
use ndarray::{Array2, ArrayD, IxDyn, SliceInfoElem, s};
use std::path::{Path, PathBuf};
use std::sync::Arc;
fn open(file: &str) -> Result<DynReader, Error> {
@@ -1610,6 +1822,51 @@ mod tests {
Ok(())
}
#[test]
fn array_cache() -> Result<(), Error> {
let cache = ArrayCache::new(2);
let key = |dtype: &'static str, index: isize, axes: &[Axis]| ArrayKey {
reader: ReaderKey {
name: "test".to_string(),
path: PathBuf::from("test.tif"),
series: 0,
position: 0,
},
dtype,
slice: vec![SliceInfoElem::Index(index)],
axes: axes.to_vec(),
operations: vec![],
};
let k0 = key("u16", 0, &[Axis::T]);
let k1 = key("u16", 1, &[Axis::T]);
cache.insert(k0.clone(), ArrayD::<u16>::zeros(IxDyn(&[2, 2])));
cache.insert(k1.clone(), ArrayD::<u16>::zeros(IxDyn(&[2, 2])));
assert!(cache.get::<u16, _>(&k0).is_some());
cache.insert(
key("u16", 2, &[Axis::T]),
ArrayD::<u16>::zeros(IxDyn(&[2, 2])),
);
assert!(cache.get::<u16, _>(&k0).is_some());
assert!(cache.get::<u16, _>(&k1).is_none());
assert_eq!(cache.len(), 2);
let ops: IndexMap<Axis, Operation> = IndexMap::new();
let borrowed = ArrayKeyRef {
reader: ReaderKeyRef {
name: "test",
path: Path::new("test.tif"),
series: 0,
position: 0,
},
dtype: "u16",
slice: &[SliceInfoElem::Index(0)],
axes: &[Axis::T],
operations: &ops,
};
assert!(cache.get::<f64, _>(&borrowed).is_none());
assert!(cache.get::<u16, _>(&borrowed).is_some());
Ok(())
}
#[test]
fn as_array_caches_frames() -> Result<(), Error> {
let file = "tiffseq/YTL1841B2-2-1_1hr_DMSO_galinduction_1";