diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index fcd5d826..dc98b3e5 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -4,6 +4,35 @@ XXX version-specific blurb XXX +### New features + +- New `blosc2.random` module: seedable, NumPy-quality random `NDArray` + constructors. Each chunk gets its own independent `SeedSequence`-spawned + stream and is generated concurrently in a thread pool, giving full `PCG64` + quality with genuinely parallel generation (measured ~3x faster than + `asarray(np.random.default_rng(...).random(...))` on a 100M-element array). + Covers 42 of `numpy.random.Generator`'s 43 public methods (full + compatibility table in `doc/reference/random.rst`): + - Core: `random`, `integers`, `normal`, `uniform`, `choice` + (`replace=True` only). + - 30 scalar distributions: `beta`, `binomial`, `chisquare`, + `exponential`, `f`, `gamma`, `geometric`, `gumbel`, `hypergeometric`, + `laplace`, `logistic`, `lognormal`, `logseries`, `negative_binomial`, + `noncentral_chisquare`, `noncentral_f`, `pareto`, `poisson`, `power`, + `rayleigh`, `standard_cauchy`, `standard_exponential`, + `standard_gamma`, `standard_normal`, `standard_t`, `triangular`, + `vonmises`, `wald`, `weibull`, `zipf`. + - 4 vector-valued distributions (output shape `shape + (k,)`, one draw + per trailing vector): `dirichlet`, `multinomial`, + `multivariate_hypergeometric`, `multivariate_normal`. + - `permutation`, `permuted`, `shuffle`: unlike the rest of the module, + these are not chunk-parallel — whole-array shuffling is inherently + sequential, so they materialize the full array and shuffle it + single-threaded. `shuffle` additionally requires its argument to + already be an `NDArray`, since it mutates in place and returns `None`, + matching numpy. + - Not implemented: `bytes` (returns raw `bytes`, not an `NDArray`). + ## Changes from 4.9.0 to 4.9.1 A small hot-fix release for the Arrow interop work in 4.9.0: a real diff --git a/bench/bench_pandas_engine.py b/bench/bench_pandas_engine.py index b73a1b08..898be4dc 100644 --- a/bench/bench_pandas_engine.py +++ b/bench/bench_pandas_engine.py @@ -8,10 +8,22 @@ # Benchmark: DataFrame.apply(f, engine=blosc2.jit) vs plain DataFrame.apply(f) # # engine=blosc2.jit calls the vectorized function once per column (the -# default axis=0), so the win comes from the Blosc2/numexpr compute engine -# (operator fusion, multi-threading) beating plain NumPy on a -# multi-operation elementwise expression over a large 1D array. This script -# measures that on a 1,000,000-row, 8-column frame. +# default axis=0), so the win comes from the Blosc2 compute engine (operator +# fusion, multi-threading) beating plain NumPy, which evaluates the function +# body one operation at a time and allocates a full-size temporary at each +# step. +# +# Two conditions must both hold for the engine to pay off, and this script +# measures each one separately: +# +# * enough rows -- below ~50k the per-call setup dominates (rows sweep) +# * enough ops -- a single operation has nothing to fuse (ops sweep) +# +# numexpr is measured alongside as a reference point: on in-memory frames it +# is somewhat faster than the Blosc2 engine, so the argument for +# engine=blosc2.jit is that you write a readable Python function instead of a +# quoted expression string, not that it wins a raw speed race. See +# doc/guides/pandas_engine.md. # # Note: axis=1 (row-wise) is NOT a good fit for this engine. It still calls # the function once per row in a Python loop either way, and for a handful @@ -22,27 +34,76 @@ # # Each measurement is the minimum of NRUNS repetitions to reduce noise. +from pathlib import Path from time import perf_counter +import numexpr import numpy as np import pandas as pd import blosc2 -NRUNS = 3 +NRUNS = 5 NROWS = 1_000_000 NCOLS = 8 +ROW_SWEEP = (1_000, 10_000, 100_000, 1_000_000, 5_000_000) + +OUT_DIR = Path(__file__).resolve().parent.parent / "doc" / "guides" / "pandas_engine" + +# dataviz reference palette, same values as bench/optim_tips/common.py +COLOR_TIP = "#1baf7a" +INK = "#0b0b0b" +MUTED = "#898781" +GRID = "#e1e0d9" -def make_df(): + +def make_df(nrows=NROWS, ncols=NCOLS): rng = np.random.default_rng(0) return pd.DataFrame( - {f"c{i}": rng.random(NROWS) for i in range(NCOLS)}, + {f"c{i}": rng.normal(size=nrows) for i in range(ncols)}, ) -def transform(col): - return np.sin(col) * np.cos(col) + col**2 - np.sqrt(np.abs(col)) + np.exp(-col) +def yeo_johnson(col, lam=0.5): + """Power transform behind sklearn's PowerTransformer, applied per column. + + np.where evaluates both arms over the whole column, so each is clamped to + its own domain: without the np.maximum calls, the unselected arm raises a + negative base to a fractional power and floods the run with NaNs and + "invalid value encountered in power" warnings. + """ + pos = (np.power(np.maximum(col, 0.0) + 1.0, lam) - 1.0) / lam + neg = -(np.power(np.maximum(-col, 0.0) + 1.0, 2.0 - lam) - 1.0) / (2.0 - lam) + return np.where(col >= 0, pos, neg) + + +# The same transform as a single numexpr expression: legal, but this is what +# the readability argument is about. +YEO_JOHNSON_NX = ( + "where(c >= 0, " + "((maximum(c, 0.0) + 1.0) ** lam - 1.0) / lam, " + "-((maximum(-c, 0.0) + 1.0) ** (2.0 - lam) - 1.0) / (2.0 - lam))" +) + + +def numexpr_apply(df, expr=YEO_JOHNSON_NX, lam=0.5): + return pd.DataFrame( + {col: numexpr.evaluate(expr, local_dict={"c": df[col].values, "lam": lam}) for col in df}, + index=df.index, + ) + + +# Prefixes of one expression, so the number of fused operations is the only +# variable: mixing in different *kinds* of operation would measure the cost of +# the operations rather than the benefit of fusing them. +OPS_SWEEP = ( + ("1", lambda col: np.sin(col)), + ("2", lambda col: np.sin(col) * np.cos(col)), + ("3", lambda col: np.sin(col) * np.cos(col) + col**2), + ("4", lambda col: np.sin(col) * np.cos(col) + col**2 - np.sqrt(np.abs(col))), + ("5", lambda col: np.sin(col) * np.cos(col) + col**2 - np.sqrt(np.abs(col)) + np.exp(-col)), +) def timeit(fn): @@ -55,18 +116,89 @@ def timeit(fn): return best, result +def speedup(df, func): + """Plain apply vs engine=blosc2.jit, returning (speedup, t_plain, t_engine).""" + t_plain, _ = timeit(lambda: df.apply(func)) + t_engine, _ = timeit(lambda: df.apply(func, engine=blosc2.jit)) + return t_plain / t_engine, t_plain, t_engine + + +def save_plot(row_speedups, ops_speedups, out_path): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, (ax_rows, ax_ops) = plt.subplots(1, 2, figsize=(8, 3.2)) + + ax_rows.semilogx(ROW_SWEEP, row_speedups, "o-", color=COLOR_TIP, linewidth=2) + ax_rows.set_xlabel("rows (log scale)", color=INK, fontsize=9) + ax_rows.set_ylabel("speedup vs plain apply", color=INK, fontsize=9) + ax_rows.set_title(f"{NCOLS} columns, Yeo-Johnson", color=MUTED, fontsize=9) + + labels = [name for name, _ in OPS_SWEEP] + ax_ops.plot(range(len(labels)), ops_speedups, "o-", color=COLOR_TIP, linewidth=2) + ax_ops.set_xticks(range(len(labels))) + ax_ops.set_xticklabels(labels) + ax_ops.set_xlabel("operations fused into one pass", color=INK, fontsize=9) + ax_ops.set_title(f"{NROWS:,} rows x {NCOLS} columns", color=MUTED, fontsize=9) + + for ax, values in ((ax_rows, row_speedups), (ax_ops, ops_speedups)): + # Break-even: below this line the engine is a net loss. + ax.axhline(1.0, color=MUTED, linestyle="--", linewidth=1) + ax.set_ylim(0, max(values) * 1.25) + ax.yaxis.set_major_formatter(lambda v, _pos: f"{v:g}x") + ax.yaxis.grid(True, color=GRID, linewidth=0.8) + ax.set_axisbelow(True) + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + ax.spines["left"].set_color(GRID) + ax.spines["bottom"].set_color(GRID) + ax.tick_params(labelsize=9, colors=MUTED) + + fig.suptitle( + "df.apply(f, engine=blosc2.jit): when it pays off", + fontsize=11, + color=INK, + ) + fig.tight_layout(rect=[0, 0, 1, 0.90]) + out_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(out_path, dpi=150) + plt.close(fig) + + def main(): df = make_df() - t_plain, result_plain = timeit(lambda: df.apply(transform)) - t_engine, result_engine = timeit(lambda: df.apply(transform, engine=blosc2.jit)) + t_plain, result_plain = timeit(lambda: df.apply(yeo_johnson)) + t_engine, result_engine = timeit(lambda: df.apply(yeo_johnson, engine=blosc2.jit)) + t_numexpr, result_numexpr = timeit(lambda: numexpr_apply(df)) pd.testing.assert_frame_equal(result_engine, result_plain) - - print(f"rows={NROWS}, cols={NCOLS}") - print(f"plain df.apply(f): {t_plain:.4f} s") - print(f"df.apply(f, engine=blosc2.jit): {t_engine:.4f} s") - print(f"speedup: {t_plain / t_engine:.1f}x") + pd.testing.assert_frame_equal(result_numexpr, result_plain) + + print(f"rows={NROWS}, cols={NCOLS}, transform=Yeo-Johnson") + print(f"plain df.apply(f): {t_plain:.4f} s") + print(f"df.apply(f, engine=blosc2.jit): {t_engine:.4f} s {t_plain / t_engine:.2f}x") + print(f"numexpr per column: {t_numexpr:.4f} s {t_plain / t_numexpr:.2f}x") + + print("\nrows sweep (speedup vs plain apply):") + row_speedups = [] + for nrows in ROW_SWEEP: + sp, tp, te = speedup(make_df(nrows=nrows), yeo_johnson) + row_speedups.append(sp) + print(f" {nrows:>9,} rows: plain {tp:.4f} s engine {te:.4f} s {sp:.2f}x") + + print("\nops sweep (speedup vs plain apply):") + ops_speedups = [] + for name, func in OPS_SWEEP: + sp, tp, te = speedup(df, func) + ops_speedups.append(sp) + print(f" {name.replace(chr(10), ' '):>20}: plain {tp:.4f} s engine {te:.4f} s {sp:.2f}x") + + out_path = OUT_DIR / "speedup.png" + save_plot(row_speedups, ops_speedups, out_path) + print(f"\nplot saved to {out_path}") if __name__ == "__main__": diff --git a/doc/guides/pandas_engine.md b/doc/guides/pandas_engine.md index fde3c366..d80dd5fd 100644 --- a/doc/guides/pandas_engine.md +++ b/doc/guides/pandas_engine.md @@ -1,50 +1,169 @@ # Using Blosc2 as a pandas Engine -pandas' `DataFrame.apply` and `Series.map` accept an `engine=` argument: a -callable exposing a `__pandas_udf__` attribute that pandas dispatches to -instead of running the Python-level per-row/per-element loop. `blosc2.jit` -is such an engine. - -The contract is different from a plain `apply`/`map` callback: the function -passed to `engine=blosc2.jit` must be **vectorized** — it is called *once* -with a full NumPy array (a column, a row, or the whole array, depending on -`axis`), not once per element. This is the same contract `@blosc2.jit` -already has everywhere else in the library; using it as a pandas engine -just changes who supplies the array. +pandas' `DataFrame.apply` and `Series.map` accept an `engine=` argument, and +`blosc2.jit` is one such engine. Instead of running your function once per +element in a Python loop, pandas hands it the **whole column at once**, and +Blosc2 evaluates the entire function body in a single multi-threaded pass +over the data. + +The result is typically **2-3x faster than a plain `apply`**, with the +function itself left exactly as you wrote it. + +## An example + +Yeo-Johnson is the power transform behind scikit-learn's `PowerTransformer`, +used to make skewed features more normally distributed. Its parameter is +fitted per column, so applying it to every column of a feature matrix is +precisely what you want: ```python import numpy as np import pandas as pd + import blosc2 -df = pd.DataFrame( - { - "a": np.arange(1_000_000, dtype=np.float64), - "b": np.arange(1_000_000, dtype=np.float64), - } -) +rng = np.random.default_rng(0) +df = pd.DataFrame({f"c{i}": rng.normal(size=1_000_000) for i in range(8)}) @blosc2.jit -def add_one(col): - return col + 1 +def yeo_johnson(col, lam=0.5): + # np.where evaluates both arms, so clamp each to its own domain + pos = (np.power(np.maximum(col, 0.0) + 1.0, lam) - 1.0) / lam + neg = -(np.power(np.maximum(-col, 0.0) + 1.0, 2.0 - lam) - 1.0) / (2.0 - lam) + return np.where(col >= 0, pos, neg) + + +result = df.apply(yeo_johnson, engine=blosc2.jit) +``` + +`result` is a DataFrame of the same shape and column names as `df`, with the +transform applied to every column — the same thing plain `df.apply(yeo_johnson)` +returns, only computed differently. On the machine below it takes 0.056 s +instead of 0.119 s, a **2.1x** speedup. + +## Why it is faster + +Evaluated by NumPy, that function body is a sequence of separate steps: raise +to a power, subtract, divide, do it again for the negative branch, then select. +Each step walks the full column and allocates a new full-size temporary array +to hold its result. + +Blosc2 does not execute the steps one at a time. It captures the whole +expression first, then makes a single pass over the data, computing every step +on one small piece while that piece is still in cache, and spreading the pieces +across cores. The intermediate arrays are never created. +## When it pays off -result = df.apply(add_one, engine=blosc2.jit) +Two things have to be true, and the plot below measures each one on its own. + +![Speedup vs rows and vs number of fused operations](pandas_engine/speedup.png) + +**Enough rows.** Setting up the compute engine costs a fixed amount per call. +At a few hundred thousand rows that setup is still larger than anything it +saves, and the engine is a net loss — at 100,000 rows it runs at 0.70x. Break-even +falls between 100,000 and 1,000,000 rows. + +**Enough arithmetic.** The more operations there are to fuse into one pass, the +more temporaries are avoided and the bigger the win — from 2.1x for a single +operation up to 3.8x for five. + +Beyond a few million rows the speedup flattens and then eases off (1.9x at 5M +above): the arrays no longer fit in cache and the whole computation becomes +limited by memory bandwidth, which fusion can reduce but not eliminate. + +## When not to use it + +**Trivial expressions.** Arithmetic intensity matters more than the raw number +of operations. A single cheap operation over a large array is limited by memory +bandwidth, not by computation, so there is nothing for the engine to win back: + +```python +result = df.apply(lambda col: col + 1, engine=blosc2.jit) # 0.53x — slower! ``` -`axis=0` (the default) calls the function once per column; `axis=1` calls it -once per row. **Use `axis=0`** (or restructure the computation so it works -column-wise): the win comes from the Blosc2/numexpr compute engine (operator -fusion, multi-threading) processing one large 1D array per call, and that -only happens for columns. `axis=1` still calls the function once per row — -same as plain pandas — and for a handful of columns, the overhead of -wrapping each tiny row array for the compute engine outweighs any benefit, -so `engine=blosc2.jit` with `axis=1` is typically *slower* than plain -`apply(axis=1)`. See the benchmark below. +That runs at roughly **half** the speed of a plain `apply`. Reach for the engine +when the function does real work per element, such as transcendental functions +or several chained operations. + +**Small frames.** See the plot above: below a few hundred thousand rows, use a +plain `apply`. + +## Compared to `pd.eval` and numexpr + +pandas' own [Enhancing performance](https://pandas.pydata.org/docs/user_guide/enhancingperf.html) +guide describes `pd.eval(..., engine="numexpr")`, which fuses expressions using +the same underlying idea. It is worth being straightforward about how the two +compare on the example above: -`Series.map(func, engine=blosc2.jit)` works the same way: `func` is called -once with the Series' full underlying array. +| approach | time | vs plain apply | +| --- | --- | --- | +| plain `df.apply(f)` | 0.1194 s | 1.00x | +| `df.apply(f, engine=blosc2.jit)` | 0.0561 s | 2.13x | +| `numexpr.evaluate(...)` per column | 0.0380 s | 3.14x | + +On an in-memory DataFrame, numexpr is somewhat faster. The reason to prefer +`engine=blosc2.jit` is not raw speed but that **you write a Python function +rather than a quoted string**. The numexpr equivalent of `yeo_johnson` has to +become a single expression: + +```python +"where(c >= 0, " +"((maximum(c, 0.0) + 1.0) ** lam - 1.0) / lam, " +"-((maximum(-c, 0.0) + 1.0) ** (2.0 - lam) - 1.0) / (2.0 - lam))" +``` + +The Blosc2 version keeps its intermediate variables and its name, can call +helper functions, can be unit-tested and reused elsewhere under `@blosc2.jit`, +and is checked by your editor and linters. That is the trade being offered. + +Note also that Blosc2's characteristic strength — computing directly over +compressed, potentially larger-than-memory arrays — does not come into play on +this path, because pandas materialises each column as a plain NumPy array +before the engine ever sees it. + +## Gotchas + +**`std()` silently changes meaning.** A plain `apply` passes your function a +pandas Series, whose `.std()` defaults to `ddof=1`. The engine passes a NumPy +array, whose `.std()` defaults to `ddof=0`. The same source code therefore +computes slightly different numbers depending on the engine. Pass `ddof` +explicitly if it matters: + +```python +z = (col - col.mean()) / col.std(ddof=0) +``` + +**No Python `if` on values.** The function is traced rather than executed +statement by statement, so branching on array contents raises +`ValueError: The truth value of an array ... is ambiguous`. Use `np.where`, +nesting it where you would have used `elif`. (This is the same restriction +numexpr has.) Branching on a scalar *parameter* is fine. + +**`np.where` evaluates both arms.** Unlike a real `if`, both branches are +computed over the whole column and only then selected between, so each one runs +on values it was never meant to see. In `yeo_johnson` above, `np.power(col + 1.0, +0.5)` would hit a negative base wherever `col < -1` — around 159,000 elements in +a million-row standard normal — producing NaNs and a +`RuntimeWarning: invalid value encountered in power`. The final answer is still +correct, because `np.where` discards them, but the warning is noise and the work +is wasted. That is why each arm is clamped to its own domain with `np.maximum`. +The same caveat applies to numexpr's `where()`. + +**`np.sign` is not supported** on the traced expressions and raises a +`TypeError`. Express it with `np.where` instead. + +## Row-wise (`axis=1`) + +`axis=0`, the default, calls the function once per column and is where the +benefit lies. `axis=1` calls it once per row — the same Python-level loop plain +pandas would run, except each call now also pays the cost of wrapping a tiny +array for the compute engine. For a handful of columns that overhead outweighs +any gain, making `engine=blosc2.jit` with `axis=1` typically *slower* than a +plain `apply(axis=1)`. + +Use `axis=0`, or restructure the computation so it works column-wise. ## Limitations @@ -61,23 +180,16 @@ once with the Series' full underlying array. the pandas-side API surface, not of the Blosc2 engine. The two entry points that do reach the engine are `DataFrame.apply` and `Series.map`. -## Benchmark +`Series.map(func, engine=blosc2.jit)` works the same way as `DataFrame.apply`: +`func` is called once with the Series' full underlying array. -`bench/bench_pandas_engine.py` compares `df.apply(f, engine=blosc2.jit)` -against plain `df.apply(f)` (`axis=0`, the default) on a 1,000,000-row, -8-column frame, for a multi-operation elementwise expression -(`sin(x)*cos(x) + x**2 - sqrt(|x|) + exp(-x)`). Measured on the development -machine (Apple M4, conda env with pandas 3.0.3): +## Reproducing these numbers -``` -rows=1000000, cols=8 -plain df.apply(f): 0.1114 s -df.apply(f, engine=blosc2.jit): 0.0260 s -speedup: 4.3x -``` - -Run the script for the numbers on your machine: +All figures on this page come from `bench/bench_pandas_engine.py`, measured on +an Apple M4 with pandas 3.0.3 and 8 threads. Run it yourself with: ``` python bench/bench_pandas_engine.py ``` + +It prints the comparison table, both sweeps, and regenerates the plot above. diff --git a/doc/guides/pandas_engine/speedup.png b/doc/guides/pandas_engine/speedup.png new file mode 100644 index 00000000..b26d37a3 Binary files /dev/null and b/doc/guides/pandas_engine/speedup.png differ diff --git a/doc/reference/array_operations.rst b/doc/reference/array_operations.rst index 2c2adb15..b40ed431 100644 --- a/doc/reference/array_operations.rst +++ b/doc/reference/array_operations.rst @@ -7,5 +7,6 @@ Operations with arrays ufuncs reduction_functions linalg + random additional_funcs index_funcs diff --git a/doc/reference/random.rst b/doc/reference/random.rst new file mode 100644 index 00000000..4fbcc32f --- /dev/null +++ b/doc/reference/random.rst @@ -0,0 +1,190 @@ +Random +------ +Chunk-parallel, NumPy-quality random :ref:`NDArray ` constructors. Each chunk +gets its own independent :class:`~numpy.random.SeedSequence`-spawned stream and is +generated concurrently in a thread pool, so generation itself parallelizes (not just +compression). + +.. currentmodule:: blosc2.random + +.. autosummary:: + + default_rng + Generator + +.. autofunction:: blosc2.random.default_rng +.. autoclass:: blosc2.random.Generator + :members: + +NumPy compatibility +~~~~~~~~~~~~~~~~~~~ +:class:`Generator` mirrors :class:`numpy.random.Generator` method-for-method: same +method names, same distribution parameters, same math. A few differences apply +consistently across the module: + +- ``shape`` replaces numpy's ``size`` and is **required** (no implicit scalar draws), + and keyword-only on every method except :meth:`~Generator.random`. +- There is no ``out=`` parameter — each call allocates its own destination + :class:`~blosc2.NDArray`, filled chunk by chunk. +- Extra keyword arguments (``chunks``, ``cparams``, ``urlpath``, ...) are forwarded to + :func:`blosc2.empty`. +- Results are reproducible for a given ``(seed, call order, shape, chunks)``; changing + the chunk layout changes the values, since chunks map one-to-one to independent + ``SeedSequence`` streams. +- :meth:`~Generator.gamma` and :meth:`~Generator.standard_gamma` rename numpy's + distribution-shape parameter (itself called ``shape``) to ``shape_param``, and make it + positional-only, to avoid colliding with the array-shape ``shape`` keyword. +- :meth:`~Generator.choice` only supports ``replace=True`` (numpy's default) and a 1-D + or scalar-int ``a``; sampling without replacement, or along an axis of a + multi-dimensional ``a``, needs whole-array coordination across chunks that this + module doesn't do. +- The vector-valued distributions (:meth:`~Generator.dirichlet`, + :meth:`~Generator.multinomial`, :meth:`~Generator.multivariate_hypergeometric`, + :meth:`~Generator.multivariate_normal`) draw one length-``k`` vector per element: + output shape is ``shape + (k,)``, and that trailing dimension is always kept whole + within a chunk. +- :meth:`~Generator.permutation`, :meth:`~Generator.permuted`, and + :meth:`~Generator.shuffle` are **not** chunk-parallel: whole-array shuffling is + inherently sequential, so these load the full array into memory and shuffle it + single-threaded rather than generating chunks independently. :meth:`~Generator.shuffle` + additionally requires its argument to already be a :class:`~blosc2.NDArray`, since it + mutates it in place and returns ``None``, matching numpy. + +Coverage of :class:`numpy.random.Generator`'s public methods: + +.. list-table:: + :header-rows: 1 + :widths: 25 10 45 + + * - Method + - Status + - Notes + * - ``beta`` + - ✅ + - + * - ``binomial`` + - ✅ + - + * - ``bytes`` + - ❌ + - returns raw ``bytes``, not an ``NDArray`` — outside this module's contract + * - ``chisquare`` + - ✅ + - + * - ``choice`` + - ✅ + - ``replace=True`` and 1-D/scalar ``a`` only + * - ``dirichlet`` + - ✅ + - vector output, see above + * - ``exponential`` + - ✅ + - + * - ``f`` + - ✅ + - + * - ``gamma`` + - ✅ + - distribution-shape parameter renamed ``shape_param`` + * - ``geometric`` + - ✅ + - + * - ``gumbel`` + - ✅ + - + * - ``hypergeometric`` + - ✅ + - + * - ``integers`` + - ✅ + - + * - ``laplace`` + - ✅ + - + * - ``logistic`` + - ✅ + - + * - ``lognormal`` + - ✅ + - + * - ``logseries`` + - ✅ + - + * - ``multinomial`` + - ✅ + - vector output, see above + * - ``multivariate_hypergeometric`` + - ✅ + - vector output, see above + * - ``multivariate_normal`` + - ✅ + - vector output, see above + * - ``negative_binomial`` + - ✅ + - + * - ``noncentral_chisquare`` + - ✅ + - + * - ``noncentral_f`` + - ✅ + - + * - ``normal`` + - ✅ + - + * - ``pareto`` + - ✅ + - + * - ``permutation`` + - ✅ + - single-threaded, full-materialization — see above + * - ``permuted`` + - ✅ + - single-threaded, full-materialization — see above + * - ``poisson`` + - ✅ + - + * - ``power`` + - ✅ + - + * - ``random`` + - ✅ + - + * - ``rayleigh`` + - ✅ + - + * - ``shuffle`` + - ✅ + - single-threaded, full-materialization; requires an ``NDArray`` — see above + * - ``standard_cauchy`` + - ✅ + - + * - ``standard_exponential`` + - ✅ + - + * - ``standard_gamma`` + - ✅ + - distribution-shape parameter renamed ``shape_param`` + * - ``standard_normal`` + - ✅ + - + * - ``standard_t`` + - ✅ + - + * - ``triangular`` + - ✅ + - + * - ``uniform`` + - ✅ + - + * - ``vonmises`` + - ✅ + - + * - ``wald`` + - ✅ + - + * - ``weibull`` + - ✅ + - + * - ``zipf`` + - ✅ + - diff --git a/examples/ndarray/random-constructor.py b/examples/ndarray/random-constructor.py new file mode 100644 index 00000000..8a12f3c6 --- /dev/null +++ b/examples/ndarray/random-constructor.py @@ -0,0 +1,73 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +# This example shows how to use `blosc2.random` to build blosc2 arrays with the same +# API shape as `numpy.random.Generator`, but generated chunk-by-chunk in parallel. + +from time import time + +import numpy as np + +import blosc2 + +# `default_rng` mirrors `numpy.random.default_rng`: same seeding, same distributions. +rng = blosc2.random.default_rng(42) + +# Every method returns a blosc2.NDArray, not a numpy array. `shape` replaces numpy's +# `size` and is required (no implicit scalar draws) and keyword-only for every +# distribution except `random`, where it stays positional to mirror numpy's own order. +a = rng.random((4, 4)) +print("*** rng.random((4, 4)) ***") +print(a[:]) + +b = rng.integers(0, 10, shape=(10,), dtype=np.int32) +print("\n*** rng.integers(0, 10, shape=(10,), dtype=np.int32) ***") +print(b[:]) + +# Reproducibility: same seed -> same array. Successive calls on the *same* generator +# draw different arrays (the "generator state advances" semantics numpy users expect). +c1 = blosc2.random.default_rng(7).normal(shape=(1000,)) +c2 = blosc2.random.default_rng(7).normal(shape=(1000,)) +same_seed = blosc2.random.default_rng(7) +d1 = same_seed.normal(shape=(1000,)) +d2 = same_seed.normal(shape=(1000,)) +print("\n*** Reproducibility ***") +print(f"Two fresh generators, same seed, first call: equal = {np.array_equal(c1[:], d1[:])}") +print(f"One generator, two successive calls: equal = {np.array_equal(d1[:], d2[:])}") + +# Extra kwargs are forwarded to blosc2.empty(): chunks, cparams, urlpath, ... +e = rng.uniform(-1, 1, shape=(1_000, 1_000), chunks=(100, 1_000), cparams={"clevel": 5}) +print(f"\n*** rng.uniform(..., chunks=(100, 1000), cparams=clevel5) -> cratio {e.schunk.cratio:.1f}x ***") + +# Most of numpy.random.Generator's other distributions are covered too. +poisson = rng.poisson(4.0, shape=(10,)) +gamma = rng.gamma(2.0, scale=1.5, shape=(10,)) # shape_param renamed, see docs +choice = rng.choice(["red", "green", "blue"], shape=(10,)) +print("\n*** A few more distributions ***") +print("poisson:", poisson[:]) +print("gamma: ", gamma[:]) +print("choice: ", choice[:]) + +# Vector-valued distributions draw a whole length-k vector per element: output shape +# is `shape + (k,)`, and that trailing dimension is always kept whole within a chunk. +dirichlet = rng.dirichlet([1.0, 2.0, 3.0], shape=(5,)) +print("\n*** rng.dirichlet([1, 2, 3], shape=(5,)) -> shape", dirichlet.shape, "***") +print(dirichlet[:]) +print("rows sum to 1:", np.allclose(dirichlet[:].sum(axis=-1), 1.0)) + +# Parallel generation vs. building the array via numpy + asarray(). +N = 100_000_000 +print(f"\n*** Timing: {N:_} float64 uniform draws ***") +t0 = time() +blosc2.asarray(np.random.default_rng(42).random(N), cparams={"clevel": 1}) +t_numpy = time() - t0 +print(f"asarray(np.random...): {t_numpy:.3f} s") + +t0 = time() +blosc2.random.default_rng(42).random((N,), cparams={"clevel": 1}) +t_blosc2 = time() - t0 +print(f"blosc2.random.default_rng...: {t_blosc2:.3f} s ({t_numpy / t_blosc2:.1f}x faster)") diff --git a/src/blosc2/__init__.py b/src/blosc2/__init__.py index 3050887b..cea161a2 100644 --- a/src/blosc2/__init__.py +++ b/src/blosc2/__init__.py @@ -595,6 +595,7 @@ def _raise(exc): from .linalg import tensordot, vecdot, permute_dims, matrix_transpose, matmul, transpose, diagonal, outer from .utils import linalg_funcs as linalg_funcs_list from . import fft +from . import random # Registry for postfilters postfilter_funcs = {} diff --git a/src/blosc2/random.py b/src/blosc2/random.py new file mode 100644 index 00000000..fd36d708 --- /dev/null +++ b/src/blosc2/random.py @@ -0,0 +1,380 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### +"""Chunk-parallel NumPy-quality random generators producing :class:`blosc2.NDArray`. + +Each chunk gets its own independent stream, spawned from a root +:class:`numpy.random.SeedSequence` via :meth:`~numpy.random.SeedSequence.spawn`, and +generation runs in a thread pool (NumPy ``Generator`` fill loops release the GIL). +Results are reproducible for a given ``(seed, call order, shape, chunks)``; changing +the chunk layout changes the values, since the chunk grid determines the stream +assignment. +""" + +from __future__ import annotations + +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from typing import Any + +import numpy as np + +import blosc2 + + +def _chunk_slices(a: blosc2.NDArray): + for info in a.iterchunks_info(): + yield tuple( + slice(c * s, min((c + 1) * s, sh)) + for c, s, sh in zip(info.coords, a.chunks, a.shape, strict=False) + ) + + +def _materialize(x) -> np.ndarray: + return x[:] if isinstance(x, blosc2.NDArray) else np.asarray(x) + + +def _bounded_map(fn, n: int, max_workers: int): + # ponytail: window bounds peak memory to O(workers * chunk); a plain + # executor.map would submit all n chunks upfront and could buffer O(n). + window = 2 * max_workers + try: + with ThreadPoolExecutor(max_workers=max_workers) as ex: + pending = {ex.submit(fn, i): i for i in range(min(window, n))} + next_i = len(pending) + while pending: + done, _ = wait(pending, return_when=FIRST_COMPLETED) + for fut in done: + del pending[fut] + yield fut.result() + if next_i < n: + pending[ex.submit(fn, next_i)] = next_i + next_i += 1 + except RuntimeError: + # ponytail: some sandboxes (e.g. wasm32/Pyodide without pthreads) report a + # real nthreads count but can't actually start OS threads at all — submit() + # then raises "can't start new thread" on the very first call, before + # anything has been yielded. Fall back to serial generation rather than + # crashing; upgrade path is none needed, this is the correct behavior. + for i in range(n): + yield fn(i) + + +class Generator: + """Chunk-parallel counterpart of :class:`numpy.random.Generator`. + + Create via :func:`default_rng`. Each method call draws a full + :class:`~blosc2.NDArray` in one shot, generating its chunks concurrently. + """ + + def __init__(self, seed=None): + self._root = seed if isinstance(seed, np.random.SeedSequence) else np.random.SeedSequence(seed) + + def _fill_chunks(self, dst, method, args, kwargs, chunk_shape_of) -> blosc2.NDArray: + slices = list(_chunk_slices(dst)) + call_ss = self._root.spawn(1)[0] + chunk_streams = call_ss.spawn(len(slices)) + + def gen(i): + rng = np.random.default_rng(chunk_streams[i]) + return i, getattr(rng, method)(*args, size=chunk_shape_of(slices[i]), **kwargs) + + if len(slices) <= 1 or blosc2.nthreads == 1: + results = (gen(i) for i in range(len(slices))) + else: + results = _bounded_map(gen, len(slices), blosc2.nthreads) + + for i, buf in results: + dst[slices[i]] = buf + return dst + + def _fill(self, method: str, args: tuple, kwargs: dict, shape, **b2_kwargs: Any) -> blosc2.NDArray: + if shape is None: + raise TypeError("shape is required") + dtype = getattr(np.random.default_rng(0), method)(*args, size=1, **kwargs).dtype + dst = blosc2.empty(shape, dtype=dtype, **b2_kwargs) + return self._fill_chunks(dst, method, args, kwargs, lambda sl: tuple(s.stop - s.start for s in sl)) + + def _fill_vector( + self, method: str, args: tuple, kwargs: dict, shape, k: int, **b2_kwargs: Any + ) -> blosc2.NDArray: + # Each draw produces a whole length-k vector, so the trailing dimension must + # never be split across chunks (a chunk can only hold complete draws). + if shape is None: + raise TypeError("shape is required") + shape = (shape,) if isinstance(shape, int | np.integer) else tuple(shape) + full_shape = shape + (k,) + dtype = getattr(np.random.default_rng(0), method)(*args, size=1, **kwargs).dtype + + chunks = b2_kwargs.get("chunks") + if chunks is not None: + if tuple(chunks)[-1] != k: + raise ValueError(f"chunks[-1] must equal the trailing vector length {k}") + else: + # ponytail: reuse empty()'s own chunk-sizing heuristic instead of + # reimplementing it, then pin the non-splittable trailing dim to its full length. + b2_kwargs["chunks"] = blosc2.empty(full_shape, dtype=dtype).chunks[:-1] + (k,) + + dst = blosc2.empty(full_shape, dtype=dtype, **b2_kwargs) + return self._fill_chunks( + dst, method, args, kwargs, lambda sl: tuple(s.stop - s.start for s in sl[:-1]) + ) + + def random(self, shape, dtype=np.float64, **kwargs: Any) -> blosc2.NDArray: + """Draw uniform floats in [0, 1). Mirrors :meth:`numpy.random.Generator.random`.""" + return self._fill("random", (), {"dtype": dtype}, shape, **kwargs) + + def integers( + self, low, high=None, *, shape, dtype=np.int64, endpoint: bool = False, **kwargs: Any + ) -> blosc2.NDArray: + """Draw random integers. Mirrors :meth:`numpy.random.Generator.integers`.""" + return self._fill("integers", (low, high), {"dtype": dtype, "endpoint": endpoint}, shape, **kwargs) + + def normal(self, loc=0.0, scale=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a normal distribution. Mirrors :meth:`numpy.random.Generator.normal`.""" + return self._fill("normal", (loc, scale), {}, shape, **kwargs) + + def uniform(self, low=0.0, high=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a uniform distribution. Mirrors :meth:`numpy.random.Generator.uniform`.""" + return self._fill("uniform", (low, high), {}, shape, **kwargs) + + def beta(self, a, b, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a beta distribution. Mirrors :meth:`numpy.random.Generator.beta`.""" + return self._fill("beta", (a, b), {}, shape, **kwargs) + + def binomial(self, n, p, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a binomial distribution. Mirrors :meth:`numpy.random.Generator.binomial`.""" + return self._fill("binomial", (n, p), {}, shape, **kwargs) + + def chisquare(self, df, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a chi-square distribution. Mirrors :meth:`numpy.random.Generator.chisquare`.""" + return self._fill("chisquare", (df,), {}, shape, **kwargs) + + def exponential(self, scale=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from an exponential distribution. Mirrors :meth:`numpy.random.Generator.exponential`.""" + return self._fill("exponential", (scale,), {}, shape, **kwargs) + + def f(self, dfnum, dfden, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from an F distribution. Mirrors :meth:`numpy.random.Generator.f`.""" + return self._fill("f", (dfnum, dfden), {}, shape, **kwargs) + + def gamma(self, shape_param, /, scale=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a gamma distribution. Mirrors :meth:`numpy.random.Generator.gamma`. + + ``shape_param`` is numpy's ``shape`` (the gamma distribution's shape parameter, + usually called *k*); renamed here, and made positional-only, to avoid colliding + with the array-shape ``shape`` keyword. + """ + return self._fill("gamma", (shape_param, scale), {}, shape, **kwargs) + + def geometric(self, p, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a geometric distribution. Mirrors :meth:`numpy.random.Generator.geometric`.""" + return self._fill("geometric", (p,), {}, shape, **kwargs) + + def gumbel(self, loc=0.0, scale=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a Gumbel distribution. Mirrors :meth:`numpy.random.Generator.gumbel`.""" + return self._fill("gumbel", (loc, scale), {}, shape, **kwargs) + + def hypergeometric(self, ngood, nbad, nsample, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a hypergeometric distribution. Mirrors :meth:`numpy.random.Generator.hypergeometric`.""" + return self._fill("hypergeometric", (ngood, nbad, nsample), {}, shape, **kwargs) + + def laplace(self, loc=0.0, scale=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a Laplace distribution. Mirrors :meth:`numpy.random.Generator.laplace`.""" + return self._fill("laplace", (loc, scale), {}, shape, **kwargs) + + def logistic(self, loc=0.0, scale=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a logistic distribution. Mirrors :meth:`numpy.random.Generator.logistic`.""" + return self._fill("logistic", (loc, scale), {}, shape, **kwargs) + + def lognormal(self, mean=0.0, sigma=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a log-normal distribution. Mirrors :meth:`numpy.random.Generator.lognormal`.""" + return self._fill("lognormal", (mean, sigma), {}, shape, **kwargs) + + def logseries(self, p, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a logarithmic series distribution. Mirrors :meth:`numpy.random.Generator.logseries`.""" + return self._fill("logseries", (p,), {}, shape, **kwargs) + + def negative_binomial(self, n, p, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a negative binomial distribution. Mirrors :meth:`numpy.random.Generator.negative_binomial`.""" + return self._fill("negative_binomial", (n, p), {}, shape, **kwargs) + + def noncentral_chisquare(self, df, nonc, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a noncentral chi-square distribution. Mirrors :meth:`numpy.random.Generator.noncentral_chisquare`.""" + return self._fill("noncentral_chisquare", (df, nonc), {}, shape, **kwargs) + + def noncentral_f(self, dfnum, dfden, nonc, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a noncentral F distribution. Mirrors :meth:`numpy.random.Generator.noncentral_f`.""" + return self._fill("noncentral_f", (dfnum, dfden, nonc), {}, shape, **kwargs) + + def pareto(self, a, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a Pareto distribution. Mirrors :meth:`numpy.random.Generator.pareto`.""" + return self._fill("pareto", (a,), {}, shape, **kwargs) + + def poisson(self, lam=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a Poisson distribution. Mirrors :meth:`numpy.random.Generator.poisson`.""" + return self._fill("poisson", (lam,), {}, shape, **kwargs) + + def power(self, a, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a power distribution. Mirrors :meth:`numpy.random.Generator.power`.""" + return self._fill("power", (a,), {}, shape, **kwargs) + + def rayleigh(self, scale=1.0, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a Rayleigh distribution. Mirrors :meth:`numpy.random.Generator.rayleigh`.""" + return self._fill("rayleigh", (scale,), {}, shape, **kwargs) + + def standard_cauchy(self, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a standard Cauchy distribution. Mirrors :meth:`numpy.random.Generator.standard_cauchy`.""" + return self._fill("standard_cauchy", (), {}, shape, **kwargs) + + def standard_exponential( + self, *, shape, dtype=np.float64, method="zig", **kwargs: Any + ) -> blosc2.NDArray: + """Draw samples from a standard exponential distribution. + + Mirrors :meth:`numpy.random.Generator.standard_exponential`. + """ + return self._fill("standard_exponential", (), {"dtype": dtype, "method": method}, shape, **kwargs) + + def standard_gamma(self, shape_param, /, *, shape, dtype=np.float64, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a standard gamma distribution. + + Mirrors :meth:`numpy.random.Generator.standard_gamma`; see :meth:`gamma` for why the + distribution's ``shape`` parameter is renamed ``shape_param`` here. + """ + return self._fill("standard_gamma", (shape_param,), {"dtype": dtype}, shape, **kwargs) + + def standard_normal(self, *, shape, dtype=np.float64, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a standard normal distribution. + + Mirrors :meth:`numpy.random.Generator.standard_normal`. + """ + return self._fill("standard_normal", (), {"dtype": dtype}, shape, **kwargs) + + def standard_t(self, df, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a standard Student's t distribution. Mirrors :meth:`numpy.random.Generator.standard_t`.""" + return self._fill("standard_t", (df,), {}, shape, **kwargs) + + def triangular(self, left, mode, right, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a triangular distribution. Mirrors :meth:`numpy.random.Generator.triangular`.""" + return self._fill("triangular", (left, mode, right), {}, shape, **kwargs) + + def vonmises(self, mu, kappa, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a von Mises distribution. Mirrors :meth:`numpy.random.Generator.vonmises`.""" + return self._fill("vonmises", (mu, kappa), {}, shape, **kwargs) + + def wald(self, mean, scale, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a Wald distribution. Mirrors :meth:`numpy.random.Generator.wald`.""" + return self._fill("wald", (mean, scale), {}, shape, **kwargs) + + def weibull(self, a, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a Weibull distribution. Mirrors :meth:`numpy.random.Generator.weibull`.""" + return self._fill("weibull", (a,), {}, shape, **kwargs) + + def zipf(self, a, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a Zipf distribution. Mirrors :meth:`numpy.random.Generator.zipf`.""" + return self._fill("zipf", (a,), {}, shape, **kwargs) + + def choice(self, a, *, shape, p=None, replace: bool = True, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from ``a`` with replacement. Mirrors :meth:`numpy.random.Generator.choice`. + + Only ``replace=True`` (numpy's default) and 1-D or scalar-int ``a`` are + supported: sampling without replacement, or along an axis of a multi-dimensional + ``a``, needs whole-array coordination across chunks that this module doesn't do. + """ + if not replace: + raise NotImplementedError( + "blosc2.random.Generator.choice only supports replace=True: sampling " + "without replacement needs whole-array coordination across chunks" + ) + if np.ndim(a) > 1: + raise NotImplementedError("blosc2.random.Generator.choice only supports 1-D or scalar `a`") + return self._fill("choice", (a,), {"p": p, "replace": True}, shape, **kwargs) + + def dirichlet(self, alpha, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a Dirichlet distribution. Mirrors :meth:`numpy.random.Generator.dirichlet`. + + Output shape is ``shape + (len(alpha),)``; see :meth:`_fill_vector` — the + trailing dimension holds one full draw and is never split across chunks. + """ + return self._fill_vector("dirichlet", (alpha,), {}, shape, len(alpha), **kwargs) + + def multinomial(self, n, pvals, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a multinomial distribution. Mirrors :meth:`numpy.random.Generator.multinomial`. + + Output shape is ``shape + (len(pvals),)``; see :meth:`dirichlet` for the + trailing-dimension note. + """ + return self._fill_vector("multinomial", (n, pvals), {}, shape, len(pvals), **kwargs) + + def multivariate_hypergeometric(self, colors, nsample, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a multivariate hypergeometric distribution. + + Mirrors :meth:`numpy.random.Generator.multivariate_hypergeometric`. Output shape is + ``shape + (len(colors),)``; see :meth:`dirichlet` for the trailing-dimension note. + """ + return self._fill_vector( + "multivariate_hypergeometric", (colors, nsample), {}, shape, len(colors), **kwargs + ) + + def multivariate_normal(self, mean, cov, *, shape, **kwargs: Any) -> blosc2.NDArray: + """Draw samples from a multivariate normal distribution. + + Mirrors :meth:`numpy.random.Generator.multivariate_normal`. Output shape is + ``shape + (len(mean),)``; see :meth:`dirichlet` for the trailing-dimension note. + """ + return self._fill_vector("multivariate_normal", (mean, cov), {}, shape, len(mean), **kwargs) + + def permutation(self, x, *, axis: int = 0, **kwargs: Any) -> blosc2.NDArray: + """Return a shuffled copy of ``x`` (or of ``arange(x)`` if ``x`` is an int). + + Mirrors :meth:`numpy.random.Generator.permutation`. + + Unlike the rest of this module, permuting is inherently sequential: it loads + ``x`` fully into memory and shuffles it single-threaded, rather than generating + chunks independently in parallel. + """ + rng = np.random.default_rng(self._root.spawn(1)[0]) + arr = np.arange(x) if isinstance(x, int | np.integer) else _materialize(x) + return blosc2.asarray(rng.permutation(arr, axis=axis), **kwargs) + + def permuted(self, x, *, axis: int | None = None, **kwargs: Any) -> blosc2.NDArray: + """Return a copy of ``x`` with elements permuted along ``axis`` (flattened if ``None``). + + Mirrors :meth:`numpy.random.Generator.permuted`. See :meth:`permutation` for the + single-threaded, full-materialization caveat. + """ + rng = np.random.default_rng(self._root.spawn(1)[0]) + return blosc2.asarray(rng.permuted(_materialize(x), axis=axis), **kwargs) + + def shuffle(self, x: blosc2.NDArray, *, axis: int = 0) -> None: + """Shuffle ``x`` in place along ``axis``. Mirrors :meth:`numpy.random.Generator.shuffle`. + + ``x`` must be a :class:`~blosc2.NDArray`: its full contents are read into memory, + shuffled single-threaded (see :meth:`permutation`), and written back. Returns + ``None``, matching numpy. + """ + if not isinstance(x, blosc2.NDArray): + raise TypeError("blosc2.random.Generator.shuffle requires a blosc2.NDArray to shuffle in place") + rng = np.random.default_rng(self._root.spawn(1)[0]) + arr = x[:] + rng.shuffle(arr, axis=axis) + x[:] = arr + + +def default_rng(seed=None) -> Generator: + """Construct a chunk-parallel :class:`Generator`, mirroring :func:`numpy.random.default_rng`. + + Parameters + ---------- + seed: int, array-like, SeedSequence or None + Seed for the root :class:`numpy.random.SeedSequence`. ``None`` (default) draws + entropy from the OS. + + Returns + ------- + out: :class:`Generator` + """ + return Generator(seed) diff --git a/tests/test_random.py b/tests/test_random.py new file mode 100644 index 00000000..b5fdf9c0 --- /dev/null +++ b/tests/test_random.py @@ -0,0 +1,303 @@ +####################################################################### +# Copyright (c) 2019-present, Blosc Development Team +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +####################################################################### + +import numpy as np +import pytest + +import blosc2 + + +def test_reproducible_same_seed(): + a = blosc2.random.default_rng(42).random((100, 100), chunks=(17, 17)) + b = blosc2.random.default_rng(42).random((100, 100), chunks=(17, 17)) + np.testing.assert_array_equal(a[:], b[:]) + + +def test_successive_calls_differ(): + rng = blosc2.random.default_rng(42) + a = rng.random((1000,)) + b = rng.random((1000,)) + assert not np.array_equal(a[:], b[:]) + + +def test_different_seeds_differ(): + a = blosc2.random.default_rng(1).random((1000,)) + b = blosc2.random.default_rng(2).random((1000,)) + assert not np.array_equal(a[:], b[:]) + + +def test_seed_none_works(): + a = blosc2.random.default_rng(None).random((10,)) + assert a.shape == (10,) + + +@pytest.mark.parametrize("nthreads", [1, 2]) +def test_random_distribution_sanity(nthreads, monkeypatch): + monkeypatch.setattr(blosc2, "nthreads", nthreads) + a = blosc2.random.default_rng(0).random((200, 200), chunks=(31, 31))[:] + assert a.dtype == np.float64 + assert a.min() >= 0.0 + assert a.max() < 1.0 + assert abs(a.mean() - 0.5) < 0.05 + + +def test_normal_distribution_sanity(): + a = blosc2.random.default_rng(0).normal(loc=5.0, scale=2.0, shape=(500, 500))[:] + assert abs(a.mean() - 5.0) < 0.1 + assert abs(a.std() - 2.0) < 0.1 + + +def test_integers_bounds_dtype_endpoint(): + a = blosc2.random.default_rng(0).integers( + 0, 10, shape=(1000,), dtype=np.int32, endpoint=True, chunks=(100,) + )[:] + assert a.dtype == np.int32 + assert a.min() >= 0 + assert a.max() <= 10 + + +def test_uniform_bounds(): + a = blosc2.random.default_rng(0).uniform(-3.0, 3.0, shape=(1000,))[:] + assert a.min() >= -3.0 + assert a.max() < 3.0 + + +def test_blosc2_kwargs_passthrough(tmp_path): + urlpath = tmp_path / "r.b2nd" + a = blosc2.random.default_rng(0).random( + (100, 100), chunks=(20, 20), cparams={"clevel": 5}, urlpath=str(urlpath) + ) + assert a.chunks == (20, 20) + assert a.schunk.cparams.clevel == 5 + assert urlpath.exists() + b = blosc2.open(str(urlpath)) + np.testing.assert_array_equal(a[:], b[:]) + + +def test_nd_shape(): + a = blosc2.random.default_rng(0).random((4, 5, 6), chunks=(2, 3, 3)) + assert a.shape == (4, 5, 6) + + +def test_shape_required(): + with pytest.raises(TypeError): + blosc2.random.default_rng(0).normal() + + +def test_scheduling_independent_of_nthreads(monkeypatch): + monkeypatch.setattr(blosc2, "nthreads", 1) + serial = blosc2.random.default_rng(7).random((300, 300), chunks=(17, 23))[:] + monkeypatch.setattr(blosc2, "nthreads", 8) + parallel = blosc2.random.default_rng(7).random((300, 300), chunks=(17, 23))[:] + np.testing.assert_array_equal(serial, parallel) + + +def test_falls_back_to_serial_when_threads_unavailable(monkeypatch): + # Some sandboxes (e.g. wasm32/Pyodide without pthreads) report a real nthreads + # count but can't actually start OS threads: ThreadPoolExecutor.submit() raises + # RuntimeError. _bounded_map must catch that and fall back to serial generation + # instead of propagating it (regression for the CI failure this triggered). + class _NoThreadsExecutor: + def __init__(self, max_workers=None): + pass + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def submit(self, fn, *args): + raise RuntimeError("can't start new thread") + + monkeypatch.setattr(blosc2.random, "ThreadPoolExecutor", _NoThreadsExecutor) + monkeypatch.setattr(blosc2, "nthreads", 8) + fallback = blosc2.random.default_rng(7).random((300, 300), chunks=(17, 23))[:] + + # nthreads=1 takes the serial fast path directly, never touching + # ThreadPoolExecutor, so this is a trustworthy baseline despite the patch above. + monkeypatch.setattr(blosc2, "nthreads", 1) + serial = blosc2.random.default_rng(7).random((300, 300), chunks=(17, 23))[:] + np.testing.assert_array_equal(fallback, serial) + + +_DIST_CASES = { + "beta": (2.0, 5.0), + "binomial": (10, 0.3), + "chisquare": (3.0,), + "exponential": (2.0,), + "f": (5.0, 2.0), + "gamma": (2.0, 3.0), + "geometric": (0.3,), + "gumbel": (0.0, 1.0), + "hypergeometric": (10, 10, 5), + "laplace": (0.0, 1.0), + "logistic": (0.0, 1.0), + "lognormal": (0.0, 1.0), + "logseries": (0.5,), + "negative_binomial": (5, 0.5), + "noncentral_chisquare": (3.0, 2.0), + "noncentral_f": (5.0, 2.0, 1.0), + "pareto": (3.0,), + "poisson": (4.0,), + "power": (2.0,), + "rayleigh": (1.0,), + "standard_cauchy": (), + "standard_exponential": (), + "standard_gamma": (2.0,), + "standard_normal": (), + "standard_t": (5.0,), + "triangular": (0.0, 0.5, 1.0), + "vonmises": (0.0, 4.0), + "wald": (1.0, 1.0), + "weibull": (1.5,), + "zipf": (2.0,), +} + + +@pytest.mark.parametrize(("method", "args"), _DIST_CASES.items(), ids=_DIST_CASES.keys()) +def test_distribution_reproducible_and_finite(method, args): + def draw(): + rng = getattr(blosc2.random.default_rng(0), method) + return rng(*args, shape=(200,), chunks=(37,))[:] + + a, b = draw(), draw() + np.testing.assert_array_equal(a, b) + assert np.all(np.isfinite(a)) + + +def test_poisson_mean_sanity(): + a = blosc2.random.default_rng(0).poisson(4.0, shape=(5000,))[:] + assert abs(a.mean() - 4.0) < 0.2 + + +def test_standard_normal_dtype(): + a = blosc2.random.default_rng(0).standard_normal(shape=(10,), dtype=np.float32) + assert a.dtype == np.float32 + + +def test_choice_int_a_bounds_and_reproducible(): + a = blosc2.random.default_rng(0).choice(5, shape=(1000,), chunks=(137,)) + b = blosc2.random.default_rng(0).choice(5, shape=(1000,), chunks=(137,)) + np.testing.assert_array_equal(a[:], b[:]) + assert a[:].min() >= 0 + assert a[:].max() < 5 + + +def test_choice_array_a(): + a = blosc2.random.default_rng(0).choice(np.array(["x", "y", "z"]), shape=(20,)) + assert set(a[:].tolist()) <= {"x", "y", "z"} + + +def test_choice_replace_false_not_implemented(): + with pytest.raises(NotImplementedError): + blosc2.random.default_rng(0).choice(5, shape=(10,), replace=False) + + +def test_choice_2d_a_not_implemented(): + with pytest.raises(NotImplementedError): + blosc2.random.default_rng(0).choice(np.zeros((3, 3)), shape=(10,)) + + +_VECTOR_DIST_CASES = { + "dirichlet": (([1.0, 2.0, 3.0],), 3), + "multinomial": ((10, [0.2, 0.3, 0.5]), 3), + "multivariate_hypergeometric": (([3, 4, 5], 6), 3), + "multivariate_normal": (([0.0, 0.0], [[1.0, 0.0], [0.0, 1.0]]), 2), +} + + +@pytest.mark.parametrize( + ("method", "args", "k"), + [(m, a, k) for m, (a, k) in _VECTOR_DIST_CASES.items()], + ids=_VECTOR_DIST_CASES.keys(), +) +def test_vector_distribution_shape_reproducible_and_finite(method, args, k): + def draw(): + rng = getattr(blosc2.random.default_rng(0), method) + return rng(*args, shape=(40,)) + + a = draw() + assert a.shape == (40, k) + assert a.chunks[-1] == k + b = draw() + np.testing.assert_array_equal(a[:], b[:]) + assert np.all(np.isfinite(a[:])) + + +def test_dirichlet_sums_to_one(): + a = blosc2.random.default_rng(0).dirichlet([1.0, 2.0, 3.0], shape=(50,), chunks=(7, 3))[:] + np.testing.assert_allclose(a.sum(axis=-1), 1.0) + + +def test_vector_dist_chunks_mismatch_raises(): + with pytest.raises(ValueError): + blosc2.random.default_rng(0).dirichlet([1.0, 2.0, 3.0], shape=(50,), chunks=(7, 2)) + + +def test_vector_dist_scalar_leading_shape(): + a = blosc2.random.default_rng(0).dirichlet([1.0, 2.0, 3.0], shape=()) + assert a.shape == (3,) + + +def test_vector_dist_numpy_integer_shape(): + a = blosc2.random.default_rng(0).dirichlet([1.0, 2.0, 3.0], shape=np.int64(5)) + assert a.shape == (5, 3) + + +def test_permutation_int_is_a_permutation_and_reproducible(): + a = blosc2.random.default_rng(0).permutation(10) + b = blosc2.random.default_rng(0).permutation(10) + np.testing.assert_array_equal(a[:], b[:]) + assert sorted(a[:].tolist()) == list(range(10)) + + +def test_permutation_successive_calls_differ(): + rng = blosc2.random.default_rng(0) + a = rng.permutation(1000) + b = rng.permutation(1000) + assert not np.array_equal(a[:], b[:]) + + +def test_permutation_array_axis(): + arr = blosc2.arange(12).reshape((3, 4)) + p = blosc2.random.default_rng(1).permutation(arr, axis=1) + assert p.shape == (3, 4) + for row_before, row_after in zip(arr[:], p[:], strict=True): + assert sorted(row_after.tolist()) == sorted(row_before.tolist()) + + +def test_permutation_kwargs_passthrough(): + p = blosc2.random.default_rng(0).permutation(1000, chunks=(100,), cparams={"clevel": 5}) + assert p.chunks == (100,) + assert p.schunk.cparams.clevel == 5 + + +def test_permuted_is_a_permutation(): + a = blosc2.random.default_rng(2).permuted(np.arange(10)) + assert sorted(a[:].tolist()) == list(range(10)) + + +def test_permuted_flattens_when_axis_none(): + a = blosc2.random.default_rng(2).permuted(blosc2.arange(12).reshape((3, 4))) + assert a.shape == (3, 4) + assert sorted(a[:].flatten().tolist()) == list(range(12)) + + +def test_shuffle_in_place_preserves_multiset(): + a = blosc2.arange(10) + before = a[:].copy() + result = blosc2.random.default_rng(3).shuffle(a) + after = a[:] + assert result is None + assert sorted(before.tolist()) == sorted(after.tolist()) + assert not np.array_equal(before, after) + + +def test_shuffle_requires_ndarray(): + with pytest.raises(TypeError): + blosc2.random.default_rng(3).shuffle(np.arange(5))