arraybridge 0.2.0__py3-none-any.whl

arraybridge/__init__.py

@@ -0,0 +1,37 @@
+"""
+arraybridge: Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto.
+
+This package provides automatic memory type conversion, declarative decorators,
+and unified utilities for working with multiple array/tensor frameworks.
+"""
+
+__version__ = "0.2.0"
+
+from .converters import convert_memory, detect_memory_type
+from .decorators import cupy, jax, memory_types, numpy, tensorflow, torch
+from .exceptions import MemoryConversionError
+from .stack_utils import stack_slices, unstack_slices
+from .types import CPU_MEMORY_TYPES, GPU_MEMORY_TYPES, SUPPORTED_MEMORY_TYPES, MemoryType
+
+__all__ = [
+    # Types
+    "MemoryType",
+    "CPU_MEMORY_TYPES",
+    "GPU_MEMORY_TYPES",
+    "SUPPORTED_MEMORY_TYPES",
+    # Converters
+    "convert_memory",
+    "detect_memory_type",
+    # Decorators
+    "memory_types",
+    "numpy",
+    "cupy",
+    "torch",
+    "tensorflow",
+    "jax",
+    # Stack utilities
+    "stack_slices",
+    "unstack_slices",
+    # Exceptions
+    "MemoryConversionError",
+]

arraybridge/conversion_helpers.py

@@ -0,0 +1,150 @@
+"""
+Memory conversion helpers for OpenHCS.
+
+This module provides the ABC and metaprogramming infrastructure for memory type conversions.
+Uses enum-driven polymorphism to eliminate 1,567 lines of duplication.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+
+from arraybridge.framework_config import _FRAMEWORK_CONFIG
+from arraybridge.types import MemoryType
+from arraybridge.utils import _supports_dlpack
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryTypeConverter(ABC):
+    """Abstract base class for memory type converters.
+
+    Each memory type (numpy, cupy, torch, etc.) has a concrete converter
+    that implements these four core operations. All to_X() methods are
+    auto-generated using polymorphism.
+    """
+
+    @abstractmethod
+    def to_numpy(self, data, gpu_id):
+        """Extract to NumPy (type-specific implementation)."""
+        pass
+
+    @abstractmethod
+    def from_numpy(self, data, gpu_id):
+        """Create from NumPy (type-specific implementation)."""
+        pass
+
+    @abstractmethod
+    def from_dlpack(self, data, gpu_id):
+        """Create from DLPack capsule (type-specific implementation)."""
+        pass
+
+    @abstractmethod
+    def move_to_device(self, data, gpu_id):
+        """Move data to specified GPU device if needed (type-specific implementation)."""
+        pass
+
+
+def _add_converter_methods():
+    """Add to_X() methods to MemoryTypeConverter ABC.
+
+    NOTE: This must be called AFTER _CONVERTERS is defined (see below).
+
+    For each target memory type, generates a method like to_cupy(), to_torch(), etc.
+    that tries GPU-to-GPU conversion via DLPack first, then falls back to CPU roundtrip.
+    """
+    for target_type in MemoryType:
+        method_name = f"to_{target_type.value}"
+
+        def make_method(tgt):
+            def method(self, data, gpu_id):
+                # Try GPU-to-GPU first (DLPack)
+                if _supports_dlpack(data):
+                    try:
+                        target_converter = _CONVERTERS[tgt]
+                        result = target_converter.from_dlpack(data, gpu_id)
+                        return target_converter.move_to_device(result, gpu_id)
+                    except Exception as e:
+                        logger.warning(f"DLPack conversion failed: {e}. Using CPU roundtrip.")
+
+                # CPU roundtrip using polymorphism
+                numpy_data = self.to_numpy(data, gpu_id)
+                target_converter = _CONVERTERS[tgt]
+                return target_converter.from_numpy(numpy_data, gpu_id)
+            return method
+
+        setattr(MemoryTypeConverter, method_name, make_method(target_type))
+
+
+# NOTE: Conversion operations now defined in framework_config.py under 'conversion_ops'
+# This eliminates the scattered _OPS dict
+_OPS = {mem_type: config['conversion_ops'] for mem_type, config in _FRAMEWORK_CONFIG.items()}
+
+# Auto-generate lambdas from strings
+def _make_not_implemented(mem_type_value, method_name):
+    """Create a lambda that raises NotImplementedError with the correct signature."""
+    def not_impl(self, data, gpu_id):
+        raise NotImplementedError(f"DLPack not supported for {mem_type_value}")
+    # Add proper names for better debugging
+    not_impl.__name__ = method_name
+    not_impl.__qualname__ = f'{mem_type_value.capitalize()}Converter.{method_name}'
+    return not_impl
+
+def _make_lambda_with_name(expr_str, mem_type, method_name):
+    """Create a lambda from expression string and add proper __name__ for debugging."""
+    # Pre-compute the module string to avoid nested f-strings
+    # with backslashes (Python 3.11 limitation)
+    module_str = f'_ensure_module("{mem_type.value}")'
+    lambda_expr = f'lambda self, data, gpu_id: {expr_str.format(mod=module_str)}'
+    lambda_func = eval(lambda_expr)
+    lambda_func.__name__ = method_name
+    lambda_func.__qualname__ = f'{mem_type.value.capitalize()}Converter.{method_name}'
+    return lambda_func
+
+_TYPE_OPERATIONS = {
+    mem_type: {
+        method_name: (
+            _make_lambda_with_name(expr, mem_type, method_name)
+            if expr is not None
+            else _make_not_implemented(mem_type.value, method_name)
+        )
+        for method_name, expr in ops.items()  # Iterate over dict items - self-documenting!
+    }
+    for mem_type, ops in _OPS.items()
+}
+
+# Auto-generate all 6 converter classes
+_CONVERTERS = {
+    mem_type: type(
+        f"{mem_type.value.capitalize()}Converter",
+        (MemoryTypeConverter,),
+        _TYPE_OPERATIONS[mem_type]
+    )()
+    for mem_type in MemoryType
+}
+
+# NOW call _add_converter_methods() after _CONVERTERS exists
+_add_converter_methods()
+
+
+# Runtime validation: ensure all converters have required methods
+def _validate_converters():
+    """Validate that all generated converters have the required methods."""
+    required_methods = ['to_numpy', 'from_numpy', 'from_dlpack', 'move_to_device']
+
+    for mem_type, converter in _CONVERTERS.items():
+        # Check ABC methods
+        for method in required_methods:
+            if not hasattr(converter, method):
+                raise RuntimeError(f"{mem_type.value} converter missing method: {method}")
+
+        # Check to_X() methods for all memory types
+        for target_type in MemoryType:
+            method_name = f'to_{target_type.value}'
+            if not hasattr(converter, method_name):
+                raise RuntimeError(f"{mem_type.value} converter missing method: {method_name}")
+
+    logger.debug(f"✅ Validated {len(_CONVERTERS)} memory type converters")
+
+# Run validation at module load time
+_validate_converters()
+

arraybridge/converters.py

@@ -0,0 +1,61 @@
+"""Memory conversion public API for OpenHCS."""
+
+from typing import Any
+
+import numpy as np
+
+from arraybridge.conversion_helpers import _CONVERTERS
+from arraybridge.framework_config import _FRAMEWORK_CONFIG
+from arraybridge.types import MemoryType
+
+
+def convert_memory(data: Any, source_type: str, target_type: str, gpu_id: int) -> Any:
+    """
+    Convert data between memory types using the unified converter infrastructure.
+
+    Args:
+        data: The data to convert
+        source_type: The source memory type (e.g., "numpy", "torch")
+        target_type: The target memory type (e.g., "cupy", "jax")
+        gpu_id: The target GPU device ID
+
+    Returns:
+        The converted data in the target memory type
+
+    Raises:
+        ValueError: If source_type or target_type is invalid
+        MemoryConversionError: If conversion fails
+    """
+    source_enum = MemoryType(source_type)
+    converter = _CONVERTERS[source_enum]
+    method = getattr(converter, f"to_{target_type}")
+    return method(data, gpu_id)
+
+
+def detect_memory_type(data: Any) -> str:
+    """
+    Detect the memory type of data using framework config.
+
+    Args:
+        data: The data to detect
+
+    Returns:
+        The detected memory type string (e.g., "numpy", "torch")
+
+    Raises:
+        ValueError: If memory type cannot be detected
+    """
+    # NumPy special case (most common, check first)
+    if isinstance(data, np.ndarray):
+        return MemoryType.NUMPY.value
+
+    # Check all frameworks using their module names from config
+    module_name = type(data).__module__
+
+    for mem_type, config in _FRAMEWORK_CONFIG.items():
+        import_name = config['import_name']
+        # Check if module name starts with or contains the import name
+        if module_name.startswith(import_name) or import_name in module_name:
+            return mem_type.value
+
+    raise ValueError(f"Unknown memory type for {type(data)} (module: {module_name})")

arraybridge/decorators.py

@@ -0,0 +1,396 @@
+"""
+Memory type declaration decorators for OpenHCS.
+
+This module provides decorators for explicitly declaring the memory interface
+of pure functions, enforcing Clause 106-A (Declared Memory Types) and supporting
+memory-type-aware dispatching and orchestration.
+
+These decorators annotate functions with input_memory_type and output_memory_type
+attributes and provide automatic thread-local CUDA stream management for GPU
+frameworks to enable true parallelization across multiple threads.
+
+REFACTORED: Uses enum-driven metaprogramming to eliminate 79% of code duplication.
+"""
+
+import functools
+import inspect
+import logging
+import threading
+from enum import Enum
+from typing import Any, Callable, Optional, TypeVar
+
+import numpy as np
+
+from arraybridge.dtype_scaling import SCALING_FUNCTIONS
+from arraybridge.framework_ops import _FRAMEWORK_OPS
+from arraybridge.oom_recovery import _execute_with_oom_recovery
+from arraybridge.slice_processing import process_slices
+from arraybridge.types import MemoryType
+from arraybridge.utils import optional_import
+
+logger = logging.getLogger(__name__)
+
+F = TypeVar('F', bound=Callable[..., Any])
+
+
+class DtypeConversion(Enum):
+    """Data type conversion modes for all memory type functions."""
+
+    PRESERVE_INPUT = "preserve"     # Keep input dtype (default)
+    NATIVE_OUTPUT = "native"        # Use framework's native output
+    UINT8 = "uint8"                # Force uint8 (0-255 range)
+    UINT16 = "uint16"              # Force uint16 (microscopy standard)
+    INT16 = "int16"                # Force int16 (signed microscopy data)
+    INT32 = "int32"                # Force int32 (large integer values)
+    FLOAT32 = "float32"            # Force float32 (GPU performance)
+    FLOAT64 = "float64"            # Force float64 (maximum precision)
+
+    @property
+    def numpy_dtype(self):
+        """Get the corresponding numpy dtype."""
+        dtype_map = {
+            self.UINT8: np.uint8,
+            self.UINT16: np.uint16,
+            self.INT16: np.int16,
+            self.INT32: np.int32,
+            self.FLOAT32: np.float32,
+            self.FLOAT64: np.float64,
+        }
+        return dtype_map.get(self, None)
+
+
+# Thread-local cache for lazy-loaded GPU frameworks
+_gpu_frameworks_cache = {}
+
+
+def _create_lazy_getter(framework_name: str):
+    """Factory function that creates a lazy import getter for a framework."""
+    def getter():
+        if framework_name not in _gpu_frameworks_cache:
+            _gpu_frameworks_cache[framework_name] = optional_import(framework_name)
+            if _gpu_frameworks_cache[framework_name] is not None:
+                logger.debug(
+                    f"🔧 Lazy imported {framework_name} in thread "
+                    f"{threading.current_thread().name}"
+                )
+        return _gpu_frameworks_cache[framework_name]
+    return getter
+
+
+# Auto-generate lazy getters for all GPU frameworks
+for mem_type in MemoryType:
+    ops = _FRAMEWORK_OPS[mem_type]
+    if ops['lazy_getter'] is not None:
+        getter_func = _create_lazy_getter(ops['import_name'])
+        globals()[f"_get_{ops['import_name']}"] = getter_func
+
+
+# Thread-local storage for GPU streams and contexts
+_thread_gpu_contexts = threading.local()
+
+class ThreadGPUContext:
+    """Thread-local GPU context manager for CUDA streams."""
+
+    def __init__(self):
+        self.cupy_stream = None
+        self.torch_stream = None
+        self.tensorflow_device = None
+        self.jax_device = None
+
+    def get_cupy_stream(self):
+        """Get or create thread-local CuPy stream."""
+        if self.cupy_stream is None:
+            cupy = globals().get('_get_cupy', lambda: None)()  # noqa: F821
+            if cupy is not None and hasattr(cupy, 'cuda'):
+                self.cupy_stream = cupy.cuda.Stream()
+                logger.debug(f"🔧 Created CuPy stream for thread {threading.current_thread().name}")
+        return self.cupy_stream
+
+    def get_torch_stream(self):
+        """Get or create thread-local PyTorch stream."""
+        if self.torch_stream is None:
+            torch = globals().get('_get_torch', lambda: None)()  # noqa: F821
+            if torch is not None and hasattr(torch, 'cuda') and torch.cuda.is_available():
+                self.torch_stream = torch.cuda.Stream()
+                logger.debug(
+                    f"🔧 Created PyTorch stream for thread "
+                    f"{threading.current_thread().name}"
+                )
+        return self.torch_stream
+
+
+def _get_thread_gpu_context():
+    """Get or create thread-local GPU context."""
+    if not hasattr(_thread_gpu_contexts, 'context'):
+        _thread_gpu_contexts.context = ThreadGPUContext()
+    return _thread_gpu_contexts.context
+
+
+def memory_types(
+    input_type: str,
+    output_type: str,
+    contract: Optional[Callable[[Any], bool]] = None
+) -> Callable[[F], F]:
+    """
+    Base decorator for declaring memory types of a function.
+
+    This is the foundation decorator that all memory-type-specific decorators build upon.
+    """
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            result = func(*args, **kwargs)
+
+            # Apply contract validation if provided
+            if contract is not None and not contract(result):
+                raise ValueError(f"Function {func.__name__} violated its output contract")
+
+            return result
+
+        # Attach memory type metadata
+        wrapper.input_memory_type = input_type
+        wrapper.output_memory_type = output_type
+
+        return wrapper
+
+    return decorator
+
+
+def _create_dtype_wrapper(func, mem_type: MemoryType, func_name: str):
+    """
+    Auto-generate dtype preservation wrapper for any memory type.
+
+    This single function replaces 6 nearly-identical dtype wrapper functions.
+    """
+    _FRAMEWORK_OPS[mem_type]
+    scale_func = SCALING_FUNCTIONS[mem_type.value]
+
+    @functools.wraps(func)
+    def dtype_wrapper(image, *args, dtype_conversion=None, slice_by_slice: bool = False, **kwargs):
+        # Set default dtype_conversion if not provided
+        if dtype_conversion is None:
+            dtype_conversion = DtypeConversion.PRESERVE_INPUT
+
+        try:
+            # Store original dtype
+            original_dtype = image.dtype
+
+            # Handle slice_by_slice processing for 3D arrays
+            if slice_by_slice and hasattr(image, 'ndim') and image.ndim == 3:
+                result = process_slices(image, func, args, kwargs)
+            else:
+                # Call the original function normally
+                result = func(image, *args, **kwargs)
+
+            # Apply dtype conversion based on enum value
+            if hasattr(result, 'dtype') and dtype_conversion is not None:
+                if dtype_conversion == DtypeConversion.PRESERVE_INPUT:
+                    # Preserve input dtype
+                    if result.dtype != original_dtype:
+                        result = scale_func(result, original_dtype)
+                elif dtype_conversion == DtypeConversion.NATIVE_OUTPUT:
+                    # Return framework's native output dtype
+                    pass  # No conversion needed
+                else:
+                    # Force specific dtype
+                    target_dtype = dtype_conversion.numpy_dtype
+                    if target_dtype is not None:
+                        result = scale_func(result, target_dtype)
+
+            return result
+        except Exception as e:
+            logger.error(
+                f"Error in {mem_type.value} dtype/slice preserving wrapper "
+                f"for {func_name}: {e}"
+            )
+            # Return original result on error
+            return func(image, *args, **kwargs)
+
+    # Update function signature to include new parameters
+    try:
+        original_sig = inspect.signature(func)
+        new_params = list(original_sig.parameters.values())
+
+        # Check if parameters already exist
+        param_names = [p.name for p in new_params]
+
+        # Add dtype_conversion parameter first (before slice_by_slice)
+        if 'dtype_conversion' not in param_names:
+            dtype_param = inspect.Parameter(
+                'dtype_conversion',
+                inspect.Parameter.KEYWORD_ONLY,
+                default=DtypeConversion.PRESERVE_INPUT,
+                annotation=Optional[DtypeConversion]
+            )
+            new_params.append(dtype_param)
+
+        # Add slice_by_slice parameter
+        if 'slice_by_slice' not in param_names:
+            slice_param = inspect.Parameter(
+                'slice_by_slice',
+                inspect.Parameter.KEYWORD_ONLY,
+                default=False,
+                annotation=bool
+            )
+            new_params.append(slice_param)
+
+        # Create new signature
+        new_sig = original_sig.replace(parameters=new_params)
+        dtype_wrapper.__signature__ = new_sig
+
+        # Update docstring
+        if dtype_wrapper.__doc__:
+            dtype_wrapper.__doc__ += (
+                f"\n\n    Additional Parameters "
+                f"(added by {mem_type.value} decorator):\n"
+            )
+            dtype_wrapper.__doc__ += (
+                "        dtype_conversion (DtypeConversion, optional): "
+                "How to handle output dtype.\n"
+            )
+            dtype_wrapper.__doc__ += (
+                "            Defaults to PRESERVE_INPUT (match input dtype).\n"
+            )
+            dtype_wrapper.__doc__ += (
+                "        slice_by_slice (bool, optional): "
+                "Process 3D arrays slice-by-slice.\n"
+            )
+            dtype_wrapper.__doc__ += (
+                "            Defaults to False. "
+                "Prevents cross-slice contamination.\n"
+            )
+
+    except Exception as e:
+        logger.warning(f"Could not update signature for {func_name}: {e}")
+
+    return dtype_wrapper
+
+
+def _create_gpu_wrapper(func, mem_type: MemoryType, oom_recovery: bool):
+    """
+    Auto-generate GPU stream/device wrapper for any GPU memory type.
+
+    This function creates the GPU-specific wrapper with stream management and OOM recovery.
+    """
+    ops = _FRAMEWORK_OPS[mem_type]
+    framework_name = ops['import_name']
+    lazy_getter = globals().get(ops['lazy_getter'])
+
+    @functools.wraps(func)
+    def gpu_wrapper(*args, **kwargs):
+        framework = lazy_getter()
+
+        # Check if GPU is available for this framework
+        if framework is not None:
+            gpu_check_expr = ops['gpu_check'].format(mod=framework_name)
+            try:
+                gpu_available = eval(gpu_check_expr, {framework_name: framework})
+            except Exception:
+                gpu_available = False
+
+            if gpu_available:
+                # Get thread-local context
+                ctx = _get_thread_gpu_context()
+
+                # Get stream if framework supports it
+                stream = None
+                if mem_type == MemoryType.CUPY:
+                    stream = ctx.get_cupy_stream()
+                elif mem_type == MemoryType.TORCH:
+                    stream = ctx.get_torch_stream()
+
+                # Define execution function that captures args/kwargs
+                def execute_with_stream():
+                    if stream is not None:
+                        with stream:
+                            return func(*args, **kwargs)
+                    else:
+                        return func(*args, **kwargs)
+
+                # Execute with OOM recovery if enabled
+                if oom_recovery and ops['has_oom_recovery']:
+                    return _execute_with_oom_recovery(execute_with_stream, mem_type.value)
+                else:
+                    return execute_with_stream()
+
+        # CPU fallback or framework not available
+        return func(*args, **kwargs)
+
+    # Preserve memory type attributes
+    gpu_wrapper.input_memory_type = func.input_memory_type
+    gpu_wrapper.output_memory_type = func.output_memory_type
+
+    return gpu_wrapper
+
+
+def _create_memory_decorator(mem_type: MemoryType):
+    """
+    Factory function that creates a decorator for a specific memory type.
+
+    This single factory replaces 6 nearly-identical decorator functions.
+    """
+    ops = _FRAMEWORK_OPS[mem_type]
+
+    def decorator(func=None, *, input_type=mem_type.value, output_type=mem_type.value,
+                  oom_recovery=True, contract=None):
+        """
+        Decorator for {mem_type} memory type functions.
+
+        Args:
+            func: Function to decorate (when used as @decorator)
+            input_type: Expected input memory type (default: {mem_type})
+            output_type: Expected output memory type (default: {mem_type})
+            oom_recovery: Enable automatic OOM recovery (default: True)
+            contract: Optional validation function for outputs
+
+        Returns:
+            Decorated function with memory type metadata and dtype preservation
+        """
+        def inner_decorator(func):
+            # Apply base memory_types decorator
+            memory_decorator = memory_types(
+                input_type=input_type,
+                output_type=output_type,
+                contract=contract
+            )
+            func = memory_decorator(func)
+
+            # Apply dtype preservation wrapper
+            func = _create_dtype_wrapper(func, mem_type, func.__name__)
+
+            # Apply GPU wrapper if this is a GPU memory type
+            if ops['gpu_check'] is not None:
+                func = _create_gpu_wrapper(func, mem_type, oom_recovery)
+
+            return func
+
+        # Handle both @decorator and @decorator() forms
+        if func is None:
+            return inner_decorator
+        return inner_decorator(func)
+
+    # Set proper function name and docstring
+    decorator.__name__ = mem_type.value
+    decorator.__doc__ = decorator.__doc__.format(mem_type=ops['display_name'])
+
+    return decorator
+
+
+# Auto-generate all 6 memory type decorators
+for mem_type in MemoryType:
+    decorator_func = _create_memory_decorator(mem_type)
+    globals()[mem_type.value] = decorator_func
+
+
+# Export all decorators
+__all__ = [
+    'memory_types',
+    'DtypeConversion',
+    'numpy',  # noqa: F822
+    'cupy',  # noqa: F822
+    'torch',  # noqa: F822
+    'tensorflow',  # noqa: F822
+    'jax',  # noqa: F822
+    'pyclesperanto',  # noqa: F822
+]
+