Merge pull request #423 from Blosc/concatenate

Concatenate

Author: lshaw8317

CMakeLists.txt

@@ -50,7 +50,7 @@ else()
     include(FetchContent)
     FetchContent_Declare(blosc2
         GIT_REPOSITORY https://github.com/Blosc/c-blosc2
-        GIT_TAG 4ef3c7440a85632e6c8b6c5d2a9e651e45569fc1  # v2.17.1 + mmap fix
+        GIT_TAG 1c2f8bb0c914c43e23b751fbcf6642cd7aec09db  # v2.18.0 (concatenate added)
     )
     FetchContent_MakeAvailable(blosc2)
     include_directories("${blosc2_SOURCE_DIR}/include")

README_DEVELOPERS.md

@@ -22,6 +22,20 @@ You are done!
     pip install .   # add -e for editable mode
 ```
 
+There are situations where you may want to build the C-Blosc2 library separately, for example, when debugging issues in the C library. In that case, let's assume you have the C-Blosc2 library installed in `/usr/local`:
+
+```bash
+CMAKE_PREFIX_PATH=/usr/local USE_SYSTEM_BLOSC2=1 pip install -e .
+```
+
+and then, you can run the tests with:
+
+```bash
+LD_LIBRARY_PATH=/usr/local/lib pytest
+```
+
+[replace `LD_LIBRARY_PATH` with the appropriate environment variable for your system, such as `DYLD_LIBRARY_PATH` on macOS or `PATH` on Windows, if necessary].
+
 That's it! You can now proceed to the testing section.
 
 ## Testing

bench/ndarray/concatenate.py

@@ -0,0 +1,271 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <[email protected]>
+# All rights reserved.
+#
+# This source code is licensed under a BSD-style license (found in the
+# LICENSE file in the root directory of this source tree)
+#######################################################################
+
+import numpy as np
+import blosc2
+import time
+import matplotlib.pyplot as plt
+import os
+from matplotlib.ticker import ScalarFormatter
+
+
+def run_benchmark(num_arrays=10, size=500, aligned_chunks=False, axis=0):
+    """
+    Benchmark blosc2.concatenate performance with different chunk alignments.
+
+    Parameters:
+    - num_arrays: Number of arrays to concatenate
+    - size: Base size for array dimensions
+    - aligned_chunks: Whether to use aligned chunk shapes
+    - axis: Axis along which to concatenate (0 or 1)
+
+    Returns:
+    - duration: Time taken in seconds
+    - result_shape: Shape of the resulting array
+    - data_size_gb: Size of data processed in GB
+    """
+    if axis == 0:
+        # For concatenating along axis 0, the second dimension must be consistent
+        shapes = [(size // num_arrays, size) for _ in range(num_arrays)]
+    elif axis == 1:
+        # For concatenating along axis 1, the first dimension must be consistent
+        shapes = [(size, size // num_arrays) for _ in range(num_arrays)]
+    else:
+        raise ValueError("Only axis 0 and 1 are supported")
+
+    # Create appropriate chunk shapes
+    if aligned_chunks:
+        # Aligned chunks: divisors of the shape dimensions
+        chunk_shapes = [(shape[0] // 4, shape[1] // 4) for shape in shapes]
+    else:
+        # Unaligned chunks: not divisors of shape dimensions
+        chunk_shapes = [(shape[0] // 4 + 1, shape[1] // 4 - 1) for shape in shapes]
+
+    # Create arrays
+    arrays = []
+    for i, (shape, chunk_shape) in enumerate(zip(shapes, chunk_shapes)):
+        arr = blosc2.arange(
+            i * np.prod(shape), (i + 1) * np.prod(shape), 1, dtype="i4", shape=shape, chunks=chunk_shape
+        )
+        arrays.append(arr)
+
+    # Calculate total data size in GB (4 bytes per int32)
+    total_elements = sum(np.prod(shape) for shape in shapes)
+    data_size_gb = total_elements * 4 / (1024**3)  # Convert bytes to GB
+
+    # Time the concatenation
+    start_time = time.time()
+    result = blosc2.concatenate(arrays, axis=axis)
+    duration = time.time() - start_time
+
+    return duration, result.shape, data_size_gb
+
+
+def run_numpy_benchmark(num_arrays=10, size=500, axis=0):
+    """
+    Benchmark numpy.concatenate performance for comparison.
+
+    Parameters:
+    - num_arrays: Number of arrays to concatenate
+    - size: Base size for array dimensions
+    - axis: Axis along which to concatenate (0 or 1)
+
+    Returns:
+    - duration: Time taken in seconds
+    - result_shape: Shape of the resulting array
+    - data_size_gb: Size of data processed in GB
+    """
+    if axis == 0:
+        # For concatenating along axis 0, the second dimension must be consistent
+        shapes = [(size // num_arrays, size) for _ in range(num_arrays)]
+    elif axis == 1:
+        # For concatenating along axis 1, the first dimension must be consistent
+        shapes = [(size, size // num_arrays) for _ in range(num_arrays)]
+    else:
+        raise ValueError("Only axis 0 and 1 are supported")
+
+    # Create arrays
+    numpy_arrays = []
+    for i, shape in enumerate(shapes):
+        arr = np.arange(
+            i * np.prod(shape),
+            (i + 1) * np.prod(shape),
+            1,
+            dtype="i4"
+        ).reshape(shape)
+        numpy_arrays.append(arr)
+
+    # Calculate total data size in GB (4 bytes per int32)
+    total_elements = sum(np.prod(shape) for shape in shapes)
+    data_size_gb = total_elements * 4 / (1024**3)  # Convert bytes to GB
+
+    # Time the concatenation
+    start_time = time.time()
+    result = np.concatenate(numpy_arrays, axis=axis)
+    duration = time.time() - start_time
+
+    return duration, result.shape, data_size_gb
+
+
+def create_combined_plot(num_arrays, sizes, numpy_speeds_axis0, unaligned_speeds_axis0, aligned_speeds_axis0,
+                        numpy_speeds_axis1, unaligned_speeds_axis1, aligned_speeds_axis1, output_dir="plots"):
+    """
+    Create a figure with two side-by-side bar plots comparing the performance for both axes.
+
+    Parameters:
+    - sizes: List of array sizes
+    - *_speeds_axis0: Lists of speeds (GB/s) for axis 0 concatenation
+    - *_speeds_axis1: Lists of speeds (GB/s) for axis 1 concatenation
+    - output_dir: Directory to save the plot
+    """
+    # Create output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Set up the figure with two subplots side by side
+    fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(20, 8), sharey=True)
+
+    # Convert sizes to strings for the x-axis
+    x_labels = [str(size) for size in sizes]
+    x = np.arange(len(sizes))
+    width = 0.25
+
+    # Create bars for axis 0 plot
+    rect1_axis0 = ax0.bar(x - width, numpy_speeds_axis0, width, label='NumPy', color='#1f77b4')
+    rect2_axis0 = ax0.bar(x, unaligned_speeds_axis0, width, label='Blosc2 Unaligned', color='#ff7f0e')
+    rect3_axis0 = ax0.bar(x + width, aligned_speeds_axis0, width, label='Blosc2 Aligned', color='#2ca02c')
+
+    # Create bars for axis 1 plot
+    rect1_axis1 = ax1.bar(x - width, numpy_speeds_axis1, width, label='NumPy', color='#1f77b4')
+    rect2_axis1 = ax1.bar(x, unaligned_speeds_axis1, width, label='Blosc2 Unaligned', color='#ff7f0e')
+    rect3_axis1 = ax1.bar(x + width, aligned_speeds_axis1, width, label='Blosc2 Aligned', color='#2ca02c')
+
+    # Add labels and titles
+    for ax, axis in [(ax0, 0), (ax1, 1)]:
+        ax.set_xlabel('Array Size (N for NxN array)', fontsize=12)
+        ax.set_title(f'Concatenation Performance for {num_arrays} arrays (axis={axis})', fontsize=14)
+        ax.set_xticks(x)
+        ax.set_xticklabels(x_labels)
+        ax.grid(True, axis='y', linestyle='--', alpha=0.7)
+        ax.yaxis.set_major_formatter(ScalarFormatter(useOffset=False))
+
+        # Add legend inside each plot
+        ax.legend(title="Concatenation Methods",
+                 loc='upper left',
+                 fontsize=12,
+                 frameon=True,
+                 facecolor='white',
+                 edgecolor='black',
+                 framealpha=0.8)
+
+    # Add y-label only to the left subplot
+    ax0.set_ylabel('Throughput (GB/s)', fontsize=12)
+
+    # Add value labels on top of the bars
+    def autolabel(rects, ax):
+        for rect in rects:
+            height = rect.get_height()
+            ax.annotate(f'{height:.2f} GB/s',
+                        xy=(rect.get_x() + rect.get_width() / 2, height),
+                        xytext=(0, 3),  # 3 points vertical offset
+                        textcoords="offset points",
+                        ha='center', va='bottom', rotation=90, fontsize=8)
+
+    autolabel(rect1_axis0, ax0)
+    autolabel(rect2_axis0, ax0)
+    autolabel(rect3_axis0, ax0)
+
+    autolabel(rect1_axis1, ax1)
+    autolabel(rect2_axis1, ax1)
+    autolabel(rect3_axis1, ax1)
+
+    # Save the plot
+    plt.tight_layout()
+    plt.savefig(os.path.join(output_dir, 'concatenate_benchmark_combined.png'), dpi=300)
+    plt.show()
+    plt.close()
+
+    print(f"Combined plot saved to {os.path.join(output_dir, 'concatenate_benchmark_combined.png')}")
+
+
+def main():
+    print(f"{'=' * 60}")
+    print(f"Blosc2 vs NumPy concatenation benchmark")
+    print(f"{'=' * 60}")
+
+    # Parameters
+    sizes = [500, 1000, 2000, 4000] #, 10000]  # must be divisible by 4 for aligned chunks
+    num_arrays = 10
+
+    # Lists to store results for both axes
+    numpy_speeds_axis0 = []
+    unaligned_speeds_axis0 = []
+    aligned_speeds_axis0 = []
+    numpy_speeds_axis1 = []
+    unaligned_speeds_axis1 = []
+    aligned_speeds_axis1 = []
+
+    for axis in [0, 1]:
+        print(f"\nConcatenating {num_arrays} arrays along axis {axis}")
+        print(f"{'Size':<10} {'NumPy (GB/s)':<14} {'Unaligned (GB/s)':<18} {'Aligned (GB/s)':<16} {'Alig vs Unalig':<16} {'Alig vs NumPy':<16}")
+        print(f"{'-' * 90}")
+
+        for size in sizes:
+            # Run the benchmarks
+            numpy_time, numpy_shape, data_size_gb = run_numpy_benchmark(num_arrays, size, axis=axis)
+            unaligned_time, shape1, _ = run_benchmark(num_arrays, size, aligned_chunks=False, axis=axis)
+            aligned_time, shape2, _ = run_benchmark(num_arrays, size, aligned_chunks=True, axis=axis)
+
+            # Calculate throughputs in GB/s
+            numpy_speed = data_size_gb / numpy_time if numpy_time > 0 else float("inf")
+            unaligned_speed = data_size_gb / unaligned_time if unaligned_time > 0 else float("inf")
+            aligned_speed = data_size_gb / aligned_time if aligned_time > 0 else float("inf")
+
+            # Store speeds in the appropriate list
+            if axis == 0:
+                numpy_speeds_axis0.append(numpy_speed)
+                unaligned_speeds_axis0.append(unaligned_speed)
+                aligned_speeds_axis0.append(aligned_speed)
+            else:
+                numpy_speeds_axis1.append(numpy_speed)
+                unaligned_speeds_axis1.append(unaligned_speed)
+                aligned_speeds_axis1.append(aligned_speed)
+
+            # Calculate speedup ratios
+            aligned_vs_unaligned = aligned_speed / unaligned_speed if unaligned_speed > 0 else float("inf")
+            aligned_vs_numpy = aligned_speed / numpy_speed if numpy_speed > 0 else float("inf")
+
+            # Print results
+            print(f"{size:<10} {numpy_speed:<14.2f} {unaligned_speed:<18.2f} {aligned_speed:<16.2f} "
+                  f"{aligned_vs_unaligned:>10.2f}x {aligned_vs_numpy:>10.2f}x")
+
+            # Quick verification of result shape
+            if axis == 0:
+                expected_shape = (size, size)  # After concatenation along axis 0
+            else:
+                expected_shape = (size, size)  # After concatenation along axis 1
+
+            # Verify shapes match
+            shapes = [numpy_shape, shape1, shape2]
+            if any(shape != expected_shape for shape in shapes):
+                for i, shape_name in enumerate(["NumPy", "Blosc2 unaligned", "Blosc2 aligned"]):
+                    if shapes[i] != expected_shape:
+                        print(f"Warning: {shape_name} shape {shapes[i]} does not match expected {expected_shape}")
+
+    print(f"{'=' * 90}")
+
+    # Create the combined plot with both axes
+    create_combined_plot(
+        num_arrays,
+        sizes,
+        numpy_speeds_axis0, unaligned_speeds_axis0, aligned_speeds_axis0,
+        numpy_speeds_axis1, unaligned_speeds_axis1, aligned_speeds_axis1
+    )
+
+
+if __name__ == "__main__":
+    main()

bench/ndarray/concatenate_benchmark_combined.png

@@ -72,6 +72,7 @@ Constructors
 
     asarray
     copy
+    concatenate
     empty
     frombuffer
     fromiter

doc/reference/ndarray.rst

@@ -238,6 +238,7 @@ class Tuner(Enum):
     sort,
     reshape,
     copy,
+    concatenate,
     empty,
     frombuffer,
     fromiter,

src/blosc2/__init__.py

@@ -507,6 +507,8 @@ cdef extern from "b2nd.h":
     int b2nd_squeeze_index(b2nd_array_t *array, const c_bool *index)
     int b2nd_resize(b2nd_array_t *array, const int64_t *new_shape, const int64_t *start)
     int b2nd_copy(b2nd_context_t *ctx, b2nd_array_t *src, b2nd_array_t **array)
+    int b2nd_concatenate(b2nd_context_t *ctx, b2nd_array_t *src1, b2nd_array_t *src2,
+                         int8_t axis, c_bool copy, b2nd_array_t **array)
     int b2nd_from_schunk(blosc2_schunk *schunk, b2nd_array_t **array)
 
     void blosc2_unidim_to_multidim(uint8_t ndim, int64_t *shape, int64_t i, int64_t *index)
@@ -2867,3 +2869,26 @@ def schunk_get_slice_nchunks(schunk: SChunk, key):
         res[i] = chunks_idx[i]
     free(chunks_idx)
     return res
+
+
+def concatenate(arr1: NDArray, arr2: NDArray, axis: int, **kwargs):
+    """
+    Concatenate two NDArray objects along a specified axis.
+    """
+    cdef c_bool copy = kwargs.pop("copy", True)
+    cdef b2nd_context_t *ctx = create_b2nd_context(arr1.shape, arr1.chunks, arr1.blocks, arr1.dtype, kwargs)
+    if ctx == NULL:
+        raise RuntimeError("Error while creating the context for concatenation")
+
+    cdef b2nd_array_t *array
+    _check_rc(b2nd_concatenate(ctx, arr1.array, arr2.array, axis, copy, &array),
+              "Error while concatenating the arrays")
+    _check_rc(b2nd_free_ctx(ctx), "Error while freeing the context")
+
+    if copy:
+        # We have copied the concatenated data into a new array
+        return blosc2.NDArray(_schunk=PyCapsule_New(array.sc, <char *> "blosc2_schunk*", NULL),
+                              _array=PyCapsule_New(array, <char *> "b2nd_array_t*", NULL))
+    else:
+        # Return the first array, which now contains the concatenated data
+        return arr1

src/blosc2/blosc2_ext.pyx

@@ -1,8 +1,7 @@
-
 class MissingOperands(ValueError):
     def __init__(self, expr, missing_ops):
         self.expr = expr
         self.missing_ops = missing_ops
 
-        message = f"Lazy expression \"{expr}\" with missing operands: {missing_ops}"
+        message = f'Lazy expression "{expr}" with missing operands: {missing_ops}'
         super().__init__(message)