[WIP] Shared memory for NBLAST

navis/nbl/base.py

@@ -13,10 +13,13 @@
 
 """Module containing base classes for BLASTING."""
 
+import atexit
+
 import numpy as np
 import pandas as pd
 
 from abc import ABC, abstractmethod
+from multiprocessing import shared_memory
 
 from .. import utils, config
 
@@ -25,7 +28,7 @@
 
 
 class Blaster(ABC):
- """Base class for blasting."""
+ """Base class for Blasting."""
 
  def __init__(self, dtype=np.float64, progress=True):
  """Initialize class."""
@@ -58,15 +61,7 @@ def dtype(self):
 
  @dtype.setter
  def dtype(self, dtype):
- try:
- self._dtype = np.dtype(dtype)
- except TypeError:
- try:
- self._dtype = FLOAT_DTYPES[dtype]
- except KeyError:
- raise ValueError(
- f'Unknown precision/dtype {dtype}. Expected on of the following: 16, 32 or 64 (default)'
- )
+ self._dtype = parse_precision(dtype)
 
  def pair_query_target(self, pairs, scores='forward'):
  """BLAST multiple pairs.
@@ -101,7 +96,7 @@ def pair_query_target(self, pairs, scores='forward'):
 
  return scr
 
- def multi_query_target(self, q_idx, t_idx, scores='forward'):
+ def multi_query_target(self, q_idx, t_idx, scores='forward', out=None):
  """BLAST multiple queries against multiple targets.
 
  Parameters
@@ -110,6 +105,8 @@ def multi_query_target(self, q_idx, t_idx, scores='forward'):
  Iterable of query/target neuron indices to BLAST.
  scores : "forward" | "mean" | "min" | "max"
  Which scores to return.
+ out : np.ndarray, optional
+ Array to write results to.
 
  """
  if utils.is_jupyter() and config.tqdm == config.tqdm_notebook:
@@ -124,18 +121,24 @@ def multi_query_target(self, q_idx, t_idx, scores='forward'):
  else:
  position = getattr(self, 'pbar_position', 0)
 
- res = np.zeros((len(q_idx), len(t_idx)),
- dtype=self.dtype)
+ shape = (len(q_idx), len(t_idx))
+ if not isinstance(out, np.ndarray):
+ out = np.zeros(shape, dtype=self.dtype)
+ elif out.shape != shape:
+ raise TypeError(f'Expected out array of shape {shape}, got {out.shape}')
+ elif out.dtype != self.dtype:
+ raise TypeError(f'Expected out array to be {self.dtype}, got {out.dtype}')
+
  for i, q in enumerate(config.tqdm(q_idx,
  desc=self.desc,
  leave=False,
  position=position,
  disable=not self.progress)):
  for k, t in enumerate(t_idx):
- res[i, k] = self.single_query_target(q, t, scores=scores)
+ out[i, k] = self.single_query_target(q, t, scores=scores)
 
  # Generate results
- res = pd.DataFrame(res)
+ res = pd.DataFrame(out)
  res.columns = [self.ids[t] for t in t_idx]
  res.index = [self.ids[q] for q in q_idx]
 
@@ -157,3 +160,109 @@ def all_by_all(self, scores='forward'):
  res.loc[:, :] = np.dstack((res, res.T)).max(axis=2)
 
  return res
+
+
+class SharedBlaster(Blaster, ABC):
+ """Version of the Blaster that works with shared memory buffers."""
+
+ def create_array(self, shm, shape, offset=None):
+ """Create array from shared memory buffer."""
+ # Generate the out array from shared memory buffer
+ arr = np.ndarray(shape, dtype=self.dtype, buffer=shm.buf)
+
+ if not isinstance(offset, type(None)):
+ arr = arr[offset]
+
+ return arr
+
+ def multi_query_target(self, q_idx, t_idx, shm, shape, offset=None, scores='forward'):
+ """BLAST multiple queries against multiple targets.
+
+ Parameters
+ ----------
+ q_idx,t_idx : iterable
+ Iterable of query/target neuron indices to BLAST.
+ shm : multiprocessing.shared_memory.SharedMemory 
+ shape : tuple (N, M)
+ Shape of the array in the memory buffer.
+ offset : tuple (slice, slice)
+ The view inside the full array to pass through as `out` to
+ `multi_query_target`.
+ scores : "forward" | "mean" | "min" | "max"
+ Which scores to produce.
+
+ """
+ out = self.create_array(shm, shape, offset=offset)
+
+ _ = super().multi_query_target(q_idx, t_idx, scores=scores, out=out)
+
+ def all_by_all(self, shm, shape, offset=None, scores='forward'):
+ """BLAST all-by-all neurons."""
+ out = self.create_array(shm, shape, offset=offset)
+
+ _ = super().multi_query_target(range(len(self.neurons)),
+ range(len(self.neurons)),
+ scores='forward',
+ out=out)
+
+ # For all-by-all BLAST we can get the mean score by
+ # transposing the scores
+ if scores == 'mean':
+ out[:, :] = (out + out.T) / 2
+ elif scores == 'min':
+ out[:, :] = np.dstack((out, out.T)).min(axis=2)
+ elif scores == 'max':
+ out[:, :] = np.dstack((out, out.T)).max(axis=2)
+
+
+def create_shared_array(shape, dtype):
+ """Create shared array.
+
+ Parameters
+ ----------
+ shape : tuple
+ dtype : str | np.dtype
+
+ Returns
+ -------
+ multiprocessing.shared_memory.SharedMemory
+ The shared memory buffer for the array.
+ np.ndarray
+ A numpy array accessing the shared memory buffer.
+
+ """
+ # Get the number of items in the requested array
+ if utils.is_iterable(shape):
+ items = np.prod(shape)
+ else:
+ items = shape
+
+ # Force dtype to numpy dtype
+ dtype = np.dtype(dtype)
+ # Calculate required size for memory buffer
+ size = dtype.itemsize * items
+
+ # Create shared memory buffer
+ shm = shared_memory.SharedMemory(create=True, size=size)
+
+ # We need to make sure that the memory is released on exit of this process
+ # Note: order is reverse -> last registered is executed first
+ atexit.register(shm.unlink)
+ atexit.register(shm.close)
+
+ # Create array based on buffer
+ arr = np.ndarray(shape, dtype=dtype, buffer=shm.buf)
+
+ return shm, arr
+
+
+def parse_precision(dtype):
+ """Parse precision into numpy dtype."""
+ try:
+ return np.dtype(dtype)
+ except TypeError:
+ try:
+ return FLOAT_DTYPES[dtype]
+ except KeyError:
+ raise ValueError(f'Unknown precision/dtype {dtype}. Expected one '
+ 'of the following: 16, 32 or 64 (default)')

navis/nbl/nblast_funcs.py

@@ -27,7 +27,7 @@
 
 from .. import utils, config
 from ..core import NeuronList, Dotprops, make_dotprops
-from .base import Blaster
+from .base import Blaster, SharedBlaster, create_shared_array, parse_precision
 
 __all__ = ['nblast', 'nblast_smart', 'nblast_allbyall', 'sim_to_dist']
 
@@ -234,6 +234,11 @@ def single_query_target(self, q_idx, t_idx, scores='forward'):
  return scr
 
 
+class SharedNBlaster(NBlaster, SharedBlaster):
+ """An NBLASTER working with shared memory."""
+ pass
+
+
 def nblast_smart(query: Union[Dotprops, NeuronList],
  target: Optional[str] = None,
  t: Union[int, float] = 90,
@@ -316,7 +321,7 @@ def nblast_smart(query: Union[Dotprops, NeuronList],
  precision : int [16, 32, 64] | str [e.g. "float64"] | np.dtype
  Precision for scores. Defaults to 64 bit (double) floats.
  This is useful to reduce the memory footprint for very large
- matrices. In real-world scenarios 32 bit (single)- and
+ matrices. In real-world scenarios 32 bit (single) - and
  depending on the purpose even 16 bit (half) - are typically
  sufficient.
  n_cores : int, optional
@@ -685,21 +690,30 @@ def nblast(query: Union[Dotprops, NeuronList],
  else:
  n_rows, n_cols = find_batch_partition(batch_size, query_dps, target_dps)
 
+ # Force precision into numpy dtype
+ dtype = parse_precision(precision)
+
+ # Create scores array
+ shm, arr = create_shared_array(shape=(len(query_dps), len(target_dps)),
+ dtype=dtype)
+
  nblasters = []
  with config.tqdm(desc='Preparing',
  total=n_rows * n_cols,
  leave=False,
  disable=not progress) as pbar:
+ offset_x = 0
  for q in np.array_split(query_dps, n_rows):
+ offset_y = 0
  for t in np.array_split(target_dps, n_cols):
  # Initialize NBlaster
- this = NBlaster(use_alpha=use_alpha,
- normalized=normalized,
- smat=smat,
- limit_dist=limit_dist,
- dtype=precision,
- approx_nn=approx_nn,
- progress=progress)
+ this = SharedNBlaster(use_alpha=use_alpha,
+  normalized=normalized,
+  smat=smat,
+  limit_dist=limit_dist,
+  dtype=dtype,
+  approx_nn=approx_nn,
+  progress=progress)
 
  # Use better description if we process in batches
  if batch_size:
@@ -713,32 +727,41 @@ def nblast(query: Union[Dotprops, NeuronList],
  # Keep track of indices of queries and targets
  this.queries = np.arange(len(q))
  this.targets = np.arange(len(t)) + len(q)
+
+ # Make sure progress bars don't overlap
  this.pbar_position = len(nblasters) % n_cores
 
+ # This is the offset into the shared array
+ this.offset_x = slice(offset_x, offset_x + len(q))
+ this.offset_y = slice(offset_y, offset_y + len(t))
+
  nblasters.append(this)
  pbar.update(1)
+ offset_y += len(t)
+ offset_x += len(q)
 
  # If only one core, we don't need to break out the multiprocessing
  if n_cores == 1:
- return this.multi_query_target(this.queries,
- this.targets,
- scores=scores)
-
- with ProcessPoolExecutor(max_workers=n_cores) as pool:
- # Each nblaster is passed to its own process
- futures = [pool.submit(this.multi_query_target,
- q_idx=this.queries,
- t_idx=this.targets,
- scores=scores) for this in nblasters]
-
- results = [f.result() for f in futures]
+ _ = this.multi_query_target(this.queries,
+ this.targets,
+ shm=shm,
+ shape=arr.shape,
+ offset=None,
+ scores=scores)
+ else:
+ with ProcessPoolExecutor(max_workers=n_cores) as pool:
+ # Each nblaster is passed to its own process
+ futures = [pool.submit(this.multi_query_target,
+ q_idx=this.queries,
+ t_idx=this.targets,
+ shm=shm,
+ shape=arr.shape,
+ offset=(this.offset_x, this.offset_y),
+ scores=scores) for this in nblasters]
 
- scores = pd.DataFrame(np.zeros((len(query_dps), len(target_dps)),
- dtype=this.dtype),
- index=query_dps.id, columns=target_dps.id)
+ _ = [f.result() for f in futures]
 
- for res in results:
- scores.loc[res.index, res.columns] = res.values
+ scores = pd.DataFrame(arr, index=query_dps.id, columns=target_dps.id)
 
  return scores