|
| 1 | +####################################################################### |
| 2 | +# Copyright (c) 2019-present, Blosc Development Team <[email protected]> |
| 3 | +# All rights reserved. |
| 4 | +# |
| 5 | +# This source code is licensed under a BSD-style license (found in the |
| 6 | +# LICENSE file in the root directory of this source tree) |
| 7 | +####################################################################### |
| 8 | + |
| 9 | +# Benchmark for comparing compute speeds of Blosc2 and Numexpr. |
| 10 | +# This version compares across different distributions of data: |
| 11 | +# constant, arange, linspace, or random |
| 12 | +# The expression can be any valid Numexpr expression. |
| 13 | + |
| 14 | +import blosc2 |
| 15 | +from time import time |
| 16 | +import numpy as np |
| 17 | +import numexpr as ne |
| 18 | +import matplotlib.pyplot as plt |
| 19 | +import seaborn as sns |
| 20 | + |
| 21 | +# Bench params |
| 22 | +N = 10_000 |
| 23 | +step = 3000 |
| 24 | +dtype = np.dtype(np.float64) |
| 25 | +persistent = False |
| 26 | +distributions = ["constant", "arange", "linspace", "random"] |
| 27 | +expr = "(a - b)" |
| 28 | +#expr = "sum(a - b)" |
| 29 | +#expr = "cos(a)**2 + sin(b)**2 - 1" |
| 30 | +#expr = "sum(cos(a)**2 + sin(b)**2 - 1)" |
| 31 | + |
| 32 | +# Set default compression params |
| 33 | +cparams = blosc2.CParams(clevel=1, codec=blosc2.Codec.BLOSCLZ) |
| 34 | +blosc2.cparams_dflts["codec"] = cparams.codec |
| 35 | +blosc2.cparams_dflts["clevel"] = cparams.clevel |
| 36 | +# Set default storage params |
| 37 | +storage = blosc2.Storage(contiguous=True, mode="w") |
| 38 | +blosc2.storage_dflts["contiguous"] = storage.contiguous |
| 39 | +blosc2.storage_dflts["mode"] = storage.mode |
| 40 | + |
| 41 | +# Create dictionaries to store results for each distribution |
| 42 | +blosc2_speeds = {dist: [] for dist in distributions} |
| 43 | +numexpr_speeds = {dist: [] for dist in distributions} |
| 44 | +ws_sizes = [] |
| 45 | + |
| 46 | +# Generate working set sizes once |
| 47 | +sizes = list(range(step, N + step, step)) |
| 48 | +for i in sizes: |
| 49 | + ws_sizes.append((i * i * 3 * np.dtype(dtype).itemsize) / 2**30) # Approximate size in GB |
| 50 | + |
| 51 | +# Loop through different distributions for benchmarking |
| 52 | +for dist in distributions: |
| 53 | + print(f"\nBenchmarking {dist} distribution...") |
| 54 | + |
| 55 | + # Evaluate using Blosc2 |
| 56 | + for i in sizes: |
| 57 | + shape = (i, i) |
| 58 | + urlpath = {name: None for name in ("a", "b", "c")} |
| 59 | + |
| 60 | + if dist == "constant": |
| 61 | + a = blosc2.ones(shape, dtype=dtype, urlpath=urlpath['a']) |
| 62 | + b = blosc2.full(shape, 2, dtype=dtype, urlpath=urlpath['b']) |
| 63 | + elif dist == "arange": |
| 64 | + a = blosc2.arange(0, i * i, dtype=dtype, shape=shape, urlpath=urlpath['a']) |
| 65 | + b = blosc2.arange(i * i, 2* i * i, dtype=dtype, shape=shape, urlpath=urlpath['b']) |
| 66 | + elif dist == "linspace": |
| 67 | + a = blosc2.linspace(0, 1, dtype=dtype, shape=shape, urlpath=urlpath['a']) |
| 68 | + b = blosc2.linspace(1, 2, dtype=dtype, shape=shape, urlpath=urlpath['b']) |
| 69 | + elif dist == "random": |
| 70 | + _ = np.random.random(shape) |
| 71 | + a = blosc2.fromiter(np.nditer(_), dtype=dtype, shape=shape, urlpath=urlpath['a']) |
| 72 | + # b = a.copy(urlpath=urlpath['b']) # faster, but output is not random |
| 73 | + _ = np.random.random(shape) |
| 74 | + b = blosc2.fromiter(np.nditer(_), dtype=dtype, shape=shape, urlpath=urlpath['b']) |
| 75 | + |
| 76 | + t0 = time() |
| 77 | + c = blosc2.lazyexpr(expr).compute(urlpath=urlpath['c']) |
| 78 | + t = time() - t0 |
| 79 | + speed = (a.schunk.nbytes + b.schunk.nbytes + c.schunk.nbytes) / 2**30 / t |
| 80 | + print(f"Blosc2 - {dist} - Size {i}x{i}: {speed:.2f} GB/s - cratio: {c.schunk.cratio:.1f}x") |
| 81 | + blosc2_speeds[dist].append(speed) |
| 82 | + |
| 83 | + # Evaluate using Numexpr |
| 84 | + for i in sizes: |
| 85 | + shape = (i, i) |
| 86 | + |
| 87 | + if dist == "constant": |
| 88 | + a = np.ones(shape, dtype=dtype) |
| 89 | + b = np.full(shape, 2, dtype=dtype) |
| 90 | + elif dist == "arange": |
| 91 | + a = np.arange(0, i * i, dtype=dtype).reshape(shape) |
| 92 | + b = np.arange(i * i, 2 * i * i, dtype=dtype).reshape(shape) |
| 93 | + elif dist == "linspace": |
| 94 | + a = np.linspace(0, 1, num=i * i, dtype=dtype).reshape(shape) |
| 95 | + b = np.linspace(1, 2, num=i * i, dtype=dtype).reshape(shape) |
| 96 | + elif dist == "random": |
| 97 | + a = np.random.random(shape) |
| 98 | + b = np.random.random(shape) |
| 99 | + |
| 100 | + t0 = time() |
| 101 | + c = ne.evaluate(expr) |
| 102 | + t = time() - t0 |
| 103 | + speed = (a.nbytes + b.nbytes + c.nbytes) / 2**30 / t |
| 104 | + print(f"Numexpr - {dist} - Size {i}x{i}: {speed:.2f} GB/s") |
| 105 | + numexpr_speeds[dist].append(speed) |
| 106 | + |
| 107 | +# Create a figure with four subplots (2x2 grid) |
| 108 | +sns.set_theme(style="whitegrid") |
| 109 | +fig, axes = plt.subplots(2, 2, figsize=(14, 10), sharex=True) |
| 110 | + |
| 111 | +# Flatten axes for easier iteration |
| 112 | +axes = axes.flatten() |
| 113 | + |
| 114 | +# Plot each distribution in its own subplot |
| 115 | +for i, dist in enumerate(distributions): |
| 116 | + axes[i].plot(ws_sizes, blosc2_speeds[dist], marker='o', linestyle='-', label="Blosc2") |
| 117 | + axes[i].plot(ws_sizes, numexpr_speeds[dist], marker='s', linestyle='--', label="Numexpr") |
| 118 | + axes[i].set_title(f"{dist.capitalize()} Distribution") |
| 119 | + axes[i].set_ylabel("Speed (GB/s)") |
| 120 | + axes[i].grid(True) |
| 121 | + axes[i].legend() |
| 122 | + if i >= 2: # Add x-label only to bottom subplots |
| 123 | + axes[i].set_xlabel("Working set size (GB)") |
| 124 | + |
| 125 | +# Add a shared title |
| 126 | +fig.suptitle(f"Blosc2 vs Numexpr Performance Across Different Data Distributions ({expr=})", fontsize=16) |
| 127 | +plt.tight_layout(rect=[0, 0, 1, 0.96]) # Adjust the rect parameter to make room for the suptitle |
| 128 | + |
| 129 | +# Save the unified plot with subplots |
| 130 | +plt.savefig("blosc2_vs_numexpr_subplots.png", dpi=300, bbox_inches='tight') |
| 131 | +plt.show() |
0 commit comments