zea 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

zea/utils.py

@@ -2,117 +2,33 @@
 
 import collections.abc
 import datetime
-import functools
-import hashlib
-import inspect
-import platform
 import time
 from functools import wraps
-from pathlib import Path
 from statistics import mean, median, stdev
 
-import numpy as np
+import keras
 import yaml
-from keras import ops
-from PIL import Image
 
 from zea import log
 
 
-def _assert_uint8_images(images: np.ndarray):
-    """
-    Asserts that the input images have the correct properties.
-
-    Args:
-        images (np.ndarray): The input images.
-
-    Raises:
-        AssertionError: If the dtype of images is not uint8.
-        AssertionError: If the shape of images is not (n_frames, height, width, channels)
-            or (n_frames, height, width) for grayscale images.
-        AssertionError: If images have anything other than 1 (grayscale),
-            3 (rgb) or 4 (rgba) channels.
-    """
-    assert images.dtype == np.uint8, f"dtype of images should be uint8, got {images.dtype}"
-
-    assert images.ndim in (3, 4), (
-        "images must have shape (n_frames, height, width, channels),"
-        f" or (n_frames, height, width) for grayscale images. Got {images.shape}"
-    )
-
-    if images.ndim == 4:
-        assert images.shape[-1] in (1, 3, 4), (
-            "Grayscale images must have 1 channel, "
-            "RGB images must have 3 channels, and RGBA images must have 4 channels. "
-            f"Got shape: {images.shape}, channels: {images.shape[-1]}"
-        )
-
+def canonicalize_axis(axis, num_dims) -> int:
+    """Canonicalize an axis in [-num_dims, num_dims) to [0, num_dims)."""
+    if not -num_dims <= axis < num_dims:
+        raise ValueError(f"axis {axis} is out of bounds for array of dimension {num_dims}")
+    if axis < 0:
+        axis = axis + num_dims
+    return axis
 
-def translate(array, range_from=None, range_to=(0, 255)):
-    """Map values in array from one range to other.
 
-    Args:
-        array (ndarray): input array.
-        range_from (Tuple, optional): lower and upper bound of original array.
-            Defaults to min and max of array.
-        range_to (Tuple, optional): lower and upper bound to which array should be mapped.
-            Defaults to (0, 255).
-
-    Returns:
-        (ndarray): translated array
-    """
-    if range_from is None:
-        left_min, left_max = ops.min(array), ops.max(array)
-    else:
-        left_min, left_max = range_from
-    right_min, right_max = range_to
-
-    # Convert the left range into a 0-1 range (float)
-    value_scaled = (array - left_min) / (left_max - left_min)
-
-    # Convert the 0-1 range into a value in the right range.
-    return right_min + (value_scaled * (right_max - right_min))
-
-
-def map_negative_indices(indices: list, length: int):
+def map_negative_indices(indices: list, num_dims: int):
     """Maps negative indices for array indexing to positive indices.
     Example:
         >>> from zea.utils import map_negative_indices
         >>> map_negative_indices([-1, -2], 5)
         [4, 3]
     """
-    return [i if i >= 0 else length + i for i in indices]
-
-
-def find_key(dictionary, contains, case_sensitive=False):
-    """Find key in dictionary that contains partly the string `contains`
-
-    Args:
-        dictionary (dict): Dictionary to find key in.
-        contains (str): String which the key should .
-        case_sensitive (bool, optional): Whether the search is case sensitive.
-            Defaults to False.
-
-    Returns:
-        str: the key of the dictionary that contains the query string.
-
-    Raises:
-        TypeError: if not all keys are strings.
-        KeyError: if no key is found containing the query string.
-    """
-    # Assert that all keys are strings
-    if not all(isinstance(k, str) for k in dictionary.keys()):
-        raise TypeError("All keys must be strings.")
-
-    if case_sensitive:
-        key = [k for k in dictionary.keys() if contains in k]
-    else:
-        key = [k for k in dictionary.keys() if contains in k.lower()]
-
-    if len(key) == 0:
-        raise KeyError(f"Key containing '{contains}' not found in dictionary.")
-
-    return key[0]
+    return [canonicalize_axis(idx, num_dims) for idx in indices]
 
 
 def print_clear_line():
@@ -139,148 +55,6 @@ def strtobool(val: str):
         raise ValueError(f"invalid truth value {val}")
 
 
-def grayscale_to_rgb(image):
-    """Converts a grayscale image to an RGB image.
-
-    Args:
-        image (ndarray): Grayscale image. Must have shape (height, width).
-
-    Returns:
-        ndarray: RGB image.
-    """
-    assert image.ndim == 2, "Input image must be grayscale."
-    # Stack the grayscale image into 3 channels (RGB)
-    return np.stack([image] * 3, axis=-1)
-
-
-def preprocess_for_saving(images):
-    """Preprocesses images for saving to GIF or MP4.
-
-    Args:
-        images (ndarray, list[ndarray]): Images. Must have shape (n_frames, height, width, channels)
-            or (n_frames, height, width).
-    """
-    images = np.array(images)
-    _assert_uint8_images(images)
-
-    # Remove channel axis if it is 1 (grayscale image)
-    if images.ndim == 4 and images.shape[-1] == 1:
-        images = np.squeeze(images, axis=-1)
-
-    # convert grayscale images to RGB
-    if images.ndim == 3:
-        images = [grayscale_to_rgb(image) for image in images]
-        images = np.array(images)
-
-    return images
-
-
-def save_to_gif(images, filename, fps=20, shared_color_palette=False):
-    """Saves a sequence of images to a GIF file.
-
-    Args:
-        images (list or np.ndarray): List or array of images. Must have shape
-            (n_frames, height, width, channels) or (n_frames, height, width).
-            If channel axis is not present, or is 1, grayscale image is assumed,
-            which is then converted to RGB. Images should be uint8.
-        filename (str or Path): Filename to which data should be written.
-        fps (int): Frames per second of rendered format.
-        shared_color_palette (bool, optional): If True, creates a global
-            color palette across all frames, ensuring consistent colors
-            throughout the GIF. Defaults to False, which is default behavior
-            of PIL.Image.save. Note: True can cause slow saving for longer sequences.
-
-    """
-    images = preprocess_for_saving(images)
-
-    if fps > 50:
-        log.warning(f"Cannot set fps ({fps}) > 50. Setting it automatically to 50.")
-        fps = 50
-
-    duration = 1 / (fps) * 1000  # milliseconds per frame
-
-    pillow_imgs = [Image.fromarray(img) for img in images]
-
-    if shared_color_palette:
-        # Apply the same palette to all frames without dithering for consistent color mapping
-        # Convert all images to RGB and combine their colors for palette generation
-        all_colors = np.vstack([np.array(img.convert("RGB")).reshape(-1, 3) for img in pillow_imgs])
-        combined_image = Image.fromarray(all_colors.reshape(-1, 1, 3))
-
-        # Generate palette from all frames
-        global_palette = combined_image.quantize(
-            colors=256,
-            method=Image.MEDIANCUT,
-            kmeans=1,
-        )
-
-        # Apply the same palette to all frames without dithering
-        pillow_imgs = [
-            img.convert("RGB").quantize(
-                palette=global_palette,
-                dither=Image.NONE,
-            )
-            for img in pillow_imgs
-        ]
-
-    pillow_img, *pillow_imgs = pillow_imgs
-
-    pillow_img.save(
-        fp=filename,
-        format="GIF",
-        append_images=pillow_imgs,
-        save_all=True,
-        loop=0,
-        duration=duration,
-        interlace=False,
-        optimize=False,
-    )
-    log.success(f"Succesfully saved GIF to -> {log.yellow(filename)}")
-
-
-def save_to_mp4(images, filename, fps=20):
-    """Saves a sequence of images to an MP4 file.
-
-    Args:
-        images (list or np.ndarray): List or array of images. Must have shape
-            (n_frames, height, width, channels) or (n_frames, height, width).
-            If channel axis is not present, or is 1, grayscale image is assumed,
-            which is then converted to RGB. Images should be uint8.
-        filename (str or Path): Filename to which data should be written.
-        fps (int): Frames per second of rendered format.
-
-    Returns:
-        str: Success message.
-
-    """
-    images = preprocess_for_saving(images)
-
-    filename = str(filename)
-
-    parent_dir = Path(filename).parent
-    if not parent_dir.exists():
-        raise FileNotFoundError(f"Directory '{parent_dir}' does not exist.")
-
-    try:
-        import cv2
-    except ImportError as exc:
-        raise ImportError(
-            "OpenCV is required to save MP4 files. "
-            "Please install it with 'pip install opencv-python' or "
-            "'pip install opencv-python-headless'."
-        ) from exc
-
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    _, height, width, _ = images.shape
-    video_writer = cv2.VideoWriter(filename, fourcc, fps, (width, height))
-
-    for image in images:
-        video_writer.write(image)
-
-    video_writer.release()
-    return log.success(f"Successfully saved MP4 to -> {filename}")
-
-
 def update_dictionary(dict1: dict, dict2: dict, keep_none: bool = False) -> dict:
     """Updates dict1 with values dict2
 
@@ -341,235 +115,6 @@ def date_string_to_readable(date_string: str, include_time: bool = False):
         return date.strftime("%B %d, %Y")
 
 
-def find_first_nonzero_index(arr, axis, invalid_val=-1):
-    """
-    Find the index of the first non-zero element along a specified axis in a NumPy array.
-
-    Args:
-        arr (numpy.ndarray): The input array to search for the first non-zero element.
-        axis (int): The axis along which to perform the search.
-        invalid_val (int, optional): The value to assign to elements where no
-            non-zero values are found along the axis.
-
-    Returns:
-        numpy.ndarray: An array of indices where the first non-zero element
-            occurs along the specified axis. Elements with no non-zero values along
-            the axis are replaced with the 'invalid_val'.
-
-    """
-    nonzero_mask = arr != 0
-    return np.where(nonzero_mask.any(axis=axis), nonzero_mask.argmax(axis=axis), invalid_val)
-
-
-def first_not_none_item(arr):
-    """
-    Finds and returns the first non-None item in the given array.
-
-    Args:
-        arr (list): The input array.
-
-    Returns:
-        The first non-None item found in the array, or None if no such item exists.
-    """
-    non_none_items = [item for item in arr if item is not None]
-    return non_none_items[0] if non_none_items else None
-
-
-def deprecated(replacement=None):
-    """Decorator to mark a function, method, or attribute as deprecated.
-
-    Args:
-        replacement (str, optional): The name of the replacement function, method, or attribute.
-
-    Returns:
-        callable: The decorated function, method, or property.
-
-    Raises:
-        DeprecationWarning: A warning is issued when the deprecated item is called or accessed.
-
-    Example:
-        >>> from zea.utils import deprecated
-        >>> class MyClass:
-        ...     @deprecated(replacement="new_method")
-        ...     def old_method(self):
-        ...         print("This is the old method.")
-        ...
-        ...     @deprecated(replacement="new_attribute")
-        ...     def __init__(self):
-        ...         self._old_attribute = "Old value"
-        ...
-        ...     @deprecated(replacement="new_property")
-        ...     @property
-        ...     def old_property(self):
-        ...         return self._old_attribute
-
-        >>> # Using the deprecated method
-        >>> obj = MyClass()
-        >>> obj.old_method()
-        This is the old method.
-        >>> # Accessing the deprecated attribute
-        >>> print(obj.old_property)
-        Old value
-        >>> # Setting value to the deprecated attribute
-        >>> obj.old_property = "New value"
-    """
-
-    def decorator(item):
-        if callable(item):
-            # If it's a function or method
-            @functools.wraps(item)
-            def wrapper(*args, **kwargs):
-                if replacement:
-                    log.deprecated(
-                        f"Call to deprecated {item.__name__}. Use {replacement} instead."
-                    )
-                else:
-                    log.deprecated(f"Call to deprecated {item.__name__}.")
-                return item(*args, **kwargs)
-
-            return wrapper
-        elif isinstance(item, property):
-            # If it's a property of a class
-            def getter(self):
-                if replacement:
-                    log.deprecated(
-                        f"Access to deprecated attribute {item.fget.__name__}, "
-                        f"use {replacement} instead."
-                    )
-                else:
-                    log.deprecated(f"Access to deprecated attribute {item.fget.__name__}.")
-                return item.fget(self)
-
-            def setter(self, value):
-                if replacement:
-                    log.deprecated(
-                        f"Setting value to deprecated attribute {item.fget.__name__}, "
-                        f"use {replacement} instead."
-                    )
-                else:
-                    log.deprecated(f"Setting value to deprecated attribute {item.fget.__name__}.")
-                item.fset(self, value)
-
-            def deleter(self):
-                if replacement:
-                    log.deprecated(
-                        f"Deleting deprecated attribute {item.fget.__name__}, "
-                        f"use {replacement} instead."
-                    )
-                else:
-                    log.deprecated(f"Deleting deprecated attribute {item.fget.__name__}.")
-                item.fdel(self)
-
-            return property(getter, setter, deleter)
-
-        else:
-            raise TypeError("Decorator can only be applied to functions, methods, or properties.")
-
-    return decorator
-
-
-def calculate_file_hash(file_path, omit_line_str=None):
-    """Calculates the hash of a file.
-
-    Args:
-        file_path (str): Path to file.
-        omit_line_str (str, optional): If this string is found in a line, the line will
-            be omitted when calculating the hash. This is useful for example
-            when the file contains the hash itself.
-
-    Returns:
-        str: The hash of the file.
-
-    """
-    hash_object = hashlib.sha256()
-    with open(file_path, "rb") as f:
-        for line in f:
-            if omit_line_str is not None and omit_line_str in str(line):
-                continue
-            hash_object.update(line)
-    return hash_object.hexdigest()
-
-
-def check_architecture():
-    """Checks the architecture of the system."""
-    return platform.uname()[-1]
-
-
-def get_function_args(func):
-    """Get the names of the arguments of a function."""
-    sig = inspect.signature(func)
-    return tuple(sig.parameters)
-
-
-def fn_requires_argument(fn, arg_name):
-    """Returns True if the function requires the argument 'arg_name'."""
-    params = get_function_args(fn)
-    return arg_name in params
-
-
-def find_methods_with_return_type(cls, return_type_hint):
-    """
-    Find all methods in a class that have the specified return type hint.
-
-    Args:
-        cls: The class to inspect.
-        return_type_hint: The return type hint to match (as a string).
-
-    Returns:
-        A list of method names that match the return type hint.
-    """
-    matching_methods = []
-    for name, member in inspect.getmembers(cls, predicate=inspect.isfunction):
-        annotations = getattr(member, "__annotations__", {})
-        if annotations.get("return") == return_type_hint:
-            matching_methods.append(name)
-    return matching_methods
-
-
-def keep_trying(fn, args=None, required_set=None):
-    """Keep trying to run a function until it succeeds.
-
-    Args:
-        fn (callable): Function to run.
-        args (dict, optional): Arguments to pass to function.
-        required_set (set, optional): Set of required outputs.
-            If output is not in required_set, function will be rerun.
-
-    Returns:
-        Any: The output of the function if successful.
-
-    """
-    while True:
-        try:
-            out = fn(**args) if args is not None else fn()
-            if required_set is not None:
-                assert out is not None
-                assert out in required_set, f"Output {out} not in {required_set}"
-            return out
-        except Exception as e:
-            print(e)
-
-
-def reduce_to_signature(func, kwargs):
-    """Reduce the kwargs to the signature of the function."""
-    # Retrieve the argument names of the function
-    sig = inspect.signature(func)
-
-    # Filter out the arguments that are not part of the function
-    reduced_params = {key: kwargs[key] for key in sig.parameters if key in kwargs}
-
-    return reduced_params
-
-
-def safe_initialize_class(cls, **kwargs):
-    """Safely initialize a class by removing any invalid arguments."""
-
-    # Filter out the arguments that are not part of the Scan class
-    reduced_params = reduce_to_signature(cls.__init__, kwargs)
-
-    return cls(**reduced_params)
-
-
 def deep_compare(obj1, obj2):
     """Recursively compare two objects for equality."""
     # Only recurse into dicts
@@ -591,24 +136,94 @@ def deep_compare(obj1, obj2):
     return obj1 == obj2
 
 
+def block_until_ready(func):
+    """Decorator that ensures asynchronous (gpu) operations complete before returning."""
+    if keras.backend.backend() == "jax":
+        import jax
+
+        def _block(value):
+            if hasattr(value, "__array__"):
+                return jax.block_until_ready(value)
+            else:
+                return value
+    else:
+
+        def _block(value):
+            if hasattr(value, "__array__"):
+                # convert to numpy but return as original type
+                _ = keras.ops.convert_to_numpy(value)
+            return value
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        result = func(*args, **kwargs)
+
+        # Handle different return types
+        if isinstance(result, (list, tuple)):
+            # For multiple outputs, block each one
+            blocked_results = [_block(r) for r in result]
+            return type(result)(blocked_results)
+        elif isinstance(result, dict):
+            # For dict outputs, block array values
+            return {k: _block(v) for k, v in result.items()}
+        else:
+            # Single output
+            return _block(result)
+
+    return wrapper
+
+
 class FunctionTimer:
     """
     A decorator class for timing the execution of functions.
 
     Example:
-        >>> from zea.utils import FunctionTimer
-        >>> timer = FunctionTimer()
-        >>> my_function = lambda: sum(range(10))
-        >>> my_function = timer(my_function)
-        >>> _ = my_function()
-        >>> print(timer.get_stats("my_function"))
+        .. doctest::
+
+            >>> from zea.utils import FunctionTimer
+            >>> timer = FunctionTimer()
+            >>> my_function = lambda: sum(range(10))
+            >>> my_function = timer(my_function, name="my_function")
+            >>> _ = my_function()
+            >>> print(timer.get_stats("my_function"))  # doctest: +ELLIPSIS
+            {'mean': ..., 'median': ..., 'std_dev': ..., 'min': ..., 'max': ..., 'count': ...}
     """
 
     def __init__(self):
         self.timings = {}
         self.last_append = 0
+        self.decorated_functions = {}  # Track decorated functions
 
     def __call__(self, func, name=None):
+        _name = name if name is not None else func.__name__
+
+        # Create a unique identifier for this function
+        func_id = id(func)
+
+        # Check if this exact function has already been decorated
+        if func_id in self.decorated_functions:
+            existing_name = self.decorated_functions[func_id]
+            raise ValueError(
+                f"Function '{func.__name__}' (id: {func_id}) has already been "
+                f"decorated with timer name '{existing_name}'. "
+                f"Cannot decorate the same function instance multiple times."
+            )
+
+        # Handle name conflicts by appending a suffix
+        original_name = _name
+        counter = 1
+        while _name in self.timings:
+            _name = f"{original_name}_{counter}"
+            counter += 1
+
+        # Initialize timing storage for this function
+        self.timings[_name] = []
+
+        # Track this decorated function
+        self.decorated_functions[func_id] = _name
+
+        func = block_until_ready(func)
+
         @wraps(func)
         def wrapper(*args, **kwargs):
             start_time = time.perf_counter()
@@ -617,27 +232,27 @@ class FunctionTimer:
             elapsed_time = end_time - start_time
 
             # Store the timing result
-            _name = name if name is not None else func.__name__
-            if _name not in self.timings:
-                self.timings[_name] = []
             self.timings[_name].append(elapsed_time)
 
             return result
 
         return wrapper
 
-    def get_stats(self, func_name, drop_first: bool | int = False):
-        """Calculate statistics for the given function."""
-        if func_name not in self.timings:
-            raise ValueError(f"No timings recorded for function '{func_name}'.")
-
+    def _parse_drop_first(self, drop_first: bool | int):
         if isinstance(drop_first, bool):
             idx = 1 if drop_first else 0
         elif isinstance(drop_first, int):
             idx = drop_first
         else:
             raise ValueError("drop_first must be a boolean or an integer.")
+        return idx
+
+    def get_stats(self, func_name, drop_first: bool | int = False):
+        """Calculate statistics for the given function."""
+        if func_name not in self.timings:
+            raise ValueError(f"No timings recorded for function '{func_name}'.")
 
+        idx = self._parse_drop_first(drop_first)
         times = self.timings[func_name][idx:]
         return {
             "mean": mean(times),
@@ -663,7 +278,7 @@ class FunctionTimer:
 
         self.last_append = len(self.timings[func_name])
 
-    def print(self, drop_first: bool | int = False):
+    def print(self, drop_first: bool | int = False, total_time: bool = False):
         """Print timing statistics for all recorded functions using formatted output."""
 
         # Print title
@@ -693,3 +308,9 @@ class FunctionTimer:
                 f"{log.magenta(str(stats['count'])):<18}"
             )
             print(row)
+
+        if total_time:
+            idx = self._parse_drop_first(drop_first)
+            total = sum(mean(times[idx:]) for times in self.timings.values())
+            print("-" * length)
+            print(f"{log.bold('Mean Total Time:')} {log.bold(log.number_to_str(total, 6))} seconds")