diff --git a/deeptrack/elementwise.py b/deeptrack/elementwise.py index b1c9fe6a..3ef7210a 100644 --- a/deeptrack/elementwise.py +++ b/deeptrack/elementwise.py @@ -1,759 +1,2154 @@ -"""Classes that apply functions to features elementwise. +"""Elementwise mathematical operations for DeepTrack features. -This module provides the `elementwise` DeepTrack2 classes -which work as a handle to apply various NumPy functions -to `Feature` objects elementwise. +This module defines a collection of `Feature` classes that apply mathematical +functions elementwise to arrays or tensors in a DeepTrack2 pipeline. These +operations are backend-agnostic and compatible with both NumPy and PyTorch. -Key Features ------------- -- **Extends NumPy Functions** +Elementwise features can be created in two ways: - The convenience of NumPy functions are extended with this module such that - they can be applied elementwise to a DeepTrack `Feature` object. +1. **Using the Factory Function (`create_elementwise_class`)** -- **Trigonometric Functions** - The elementary trigonometric functions: Sin, Cos, Tan. + For most functions that are available in both NumPy and PyTorch and are + supported by the `array-api-compat` backend abstraction (`xp`), the class + can be generated dynamically using the `create_elementwise_class` factory. -- **Hyperbolic Functions** - The trigonometric hyperbolic functions: Sinh, Cosh, Tanh. + For example: -- **Rounding Functions** - Common rounding functions: nearest integer rounding `Round`, - nearest lowest integer `Floor`, nearest highest integer `Ceil`. + >>> from deeptrack.backend import xp + >>> from deeptrack.elementwise import create_elementwise_class + >>> Abs = create_elementwise_class("Abs", xp.abs) -- **Exponents And Logarithm Functions** - Includes Exponential (exp) function, Natural Logarithm function, - Logarithm function with base 10, and Logarithm function with base 2. + This creates a `Feature` class named `Abs` that applies `abs()` to the + input elementwise, supporting both direct and pipeline usage. -- **Complex Number Functions** - Functions to get various values from a complex number: - Angle, Absolute value, Real value, Imaginary value, Conjugate +2. **Defining a Custom Subclass of `ElementwiseFeature`** -- **Miscellaneous Functions** - Contains Square root, Square, Sign function. + In cases where the `array-api-compat` implementation fails or does not + support the required function for a backend (e.g., `torch.float32` with + `xp.floor`), a custom subclass of `ElementwiseFeature` can be defined + explicitly. -Module Structure ----------------- -Classes: + These subclasses manually dispatch to the appropriate backend function + (e.g., `torch.floor`, `np.floor`) depending on the input type and device, + ensuring robust and backend-safe behavior. -- `ElementwiseFeature` - Forms the base from which other classes inherit from. + For example, `Floor`, `Ceil`, `Imag`, and `Sign` are implemented this way. -- `Sin` +This dual mechanism provides both flexibility and robustness for applying +mathematical operations in a pipeline-agnostic, extensible, and modular way. -- `Cos` - -- `Tan` +Key Features +------------ +- **Seamless Backend Compatibility** -- `ArcSin` + Most functions use `array-api-compat` (`xp`) to ensure compatibility with + both NumPy and PyTorch backends. Manual dispatch is used for cases where + `xp` fails (e.g., `ceil`, `floor`, `imag`, `sign`). -- `Arccos` +- **Factory-Generated and Manual Implementations** -- `ArcTan` + Elementwise operations are implemented using either: + - `create_elementwise_class()` for standard backend-agnostic functions. + - Dedicated subclasses of `ElementwiseFeature` for operations requiring + manual backend dispatch. -- `Sinh` +- **Supports Direct and Pipeline Composition** -- `Cosh` + These features can be applied directly to NumPy arrays or PyTorch tensors, + or used in combination with other DeepTrack `Feature` objects in pipelines. -- `Tanh` +- **Extensive Documentation and Examples** -- `ArcSinh` + Each class includes detailed docstrings with usage examples for both + backends and for pipeline integration. -- `Arccosh` +Module Structure +---------------- +Classes: -- `ArcTanh` +- `ElementwiseFeature` + + Base class for features that apply mathematical operations elementwise to + NumPy arrays or PyTorch tensors. Accepts a function and an optional + input `Feature`. -- `Round` +Functions: -- `Floor` +- `create_elementwise_class(name, function, docstring="")` -- `Ceil` + def create_elementwise_class( + name: str, + function: Callable[[NDArray | torch.Tensor], NDArray | torch.Tensor], + docstring: str = "", + ) -> type -- `Exp` + Factory function that returns a new subclass of `ElementwiseFeature` with + the given `name` and `function`. Automatically sets the class name, + module, and docstring for full introspection and documentation support. -- `Log` +Elementwise Features +-------------------- -- `Log10` +All elementwise features inherit from `ElementwiseFeature` and apply a +mathematical operation elementwise to the output of another `Feature`, or +directly to an input array or tensor. -- `Log2` +The following features are available: -- `Angle` +Trigonometric Functions: +- `Sin`: Applies the sine function `sin(x)` elementwise. +- `Cos`: Applies the cosine function `cos(x)` elementwise. +- `Tan`: Applies the tangent function `tan(x)` elementwise. -- `Real` +Inverse Trigonometric Functions: +- `Arcsin`: Applies the arcsine function `arcsin(x)` elementwise. +- `Arctan`: Applies the arctangent function `arctan(x)` elementwise. -- `Imag` +Hyperbolic Functions: +- `Sinh`: Applies the hyperbolic sine function `sinh(x)` elementwise. +- `Cosh`: Applies the hyperbolic cosine function `cosh(x)` elementwise. +- `Tanh`: Applies the hyperbolic tangent function `tanh(x)` elementwise. -- `Abs` +Inverse Hyperbolic Functions: +- `Arcsinh`: Applies the inverse hyperbolic sine `arcsinh(x)` elementwise. +- `Arccosh`: Applies the inverse hyperbolic cosine `arccosh(x)` elementwise. +- `Arctanh`: Applies the inverse hyperbolic tangent `arctanh(x)` elementwise. -- `Conjugate` +Rounding Functions: +- `Round`: Applies nearest-integer rounding elementwise. +- `Floor`: Applies floor function `floor(x)` elementwise. +- `Ceil`: Applies ceil function `ceil(x)` elementwise. -- `Sqrt` +Exponential and Logarithmic Functions: +- `Exp`: Applies the exponential function `exp(x)` elementwise. +- `Log`: Applies the natural logarithm `log(x)` elementwise. +- `Log10`: Applies the base-10 logarithm `log10(x)` elementwise. +- `Log2`: Applies the base-2 logarithm `log2(x)` elementwise. -- `Square` +Complex Number Functions: +- `Angle`: Returns the phase angle `angle(x)` of complex inputs. +- `Real`: Extracts the real part `real(x)` of complex inputs. +- `Imag`: Extracts the imaginary part `imag(x)`; returns zero for real tensors. +- `Abs`: Returns the magnitude or absolute value `abs(x)`. +- `Conj`, `Conjugate`: Returns the complex conjugate `conj(x)`. -- `Sign` +Miscellaneous Mathematical Functions: +- `Sqrt`: Applies the square root `sqrt(x)` elementwise. +- `Square`: Applies squaring operation `x**2` elementwise. +- `Sign`: Applies the sign function `sign(x)`; returns -1, 0, or 1. Examples -------- -Apply cosine elementwise to a Feature: +import deeptrack as dt + +Import the backend-agnostic functionality from DeepTrack2: +>>> from deeptrack.backend import xp + +Create a elementwise feature to execute a backend-agnostic function: +>>> from deeptrack.elementwise import create_elementwise_class +>>> +>>> Abs = create_elementwise_class( +... name="Abs", +... function=xp.abs, +... docstring="Elementwise abs function." +... ) + +**NumPy backend with direct resolved input** >>> import numpy as np +>>> +>>> array = np.array([-1.0, 0.0, 2.5]) +>>> result = Abs()(array) +>>> result +array([1. , 0. , 2.5]) ->>> from deeptrack import Feature, elementwise +**PyTorch backend with direct resolved input** ->>> class TestFeature(Feature): ->>> __distributed__ = False ->>> def get(self, image, **kwargs): ->>> output = np.array([[np.pi, 0], -... [np.pi / 4, 0]]) ->>> return output +>>> import torch +>>> +>>> tensor = torch.tensor([-1.0, 0.0, 2.5]) +>>> result = Abs()(tensor) +>>> result +tensor([1.0000, 0.0000, 2.5000]) ->>> test_feature = TestFeature() ->>> elementwise_cosine = test_feature >> elementwise.Cos() -[[-1. 1. ] - [ 0.70710678 1. ]] +**NumPy pipeline** -""" +>>> value = dt.Value(value=np.array([-3.0, 0.0, 3.0])) +>>> pipeline = value >> Abs() +>>> result = pipeline() +>>> result +array([3., 0., 3.]) + +This is equivalent to: +>>> pipeline = Abs(value) + +**PyTorch pipeline** -#TODO ***??*** revise class docstring -#TODO ***??*** revise DTAT389 +>>> value = dt.Value(value=torch.tensor([-3.0, 0.0, 3.0])) +>>> pipeline = value >> Abs() +>>> result = pipeline() +>>> result +tensor([3., 0., 3.]) + +This is equivalent to: +>>> pipeline = Abs(value) + +""" from __future__ import annotations -from typing import Any, Callable +from typing import Any, Callable, TYPE_CHECKING import numpy as np - -from deeptrack.features import Feature +from numpy.typing import NDArray + +from deeptrack import Feature, TORCH_AVAILABLE, xp + +if TORCH_AVAILABLE: + import torch + +__all__ = [ + "ElementwiseFeature", + "create_elementwise_class", + "Sin", + "Cos", + "Tan", + "Arcsin", + "Arctan", + "Sinh", + "Cosh", + "Tanh", + "Arcsinh", + "Arccosh", + "Arctanh", + "Round", + "Floor", + "Ceil", + "Exp", + "Log", + "Log10", + "Log2", + "Angle", + "Real", + "Imag", + "Abs", + "Conj", + "Conjugate", + "Sqrt", + "Square", + "Sign", +] + + +if TYPE_CHECKING: + import torch -#TODO ***??*** revise ElementwiseFeature - torch, typing, docstring, unit test class ElementwiseFeature(Feature): - """ - Base class for applying NumPy functions elementwise. + """Base class for applying NumPy or PyTorch functions elementwise. - This class provides the foundation for subclasses that apply specific - NumPy functions (e.g., sin, cos, exp) to the elements of an input array. + This class wraps a backend function (e.g., `np.sin`, `torch.exp`) and + applies it elementwise to the output of another `Feature`. + + If no input feature is provided, the function is applied directly to the + input passed during resolution. Parameters ---------- - function : Callable[[np.ndarray], np.ndarray] - The NumPy function to be applied elementwise. - feature : Feature or None, optional - The input feature to which the function will be applied. If None, - the function will be applied to the input array directly. + function: Callable[[array], array] + A backend-specific function (e.g., `np.sin`, `torch.abs`) or a + backend-agnostic function (e.g., `xp.sin`, `xp.abs`) that will be + applied elementwise to the input NumPy array or PyTorch tensor. + feature: Feature or None, optional + The input feature to be transformed. If provided, the function is + applied to the output of this feature. If None, the function is applied + directly to the input passed to `resolve`. + + Attributes + ---------- + __distributed__: bool + It overrides the default behavior to disable distributed resolution if + a feature is chained. This ensures the transformation is computed + locally. Methods ------- - `get(image: np.ndarray, **kwargs: Any) -> np.ndarray` - Returns the result of applying the function to the input array. + get(image: array, **kwargs: Any) -> array + It applies the stored function to the input, optionally resolving the + wrapped feature first. """ + __distributed__: bool + function: Callable[[NDArray[Any] | torch.Tensor], + NDArray[Any] | torch.Tensor] | None + feature: Feature + def __init__( self: ElementwiseFeature, - function: Callable[[np.ndarray], np.ndarray], + function: Callable[ + [NDArray[Any] | torch.Tensor], + NDArray[Any] | torch.Tensor + ], feature: Feature | None = None, + **kwargs: Any, + ): + """Initialize ElementwiseFeature. + + It initializes ElementwiseFeature with function and optional input + feature. + + Parameters + ---------- + function: Callable[[array], array] + The function to apply elementwise to the input NumPy array or + PyTorch tensor. + feature: Feature or None, optional + The feature whose output will be transformed. If None, the + function is applied to the direct input. **kwargs: Any - ) -> None: + Additional keyword arguments passed to the `Feature` base class. + + """ + super().__init__(**kwargs) + + # Store the function to be applied elementwise self.function = function + + # Add the feature dependency if provided self.feature = ( self.add_feature(feature) if feature is not None else None ) + + # If the feature is set, prevent distributed resolution if feature: self.__distributed__ = False def get( self: ElementwiseFeature, - image: np.ndarray, + image: NDArray[Any] | torch.Tensor, + **kwargs: Any, + ) -> NDArray[Any] | torch.Tensor: + """Apply the stored function. + + It applies the stored function to the input or the result of the + wrapped feature. + + Parameters + ---------- + image: array + The input data to process, or a placeholder if a feature is + chained. **kwargs: Any - ) -> np.ndarray: + Additional keyword arguments for compatibility. + + Returns + ------- + array + The result of applying the elementwise function. + + """ + + # Resolve the input from the chained feature if present if self.feature: image = self.feature() + + # Apply the function elementwise return self.function(image) -#TODO ***??*** revise Sin - torch, typing, docstring, unit test -class Sin(ElementwiseFeature): - """ - Applies the sine function elementwise. +def create_elementwise_class( + name: str, + function: Callable[ + [NDArray[Any] | torch.Tensor], + NDArray[Any] | torch.Tensor + ], + docstring: str = "", +) -> type: + """Factory function to create subclasses of ElementwiseFeature. + + This function generates a new subclass of `ElementwiseFeature` that + applies the given function elementwise to a NumPy array or PyTorch tensor. + It dynamically sets the class name, qualified name, module, and docstring + to make the generated class fully compatible with IDEs and documentation + tools such as Sphinx. Parameters ---------- - feature : Feature or None, optional - The input feature to which the sine function will be applied. - If None, the function is applied to the input array directly. - + name: str + Name of the new class to be created (e.g., "Sin", "Exp"). + function: Callable[[array], array] + The elementwise function to apply, such as `np.sin`, `torch.exp`, or + `xp.abs`. + docstring: str, optional + The docstring for the generated class. This string will be visible + in IDE tooltips and Sphinx documentation. + + Returns + ------- + type + A dynamically generated subclass of `ElementwiseFeature` that wraps + the given function. + + Examples + -------- + import deeptrack as dt + + Import the backend-agnostic functionality from DeepTrack2: + >>> from deeptrack.backend import xp + + Create a elementwise feature to execute a backend-agnostic function: + >>> from deeptrack.elementwise import create_elementwise_class + >>> + >>> Abs = create_elementwise_class( + ... name="Abs", + ... function=xp.abs, + ... docstring="Elementwise abs function." + ... ) + + **NumPy backend with direct resolved input** + + >>> import numpy as np + >>> + >>> array = np.array([-1.0, 0.0, 2.5]) + >>> result = Abs()(array) + >>> result + array([1. , 0. , 2.5]) + + **PyTorch backend with direct resolved input** + + >>> import torch + >>> + >>> tensor = torch.tensor([-1.0, 0.0, 2.5]) + >>> result = Abs()(tensor) + >>> result + tensor([1.0000, 0.0000, 2.5000]) + + **NumPy pipeline** + + >>> value = dt.Value(value=np.array([-3.0, 0.0, 3.0])) + >>> pipeline = value >> Abs() + >>> result = pipeline() + >>> result + array([3., 0., 3.]) + + This is equivalent to: + >>> pipeline = Abs(value) + + **PyTorch pipeline** + + >>> value = dt.Value(value=torch.tensor([-3.0, 0.0, 3.0])) + >>> pipeline = value >> Abs() + >>> result = pipeline() + >>> result + tensor([3., 0., 3.]) + + This is equivalent to: + >>> pipeline = Abs(value) + """ - - def __init__( - self: Sin, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.sin, feature=feature, **kwargs) + class _GeneratedElementwise(ElementwiseFeature): + """Dynamically generated subclass of ElementwiseFeature.""" -#TODO ***??*** revise Cos - torch, typing, docstring, unit test -class Cos(ElementwiseFeature): - """ - Applies the cosine function elementwise. + def __init__( + self: _GeneratedElementwise, + feature: Feature | None = None, + **kwargs: Any, + ): + # Initialize the ElementwiseFeature with the fixed function + super().__init__(function=function, feature=feature, **kwargs) + + # Set the class name to match the intended name + _GeneratedElementwise.__name__ = name + + # Set the qualified name to improve introspection and traceback readability + _GeneratedElementwise.__qualname__ = name + + # Attach the user-specified docstring to enable documentation + _GeneratedElementwise.__doc__ = docstring + + # Set correct module to ensure proper Sphinx indexing and import tracing + _GeneratedElementwise.__module__ = __name__ + + return _GeneratedElementwise + + +Sin = create_elementwise_class( + name="Sin", + function=xp.sin, + docstring=""" + Apply the sine function elementwise. + + This feature applies `xp.sin` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. Parameters ---------- - feature : Feature or None, optional - The input feature to which the cosine function will be applied. + feature: Feature or None, optional + The input feature to which the sine function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Sin + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Sin()(np.array([0, np.pi / 2, np.pi])) + >>> result + array([0.0000000e+00, 1.0000000e+00, 1.2246468e-16]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Sin()(torch.tensor([0, torch.pi / 2, torch.pi])) + >>> result + tensor([ 0.0000e+00, 1.0000e+00, -8.7423e-08]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([0, np.pi / 2, np.pi])) + >>> pipeline = value >> Sin() + >>> result = pipeline() + >>> result + array([0.0000000e+00, 1.0000000e+00, 1.2246468e-16]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([0, torch.pi / 2, torch.pi])) + >>> pipeline = value >> Sin() + >>> result = pipeline() + >>> result + tensor([ 0.0000e+00, 1.0000e+00, -8.7423e-08]) + + These are equivalent to: + >>> pipeline = Sin(value) + """ +) - def __init__( - self: Cos, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.cos, feature=feature, **kwargs) +Cos = create_elementwise_class( + name="Cos", + function=xp.cos, + docstring=""" + Apply the cosine function elementwise. -#TODO ***??*** revise Tan - torch, typing, docstring, unit test -class Tan(ElementwiseFeature): - """ - Applies the tangent function elementwise. + This feature applies `xp.cos` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. Parameters ---------- - feature : Feature or None, optional - The input feature to which the tangent function will be applied. + feature: Feature or None, optional + The input feature to which the cosine function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Cos + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Cos()(np.array([0, np.pi / 2, np.pi])) + >>> result + array([ 1.000000e+00, 6.123234e-17, -1.000000e+00]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Cos()(torch.tensor([0, torch.pi / 2, torch.pi])) + >>> result + tensor([ 1.0000e+00, -4.3711e-08, -1.0000e+00]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([0, np.pi / 2, np.pi])) + >>> pipeline = value >> Cos() + >>> result = pipeline() + >>> result + array([ 1.000000e+00, 6.123234e-17, -1.000000e+00]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([0, torch.pi / 2, torch.pi])) + >>> pipeline = value >> Cos() + >>> result = pipeline() + >>> result + tensor([ 1.0000e+00, -4.3711e-08, -1.0000e+00]) + + These are equivalent to: + >>> pipeline = Cos(value) + """ +) - def __init__( - self: Tan, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.tan, feature=feature, **kwargs) +Tan = create_elementwise_class( + name="Tan", + function=xp.tan, + docstring=""" + Apply the tangent function elementwise. -#TODO ***??*** revise Arcsin - torch, typing, docstring, unit test -class Arcsin(ElementwiseFeature): - """ - Applies the arcsine function elementwise. + This feature applies `xp.tan` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. Parameters ---------- - feature : Feature or None, optional - The input feature to which the arcsine function will be applied. + feature: Feature or None, optional + The input feature to which the tangent function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Tan + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Tan()(np.array([0, np.pi / 4, np.pi / 2])) + >>> result + array([0.00000000e+00, 1.00000000e+00, 1.63312394e+16]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Tan()(torch.tensor([0, torch.pi / 4, torch.pi / 2])) + >>> result + tensor([ 0.0000e+00, 1.0000e+00, -2.2877e+07]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([0, np.pi / 4, np.pi / 2])) + >>> pipeline = value >> Tan() + >>> result = pipeline() + >>> result + array([0.00000000e+00, 1.00000000e+00, 1.63312394e+16]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([0, torch.pi / 4, torch.pi / 2])) + >>> pipeline = value >> Tan() + >>> result = pipeline() + >>> result + tensor([ 0.0000e+00, 1.0000e+00, -2.2877e+07]) + + These are equivalent to: + >>> pipeline = Tan(value) + """ +) - def __init__( - self: Arcsin, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.arcsin, feature=feature, **kwargs) +Arcsin = create_elementwise_class( + name="Arcsin", + function=xp.arcsin, + docstring=""" + Apply the arcsine function elementwise. -#TODO ***??*** revise Arccos - torch, typing, docstring, unit test -class Arccos(ElementwiseFeature): - """ - Applies the arccosine function elementwise. + This feature applies `xp.arcsin` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + The input must be in the domain [-1, 1]. Values outside this range will + produce NaNs or raise runtime warnings or errors, depending on the backend. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the arccosine function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Arcsin + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Arcsin()(np.array([0.0, 0.5, 1.0])) + >>> result + array([0. , 0.52359878, 1.57079633]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Arcsin()(torch.tensor([0.0, 0.5, 1.0])) + >>> result + tensor([0.0000, 0.5236, 1.5708]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([0.0, 0.5, 1.0])) + >>> pipeline = value >> Arcsin() + >>> result = pipeline() + >>> result + array([0. , 0.52359878, 1.57079633]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([0.0, 0.5, 1.0])) + >>> pipeline = value >> Arcsin() + >>> result = pipeline() + >>> result + tensor([0.0000, 0.5236, 1.5708]) + + These are equivalent to: + >>> pipeline = Arcsin(value) + """ +) - def __init__( - self: Arccos, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.arccos, feature=feature, **kwargs) +Arctan = create_elementwise_class( + name="Arctan", + function=xp.arctan, + docstring=""" + Apply the arctangent function elementwise. -#TODO ***??*** revise Arctan - torch, typing, docstring, unit test -class Arctan(ElementwiseFeature): - """ - Applies the arctangent function elementwise. + This feature applies `xp.arctan` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the arctangent function will be applied. If None, the function is applied to the input array directly. - - """ - def __init__( - self: Arctan, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.arctan, feature=feature, **kwargs) + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Arctan + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Arctan()(np.array([-1.0, 0.0, 1.0])) + >>> result + array([-0.78539816, 0. , 0.78539816]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Arctan()(torch.tensor([-1.0, 0.0, 1.0])) + >>> result + tensor([-0.7854, 0.0000, 0.7854]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Arctan() + >>> result = pipeline() + >>> result + array([-0.78539816, 0. , 0.78539816]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Arctan() + >>> result = pipeline() + >>> result + tensor([-0.7854, 0.0000, 0.7854]) + + These are equivalent to: + >>> pipeline = Arctan(value) -#TODO ***??*** revise Sinh - torch, typing, docstring, unit test -class Sinh(ElementwiseFeature): """ - Applies the hyperbolic sine function elementwise. +) + + +Sinh = create_elementwise_class( + name="Sinh", + function=xp.sinh, + docstring=""" + Apply the hyperbolic sine function elementwise. + + This feature applies `xp.sinh` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the hyperbolic sine function will be applied. If None, the function is applied to the input array directly. - - """ - def __init__( - self: Sinh, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.sinh, feature=feature, **kwargs) + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Sinh + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Sinh()(np.array([-1.0, 0.0, 1.0])) + >>> result + array([-1.17520119, 0. , 1.17520119]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Sinh()(torch.tensor([-1.0, 0.0, 1.0])) + >>> result + tensor([-1.1752, 0.0000, 1.1752]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Sinh() + >>> result = pipeline() + >>> result + array([-1.17520119, 0. , 1.17520119]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Sinh() + >>> result = pipeline() + >>> result + tensor([-1.1752, 0.0000, 1.1752]) + + These are equivalent to: + >>> pipeline = Sinh(value) -#TODO ***??*** revise Cosh - torch, typing, docstring, unit test -class Cosh(ElementwiseFeature): """ - Applies the hyperbolic cosine function elementwise. +) + + +Cosh = create_elementwise_class( + name="Cosh", + function=xp.cosh, + docstring=""" + Apply the hyperbolic cosine function elementwise. + + This feature applies `xp.cosh` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the hyperbolic cosine function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Cosh + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Cosh()(np.array([-1.0, 0.0, 1.0])) + >>> result + array([1.54308063, 1. , 1.54308063]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Cosh()(torch.tensor([-1.0, 0.0, 1.0])) + >>> result + tensor([1.5431, 1.0000, 1.5431]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Cosh() + >>> result = pipeline() + >>> result + array([1.54308063, 1. , 1.54308063]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Cosh() + >>> result = pipeline() + >>> result + tensor([1.5431, 1.0000, 1.5431]) + + These are equivalent to: + >>> pipeline = Cosh(value) + """ +) - def __init__( - self: Cosh, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.cosh, feature=feature, **kwargs) +Tanh = create_elementwise_class( + name="Tanh", + function=xp.tanh, + docstring=""" + Apply the hyperbolic tangent function elementwise. -#TODO ***??*** revise Tanh - torch, typing, docstring, unit test -class Tanh(ElementwiseFeature): - """ - Applies the hyperbolic tangent function elementwise. + This feature applies `xp.tanh` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the hyperbolic tangent function will be applied. If None, the function is applied to the input array directly. - - """ - - def __init__( - self: Tanh, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.tanh, feature=feature, **kwargs) + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Tanh + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Tanh()(np.array([-1.0, 0.0, 1.0])) + >>> result + array([-0.76159416, 0. , 0.76159416]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Tanh()(torch.tensor([-1.0, 0.0, 1.0])) + >>> result + tensor([-0.7616, 0.0000, 0.7616]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Tanh() + >>> result = pipeline() + >>> result + array([-0.76159416, 0. , 0.76159416]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Tanh() + >>> result = pipeline() + >>> result + tensor([-0.7616, 0.0000, 0.7616]) + + These are equivalent to: + >>> pipeline = Tanh(value) -#TODO ***??*** revise Arcsinh - torch, typing, docstring, unit test -class Arcsinh(ElementwiseFeature): """ - Applies the hyperbolic arcsine function elementwise. +) + + +Arcsinh = create_elementwise_class( + name="Arcsinh", + function=xp.arcsinh, + docstring=""" + Apply the inverse hyperbolic sine function elementwise. + + This feature applies `xp.arcsinh` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the hyperbolic arcsine function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Arcsinh + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Arcsinh()(np.array([-1.0, 0.0, 1.0])) + >>> result + array([-0.88137359, 0. , 0.88137359]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Arcsinh()(torch.tensor([-1.0, 0.0, 1.0])) + >>> result + tensor([-0.8814, 0.0000, 0.8814]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Arcsinh() + >>> result = pipeline() + >>> result + array([-0.88137359, 0. , 0.88137359]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Arcsinh() + >>> result = pipeline() + >>> result + tensor([-0.8814, 0.0000, 0.8814]) + + These are equivalent to: + >>> pipeline = Arcsinh(value) + """ +) - def __init__( - self: Arcsinh, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.arcsinh, feature=feature, **kwargs) +Arccosh = create_elementwise_class( + name="Arccosh", + function=xp.arccosh, + docstring=""" + Apply the inverse hyperbolic cosine function elementwise. -#TODO ***??*** revise Arccosh - torch, typing, docstring, unit test -class Arccosh(ElementwiseFeature): - """ - Applies the hyperbolic arccosine function elementwise. + This feature applies `xp.arccosh` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + The input must be greater than or equal to 1. Values below this will + return NaN or raise errors depending on the backend. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the hyperbolic arccosine function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Arccosh + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Arccosh()(np.array([1.0, 2.0, 3.0])) + >>> result + array([0. , 1.3169579 , 1.76274717]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Arccosh()(torch.tensor([1.0, 2.0, 3.0])) + >>> result + tensor([0.0000, 1.3170, 1.7627]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([1.0, 2.0, 3.0])) + >>> pipeline = value >> Arccosh() + >>> result = pipeline() + >>> result + array([0. , 1.3169579 , 1.76274717]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([1.0, 2.0, 3.0])) + >>> pipeline = value >> Arccosh() + >>> result = pipeline() + >>> result + tensor([0.0000, 1.3170, 1.7627]) + + These are equivalent to: + >>> pipeline = Arccosh(value) + """ +) - def __init__( - self: Arccosh, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.arccosh, feature=feature, **kwargs) +Arctanh = create_elementwise_class( + name="Arctanh", + function=xp.arctanh, + docstring=""" + Apply the inverse hyperbolic tangent function elementwise. -#TODO ***??*** revise Arctanh - torch, typing, docstring, unit test -class Arctanh(ElementwiseFeature): - """ - Applies the hyperbolic arctangent function elementwise. + This feature applies `xp.arctanh` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + The input must be within the open interval (-1, 1). Values outside this + range will produce NaNs or raise domain errors depending on the backend. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the hyperbolic arctangent function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Arctanh + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Arctanh()(np.array([-0.5, 0.0, 0.5])) + >>> result + array([-0.54930614, 0. , 0.54930614]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Arctanh()(torch.tensor([-0.5, 0.0, 0.5])) + >>> result + tensor([-0.5493, 0.0000, 0.5493]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-0.5, 0.0, 0.5])) + >>> pipeline = value >> Arctanh() + >>> result = pipeline() + >>> result + array([-0.54930614, 0. , 0.54930614]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-0.5, 0.0, 0.5])) + >>> pipeline = value >> Arctanh() + >>> result = pipeline() + >>> result + tensor([-0.5493, 0.0000, 0.5493]) + + These are equivalent to: + >>> pipeline = Arctanh(value) + """ +) - def __init__( - self: Arctanh, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.arctanh, feature=feature, **kwargs) +Round = create_elementwise_class( + name="Round", + function=xp.round, + docstring=""" + Apply the rounding function elementwise. -#TODO ***??*** revise Round - torch, typing, docstring, unit test -class Round(ElementwiseFeature): - """ - Applies the round function elementwise. + This feature applies `xp.round` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + This function rounds to the nearest integer. For NumPy, ties round to + the even number (bankers' rounding). For PyTorch, ties round away from + zero. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the round function will be applied. If None, the function is applied to the input array directly. - - """ - def __init__( - self: Round, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.around, feature=feature, **kwargs) + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Round + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Round()(np.array([-1.5, -0.5, 0.5, 1.5])) + >>> result + array([-2., -0., 0., 2.]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Round()(torch.tensor([-1.5, -0.5, 0.5, 1.5])) + >>> result + tensor([-2., -1., 1., 2.]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-1.5, -0.5, 0.5, 1.5])) + >>> pipeline = value >> Round() + >>> result = pipeline() + >>> result + array([-2., -0., 0., 2.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-1.5, -0.5, 0.5, 1.5])) + >>> pipeline = value >> Round() + >>> result = pipeline() + >>> result + tensor([-2., -1., 1., 2.]) + + These are equivalent to: + >>> pipeline = Round(value) + + """ +) -#TODO ***??*** revise Floor - torch, typing, docstring, unit test class Floor(ElementwiseFeature): - """ - Applies the floor function elementwise. + """Apply the floor function elementwise. + + This feature applies `xp.floor` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + The floor function returns the greatest integer less than or equal to + each element of the input. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the floor function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Floor + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Floor()(np.array([-1.7, -0.5, 0.0, 0.5, 1.7])) + >>> result + array([-2., -1., 0., 0., 1.]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Floor()(torch.tensor([-1.7, -0.5, 0.0, 0.5, 1.7])) + >>> result + tensor([-2., -1., 0., 0., 1.]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-1.7, -0.5, 0.0, 0.5, 1.7])) + >>> pipeline = value >> Floor() + >>> result = pipeline() + >>> result + array([-2., -1., 0., 0., 1.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-1.7, -0.5, 0.0, 0.5, 1.7])) + >>> pipeline = value >> Floor() + >>> result = pipeline() + >>> result + tensor([-2., -1., 0., 0., 1.]) + + These are equivalent to: + >>> pipeline = Floor(value) + """ def __init__( self: Floor, feature: Feature | None = None, - **kwargs: Any + **kwargs: Any, ) -> None: - super().__init__(np.floor, feature=feature, **kwargs) + """Initialize the Floor feature. + + Parameters + ---------- + feature: Feature or None, optional + The input feature whose output will be transformed. + **kwargs: Any + Additional keyword arguments passed to the base Feature class. + + """ + + super().__init__( + function=self._floor_dispatch, + feature=feature, + **kwargs, + ) + + @staticmethod + def _floor_dispatch(x): + """Dispatch floor function based on backend. + + This method applies `torch.floor` if the input is a PyTorch tensor, + and `np.floor` if it is a NumPy array. It ensures compatibility + across both backends and avoids errors caused by `array-api-compat`. + + This explicit dispatch is necessary because `array-api-compat`'s + `xp.floor` internally calls `xp.issubdtype(x.dtype, xp.integer)`, + which fails when `x` is a `torch.Tensor`. + + As a result, this class cannot safely use the + `create_elementwise_class()` factory, and must be defined manually + using this backend-aware method. + + Parameters + ---------- + x: np.ndarray or torch.Tensor + The input to transform. + + Returns + ------- + np.ndarray or torch.Tensor + The result after applying floor elementwise. + + """ + + if TORCH_AVAILABLE and isinstance(x, torch.Tensor): + return torch.floor(x) + + return np.floor(x) -#TODO ***??*** revise Ceil - torch, typing, docstring, unit test class Ceil(ElementwiseFeature): - """ - Applies the ceil function elementwise. + """Apply the ceiling function elementwise. + + This feature applies `xp.ceil` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + The ceiling function returns the smallest integer greater than or equal to + each element of the input. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the ceil function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Ceil + + Use with NumPy directly: + >>> import numpy as np + >>> + >>> result = Ceil()(np.array([-1.7, -0.5, 0.0, 0.5, 1.7])) + >>> result + array([-1., -0., 0., 1., 2.]) + + Use with PyTorch directly: + >>> import torch + >>> + >>> result = Ceil()(torch.tensor([-1.7, -0.5, 0.0, 0.5, 1.7])) + >>> result + tensor([-1., -0., 0., 1., 2.]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-1.7, -0.5, 0.0, 0.5, 1.7])) + >>> pipeline = value >> Ceil() + >>> result = pipeline() + >>> result + array([-1., -0., 0., 1., 2.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-1.7, -0.5, 0.0, 0.5, 1.7])) + >>> pipeline = value >> Ceil() + >>> result = pipeline() + >>> result + tensor([-1., -0., 0., 1., 2.]) + + These are equivalent to: + >>> pipeline = Ceil(value) + """ def __init__( self: Ceil, feature: Feature | None = None, - **kwargs: Any + **kwargs: Any, ) -> None: - super().__init__(np.ceil, feature=feature, **kwargs) + """Initialize the Ceil feature. + Parameters + ---------- + feature: Feature or None, optional + The input feature whose output will be transformed. + **kwargs: Any + Additional keyword arguments passed to the base Feature class. -#TODO ***??*** revise Exp - torch, typing, docstring, unit test -class Exp(ElementwiseFeature): - """ - Applies the exponential function elementwise. + """ + + super().__init__( + function=self._ceil_dispatch, + feature=feature, + **kwargs + ) + + @staticmethod + def _ceil_dispatch(x): + """Dispatch ceiling function based on backend. + + This method applies `torch.ceil` if the input is a PyTorch tensor, + and `np.ceil` if it is a NumPy array. It ensures compatibility + across both backends and avoids errors caused by `array-api-compat`. + + This explicit dispatch is necessary because `array-api-compat`'s + `xp.ceil` internally calls `xp.issubdtype(x.dtype, xp.integer)`, + which fails when `x` is a `torch.Tensor`. + + As a result, this class cannot safely use the + `create_elementwise_class()` factory, and must be defined manually + using this backend-aware method. + + Parameters + ---------- + x: np.ndarray or torch.Tensor + The input to transform. + + Returns + ------- + np.ndarray or torch.Tensor + The result after applying ceil elementwise. + + """ + + if TORCH_AVAILABLE and isinstance(x, torch.Tensor): + return torch.ceil(x) + + return np.ceil(x) + + +Exp = create_elementwise_class( + name="Exp", + function=xp.exp, + docstring=""" + Apply the exponential function elementwise. + + This feature applies `xp.exp` (NumPy or PyTorch) to each element in the + input. It supports both direct input and pipeline composition. + + The exponential function computes `e**x` elementwise. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the exponential function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Exp + + Use with NumPy directly: + >>> import numpy as np + >>> result = Exp()(np.array([-1.0, 0.0, 1.0])) + >>> result + array([0.36787944, 1. , 2.71828183]) + + Use with PyTorch directly: + >>> import torch + >>> result = Exp()(torch.tensor([-1.0, 0.0, 1.0])) + >>> result + tensor([0.3679, 1.0000, 2.7183]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Exp() + >>> result = pipeline() + >>> result + array([0.36787944, 1. , 2.71828183]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-1.0, 0.0, 1.0])) + >>> pipeline = value >> Exp() + >>> result = pipeline() + >>> result + tensor([0.3679, 1.0000, 2.7183]) + + These are equivalent to: + >>> pipeline = Exp(value) + """ +) - def __init__( - self: Exp, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.exp, feature=feature, **kwargs) +Log = create_elementwise_class( + name="Log", + function=xp.log, + docstring=""" + Apply the natural logarithm function elementwise. -#TODO ***??*** revise Log - torch, typing, docstring, unit test -class Log(ElementwiseFeature): - """ - Applies the natural logarithm function elementwise. + This feature applies `xp.log` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + The input must be strictly positive. Passing zero or negative values will + return `-inf` or `NaN`, and may raise warnings or errors depending on + the backend. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the natural logarithm function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Log + + Use with NumPy directly: + >>> import numpy as np + >>> result = Log()(np.array([1.0, np.e, 10.0])) + >>> result + array([0. , 1. , 2.30258509]) + + Use with PyTorch directly: + >>> import torch + >>> result = Log()(torch.tensor([1.0, torch.exp(torch.tensor(1.0)), 10.0])) + >>> result + tensor([0.0000, 1.0000, 2.3026]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([1.0, np.e, 10.0])) + >>> pipeline = value >> Log() + >>> result = pipeline() + >>> result + array([0. , 1. , 2.30258509]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor( + ... [1.0, torch.exp(torch.tensor(1.0)), 10.0]) + ... ) + >>> pipeline = value >> Log() + >>> result = pipeline() + >>> result + tensor([0.0000, 1.0000, 2.3026]) + + These are equivalent to: + >>> pipeline = Log(value) + """ +) - def __init__( - self: Log, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.log, feature=feature, **kwargs) +Log10 = create_elementwise_class( + name="Log10", + function=xp.log10, + docstring=""" + Apply the base-10 logarithm function elementwise. -#TODO ***??*** revise Log10 - torch, typing, docstring, unit test -class Log10(ElementwiseFeature): - """ - Applies the logarithm function with base 10 elementwise. + This feature applies `xp.log10` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + The input must be strictly positive. Passing zero or negative values will + return `-inf` or `NaN`, and may raise warnings or errors depending on + the backend. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the logarithm function with base 10 will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Log10 + + Use with NumPy directly: + >>> import numpy as np + >>> result = Log10()(np.array([1.0, 10.0, 100.0])) + >>> result + array([0., 1., 2.]) + + Use with PyTorch directly: + >>> import torch + >>> result = Log10()(torch.tensor([1.0, 10.0, 100.0])) + >>> result + tensor([0., 1., 2.]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([1.0, 10.0, 100.0])) + >>> pipeline = value >> Log10() + >>> result = pipeline() + >>> result + array([0., 1., 2.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([1.0, 10.0, 100.0])) + >>> pipeline = value >> Log10() + >>> result = pipeline() + >>> result + tensor([0., 1., 2.]) + + These are equivalent to: + >>> pipeline = Log10(value) + """ +) - def __init__( - self: Log10, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.log10, feature=feature, **kwargs) +Log2 = create_elementwise_class( + name="Log2", + function=xp.log2, + docstring=""" + Apply the base-2 logarithm function elementwise. -#TODO ***??*** revise Log2 - torch, typing, docstring, unit test -class Log2(ElementwiseFeature): - """ - Applies the logarithm function with base 2 elementwise. + This feature applies `xp.log2` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + The input must be strictly positive. Passing zero or negative values will + return `-inf` or `NaN`, and may raise warnings or errors depending on + the backend. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the logarithm function with base 2 will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Log2 + + Use with NumPy directly: + >>> import numpy as np + >>> result = Log2()(np.array([1.0, 2.0, 4.0, 8.0])) + >>> result + array([0., 1., 2., 3.]) + + Use with PyTorch directly: + >>> import torch + >>> result = Log2()(torch.tensor([1.0, 2.0, 4.0, 8.0])) + >>> result + tensor([0., 1., 2., 3.]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([1.0, 2.0, 4.0, 8.0])) + >>> pipeline = value >> Log2() + >>> result = pipeline() + >>> result + array([0., 1., 2., 3.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([1.0, 2.0, 4.0, 8.0])) + >>> pipeline = value >> Log2() + >>> result = pipeline() + >>> result + tensor([0., 1., 2., 3.]) + + These are equivalent to: + >>> pipeline = Log2(value) + """ +) - def __init__( - self: Log2, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.log2, feature=feature, **kwargs) +Angle = create_elementwise_class( + name="Angle", + function=xp.angle, + docstring=""" + Apply the angle (phase) function elementwise. -#TODO ***??*** revise Angle - torch, typing, docstring, unit test -class Angle(ElementwiseFeature): - """ - Applies the angle function elementwise. + This feature applies `xp.angle` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + The angle function returns the phase angle (in radians) of a complex + number. For real-valued inputs, it returns 0 for positive and π for + negative values. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the angle function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Angle + + Use with NumPy directly: + >>> import numpy as np + >>> result = Angle()(np.array([1+0j, 0+1j, -1+0j, 1+1j])) + >>> result + array([0. , 1.57079633, 3.14159265, 0.78539816]) + + Use with PyTorch directly: + >>> import torch + >>> result = Angle()(torch.tensor([1+0j, 0+1j, -1+0j, 1+1j])) + >>> result + tensor([0.0000, 1.5708, 3.1416, 0.7854]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([1+0j, 0+1j, -1+0j, 1+1j])) + >>> pipeline = value >> Angle() + >>> result = pipeline() + >>> result + array([0. , 1.57079633, 3.14159265, 0.78539816]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([1+0j, 0+1j, -1+0j, 1+1j])) + >>> pipeline = value >> Angle() + >>> result = pipeline() + >>> result + tensor([0.0000, 1.5708, 3.1416, 0.7854]) + + These are equivalent to: + >>> pipeline = Angle(value) + """ +) - def __init__( - self: Angle, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.angle, feature=feature, **kwargs) +Real = create_elementwise_class( + name="Real", + function=xp.real, + docstring=""" + Apply the real-part function elementwise. -#TODO ***??*** revise Real - torch, typing, docstring, unit test -class Real(ElementwiseFeature): - """ - Applies the real function elementwise. + This feature applies `xp.real` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + For real-valued inputs, it returns the input unchanged. + For complex-valued inputs, it returns the real part. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the real function will be applied. If None, the function is applied to the input array directly. - - """ - def __init__( - self: Real, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.real, feature=feature, **kwargs) + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Real + + Use with NumPy directly: + >>> import numpy as np + >>> result = Real()(np.array([1+2j, 3+0j, -4.5])) + >>> result + array([ 1. , 3. , -4.5]) + + Use with PyTorch directly: + >>> import torch + >>> result = Real()(torch.tensor([1+2j, 3+0j, -4.5+0j])) + >>> result + tensor([ 1.0000, 3.0000, -4.5000]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([1+2j, 3+0j, -4.5])) + >>> pipeline = value >> Real() + >>> result = pipeline() + >>> result + array([ 1. , 3. , -4.5]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([1+2j, 3+0j, -4.5+0j])) + >>> pipeline = value >> Real() + >>> result = pipeline() + >>> result + tensor([ 1.0000, 3.0000, -4.5000]) + + These are equivalent to: + >>> pipeline = Real(value) + + """ +) -#TODO ***??*** revise Imag - torch, typing, docstring, unit test class Imag(ElementwiseFeature): - """ - Applies the imaginary function elementwise. + """Apply the imaginary-part function elementwise. + + This class handles real and complex inputs for both NumPy and PyTorch + backends. For real inputs, it returns 0; for complex inputs, it extracts + the imaginary part. Parameters ---------- - feature : Feature or None, optional - The input feature to which the imaginary function will be applied. - If None, the function is applied to the input array directly. - + feature: Feature or None, optional + The input feature to which the imaginary-part function will be applied. + If None, the function is applied directly to the input. + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Imag + + Use with NumPy directly: + >>> import numpy as np + >>> result = Imag()(np.array([1+2j, 3+0j, -4.5])) + >>> result + array([ 2., 0., 0.]) + + Use with PyTorch directly: + >>> import torch + >>> result = Imag()(torch.tensor([1+2j, 3+0j, -4.5+0j])) + >>> result + tensor([2., 0., 0.]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([1+2j, 3+0j, -4.5])) + >>> pipeline = value >> Imag() + >>> result = pipeline() + >>> result + array([ 2., 0., 0.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([1+2j, 3+0j, -4.5+0j])) + >>> pipeline = value >> Imag() + >>> result = pipeline() + >>> result + tensor([2., 0., 0.]) + + These are equivalent to: + >>> pipeline = Imag(value) + """ def __init__( self: Imag, feature: Feature | None = None, - **kwargs: Any + **kwargs: Any, ) -> None: - super().__init__(np.imag, feature=feature, **kwargs) + """Initialize the Imag feature. + Parameters + ---------- + feature: Feature or None, optional + The input feature whose output will be transformed. + **kwargs: Any + Additional keyword arguments passed to the base Feature class. -#TODO ***??*** revise Abs - torch, typing, docstring, unit test -class Abs(ElementwiseFeature): - """ - Applies the absolute value function elementwise. + """ + super().__init__( + function=self._imag_dispatch, + feature=feature, + **kwargs + ) - Parameters - ---------- - feature : Feature or None, optional - The input feature to which the absolute value function will be applied. - If None, the function is applied to the input array directly. - - """ + @staticmethod + def _imag_dispatch(x): + """Dispatch imag function based on backend and dtype. + + This method extracts the imaginary part of the input. For NumPy arrays, + `np.imag` always returns an array, returning zeros for real-valued inputs. + However, PyTorch's `torch.imag()` raises a `RuntimeError` when called on + real tensors. + + To ensure compatibility with both backends, this function checks whether + the input is a complex tensor before calling `torch.imag`. If it is not + complex, it returns a zero tensor of the same shape and dtype. + + This logic is necessary because `xp.imag` (from array-api-compat) does not + handle real PyTorch tensors safely, and thus this function cannot be + created using the factory-based method. + + Parameters + ---------- + x: np.ndarray or torch.Tensor + + Returns + ------- + np.ndarray or torch.Tensor + Imaginary part of the input, or zero if real. + + """ + + if TORCH_AVAILABLE and isinstance(x, torch.Tensor): + if x.is_complex(): + return torch.imag(x) + return torch.zeros_like(x) + + return np.imag(x) - def __init__( - self: Abs, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.abs, feature=feature, **kwargs) +Abs = create_elementwise_class( + name="Abs", + function=xp.abs, + docstring=""" + Apply the absolute value function elementwise. + + This feature applies `xp.abs` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + For real-valued inputs, this is the absolute value. + For complex-valued inputs, this is the magnitude. + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Abs + + Use with NumPy directly: + >>> import numpy as np + >>> result = Abs()(np.array([-1.0, 0.0, 2.5])) + >>> result + array([1. , 0. , 2.5]) + + Use with PyTorch directly: + >>> import torch + >>> result = Abs()(torch.tensor([-1.0, 0.0, 2.5])) + >>> result + tensor([1.0000, 0.0000, 2.5000]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-3.0, 0.0, 3.0])) + >>> pipeline = value >> Abs() + >>> result = pipeline() + >>> result + array([3., 0., 3.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-3.0, 0.0, 3.0])) + >>> pipeline = value >> Abs() + >>> result = pipeline() + >>> result + tensor([3., 0., 3.]) + + These are equivalent to: + >>> pipeline = Abs(value) -#TODO ***??*** revise Conjugate - torch, typing, docstring, unit test -class Conjugate(ElementwiseFeature): """ - Applies the conjugate function elementwise. +) + + +Conj = create_elementwise_class( + name="Conj", + function=xp.conj, + docstring=""" + Apply the complex conjugate function elementwise. + + This feature applies `xp.conj` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + For real-valued inputs, the result is unchanged. For complex-valued inputs, + it returns the complex conjugate (i.e., `a + bj → a - bj`). Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the conjugate function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Conj + + Use with NumPy directly: + >>> import numpy as np + >>> result = Conj()(np.array([1+2j, 3+0j, -4.5])) + >>> result + array([ 1.-2.j, 3.-0.j, -4.5+0.j]) + + Use with PyTorch directly: + >>> import torch + >>> result = Conj()(torch.tensor([1+2j, 3+0j, -4.5+0j])) + >>> result + tensor([ 1.-2.j, 3.-0.j, -4.5+0.j]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([1+2j, 3+0j, -4.5])) + >>> pipeline = value >> Conj() + >>> result = pipeline() + >>> result + array([ 1.-2.j, 3.-0.j, -4.5+0.j]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([1+2j, 3+0j, -4.5+0j])) + >>> pipeline = value >> Conj() + >>> result = pipeline() + >>> result + tensor([ 1.-2.j, 3.-0.j, -4.5+0.j]) + + These are equivalent to: + >>> pipeline = Conj(value) + """ +) - def __init__( - self: Conjugate, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.conjugate, feature=feature, **kwargs) +Conjugate = Conj -#TODO ***??*** revise Sqrt - torch, typing, docstring, unit test -class Sqrt(ElementwiseFeature): - """ - Applies the square root function elementwise. + +Sqrt = create_elementwise_class( + name="Sqrt", + function=xp.sqrt, + docstring=""" + Apply the square root function elementwise. + + This feature applies `xp.sqrt` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + For non-negative real values, it returns the usual square root. + For negative inputs, the behavior depends on the backend: + - NumPy may return `nan` or a complex result depending on dtype. + - PyTorch raises an error unless the input is explicitly complex. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the square root function will be applied. If None, the function is applied to the input array directly. - + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Sqrt + + Use with NumPy directly: + >>> import numpy as np + >>> result = Sqrt()(np.array([0.0, 1.0, 4.0])) + >>> result + array([0., 1., 2.]) + + Use with PyTorch directly: + >>> import torch + >>> result = Sqrt()(torch.tensor([0.0, 1.0, 4.0])) + >>> result + tensor([0., 1., 2.]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([0.0, 1.0, 4.0])) + >>> pipeline = value >> Sqrt() + >>> result = pipeline() + >>> result + array([0., 1., 2.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([0.0, 1.0, 4.0])) + >>> pipeline = value >> Sqrt() + >>> result = pipeline() + >>> result + tensor([0., 1., 2.]) + + These are equivalent to: + >>> pipeline = Sqrt(value) + """ +) - def __init__( - self: Sqrt, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.sqrt, feature=feature, **kwargs) +Square = create_elementwise_class( + name="Square", + function=xp.square, + docstring=""" + Apply the square function elementwise. -#TODO ***??*** revise Square - torch, typing, docstring, unit test -class Square(ElementwiseFeature): - """ - Applies the square function elementwise. + This feature applies `xp.square` to each element in a NumPy array or a + PyTorch tensor. It supports both direct input and pipeline composition. + + This operation computes `x ** 2` for each element. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the square function will be applied. If None, the function is applied to the input array directly. - - """ - def __init__( - self: Square, - feature: Feature | None = None, - **kwargs: Any - ) -> None: - super().__init__(np.square, feature=feature, **kwargs) + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Square + + Use with NumPy directly: + >>> import numpy as np + >>> result = Square()(np.array([-2.0, 0.0, 3.0])) + >>> result + array([4., 0., 9.]) + + Use with PyTorch directly: + >>> import torch + >>> result = Square()(torch.tensor([-2.0, 0.0, 3.0])) + >>> result + tensor([4., 0., 9.]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-2.0, 0.0, 3.0])) + >>> pipeline = value >> Square() + >>> result = pipeline() + >>> result + array([4., 0., 9.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-2.0, 0.0, 3.0])) + >>> pipeline = value >> Square() + >>> result = pipeline() + >>> result + tensor([4., 0., 9.]) + + These are equivalent to: + >>> pipeline = Square(value) + + """ +) -#TODO ***??*** revise Sign - torch, typing, docstring, unit test class Sign(ElementwiseFeature): - """ - Applies the sign function elementwise. + """Apply the sign function elementwise. + + This class uses a backend-aware dispatch to handle NumPy and PyTorch + tensors safely. It returns: + - -1 for negative values, + - 0 for zero, + - +1 for positive values. + + For complex numbers, it returns `x / abs(x)` when `x != 0`. Parameters ---------- - feature : Feature or None, optional + feature: Feature or None, optional The input feature to which the sign function will be applied. - If None, the function is applied to the input array directly. - + If None, the function is applied directly to the input array. + + Examples + -------- + >>> import deeptrack as dt + >>> from deeptrack.elementwise import Sign + + Use with NumPy directly: + >>> import numpy as np + >>> result = Sign()(np.array([-5.0, 0.0, 2.0])) + >>> result + array([-1., 0., 1.]) + + Use with PyTorch directly: + >>> import torch + >>> result = Sign()(torch.tensor([-5.0, 0.0, 2.0])) + >>> result + tensor([-1., 0., 1.]) + + Use in a pipeline with a NumPy value: + >>> value = dt.Value(value=np.array([-5.0, 0.0, 2.0])) + >>> pipeline = value >> Sign() + >>> result = pipeline() + >>> result + array([-1., 0., 1.]) + + Use in a pipeline with a Torch value: + >>> value = dt.Value(value=torch.tensor([-5.0, 0.0, 2.0])) + >>> pipeline = value >> Sign() + >>> result = pipeline() + >>> result + tensor([-1., 0., 1.]) + + These are equivalent to: + >>> pipeline = Sign(value) + """ def __init__( self: Sign, feature: Feature | None = None, - **kwargs: Any + **kwargs: Any, ) -> None: - super().__init__(np.sign, feature=feature, **kwargs) + """Initialize the Sign feature. + + This constructor sets up the elementwise sign operation using a + backend-aware dispatch. It optionally accepts another Feature whose output + will be processed by the sign function. + + Parameters + ---------- + feature : Feature or None, optional + An optional input feature to be wrapped. If provided, the sign + operation will be applied to the result of this feature. + If None, the sign function is applied to the input directly. + **kwargs : Any + Additional keyword arguments passed to the base Feature class. + + """ + + super().__init__( + function=self._sign_dispatch, + feature=feature, + **kwargs, + ) + + @staticmethod + def _sign_dispatch(x): + """Dispatch the sign operation depending on backend and input type. + + This method returns the sign of each element in the input: + - -1 for negative values, + - 0 for zero, + - +1 for positive values. + + For complex inputs, it returns `x / abs(x)` if `x ≠ 0`. + + This function uses `torch.sign()` when the input is a `torch.Tensor`, + and `np.sign()` otherwise. It avoids using `xp.sign()` from + `array-api-compat`, which internally calls `xp.issubdtype(...)` and + fails when provided with a `torch.float32` tensor. + + This method enables full compatibility with both NumPy and PyTorch + backends, and supports both real and complex-valued inputs. + Parameters + ---------- + x : np.ndarray or torch.Tensor + The input array or tensor whose elementwise signs will be computed. -#TODO ***GV*** Consider creating classes dynamically -#TODO ***BM*** Does this eliminate the need for the classes above? + Returns + ------- + np.ndarray or torch.Tensor + The elementwise sign values of the input. -# def create_elementwise_class(name: str, np_function: Callable) -> type: -# """Factory function to create an ElementwiseFeature subclass.""" -# return type( -# name, -# (ElementwiseFeature,), -# { -# "__init__": lambda self, feature=None, **kwargs: ElementwiseFeature.__init__(self, np_function, feature, **kwargs), -# }, -# ) + """ + if TORCH_AVAILABLE and isinstance(x, torch.Tensor): + return torch.sign(x) -# Sin = create_elementwise_class("Sin", np.sin) -# Cos = create_elementwise_class("Cos", np.cos) -# Tan = create_elementwise_class("Tan", np.tan) \ No newline at end of file + return np.sign(x) diff --git a/deeptrack/tests/test_elementwise.py b/deeptrack/tests/test_elementwise.py index 6c278a3b..edb5a32e 100644 --- a/deeptrack/tests/test_elementwise.py +++ b/deeptrack/tests/test_elementwise.py @@ -1,86 +1,93 @@ -import sys +# pylint: disable=C0115:missing-class-docstring +# pylint: disable=C0116:missing-function-docstring +# pylint: disable=C0103:invalid-name -# sys.path.append(".") # Adds the module to path +# Use this only when running the test locally. +# import sys +# sys.path.append(".") # Adds the module to path. +import inspect import unittest -import operator -import itertools -from numpy.core.numeric import array_equal - -from numpy.testing._private.utils import assert_almost_equal - -from deeptrack import elementwise, features, Image import numpy as np -import numpy.testing -import inspect +from deeptrack import elementwise, features, TORCH_AVAILABLE, xp +if TORCH_AVAILABLE: + import torch def grid_test_features( tester, - feature, + elementwise_class, feature_inputs, - expected_result_function, + function_name, ): - - for f_a_input in feature_inputs: - - inp = features.Value(f_a_input) - - f_a = feature(inp) - f_b = inp >> feature() - - for f in [f_a, f_b]: - try: - output = f() - except Exception as e: - tester.assertRaises( - type(e), - lambda: expected_result_function(f_a_input), + for feature_input in feature_inputs: + pip_a = elementwise_class(features.Value(feature_input)) + pip_b = features.Value(feature_input) >> elementwise_class() + + for pip in [pip_a, pip_b]: + result = pip() + + if TORCH_AVAILABLE and isinstance(result, torch.Tensor): + function = torch.__dict__[function_name] + if function == torch.imag: + # Torch workaround: handle real vs complex manually + if feature_input.is_complex(): + expected_result = torch.imag(feature_input) + else: + expected_result = torch.zeros_like(feature_input) + else: + expected_result = function(feature_input) + + # In PyTorch, NaNs are unequal by default + valid_mask = ~(torch.isnan(result) + | torch.isnan(expected_result)) + + torch.testing.assert_close( + result[valid_mask], + expected_result[valid_mask], + rtol=1e-5, + atol=1e-8, + msg=f"{elementwise_class.__name__} failed with PyTorch.", ) - continue - - expected_result = expected_result_function(f_a_input) - output = np.array(output) - try: - expected_result = np.array(expected_result) - except TypeError: - expected_result = expected_result.get() - - if isinstance(output, list) and isinstance(expected_result, list): - [ - np.testing.assert_almost_equal(np.array(a), np.array(b)) - for a, b in zip(output, expected_result) - ] - else: - is_equal = np.allclose( - np.array(output), np.array(expected_result), equal_nan=True - ) + function = np.__dict__[function_name] + expected_result = function(feature_input) + + # In NumPy, NaNs are ignored - tester.assertFalse( - not is_equal, - "Feature output {} is not equal to expect result {}.\n Using arguments {}".format( - output, expected_result, f_a_input - ), + np.testing.assert_allclose( + result, + expected_result, + rtol=1e-5, + atol=1e-8, + err_msg=f"{elementwise_class.__name__} failed with NumPy.", ) -def create_test(cl): - testname = "test_{}".format(cl.__name__) +def create_test(elementwise_class): + testname = f"test_{elementwise_class.__name__}" def test(self): + inputs = [ + np.array(-1.0), + np.array(0.0), + np.array(1.0), + (np.random.rand(8, 15) - 0.5) * 100, + ] + + if TORCH_AVAILABLE: + inputs.extend([ + torch.tensor([-1.0, 0.0, 1.0]), + (torch.rand(8, 15) - 0.5) * 100, + ]) + grid_test_features( self, - cl, - [ - -1, - 0, - 1, - (np.random.rand(50, 500) - 0.5) * 100, - ], - np.__dict__[cl.__name__.lower()], + elementwise_class, + inputs, + elementwise_class.__name__.lower(), ) test.__name__ = testname @@ -88,21 +95,23 @@ def test(self): return testname, test -class TestFeatures(unittest.TestCase): +class TestElementwiseFeatures(unittest.TestCase): pass -classes = inspect.getmembers(elementwise, inspect.isclass) +elementwise_classes = inspect.getmembers(elementwise, inspect.isclass) -for clname, cl in classes: +for class_name, elementwise_class in elementwise_classes: - if not issubclass(cl, elementwise.ElementwiseFeature) or ( - cl is elementwise.ElementwiseFeature + if ( + elementwise_class is elementwise.ElementwiseFeature + or + not issubclass(elementwise_class, elementwise.ElementwiseFeature) ): continue - testname, test_method = create_test(cl) - setattr(TestFeatures, testname, test_method) + test_name, test_method = create_test(elementwise_class) + setattr(TestElementwiseFeatures, test_name, test_method) if __name__ == "__main__": diff --git a/tutorials/3-advanced-topics/DTAT389_elementwise.ipynb b/tutorials/3-advanced-topics/DTAT389_elementwise.ipynb index 309deafe..b3ab0dc1 100644 --- a/tutorials/3-advanced-topics/DTAT389_elementwise.ipynb +++ b/tutorials/3-advanced-topics/DTAT389_elementwise.ipynb @@ -4,7 +4,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# deeptrack.elementwise\n", + "# DTAT389. deeptrack.elementwise\n", "\n", "\"Open" ] @@ -29,26 +29,70 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## 1. What is elementwise?\n", + "## 1. What is `elementwise.py`?\n", "\n", - "The elementwise module introduces utility functions which lets the user apply Numpy functions to `Feature` objects elementwise, that is element by element.\n", + "The `elementwise.py` module introduces utility functions which lets the user apply NumPy or PyTorch functions to `Feature` objects elementwise, that is, element by element.\n", "\n", - "Some functions included in elementwise are:\n", - "\n", - "- Trigonometric\n", - "- Hyperbolic\n", - "- Rounding \n", - "- Exponents and Logarithms\n", - "- Complex " + "Supported categories include:\n", + "- Trigonometric functions\n", + "- Hyperbolic functions\n", + "- Rounding operations\n", + "- Exponents and logarithms\n", + "- Complex number operations (real, imag, conj, angle)\n", + "- Miscellaneous (e.g., abs, square, sqrt, sign)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "## 2. Use example\n", + "The key roles of `elementwise.py` are:\n", + "\n", + "- **Backend-Agnostic Mathematical Transformations:**\n", + " The module provides a unified interface to apply mathematical functions\n", + " (e.g., `sin`, `exp`, `sqrt`, `abs`) elementwise across both NumPy and \n", + " PyTorch backends. This ensures compatibility without requiring users to \n", + " manage backend-specific logic.\n", + "\n", + "- **Composable Elementwise Features:**\n", + " Each function is wrapped in a `Feature` class, making it fully compatible \n", + " with DeepTrack2 pipelines. Users can plug these features directly into \n", + " image-processing or data-generation workflows.\n", + "\n", + "- **Factory-Based Extensibility:**\n", + " The `create_elementwise_class()` function allows dynamically generating \n", + " new elementwise features for backend functions, simplifying extension and \n", + " ensuring consistency across the API.\n", "\n", - "Create a feature that subtracts values from an image ..." + "- **Automatic Support for Static and Dynamic Inputs:**\n", + " Elementwise features automatically handle both static inputs (e.g., \n", + " `np.array`) and dynamic `Feature`-based inputs, preserving full support \n", + " for lazy evaluation, parameterization, and randomization.\n", + "\n", + "- **Robust Compatibility Layer for Edge Cases:**\n", + " For functions like `ceil`, `floor`, `sign`, or `imag` that are not safely \n", + " handled by the `array-api-compat` layer, the module provides manually \n", + " defined classes with backend-specific dispatching to ensure stability and \n", + " correctness.\n", + "\n", + "- **Operator Chaining and Feature Composition:**\n", + " Elementwise features support chaining via the `>>` operator or nesting in \n", + " constructors, allowing concise and readable expression of transformation \n", + " pipelines." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Using Elementwise Features" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.1. Using with NumPy Array with Direct Input" ] }, { @@ -57,61 +101,424 @@ "metadata": {}, "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "[-9 -8 -7]\n" - ] + "data": { + "text/plain": [ + "array([1. , 0. , 2.5])" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" } ], "source": [ + "from deeptrack.elementwise import Abs\n", "import numpy as np\n", - "from deeptrack.features import Feature\n", "\n", - "class Subtract(Feature):\n", - " def get(self, image, value, **kwargs):\n", - " return image - value\n", + "array = np.array([-1.0, 0.0, 2.5])\n", + "result = Abs()(array)\n", + "result" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.2. Using with PyTorch Tensor with Direct Input" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([1.0000, 0.0000, 2.5000])" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import torch\n", + "\n", + "tensor = torch.tensor([-1.0, 0.0, 2.5])\n", + "result = Abs()(tensor)\n", + "result" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.3. Using in a NumPy Pipeline\n", "\n", - "subtract_10 = Subtract(value=10)\n", + "You can also wrap the NumPy array using `dt.Value` and apply the feature\n", + "as part of a pipeline using the `>>` operator." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([3., 0., 3.])" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import deeptrack as dt\n", "\n", - "input_image = np.array([1, 2, 3])\n", - "output_image = subtract_10(input_image)\n", - "print(output_image)" + "value = dt.Value(value=np.array([-3.0, 0.0, 3.0]))\n", + "pipeline = value >> Abs()\n", + "result = pipeline()\n", + "result" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "... and compute the absolute value of the feature (in sequence)." + "Or equivalently:" ] }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([3., 0., 3.])" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "value = dt.Value(value=np.array([-3.0, 0.0, 3.0]))\n", + "pipeline = Abs(value)\n", + "result = pipeline()\n", + "result" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.4. Using in a PyTorch Pipeline\n", + "\n", + "Same idea using a PyTorch tensor." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([3., 0., 3.])" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "value = dt.Value(value=torch.tensor([-3.0, 0.0, 3.0]))\n", + "pipeline = value >> Abs()\n", + "result = pipeline()\n", + "result" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Or equivalently:" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "tensor([3., 0., 3.])" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "value = dt.Value(value=torch.tensor([-3.0, 0.0, 3.0]))\n", + "pipeline = Abs(value)\n", + "result = pipeline()\n", + "result" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Creating an Elementwise Feature" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.1. Creating an Elementwise Feature Using the Factory Function\n", + "\n", + "The most convenient way to define new elementwise operations in DeepTrack2\n", + "is through the `create_elementwise_class` factory function.\n", + "\n", + "This function takes three arguments:\n", + "- `name`: A string defining the name of the class.\n", + "- `function`: A callable (usually from `xp`, `np`, or `torch`) to be applied\n", + " elementwise.\n", + "- `docstring`: A documentation string (optional but recommended).\n", + "\n", + "The factory function dynamically generates a subclass of `ElementwiseFeature`\n", + "that wraps the provided function and enables it to be used in any DeepTrack\n", + "pipeline. This avoids the need to manually define the `__init__` method or\n", + "dispatch between backends like NumPy and PyTorch.\n", + "\n", + "Below we define a custom elementwise feature called `Abs` that applies the\n", + "absolute value function to its input." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "from deeptrack.backend import xp\n", + "from deeptrack.elementwise import create_elementwise_class\n", + "\n", + "Abs = create_elementwise_class(\n", + " name=\"Abs\",\n", + " function=xp.abs,\n", + " docstring=\"Elementwise absolute value feature.\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To check that this class behaves as expected, create an instance of the Abs feature:" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "abs_feature = Abs()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Evaluate the feature with NumPy input:" + ] + }, + { + "cell_type": "code", + "execution_count": 10, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "[9 8 7]\n", - "[9 8 7]\n" + "NumPy input: [-3. 0. 2.5]\n", + "NumPy output: [3. 0. 2.5]\n" ] } ], "source": [ - "from deeptrack.elementwise import Abs\n", + "numpy_input = np.array([-3.0, 0.0, 2.5])\n", + "numpy_output = abs_feature(numpy_input)\n", + "print(\"NumPy input:\", numpy_input)\n", + "print(\"NumPy output:\", numpy_output)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Evaluate the same feature with PyTorch input:" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "PyTorch input: tensor([-3.0000, 0.0000, 2.5000])\n", + "PyTorch output: tensor([3.0000, 0.0000, 2.5000])\n" + ] + } + ], + "source": [ + "torch_input = torch.tensor([-3.0, 0.0, 2.5])\n", + "torch_output = abs_feature(torch_input)\n", + "print(\"PyTorch input:\", torch_input)\n", + "print(\"PyTorch output:\", torch_output)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use Abs in a DeepTrack pipeline with NumPy:" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Pipeline output (NumPy): [3. 0. 2.5]\n" + ] + } + ], + "source": [ + "value_np = dt.Value(value=np.array([-3.0, 0.0, 2.5]))\n", + "pipeline_np = value_np >> Abs()\n", + "print(\"Pipeline output (NumPy):\", pipeline_np())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use Abs in a DeepTrack pipeline with PyTorch:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Pipeline output (PyTorch): tensor([3.0000, 0.0000, 2.5000])\n" + ] + } + ], + "source": [ + "value_torch = dt.Value(value=torch.tensor([-3.0, 0.0, 2.5]))\n", + "pipeline_torch = value_torch >> Abs()\n", + "print(\"Pipeline output (PyTorch):\", pipeline_torch())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.2. Creating an Elementwise Feature Inheriting from `ElementWise`\n", + "\n", + "Some mathematical operations, such as `ceil`, require **manual subclassing**\n", + "from `ElementwiseFeature`. This is necessary when the corresponding\n", + "`array-api-compat` function (e.g., `xp.ceil`) does not safely handle all\n", + "backends.\n", + "\n", + "In particular, calling `xp.ceil` on a `torch.Tensor` fails because it internally\n", + "uses `xp.issubdtype`, which cannot interpret PyTorch dtypes like\n", + "`torch.float32`. To avoid this issue, we subclass `ElementwiseFeature` and\n", + "manually dispatch to the appropriate backend-specific function.\n", + "\n", + "The example below demonstrates how to implement a `Ceil` feature:" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "from deeptrack import Feature, TORCH_AVAILABLE\n", + "from deeptrack.elementwise import ElementwiseFeature\n", + "import numpy as np\n", + "\n", + "if TORCH_AVAILABLE:\n", + " import torch\n", "\n", - "# Sequentially take the absolute value after subtraction.\n", - "pipeline = Abs(subtract_10)\n", - "output_image = pipeline(input_image)\n", - "print(output_image)\n", + "class Ceil(ElementwiseFeature):\n", + " \"\"\"Apply the ceiling function elementwise for NumPy and PyTorch arrays.\n", + "\n", + " This implementation dispatches to `torch.ceil` for PyTorch tensors and\n", + " `np.ceil` for NumPy arrays, ensuring compatibility across backends.\n", + "\n", + " \"\"\"\n", + "\n", + " def __init__(self, feature, **kwargs):\n", + " super().__init__(\n", + " function=self._ceil_dispatch,\n", + " feature=feature,\n", + " **kwargs,\n", + " )\n", + "\n", + " @staticmethod\n", + " def _ceil_dispatch(x):\n", + " \"\"\"Dispatch the ceiling function depending on the backend.\n", + "\n", + " Returns\n", + " -------\n", + " array or tensor with elementwise ceil applied.\n", + "\n", + " \"\"\"\n", + "\n", + " if TORCH_AVAILABLE and isinstance(x, torch.Tensor):\n", + " return torch.ceil(x)\n", + "\n", + " return np.ceil(x)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This implementation overrides the `__init__()` method to pass a custom dispatch function _ceil_dispatch to ElementwiseFeature. The dispatch function checks whether the input is a `torch.Tensor` or a `np.ndarray` and applies the appropriate backend function.\n", "\n", - "# Or equivalently:\n", - "pipeline = subtract_10 >> Abs()\n", - "output_image = pipeline(input_image)\n", - "print(output_image)" + "This approach ensures correct operation on both NumPy and PyTorch backends, and avoids failures that would arise from using the factory function." ] } ],