Add a plot for the different concatenate methods

Author: FrancescAlted

bench/ndarray/concatenate.py

@@ -9,6 +9,9 @@
 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):
@@ -20,6 +23,11 @@ def run_benchmark(num_arrays=10, size=500, aligned_chunks=False, axis=0):
     - 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
@@ -46,12 +54,16 @@ def run_benchmark(num_arrays=10, size=500, aligned_chunks=False, axis=0):
         )
         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
+    return duration, result.shape, data_size_gb
 
 
 def run_numpy_benchmark(num_arrays=10, size=500, axis=0):
@@ -62,6 +74,11 @@ def run_numpy_benchmark(num_arrays=10, size=500, axis=0):
     - 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
@@ -83,12 +100,96 @@ def run_numpy_benchmark(num_arrays=10, size=500, axis=0):
         ).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
+    return duration, result.shape, data_size_gb
+
+
+def create_combined_plot(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 (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():
@@ -100,24 +201,47 @@ def main():
     sizes = [500, 1000, 2000, 4000]  # 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':<10} {'Unaligned':<10} {'Aligned':<10} {'Alig vs Unalig':<16} {'Speedup vs NumPy':<16}")
-        print(f"{'-' * 80}")
+        print(f"{'Size':<10} {'NumPy (GB/s)':<14} {'Unaligned (GB/s)':<18} {'Aligned (GB/s)':<16} {'Alig vs Unalig':<16} {'Aligned vs NumPy':<16}")
+        print(f"{'-' * 90}")
 
         for size in sizes:
             # Run the benchmarks
-            numpy_time, numpy_shape = 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)
+            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 speedups
-            numpy_vs_aligned = numpy_time / aligned_time if aligned_time > 0 else float("inf")
-            aligned_vs_unaligned = unaligned_time / aligned_time if aligned_time > 0 else float("inf")
+            # 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_time:<10.4f} {unaligned_time:<10.4f} {aligned_time:<10.4f} "
-                  f"{aligned_vs_unaligned:>10.2f}x {numpy_vs_aligned:>10.2f}x")
+            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:
@@ -132,7 +256,14 @@ def main():
                     if shapes[i] != expected_shape:
                         print(f"Warning: {shape_name} shape {shapes[i]} does not match expected {expected_shape}")
 
-    print(f"{'=' * 80}")
+    print(f"{'=' * 90}")
+
+    # Create the combined plot with both axes
+    create_combined_plot(
+        sizes,
+        numpy_speeds_axis0, unaligned_speeds_axis0, aligned_speeds_axis0,
+        numpy_speeds_axis1, unaligned_speeds_axis1, aligned_speeds_axis1
+    )
 
 
 if __name__ == "__main__":