Feature/base classes

.github/workflows/ci.yml

@@ -21,7 +21,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: '3.x'
+          python-version: '3.11'
       - name: Install docs dependencies
         run: |
           pip install -U pip
@@ -38,7 +38,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: '3.x'
+          python-version: '3.11'
       - name: Install black
         run: pip install black
       - name: Check formatting
@@ -116,7 +116,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.x"
+          python-version: "3.11"
       - name: Install build tools
         run: pip install -U pip hatch twine
       - name: Build sdist and wheel

.gitignore

@@ -109,4 +109,10 @@ scratch/
 # PyPi builds
 dist/
 build/
-clean/
+clean/
+
+# VS Code settings, etc.
+.vscode/
+.pytest_cache/
+__pycache__/
+mypy_cache/

.pre-commit-config.yaml

@@ -12,11 +12,11 @@ repos:
   #     - id: conventional-pre-commit
   #       stages: [commit-msg]
 
-  - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.3.0
-    hooks:
-      - id: ruff
-        args: [--fix]
+  # - repo: https://github.com/charliermarsh/ruff-pre-commit
+  #   rev: v0.3.0
+  #   hooks:
+  #     - id: ruff
+  #       args: [--fix]
 
   - repo: https://github.com/psf/black
     rev: 24.2.0
@@ -28,11 +28,11 @@ repos:
     hooks:
       - id: validate-pyproject
 
-  - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.8.0
-    hooks:
-      - id: mypy
-        files: "^src/"
-        # # you have to add the things you want to type check against here
-        # additional_dependencies:
-        #   - numpy
+  # - repo: https://github.com/pre-commit/mirrors-mypy
+  #   rev: v1.8.0
+  #   hooks:
+  #     - id: mypy
+  #       files: "^src/"
+  #       # # you have to add the things you want to type check against here
+  #       # additional_dependencies:
+  #       #   - numpy

README.md

@@ -211,28 +211,34 @@ sampler = multi_dataset.get_weighted_sampler(batch_size=4)
 
 ### CellMapDataLoader
 
-High-performance data loader with optimization features:
+High-performance data loader built on PyTorch's optimized DataLoader:
 
 ```python
 loader = CellMapDataLoader(
     dataset,
-    batch_size=16,
+    batch_size=32,
     num_workers=12,
     weighted_sampler=True,
     device="cuda",
+    prefetch_factor=4,        # Preload batches for better GPU utilization
+    persistent_workers=True,  # Keep workers alive between epochs
+    pin_memory=True,          # Fast CPU-to-GPU transfer
     iterations_per_epoch=1000  # For large datasets
 )
 
 # Optimized GPU memory transfer
 loader.to("cuda", non_blocking=True)
 ```
 
-**Optimizations**:
+**Optimizations** (powered by PyTorch DataLoader):
+
+- **Prefetch Factor**: Background data loading to maximize GPU utilization
+- **Pin Memory**: Fast CPU-to-GPU transfers via pinned memory (auto-enabled on CUDA)
+- **Persistent Workers**: Reduced overhead by keeping workers alive between epochs
+- **PyTorch's Optimized Multiprocessing**: Battle-tested parallel data loading
+- **Smart Defaults**: Automatic optimization based on hardware configuration
 
-- CUDA streams for parallel GPU transfer
-- Persistent workers for reduced overhead  
-- Automatic memory estimation and optimization
-- Thread-safe multiprocessing
+See [DataLoader Optimization Guide](docs/DATALOADER_OPTIMIZATION.md) for performance tuning tips.
 
 ### CellMapDataSplit
 

src/cellmap_data/__init__.py

@@ -1,29 +1,41 @@
-"""
-CellMap Data Loading Module.
+"""Utility for loading CellMap data for machine learning training."""
 
-Utility for loading CellMap data for machine learning training,
-utilizing PyTorch, TensorStore, XArray, and PyDantic.
-"""
-
-from importlib.metadata import PackageNotFoundError, version
+try:
+    from importlib.metadata import PackageNotFoundError, version
+except ImportError:
+    from importlib_metadata import PackageNotFoundError, version
 
 try:
-    __version__ = version("cellmap_data")
+    __version__ = version("cellmap-data")
 except PackageNotFoundError:
-    __version__ = "0.1.1"
-
+    __version__ = "uninstalled"
 __author__ = "Jeff Rhoades"
 __email__ = "[email protected]"
 
-from .multidataset import CellMapMultiDataset
+from .base_dataset import CellMapBaseDataset
+from .base_image import CellMapImageBase
 from .dataloader import CellMapDataLoader
-from .datasplit import CellMapDataSplit
 from .dataset import CellMapDataset
 from .dataset_writer import CellMapDatasetWriter
-from .image import CellMapImage
+from .datasplit import CellMapDataSplit
 from .empty_image import EmptyImage
+from .image import CellMapImage
 from .image_writer import ImageWriter
+from .multidataset import CellMapMultiDataset
 from .subdataset import CellMapSubset
 from .mutable_sampler import MutableSubsetRandomSampler
-from . import transforms
-from . import utils
+
+__all__ = [
+    "CellMapBaseDataset",
+    "CellMapImageBase",
+    "CellMapDataLoader",
+    "CellMapDataset",
+    "CellMapDatasetWriter",
+    "CellMapDataSplit",
+    "CellMapImage",
+    "ImageWriter",
+    "CellMapMultiDataset",
+    "CellMapSubset",
+    "EmptyImage",
+    "MutableSubsetRandomSampler",
+]

src/cellmap_data/base_dataset.py

@@ -0,0 +1,108 @@
+"""Abstract base class for CellMap dataset objects."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Mapping, Sequence
+
+import torch
+
+
+class CellMapBaseDataset(ABC):
+    """
+    Abstract base class for CellMap dataset objects.
+
+    This class defines the common interface that all CellMap dataset objects
+    must implement, ensuring consistency across different dataset types.
+
+    Note: `classes`, `input_arrays`, and `target_arrays` are not abstract
+    properties because implementing classes define them as instance attributes
+    in __init__, not as properties.
+    """
+
+    # These are instance attributes set in __init__, not properties
+    classes: Sequence[str] | None
+    input_arrays: Mapping[str, Mapping[str, Any]]
+    target_arrays: Mapping[str, Mapping[str, Any]] | None
+
+    @property
+    @abstractmethod
+    def class_counts(self) -> dict[str, float]:
+        """
+        Return the number of samples in each class, normalized by resolution.
+
+        Returns
+        -------
+        dict[str, float]
+            Dictionary mapping class names to their counts.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def class_weights(self) -> dict[str, float]:
+        """
+        Return the class weights based on the number of samples in each class.
+
+        Returns
+        -------
+        dict[str, float]
+            Dictionary mapping class names to their weights.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def validation_indices(self) -> Sequence[int]:
+        """
+        Return the indices for the validation set.
+
+        Returns
+        -------
+        Sequence[int]
+            List of validation indices.
+        """
+        pass
+
+    @abstractmethod
+    def to(
+        self, device: str | torch.device, non_blocking: bool = True
+    ) -> "CellMapBaseDataset":
+        """
+        Move the dataset to the specified device.
+
+        Parameters
+        ----------
+        device : str | torch.device
+            The target device.
+        non_blocking : bool, optional
+            Whether to use non-blocking transfer, by default True.
+
+        Returns
+        -------
+        CellMapBaseDataset
+            Self for method chaining.
+        """
+        pass
+
+    @abstractmethod
+    def set_raw_value_transforms(self, transforms: Callable) -> None:
+        """
+        Set the value transforms for raw input data.
+
+        Parameters
+        ----------
+        transforms : Callable
+            Transform function to apply to raw data.
+        """
+        pass
+
+    @abstractmethod
+    def set_target_value_transforms(self, transforms: Callable) -> None:
+        """
+        Set the value transforms for target data.
+
+        Parameters
+        ----------
+        transforms : Callable
+            Transform function to apply to target data.
+        """
+        pass

src/cellmap_data/base_image.py

@@ -0,0 +1,100 @@
+"""Abstract base class for CellMap image objects."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Mapping
+
+import torch
+
+
+class CellMapImageBase(ABC):
+    """
+    Abstract base class for CellMap image objects.
+
+    This class defines the common interface that all CellMap image objects
+    must implement, ensuring consistency across different image types.
+    """
+
+    @abstractmethod
+    def __getitem__(self, center: Mapping[str, float]) -> torch.Tensor:
+        """
+        Return image data centered around the given point.
+
+        Parameters
+        ----------
+        center : Mapping[str, float]
+            The center coordinates in world units.
+
+        Returns
+        -------
+        torch.Tensor
+            The image data as a PyTorch tensor.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def bounding_box(self) -> Mapping[str, tuple[float, float]] | None:
+        """
+        Return the bounding box of the image in world units.
+
+        Returns
+        -------
+        Mapping[str, tuple[float, float]] | None
+            Dictionary mapping axis names to (min, max) tuples, or None.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def sampling_box(self) -> Mapping[str, tuple[float, float]] | None:
+        """
+        Return the sampling box of the image in world units.
+
+        The sampling box is the region where centers can be drawn from and
+        still have full samples drawn from within the bounding box.
+
+        Returns
+        -------
+        Mapping[str, tuple[float, float]] | None
+            Dictionary mapping axis names to (min, max) tuples, or None.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def class_counts(self) -> float | dict[str, float]:
+        """
+        Return the number of voxels for each class in the image.
+
+        Returns
+        -------
+        float | dict[str, float]
+            Class counts, either as a single float or dictionary.
+        """
+        pass
+
+    @abstractmethod
+    def to(self, device: str | torch.device, non_blocking: bool = True) -> None:
+        """
+        Move the image data to the specified device.
+
+        Parameters
+        ----------
+        device : str | torch.device
+            The target device.
+        non_blocking : bool, optional
+            Whether to use non-blocking transfer, by default True.
+        """
+        pass
+
+    @abstractmethod
+    def set_spatial_transforms(self, transforms: Mapping[str, Any] | None) -> None:
+        """
+        Set spatial transformations for the image data.
+
+        Parameters
+        ----------
+        transforms : Mapping[str, Any] | None
+            Dictionary of spatial transformations to apply.
+        """
+        pass