Monitor hook dtypes

* Customize in-memory hook data types, allowing spike tracking for very large models (O(100k)).
* `Model` method signature changes.

sapicore/engine/network/__init__.py

@@ -1,10 +1,10 @@
 """ Networks are graph representations of neuron ensembles connected by synapses. """
-from typing import Sequence
+from typing import Sequence, Type
 
 import os
 from itertools import compress
 
-import matplotlib.pyplot as plt
+import torch
 
 import networkx as nx
 from networkx import DiGraph
@@ -19,6 +19,8 @@
 from sapicore.utils.constants import SYNAPSE_SPLITTERS
 from sapicore.utils.io import DataAccumulatorHook, flatten, load_yaml, MonitorHook
 
+import matplotlib.pyplot as plt
+
 __all__ = ("Network",)
 
 
@@ -290,7 +292,11 @@ def _build(self):
             self.root_lock = False
 
     def add_monitor_hook(
-        self, steps: int = None, attrs: Sequence[str] = None, comps: Sequence[Component] = None
+        self,
+        steps: int = None,
+        attrs: Sequence[str] = None,
+        comps: Sequence[Component] = None,
+        dtype: Type = torch.float,
     ) -> dict:
         """Attach a forward hook to some or all network components, buffering accumulated output in memory.
 
@@ -313,12 +319,13 @@ def add_monitor_hook(
                 hooks[ensemble.identifier] = MonitorHook(
                     ensemble, ensemble.loggable_props if not attrs else attrs, steps
                 )
-
             for synapse in self.get_synapses():
-                hooks[synapse.identifier] = MonitorHook(synapse, synapse.loggable_props if not attrs else attrs, steps)
+                hooks[synapse.identifier] = MonitorHook(
+                    synapse, synapse.loggable_props if not attrs else attrs, steps, dtype
+                )
         else:
             for comp in comps:
-                hooks[comp.identifier] = MonitorHook(comp, comp.loggable_props if not attrs else attrs, steps)
+                hooks[comp.identifier] = MonitorHook(comp, comp.loggable_props if not attrs else attrs, steps, dtype)
 
         return hooks
 

sapicore/model/__init__.py

@@ -96,21 +96,13 @@ def fit(
         for synapse in self.network.get_synapses():
             synapse.set_learning(False)
 
-    def predict(self, data: Tensor | Sequence[Tensor], **kwargs) -> Sequence:
+    def predict(self, data: Tensor, **kwargs) -> Sequence:
         """Predicts the labels of `data`.
 
         Parameters
         ----------
-        data: Data or Tensor
-            Sapicore dataset or a standalone 2D tensor of data buffer, formatted sample X feature.
-
-        duration: int or Sequence of int
-            Duration of sample presentation. Simulates duration of exposure to a particular input.
-            If a list or a tensor is provided, the i-th sample in the batch is maintained for `duration[i]` steps.
-
-        rinse: int or Sequence of int
-            Null stimulation steps (0s in-between samples).
-            If a list or a tensor is provided, the i-th sample is followed by `rinse[i]` rinse steps.
+        data: Tensor
+            Standalone 2D tensor of data buffer, sample X feature.
 
         Returns
         -------
@@ -121,8 +113,7 @@ def predict(self, data: Tensor | Sequence[Tensor], **kwargs) -> Sequence:
         raise NotImplementedError
 
     def similarity(self, data: Tensor, metric: str | Callable, **kwargs) -> Tensor:
-        """Performs rudimentary similarity analysis on the network's responses to `data`,
-        yielding a pairwise distance matrix.
+        """Performs a similarity analysis on network responses to `data`, yielding a pairwise distance matrix.
 
         Parameters
         ----------

sapicore/utils/io.py

@@ -53,9 +53,12 @@ class MonitorHook(Module):
         Expected number of simulation steps. Used to preallocate the buffer.
         If unknown, `torch.vstack()` is used, which may result in slower performance.
 
+    dtype:
+        Data type.
+
     """
 
-    def __init__(self, component: torch.nn.Module, attributes: Sequence[str], entries: int = None):
+    def __init__(self, component: torch.nn.Module, attributes: Sequence[str], entries: int = None, dtype=torch.float):
         super().__init__()
 
         self.component = component
@@ -65,6 +68,7 @@ def __init__(self, component: torch.nn.Module, attributes: Sequence[str], entrie
         self.entries = entries
 
         self.cache = {}
+        self.dtype = dtype
 
         # (de)activate hook without removing it.
         self.active = True
@@ -98,7 +102,7 @@ def fn(_, __, output):
                         else:
                             # preallocate if number of steps is known.
                             dim = [self.entries, len(output[attr])] + ([output[attr].shape[1]] if odim == 2 else [])
-                            self.cache[attr] = torch.empty(dim)
+                            self.cache[attr] = torch.empty(dim, dtype=self.dtype)
 
                     else:
                         if self.entries is None: