From c781dd8103f72a922c4ca4f1d1294af8f1ec7510 Mon Sep 17 00:00:00 2001 From: Marcus Date: Wed, 5 Aug 2026 18:15:03 -0700 Subject: [PATCH] Unify the gp2Scale covariance computation behind one primitive The prior covariance, the blocks added by an append, and the posterior's cross-covariance were three different code paths at three levels of support: two near-duplicate distributed functions in GPprior, and nothing at all for the posterior, which called the kernel directly and so built a dense (N x n_pred) array on the client. They are all the same operation over a block grid, differing only in whether the two point sets are the same. distributed_covariance in the new gp2Scale_covariance module is that operation; GPprior._gp2Scale_covariance is its only caller and owns scatter lifetime and nothing else. Falling out of the unification: * A four-argument, args-taking kernel now works under gp2Scale. The old workers called kernel(x1, x2, hps) unconditionally, so a signature supported everywhere else in fvGP raised TypeError on the worker. * The posterior cross-covariance is distributed and stays sparse, so posterior_mean never materializes (N x n_pred). posterior_covariance cannot avoid a dense solve, KV^-1 being dense whatever KV is, so it chunks over prediction points to cap the intermediate at (N x chunk). * The joint-covariance methods (joint_gp_prior, gp_mutual_information, gp_total_correlation) are dense in N by construction; they now densify K explicitly and warn, rather than failing somewhere inside numpy. Add a row-wise distribution alongside the existing block-wise one, selected with gp2Scale_distribution. Row-wise has each worker return a finished CSR row strip, so the COO-to-CSR sort happens in parallel and host assembly is a concatenation with no global COO and no mirroring; it cannot exploit symmetry and so doubles the kernel evaluations. Block-wise remains the default. Assembly speedups: the per-block diagonal mask was a Python list comprehension over every nonzero, now vectorized (97.6 ms -> 0.34 ms on 1.1M nonzeros, and it ran once per diagonal block per likelihood evaluation); empty blocks are skipped rather than shipped; indices are int32 where the matrix allows; the mirrored triplets are written into a single preallocation instead of two np.hstack passes, freeing each gathered block as it is copied; and the augmented matrix is assembled from all-CSR blocks so scipy takes its fast path instead of converting the lot to COO. Fix the scatter race documented as unavoidable. Dask keys scattered data by a hash of its content, so a second GP on the same data, or a second prediction at the same points, lands on one key and the first copy's release races the second scatter inside the scheduler; the tasks then return CancelledError or KeyError. Every scatter now uses hash=False and so gets its own key. _harvest also raises on an exception result rather than letting it reach the assembly and fail there as an unrelated type error. Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 54 +++++- fvgp/fvgp.py | 12 ++ fvgp/gp.py | 12 ++ fvgp/gp2Scale_covariance.py | 324 ++++++++++++++++++++++++++++++++++++ fvgp/gp_posterior.py | 83 +++++++-- fvgp/gp_prior.py | 295 +++++++++++--------------------- tests/test_fvgp.py | 166 ++++++++++++++++++ 7 files changed, 725 insertions(+), 221 deletions(-) create mode 100644 fvgp/gp2Scale_covariance.py diff --git a/CLAUDE.md b/CLAUDE.md index ff09e68..48829cc 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -14,6 +14,9 @@ pytest tests/ # Run a single test (tests are top-level functions, not class methods) pytest tests/test_fvgp.py::test_single_task_init_basic +# Skip the slow distributed test while iterating +pytest tests/ -q --deselect tests/test_fvgp.py::test_gp2Scale + # Run tests with coverage pytest tests --cov=./ --cov-report=xml @@ -24,6 +27,10 @@ flake8 fvgp tests make docs ``` +There is no `setup.py` (build backend is hatchling + hatch-vcs, version written to [fvgp/_version.py](fvgp/_version.py)), so the `test`, `coverage`, `install`, and `dist` Makefile targets and `tox.ini` are stale — only `make docs` / `make lint` work. Use `pytest` and `hatch build` directly. + +There is no `conftest.py`: [tests/test_fvgp.py](tests/test_fvgp.py) imports the `client` / `loop` / `cluster_fixture` fixtures from `distributed.utils_test` at module level, so any test taking a `client` argument spins up a real local Dask cluster. + ## Architecture fvGP is a Gaussian Process library optimized for large-scale and multi-task settings. The two public-facing classes are: @@ -41,7 +48,7 @@ Both classes are composed of internal specialist objects created at `__init__` t | `GPkv` | [gp_kv.py](fvgp/gp_kv.py) | Owns K+V matrix state and all factorizations; dispatches solves/logdets across linalg modes | | `GPMarginalLikelihood` | [gp_marginal_likelihood.py](fvgp/gp_marginal_likelihood.py) | Log marginal likelihood and its gradient; delegates factorization to `GPkv`. Maintains `_warm_start_KVinvY` for iterative training solves when `args["sparse_krylov_warm_start"]=True`. | | `GPposterior` | [gp_posterior.py](fvgp/gp_posterior.py) | Posterior mean/covariance; information-theoretic quantities | -| `GPtraining` | [gp_training.py](fvgp/gp_training.py) | Hyperparameter optimization (scipy, hgdl async, MCMC, Adam) | +| `GPtraining` | [gp_training.py](fvgp/gp_training.py) | Hyperparameter optimization; owns both the sync (`train`) and async (`train_async`) dispatch over `global` / `local` / `hgdl` / `mcmc` / `adam` / `bo` | ### State propagation @@ -61,25 +68,58 @@ Gotchas: - **`update_gp_data(append=False, rank_n_update=True)`** is invalid (the previous factorization is for data that no longer exists); `GP.update_gp_data` emits a `UserWarning` and forces `rank_n_update=False`. - **`kv.solve(b, x0=...)`** zero-pads `x0` along axis 0 when shapes don't match, so a pre-append `KVinvY` can warm-start the post-append solve in iterative modes (sparseCG/MINRES/preconditioned variants). See [gp_kv.py:333-342](fvgp/gp_kv.py#L333-L342). +### Hyperparameter training (`GP.train` / `train(asynchronous=True)`) + +`GP.train` validates and normalizes everything (bounds, initial hyperparameters, objective + gradient + Hessian, gp2Scale and async restrictions) and then hands off to `GPtraining.train` or `GPtraining.train_async`. `method` is `"global"`, `"local"`, `"hgdl"`, `"mcmc"` (default), `"adam"`, `"bo"`, or a callable taking the `GP` and returning a hyperparameter vector. + +- Async is supported for `hgdl`, `mcmc`, `adam`, `bo` and needs a `dask_client`; everything else warns and falls back to sync. +- gp2Scale trains synchronously and only allows `mcmc` or `bo` (other methods are silently switched to `mcmc`). +- `mcmc` ignores a user `objective_function` (it always maximizes the log marginal likelihood). A user objective with `local`/`hgdl` must come with a gradient. +- Diagnostics land on read-only `GP` properties: `mcmc_info` after `mcmc`, `bo_info` after a sync `bo` run. + +**`method='bo'` ([gp_bo.py](fvgp/gp_bo.py))** — Bayesian optimization over the hyperparameters, for when the marginal likelihood is expensive, noisy, and effectively gradient-free (the gp2Scale/mBCG regime: stochastic-Lanczos log-determinant plus truncated CG). Points to know before touching it: + +- **The surrogate is an `fvgp.GP`.** It deliberately uses fvGP and not gpCAM, which depends on fvGP — importing gpCAM here would be circular. The recursion bottoms out: the inner GP sees only the tens-to-hundreds of θ points evaluated, uses an ARD Matérn-5/2 kernel with analytic gradients, and trains with `method='local'` — never `'bo'`. +- **`max_iter` changes meaning**: for `bo` it is a *cap on objective-function evaluations*, not iterations. The run normally stops earlier on the `patience` / `f_rtol` / `x_tol` criteria in `bo_args`. +- **Observation noise is wired in automatically.** When the objective is the default, `GP.train` injects a `bo_args["noise_function"]` reading `GPMarginalLikelihood.log_likelihood_variance()` — the SLQ log-determinant's own precision (`0.25 * kv.last_logdet_variance`), or `None` in exact modes, where the surrogate instead learns a single homoscedastic noise level whose lower bound acts as a nugget. +- **The search space is log-transformed per dimension** (`_LogAffineTransform`): log where both bounds are strictly positive, linear otherwise, then rescaled to the unit cube. Positivity is only a *proxy* for being scale-like — a positive hyperparameter that enters additively (a center in a non-stationary/Gibbs kernel, a mixing weight) is hurt by it. Override with `bo_args['log_scale']`. +- `GP._warn_about_bo_suitability` warns before the run when the budget is too small for the initial design or the hyperparameter count is too high for BO. +- `bo_info` carries the payoff beyond the optimum: `sensitivity` (curvature-based ranking of which hyperparameters matter) and `posterior covariance` (Laplace approximation at the mode, in searched coordinates), both free of extra likelihood evaluations. + ### Key supporting modules - **[gp_lin_alg.py](fvgp/gp_lin_alg.py)** — CPU/GPU linear algebra primitives; Cholesky, LU, sparse solvers; defines `NonPositiveDefiniteError` - **[gp_kv.py](fvgp/gp_kv.py)** — `GPkv` manages all K+V state across linalg modes: `"Chol"`, `"CholInv"`, `"Inv"`, `"sparseLU"`, `"sparseCG"`, `"sparseMINRES"`, and preconditioned variants. The mode is set at init and determines which factorization is updated when data or hyperparameters change. Custom solvers can be injected as a 3-tuple of callables. For `sparseMINRESpre`/`sparseCGpre`, `GPkv` caches the preconditioner across `update_KV` / `compute_new_*` calls and rebuilds when `Preconditioner_reuse_counter` ≥ `args["sparse_preconditioner_refresh_interval"] - 1` or when the shape/`sparse_preconditioner_*` args fingerprint changes. `set_KV` always force-refreshes. Aliases like `"sparseCGpre_amg"` are resolved at `__init__` into the canonical mode plus `args["sparse_preconditioner_type"]`. - **[kernels.py](fvgp/kernels.py)** — 15+ built-in kernels including Matérn, squared exponential, Wendland (compactly supported) - **[gp_mcmc.py](fvgp/gp_mcmc.py)** — Adaptive Metropolis–Hastings sampler used for Bayesian hyperparameter inference -- **[gp_actor.py](fvgp/gp_actor.py)** — `AsyncOptimizer` wraps `_MCMCActor` and `_AdamActor` for non-blocking background training; used by `GPtraining` for async MCMC and Adam modes +- **[gp_actor.py](fvgp/gp_actor.py)** — `AsyncOptimizer` wraps `_MCMCActor`, `_AdamActor`, and `_BOActor` for non-blocking background training; used by `GPtraining` for the async MCMC, Adam, and BO modes +- **[gp_bo.py](fvgp/gp_bo.py)** — `bayesian_optimize` and the noisy-EI machinery behind `method='bo'`; see the training section above. Its module docstring is the design rationale and is worth reading before changing anything here. +- **[utils.py](fvgp/utils.py)** — `log_time` context manager for cumulative timing via loguru. Note fvGP calls `logger.disable('fvgp')` at import ([`__init__.py`](fvgp/__init__.py)); re-enable it to get the debug stream. ### Scaling to large datasets (`gp2Scale`) When `gp2Scale=True`, `GP` switches to a Wendland (compactly supported) kernel producing sparse covariance matrices and uses Dask for distributed computation. This path requires a Dask client to be passed in and uses sparse linear solvers instead of dense Cholesky. +**One distributed covariance primitive.** Every kernel evaluation gp2Scale distributes goes through `distributed_covariance` in [gp2Scale_covariance.py](fvgp/gp2Scale_covariance.py) — the symmetric prior covariance, the rectangular `B` and symmetric `D` blocks of an append, and the posterior's `k(x_data, x_pred)`. They differ only in the `symmetric` flag. `GPprior._gp2Scale_covariance` is the sole caller; it owns scatter lifetime and nothing else. The module's docstring is the design rationale. + +Two ways of cutting the work, via `GP(..., gp2Scale_distribution=...)`: + +| | tasks | kernel evaluations | host assembly | +|---|---|---|---| +| `"blockwise"` (default) | (row block, col block) pairs, upper triangle only when symmetric | half | global COO + mirror, one preallocation | +| `"rowwise"` | row strips; workers return finished CSR | double (no symmetry) | `vstack` of strips — concatenation only | + +Row-wise is the choice when host assembly, not kernel evaluation, is the bottleneck; it also caps host peak memory at the finished matrix plus one strip. + +**Posterior at scale:** `GPposterior.cross_covariance` → `GPprior.compute_data_cross_covariance` returns a **sparse** `k`, so `posterior_mean` never materializes an `(N × n_pred)` array. `posterior_covariance` cannot avoid a dense solve (`KV⁻¹` is dense regardless), so `GPposterior._cross_solve_product` chunks over prediction points at `gp2Scale_batch_size` to cap the intermediate at `(N × chunk)`. `joint_gp_prior`, `joint_gp_prior_grad`, `gp_mutual_information` and `gp_total_correlation` build a joint `(N + n_pred)²` matrix and are dense-in-N by construction; they route `K` through `_dense_K()`, which warns under gp2Scale. `posterior_covariance_grad` is likewise dense-in-N via `d_kernel_dx` and was left on the direct path. + **Scatter ownership and lifecycle:** -- `GPprior.x_data_scatter_future` is the single persistent dask scatter of the current `x_data`. Scattered once at `GPprior.__init__` (see [gp_prior.py:93-96](fvgp/gp_prior.py#L93-L96)). +- `GPprior.x_data_scatter_future` is the single persistent dask scatter of the current `x_data`. Scattered once at `GPprior.__init__`. - `GPdata` does NOT scatter — it's pure-Python data only. -- `_compute_prior_covariance_gp2Scale` reads `self.x_data_scatter_future` directly; **no scatter per call**, so training stays dask-quiet. -- On data changes, `augment_state_data` / `update_state_data` refresh the scatter by **overwriting** `self.x_data_scatter_future` (no explicit `release()`). The old future loses its only Python ref and is cleaned up via `__del__`. Calling `release()` explicitly schedules a `_dec_ref` that races against subsequent scatter `replicate` operations in the scheduler — don't do it. -- `_update_prior_covariance_gp2Scale` (the augment path) uses `self.x_data_scatter_future` for the `x_old` side (no content-hash collision since it shares the existing key) and scatters only `x_new` locally, releasing that local future at the end. +- The prior covariance reads `self.x_data_scatter_future` directly; **no scatter per call**, so training stays dask-quiet. The append path additionally scatters `x_new`, and the posterior path `x_pred`; both release what they created, in a `finally`. +- **All scatters go through `GPprior._scatter`, which passes `hash=False`.** Dask otherwise keys scattered data by a content hash, so scattering the same array twice — a new GP on the same data, a second prediction at the same points — lands on one key, and the first copy's release races the second scatter inside the scheduler. The tasks then come back as `CancelledError`/`KeyError`. A unique key per scatter removes that collision at the source; `_harvest` in [gp2Scale_covariance.py](fvgp/gp2Scale_covariance.py) additionally raises on an exception result rather than letting it reach the assembly. +- On data changes, `augment_state_data` / `update_state_data` refresh the persistent scatter by **overwriting** it (no explicit `release()`); the old future loses its only Python ref and is cleaned up via `__del__`. **Cross-instance race guard:** [gp.py:14-21](fvgp/gp.py#L14-L21) defines `_GP_INSTANCES_PER_CLIENT`, a `WeakValueDictionary` keyed by `dask_client.id`. `GP.__init__` ([gp.py:285-303](fvgp/gp.py#L285-L303)) raises with a descriptive remediation message if you try to construct a second gp2Scale `GP` on a client that already has a live one — that pattern reliably triggers `FutureCancelledError`/`KeyError` from the scheduler. To reuse a client for a sequence of GPs: @@ -101,6 +141,8 @@ For `sparseCG`, `sparseMINRES`, `sparseCGpre`, and `sparseMINRESpre`, the user c Both default off so existing behavior is preserved. +**Both are gated to `method='mcmc'` during training.** `GP.train` (sync path) wraps the call in `sequential_linalg_state(self.args, method)` from [gp_kv.py](fvgp/gp_kv.py#L30), which temporarily forces `sparse_krylov_warm_start=False` and `sparse_preconditioner_refresh_interval=1` for every other method, warning if that overrides an explicit user setting and restoring it afterwards. Rationale: both mechanisms carry state between likelihood evaluations and are only sound when successive evaluations are close. For a non-local method the leftover residual of a stale-seeded truncated solve makes the likelihood *order-dependent* — a bias, not zero-mean noise, which is exactly what a Bayesian optimizer's noise model cannot absorb. `_SEQUENTIAL_STATE_METHODS` / `_SEQUENTIAL_STATE_DEFAULTS` at the top of `gp_kv.py` hold the policy; the finer per-evaluation checks are `GPkv._validated_warm_start` and `GPkv._can_reuse_sparse_preconditioner`, which discard cached state whenever K+V has actually drifted, whatever the method. + ### Customization API Kernels, mean functions, and noise models are all plain Python callables with standardized signatures. Users pass them as arguments to `GP`/`fvGP` constructors. The full hyperparameter vector is shared across kernel, mean, and noise callables, but each callable must only read its reserved index range. Kernel gradients can be user-supplied or computed via finite differences. diff --git a/fvgp/fvgp.py b/fvgp/fvgp.py index 91a91ed..ce38407 100755 --- a/fvgp/fvgp.py +++ b/fvgp/fvgp.py @@ -156,6 +156,16 @@ class fvGP(GP): The default is False. gp2Scale_batch_size : int, optional Matrix batch size for distributed computing in gp2Scale. The default is 10000. + gp2Scale_distribution : str, optional + How the covariance computation is cut across the workers in gp2Scale. + ``"blockwise"`` (default) sends (row block, column block) pairs and, for the + symmetric prior covariance, schedules only the upper triangle, so the cluster + performs half the kernel evaluations and the host mirrors the result. + ``"rowwise"`` sends whole row strips and has each worker return a finished sparse + strip, so the assembly is a concatenation rather than a global re-sort on the + host. Row-wise cannot exploit symmetry and so doubles the kernel evaluations, but + it removes the host as a bottleneck and lowers its peak memory considerably; it is + the better choice when assembly, not kernel evaluation, dominates the run time. dask_client : dask.distributed.Client, optional A dask client for gp2Scale, asynchronous training,a nd certain linear algebra operations. On HPC architecture, this client is provided by the job script. Please have a look at the examples. @@ -438,6 +448,7 @@ def __init__( gp2Scale=False, dask_client=None, gp2Scale_batch_size=10000, + gp2Scale_distribution="blockwise", linalg_mode=None, ram_economy=False, args=None @@ -470,6 +481,7 @@ def __init__( gp2Scale=gp2Scale, dask_client=dask_client, gp2Scale_batch_size=gp2Scale_batch_size, + gp2Scale_distribution=gp2Scale_distribution, linalg_mode=linalg_mode, ram_economy=ram_economy, args=args) diff --git a/fvgp/gp.py b/fvgp/gp.py index a907924..64e90f4 100755 --- a/fvgp/gp.py +++ b/fvgp/gp.py @@ -156,6 +156,16 @@ class GP: The default is False. gp2Scale_batch_size : int, optional Matrix batch size for distributed computing in gp2Scale. The default is 10000. + gp2Scale_distribution : str, optional + How the covariance computation is cut across the workers in gp2Scale. + ``"blockwise"`` (default) sends (row block, column block) pairs and, for the + symmetric prior covariance, schedules only the upper triangle, so the cluster + performs half the kernel evaluations and the host mirrors the result. + ``"rowwise"`` sends whole row strips and has each worker return a finished sparse + strip, so the assembly is a concatenation rather than a global re-sort on the + host. Row-wise cannot exploit symmetry and so doubles the kernel evaluations, but + it removes the host as a bottleneck and lowers its peak memory considerably; it is + the better choice when assembly, not kernel evaluation, dominates the run time. dask_client : dask.distributed.Client, optional A dask client for gp2Scale, asynchronous training,a nd certain linear algebra operations. On HPC architecture, this client is provided by the job script. Please have a look at the examples. @@ -377,6 +387,7 @@ def __init__( gp2Scale=False, dask_client=None, gp2Scale_batch_size=10000, + gp2Scale_distribution="blockwise", linalg_mode=None, ram_economy=False, args=None @@ -468,6 +479,7 @@ def __init__( kernel_grad=kernel_function_grad, prior_mean_function_grad=prior_mean_function_grad, gp2Scale_batch_size=gp2Scale_batch_size, + gp2Scale_distribution=gp2Scale_distribution, ) ######################################## ###init likelihood instance [tier 3]#### diff --git a/fvgp/gp2Scale_covariance.py b/fvgp/gp2Scale_covariance.py new file mode 100644 index 0000000..f95b7c0 --- /dev/null +++ b/fvgp/gp2Scale_covariance.py @@ -0,0 +1,324 @@ +"""Distributed assembly of sparse covariance matrices for gp2Scale. + +Every kernel evaluation gp2Scale distributes goes through +:py:func:`distributed_covariance`: the symmetric prior covariance, the rectangular and +symmetric blocks needed when data is appended, and the training-set cross-covariance the +posterior needs. They differ only in whether the two point sets are the same +(``symmetric``), which is a flag rather than a separate code path. + +Two ways of cutting the work across the cluster are available, chosen with +``distribution``: + +``"blockwise"`` + Tasks are (row block, column block) pairs. When ``symmetric`` only the upper + triangle is scheduled, so the cluster does half the kernel evaluations, and the host + mirrors the result. This is the historical behavior and remains the default. + +``"rowwise"`` + Tasks are row strips, and each worker returns a *finished* CSR strip. The COO-to-CSR + sort therefore happens in parallel on the workers, and host assembly is a plain + ``vstack`` -- a concatenation of ``data``/``indices`` with an offset ``indptr``, with + no global COO and no mirroring. Symmetry cannot be exploited, so the cluster does + twice the kernel evaluations; in exchange the host stops being the bottleneck and its + peak memory drops to the finished matrix plus one strip. + +Callers own their scatter futures. Nothing here scatters or releases, which keeps the +scatter-lifecycle rules in :py:class:`fvgp.gp_prior.GPprior` where they are documented. +""" + +import itertools +import time +from functools import partial + +import dask.distributed as distributed +import numpy as np +import scipy.sparse as sparse +from loguru import logger + +_DISTRIBUTIONS = ("blockwise", "rowwise") + + +def ranges(N, nb): + """Split ``range(N)`` into ``nb`` chunks given as ``(start, end)`` tuples.""" + if nb == 0: nb = 1 + step = N / nb + return [(round(step * i), round(step * (i + 1))) for i in range(nb)] + + +def num_blocks(n, batch_size): + """Number of chunks ``n`` points are cut into at ``batch_size`` points per chunk.""" + return max(1, n // batch_size) + + +def index_dtype_for(n1, n2): + """int32 indices whenever the matrix is small enough for them. + + Halves both the bytes serialized back from every worker and the size of the host + index arrays, which for a matrix with billions of nonzeros is the difference between + fitting in memory and not. + """ + return np.int32 if max(n1, n2) < 2 ** 31 else np.int64 + + +########################################################################## +###################### worker-side functions ############################# +########################################################################## +def evaluate_kernel(kernel, x1, x2, hyperparameters, k_n_params, args): + """Call the kernel with whichever signature it declares. + + Mirrors :py:meth:`fvgp.gp_prior.GPprior.compute_covariances`. The historical + gp2Scale workers called ``kernel(x1, x2, hps)`` unconditionally, so a four-argument + ``args``-taking kernel -- supported everywhere else -- raised ``TypeError`` on the + worker. + """ + if k_n_params == 4: + return kernel(x1, x2, hyperparameters, args) + elif k_n_params == 3: + return kernel(x1, x2, hyperparameters) + else: + raise Exception("No valid kernel function signature") + + +def block_to_coo(k, index_dtype): + """A dense or sparse kernel block as ``(data, rows, cols)`` in block-local indices. + + A kernel that is already support-aware (see + :py:func:`fvgp.kernels.wendland_anisotropic_gp2Scale_cpu_sparse`) hands us a sparse + block, and is passed straight through rather than round-tripped via dense. + """ + if sparse.issparse(k): + k = k.tocoo() + return k.data, k.row.astype(index_dtype, copy=False), k.col.astype(index_dtype, copy=False) + k = np.asarray(k) + rows, cols = np.nonzero(k) + return k[rows, cols], rows.astype(index_dtype, copy=False), cols.astype(index_dtype, copy=False) + + +def block_triplets(range_ij, x1, x2, hyperparameters, kernel, + k_n_params, args, symmetric, index_dtype): + """COO triplets of one block, in *global* matrix coordinates. + + ``x1``/``x2`` arrive as the materialized values of the caller's scatter futures. + Returning triplets rather than a sparse block keeps the wire format independent of + where the block sits in the matrix. + """ + (i_start, i_end), (j_start, j_end) = range_ij + k = evaluate_kernel(kernel, x1[i_start:i_end], x2[j_start:j_end], + hyperparameters, k_n_params, args) + data, rows, cols = block_to_coo(k, index_dtype) + + # Blocks straddling the diagonal of a symmetric matrix are computed once and mirrored + # by the host, so only their upper triangle may be reported. + if symmetric and i_start == j_start and data.size: + mask = rows <= cols + data, rows, cols = data[mask], rows[mask], cols[mask] + + return data, rows + index_dtype(i_start), cols + index_dtype(j_start) + + +def row_strip_csr(range_i, x1, x2, hyperparameters, kernel, + k_n_params, args, n2, col_batch_size, index_dtype): + """One finished CSR row strip, tagged with its first row index. + + The strip is evaluated in column chunks so peak worker memory stays at a single dense + block, and converted to CSR here so the sort is done by the workers in parallel + rather than by the host on the assembled whole. + """ + i_start, i_end = range_i + x1_block = x1[i_start:i_end] + data_parts, row_parts, col_parts = [], [], [] + + for j_start, j_end in ranges(n2, num_blocks(n2, col_batch_size)): + k = evaluate_kernel(kernel, x1_block, x2[j_start:j_end], + hyperparameters, k_n_params, args) + data, rows, cols = block_to_coo(k, index_dtype) + if data.size == 0: continue + data_parts.append(data) + row_parts.append(rows) + col_parts.append(cols + index_dtype(j_start)) + + shape = (i_end - i_start, n2) + if not data_parts: + return i_start, sparse.csr_matrix(shape) + + strip = sparse.coo_matrix((np.concatenate(data_parts), + (np.concatenate(row_parts), np.concatenate(col_parts))), + shape=shape) + return i_start, strip.tocsr() + + +########################################################################## +###################### host-side assembly ################################ +########################################################################## +def _harvest(future_result): + """Take a result off the wire and drop the client's reference to its future. + + A cancelled task is handed back *as* its exception rather than raised, so without + this check a scheduler-side cancellation would flow into the assembly and fail there + as an unrelated shape or type error, far from its cause. + """ + future, result = future_result + future.release() + if isinstance(result, BaseException): + raise Exception( + f"A gp2Scale covariance block failed on the cluster: " + f"{type(result).__name__}: {result}") from result + return result + + +def assemble_triplets(harvest, n1, n2, symmetric, index_dtype): + """Assemble global COO triplets into CSR with a single allocation. + + The parts are sized first, then copied into one preallocated set of arrays, each part + being written together with its mirror image and dropped immediately afterwards. COO + does not care about ordering, so interleaving a block and its mirror costs nothing + and lets the gathered results be freed as we go. The alternative -- ``np.hstack`` + over every block and then a second ``np.hstack`` for the mirrored half -- holds three + to four copies of the matrix at peak. + """ + parts, total, dtypes = [], 0, [] + for data, rows, cols in harvest: + if data.size == 0: continue + parts.append((data, rows, cols)) + dtypes.append(data.dtype) + total += data.size + # Only entries off the diagonal get mirrored. Diagonal entries exist solely in + # blocks that straddle the diagonal, so this counts zero for every other block. + if symmetric: total += data.size - int(np.count_nonzero(rows == cols)) + + if not parts: + return sparse.csr_matrix((n1, n2)) + + out_data = np.empty(total, dtype=np.result_type(*dtypes)) + out_rows = np.empty(total, dtype=index_dtype) + out_cols = np.empty(total, dtype=index_dtype) + + position = 0 + while parts: + data, rows, cols = parts.pop() + n = data.size + out_data[position:position + n] = data + out_rows[position:position + n] = rows + out_cols[position:position + n] = cols + position += n + if symmetric: + off_diagonal = rows != cols + n = int(np.count_nonzero(off_diagonal)) + if n: + out_data[position:position + n] = data[off_diagonal] + out_rows[position:position + n] = cols[off_diagonal] + out_cols[position:position + n] = rows[off_diagonal] + position += n + del data, rows, cols + + K = sparse.coo_matrix((out_data, (out_rows, out_cols)), shape=(n1, n2)) + del out_data, out_rows, out_cols + return K.tocsr() + + +def assemble_row_strips(harvest, n1, n2): + """Assemble finished CSR row strips in row order.""" + strips = dict(harvest) + if not strips: + return sparse.csr_matrix((n1, n2)) + return sparse.vstack([strips[key] for key in sorted(strips)], format="csr") + + +########################################################################## +###################### the single entry point ############################ +########################################################################## +def distributed_covariance(client, kernel, hyperparameters, + x1_future, n1, x2_future, n2, + batch_size, symmetric=False, distribution="blockwise", + k_n_params=3, args=None): + """Compute ``k(x1, x2)`` across a dask cluster and return it as CSR. + + Parameters + ---------- + client : distributed.Client + The cluster the kernel blocks are mapped over. + kernel : callable + ``f(x1, x2, hyperparameters)`` or ``f(x1, x2, hyperparameters, args)``, selected + by ``k_n_params``. May return a dense array or a sparse block. + hyperparameters : np.ndarray + x1_future, x2_future : distributed.Future + Scattered point sets. For ``symmetric=True`` these must be the same future, so + that the workers slice one broadcast copy. + n1, n2 : int + Number of points behind each future; the shape of the result. + batch_size : int + Target points per chunk along each axis. + symmetric : bool + Whether the result is ``k(x, x)``, which lets ``"blockwise"`` schedule only the + upper triangle. + distribution : str + ``"blockwise"`` or ``"rowwise"``; see the module docstring. + k_n_params : int + 3 or 4, the kernel's arity. + args : dict or None + Passed to a four-argument kernel. + + Return + ------ + Covariance matrix : scipy.sparse.csr_matrix + """ + if distribution not in _DISTRIBUTIONS: + raise Exception(f"Unknown gp2Scale distribution `{distribution}`. " + f"Choose from: {list(_DISTRIBUTIONS)}") + if symmetric: + assert n1 == n2, "a symmetric covariance must be square" + assert x1_future is x2_future, \ + "a symmetric covariance must be computed from a single scattered point set" + + st = time.time() + index_dtype = index_dtype_for(n1, n2) + logger.debug("gp2Scale covariance ({}, symmetric={}) on client {}", + distribution, symmetric, client.id) + + if distribution == "blockwise": + row_ranges = ranges(n1, num_blocks(n1, batch_size)) + col_ranges = row_ranges if symmetric else ranges(n2, num_blocks(n2, batch_size)) + tasks = list(itertools.product(row_ranges, col_ranges)) + # filter the lower triangle; the host mirrors instead + if symmetric: tasks = [task for task in tasks if task[0][0] <= task[1][0]] + worker = partial(block_triplets, + hyperparameters=hyperparameters, kernel=kernel, + k_n_params=k_n_params, args=args, + symmetric=symmetric, index_dtype=index_dtype) + else: + tasks = ranges(n1, num_blocks(n1, batch_size)) + worker = partial(row_strip_csr, + hyperparameters=hyperparameters, kernel=kernel, + k_n_params=k_n_params, args=args, n2=n2, + col_batch_size=batch_size, index_dtype=index_dtype) + + logger.debug(" gp2Scale covariance init done after {} seconds ({} tasks).", + time.time() - st, len(tasks)) + + futures = client.map(worker, tasks, [x1_future] * len(tasks), [x2_future] * len(tasks)) + harvest = map(_harvest, distributed.as_completed(futures, with_results=True)) + + if distribution == "blockwise": + K = assemble_triplets(harvest, n1, n2, symmetric, index_dtype) + else: + K = assemble_row_strips(harvest, n1, n2) + + logger.debug(" gp2Scale covariance assembled after {} seconds.", time.time() - st) + logger.debug(" gp2Scale covariance sparsity = {}.", float(K.nnz) / float(n1 * n2)) + return K + + +def stack_augmented_covariance(K, B, D): + """Assemble ``[[K, B], [B.T, D]]`` the way scipy can do fastest. + + ``scipy.sparse`` has a fast path for block assembly when *every* block is already + CSR: it concatenates along each axis instead of converting the lot to COO and + rebuilding. ``B.T`` is CSC and a freshly built block is COO, so without the explicit + conversions here the whole augmented matrix takes the slow path. + """ + K, B, D = _as_csr(K), _as_csr(B), _as_csr(D) + return sparse.block_array([[K, B], [_as_csr(B.transpose()), D]], format="csr") + + +def _as_csr(matrix): + return matrix.tocsr() if sparse.issparse(matrix) else sparse.csr_matrix(matrix) diff --git a/fvgp/gp_posterior.py b/fvgp/gp_posterior.py index 012860d..29ba154 100755 --- a/fvgp/gp_posterior.py +++ b/fvgp/gp_posterior.py @@ -1,6 +1,7 @@ import numpy as np import warnings from loguru import logger +from scipy.sparse import issparse from .gp_lin_alg import * @@ -20,8 +21,13 @@ def __init__(self, self.noise_function_available = callable(self.likelihood.noise_function) def compute_covariances(self, x1, x2, hps): + """Direct, dense kernel evaluation. For the small (n_pred x n_pred) blocks.""" return self.prior.compute_covariances(x1, x2, hps) + def cross_covariance(self, x_pred, hps): + """k(x_data, x_pred). Distributed and sparse under gp2Scale, dense otherwise.""" + return self.prior.compute_data_cross_covariance(x_pred, hps) + def compute_mean(self, x, hps): return self.prior.compute_mean(x, hps) @@ -86,6 +92,48 @@ def K(self): @property def m(self): return self.prior.m + + @property + def gp2Scale(self): + return self.data.gp2Scale + + @staticmethod + def _dense(matrix): + """A dense view of a covariance block that may have arrived sparse from gp2Scale.""" + return matrix.toarray() if issparse(matrix) else matrix + + def _dense_K(self): + """K as a dense array, for the methods that build a joint (N + n_pred)^2 matrix. + + Those methods are dense in N by construction, so under gp2Scale they are only + usable on small problems and the user should know that the sparsity they asked + for is being thrown away here. + """ + if not issparse(self.K): return self.K + warnings.warn( + "This method assembles a joint covariance over data and prediction points, " + "which is dense in the number of data points. Under gp2Scale that discards " + "the sparse representation and costs O(N^2) memory; it is only usable on " + "small problems. Consider posterior_covariance instead.") + return self.K.toarray() + + def _cross_solve_product(self, k, chunk_size=None): + """``k.T @ KV^-1 @ k``, evaluated in column chunks. + + The solve is dense whatever ``k`` is -- ``KV^-1`` is dense even when ``KV`` is not + -- so chunking over prediction points is what keeps the intermediate at + (N x chunk) rather than (N x n_pred). ``k`` itself stays sparse throughout under + gp2Scale; only the chunk handed to the solver is densified. + """ + n_pred = k.shape[1] + if chunk_size is None: chunk_size = n_pred if not self.gp2Scale else self.prior.batch_size + chunk_size = max(1, min(int(chunk_size), n_pred)) + product = np.empty((n_pred, n_pred)) + for start in range(0, n_pred, chunk_size): + end = min(start + chunk_size, n_pred) + solved = self.KVsolve(self._dense(k[:, start:end])) + product[:, start:end] = np.asarray(k.T @ solved) + return product ########################################################## def posterior_mean(self, x_pred, hyperparameters=None, x_out=None): @@ -104,8 +152,10 @@ def posterior_mean(self, x_pred, hyperparameters=None, x_out=None): x_orig = x_pred.copy() if isinstance(x_out, np.ndarray): x_pred = self.cartesian_product(x_pred, x_out) - k = self.compute_covariances(x_data, x_pred, hyperparameters) - A = k.T @ KVinvY + # Sparse under gp2Scale, and it stays sparse: the product with KVinvY is the whole + # use of k here, so the posterior mean never materializes an (N x n_pred) array. + k = self.cross_covariance(x_pred, hyperparameters) + A = np.asarray(k.T @ KVinvY) prior_mean = self.compute_mean(x_pred, hyperparameters) posterior_mean = prior_mean[:, None] + A if isinstance(x_out, np.ndarray): posterior_mean_re = posterior_mean.reshape(len(x_orig), len(x_out), order='F') @@ -175,26 +225,23 @@ def posterior_mean_grad(self, x_pred, hyperparameters=None, x_out=None, directio ########################################################################### def posterior_covariance(self, x_pred, x_out=None, variance_only=False, add_noise=False): - x_data = self.x_data.copy() if x_out is None: x_out = self.x_out self._perform_input_checks(x_pred, x_out) x_orig = x_pred.copy() if isinstance(x_out, np.ndarray): x_pred = self.cartesian_product(x_pred, x_out) - k = self.compute_covariances(x_data, x_pred, self.hyperparameters) + k = self.cross_covariance(x_pred, self.hyperparameters) kk = self.compute_covariances(x_pred, x_pred, self.hyperparameters) - if self.KVinv is not None: - if variance_only and self.y_data.shape[1] == 1: - S = None - v = np.diag(kk) - np.einsum('ij,jk,ki->i', k.T, - self.KVinv, k, optimize=True) - else: - S = kk - (k.T @ self.KVsolve(k)) - v = np.array(np.diag(S)) + if self.KVinv is not None and variance_only and self.y_data.shape[1] == 1: + # The explicit-inverse modes are dense in N anyway, so densifying k costs + # nothing here, and the einsum gets the variances without ever forming S. + k_dense = self._dense(k) + S = None + v = np.diag(kk) - np.einsum('ij,jk,ki->i', k_dense.T, + self.KVinv, k_dense, optimize=True) else: - k_cov_prod = self.KVsolve(k) - S = kk - (k_cov_prod.T @ k) + S = kk - self._cross_solve_product(k) v = np.array(np.diag(S)) if np.any(v < -0.0001): warnings.warn( @@ -284,7 +331,7 @@ def posterior_covariance_grad(self, x_pred, x_out=None, direction=None): ########################################################################### def joint_gp_prior(self, x_pred, x_out=None): x_data, K, prior_mean_vec = (self.x_data.copy(), - self.K.copy() + (np.identity(len(self.K)) * 1e-9), + self._dense_K() + (np.identity(len(self.x_data)) * 1e-9), self.m.copy()) if x_out is None: x_out = self.x_out self._perform_input_checks(x_pred, x_out) @@ -306,7 +353,7 @@ def joint_gp_prior(self, x_pred, x_out=None): ########################################################################### def joint_gp_prior_grad(self, x_pred, direction, x_out=None): x_data, K, prior_mean_vec = (self.x_data.copy(), - self.K.copy() + (np.identity(len(self.K)) * 1e-9), + self._dense_K() + (np.identity(len(self.x_data)) * 1e-9), self.m.copy()) if x_out is None: x_out = self.x_out self._perform_input_checks(x_pred, x_out) @@ -413,7 +460,7 @@ def mutual_information(self, joint, m1, m2): ########################################################################### def gp_mutual_information(self, x_pred, x_out=None, add_noise=False): - x_data, K = self.x_data.copy(), self.K.copy() + (np.identity(len(self.K)) * 1e-9) + x_data, K = self.x_data.copy(), self._dense_K() + (np.identity(len(self.x_data)) * 1e-9) if x_out is None: x_out = self.x_out self._perform_input_checks(x_pred, x_out) x_orig = x_pred.copy() @@ -429,7 +476,7 @@ def gp_mutual_information(self, x_pred, x_out=None, add_noise=False): ########################################################################### def gp_total_correlation(self, x_pred, x_out=None, add_noise=False): - x_data, K = self.x_data.copy(), self.K.copy() + (np.identity(len(self.K)) * 1e-9) + x_data, K = self.x_data.copy(), self._dense_K() + (np.identity(len(self.x_data)) * 1e-9) if x_out is None: x_out = self.x_out self._perform_input_checks(x_pred, x_out) x_orig = x_pred.copy() diff --git a/fvgp/gp_prior.py b/fvgp/gp_prior.py index c2c5c7a..063daba 100755 --- a/fvgp/gp_prior.py +++ b/fvgp/gp_prior.py @@ -1,15 +1,9 @@ import numpy as np import inspect -import dask.distributed as distributed import warnings -import itertools -import time -import scipy.sparse as sparse from .kernels import * -from functools import partial -from scipy.sparse import block_array +from .gp2Scale_covariance import distributed_covariance, stack_augmented_covariance from loguru import logger -from scipy.sparse import coo_matrix, vstack warnings.simplefilter("once", UserWarning) @@ -22,14 +16,19 @@ def __init__(self, kernel_grad=None, prior_mean_function_grad=None, gp2Scale_batch_size=10000, + gp2Scale_distribution="blockwise", ): self.kernel_function = kernel self.prior_mean_function = prior_mean_function self.batch_size = gp2Scale_batch_size + self.gp2Scale_distribution = gp2Scale_distribution self.data = data self.trainer = trainer + assert gp2Scale_distribution in ("blockwise", "rowwise"), \ + "gp2Scale_distribution must be `blockwise` or `rowwise`" + assert callable(kernel) or kernel is None, "kernel must be callable or None" assert callable(prior_mean_function) or prior_mean_function is None, \ "prior_mean_function must be callable or None" @@ -94,8 +93,7 @@ def __init__(self, self.x_data_scatter_future = None if self.gp2Scale and self.client is not None: - self.x_data_scatter_future = self.client.scatter( - self.x_data, workers=self.compute_workers, broadcast=True, direct=True) + self.x_data_scatter_future = self._scatter(self.x_data) self.m, self.K = self._compute_prior(self.x_data, self.hyperparameters) logger.debug("Prior successfully initialized.") @@ -164,19 +162,18 @@ def augment_state_data(self): # Python ref and is cleaned up via __del__. This is race-free within a # single GP's lifetime; do NOT churn many GP instances back-to-back without # a `del gp; gc.collect(); client.run(lambda: None)` between them. - self.x_data_scatter_future = self.client.scatter( - self.x_data, workers=self.compute_workers, broadcast=True, direct=True) + self.x_data_scatter_future = self._scatter(self.x_data) logger.debug("Prior mean and covariance updated after data augmentation.") def update_state_data(self): """ - This is for the case that the data has changed, but not just been augmented. For example, in an online learning setting where old data points are replaced by new ones. + This is for the case that the data has changed, but not just been augmented. + For example, in an online learning setting where old data points are replaced by new ones. """ if self.gp2Scale and self.client is not None: # Full data change: refresh the persistent scatter before rebuilding K. # Overwrite (no explicit release); the old future is GC'd at a quiet moment. - self.x_data_scatter_future = self.client.scatter( - self.x_data, workers=self.compute_workers, broadcast=True, direct=True) + self.x_data_scatter_future = self._scatter(self.x_data) self.m, self.K = self._compute_prior(self.x_data, self.hyperparameters) logger.debug("Prior mean and covariance updated after data change.") @@ -186,10 +183,31 @@ def update_state_hyperparameters(self): def compute_prior_covariance_matrix(self, x, hyperparameters): """computes the prior covariance matrix from the kernel""" - if self.gp2Scale: K = self._compute_prior_covariance_gp2Scale(x, hyperparameters) - else: K = self.compute_covariances(x, x, hyperparameters) + if self.gp2Scale: + # Every caller of this method hands over the current training set, so the + # persistent scatter is the right one to slice -- but only reuse it when the + # shapes actually agree, in case a caller ever does otherwise. + future = self.x_data_scatter_future + if future is not None and np.shape(x) != np.shape(self.x_data): future = None + K = self._gp2Scale_covariance(x, x, hyperparameters, symmetric=True, x1_future=future) + else: + K = self.compute_covariances(x, x, hyperparameters) return K + def compute_data_cross_covariance(self, x_pred, hyperparameters): + """computes k(x_data, x_pred), the cross-covariance the posterior needs. + + Under gp2Scale this is the one covariance the posterior cannot simply evaluate on + the client: it has as many rows as there are data points. It goes through the + same distributed assembler as the prior and comes back sparse, so a posterior + mean never materializes an (N x n_pred) dense array. Below one batch of data + there is nothing to gain from the cluster, so the kernel is called directly. + """ + if self.gp2Scale and self.client is not None and len(self.x_data) > self.batch_size: + return self._gp2Scale_covariance(self.x_data, x_pred, hyperparameters, + x1_future=self.x_data_scatter_future) + return self.compute_covariances(self.x_data, x_pred, hyperparameters) + def compute_covariances(self, x1, x2, hps): """computes the covariances via k(x,x')""" if self.k_n_params == 3: @@ -240,7 +258,19 @@ def _update_prior(self, x_old, x_new, hyperparameters): def _update_prior_covariance_matrix(self, x_old, x_new, hyperparameters): """This updated K based on new data""" if self.gp2Scale: - K = self._update_prior_covariance_gp2Scale(x_old, x_new, hyperparameters) + # self.x_data_scatter_future still holds x_old at this point; augment_state_data + # refreshes it to the full dataset only after this call returns. The x_new + # scatter is ours, so it is ours to release. + x_new_future = self._scatter(x_new) + try: + B = self._gp2Scale_covariance(x_old, x_new, hyperparameters, + x1_future=self.x_data_scatter_future, + x2_future=x_new_future) + D = self._gp2Scale_covariance(x_new, x_new, hyperparameters, + symmetric=True, x1_future=x_new_future) + finally: + x_new_future.release() + K = stack_augmented_covariance(self.K, B, D) else: k = self.compute_covariances(x_old, x_new, hyperparameters) kk = self.compute_covariances(x_new, x_new, hyperparameters) @@ -259,142 +289,58 @@ def _update_mean(self, x_new, hyperparameters): raise Exception("Prior mean in wrong format") return m - @staticmethod - def _ranges(N, nb): - """ splits a range(N) into nb chunks defined by chunk_start, chunk_end """ - if nb == 0: nb = 1 - step = N / nb - return [(round(step * i), round(step * (i + 1))) for i in range(nb)] - - def _compute_prior_covariance_gp2Scale(self, x_data, hyperparameters): - """computes the covariance matrix from the kernel on HPC in sparse format""" - st = time.time() - point_number = len(x_data) - num_batches = point_number // self.batch_size - NUM_RANGES = num_batches - logger.debug("client id: {}", self.client.id) - - ranges = self._ranges(len(x_data), NUM_RANGES) # the chunk ranges, as (start, end) tuples - ranges_ij = list( - itertools.product(ranges, ranges)) # all i/j ranges as ((i_start, i_end), (j_start, j_end)) pairs of tuples - ranges_ij = [range_ij for range_ij in ranges_ij if range_ij[0][0] <= range_ij[1][0]] # filter lower diagonal - logger.debug(" gp2Scale covariance matrix init done after {} seconds.", time.time() - st) - - results = list(map(self._harvest_result, distributed.as_completed(self.client.map( - partial(kernel_function, - hyperparameters=hyperparameters, - kernel=self.kernel), - ranges_ij, - [self.x_data_scatter_future] * len(ranges_ij), - [self.x_data_scatter_future] * len(ranges_ij)), - with_results=True))) - - - logger.debug(" gp2Scale covariance matrix result written after {} seconds.", time.time() - st) - - # reshape the result set into COO components - data, i_s, j_s = map(np.hstack, zip(*results)) - logger.debug(" gp2Scale covariance matrix result stacked after {} seconds.", time.time() - st) - del results - # mirror across diagonal - diagonal_mask = i_s != j_s - data, i_s, j_s = np.hstack([data, data[diagonal_mask]]), \ - np.hstack([i_s, j_s[diagonal_mask]]), \ - np.hstack([j_s, i_s[diagonal_mask]]) - K = sparse.coo_matrix((data, (i_s, j_s)), shape=(len(x_data), len(x_data))) - del data - logger.debug(" gp2Scale covariance matrix assembled after {} seconds.", time.time() - st) - #K = self._coo_to_csr_chunked(i_s, j_s, data, (len(data), len(data)), int(len(data)/2)) - K = K.tocsr() - logger.debug(" gp2Scale covariance matrix in CSR after {} seconds.", time.time() - st) - logger.debug(" gp2Scale covariance matrix sparsity = {}.", float(K.nnz) / float(K.shape[0] ** 2)) - return K - - @staticmethod - def _coo_to_csr_chunked(row, col, data, shape, chunk_size): #pragma: no cover - n_rows = shape[0] - chunks = [] - for start in range(0, n_rows, chunk_size): - end = min(start + chunk_size, n_rows) - mask = (row >= start) & (row < end) - r = row[mask] - start # Normalize to chunk - c = col[mask] - d = data[mask] - coo_chunk = coo_matrix((d, (r, c)), shape=(end - start, shape[1])) - csr_chunk = coo_chunk.tocsr() - chunks.append(csr_chunk) - return vstack(chunks, format='csr') - - def _update_prior_covariance_gp2Scale(self, x_old, x_new, hyperparameters): - """computes the covariance matrix from the kernel on HPC in sparse format. - - Uses self.x_data_scatter_future for the x_old side (pre-augment scatter, still - valid at entry) and a fresh local scatter for x_new. Only x_new's local future - is released here; the persistent self.x_data_scatter_future is refreshed by - augment_state_data after this call. + def _scatter(self, x): + """Broadcast a point set to the compute workers, under a key of its own. + + ``hash=False`` is what makes the key unique. By default dask keys scattered data + by a hash of its content, so scattering the same array twice -- a new GP on the + same data, a second prediction at the same points, a re-scatter after an append + that did not change x -- lands on one key. The first copy's release then + schedules a ``_dec_ref`` that races the second scatter inside the scheduler, and + the tasks depending on it come back as ``KeyError`` or ``CancelledError``. A + unique key per scatter removes the collision at its source; the only thing given + up is a de-duplication we never wanted, since each copy is released with the + object that made it. """ - x_new_scatter_future = self.client.scatter( - x_new, workers=self.compute_workers, broadcast=True, direct=True) - x_old_scatter_future = self.x_data_scatter_future - - point_number = len(x_old) - num_batches = point_number // self.batch_size - NUM_RANGES = num_batches - ranges_data = self._ranges(len(x_old), NUM_RANGES) # the chunk ranges, as (start, end) tuples - num_batches2 = len(x_new) // self.batch_size - ranges_input = self._ranges(len(x_new), num_batches2) - ranges_ij = list(itertools.product(ranges_data, ranges_input)) - - # K = np.block([[self.K, B], - # [B, C]]) - # Calculate B - - results = list(map(self._harvest_result, - distributed.as_completed(self.client.map( - partial(kernel_function_update, - hyperparameters=hyperparameters, - kernel=self.kernel), - ranges_ij, - [x_old_scatter_future] * len(ranges_ij), - [x_new_scatter_future] * len(ranges_ij)), - with_results=True))) - - data, i_s, j_s = map(np.hstack, zip(*results)) - B = sparse.coo_matrix((data, (i_s, j_s)), shape=(len(x_old), len(x_new))) - - # mirror across diagonal - ranges_ij2 = list(itertools.product(ranges_input, ranges_input)) - ranges_ij2 = [range_ij2 for range_ij2 in ranges_ij2 if - range_ij2[0][0] <= range_ij2[1][0]] # filter lower diagonal - - results = list(map(self._harvest_result, - distributed.as_completed(self.client.map( - partial(kernel_function, - hyperparameters=hyperparameters, - kernel=self.kernel), - ranges_ij2, - [x_new_scatter_future] * len(ranges_ij2), - [x_new_scatter_future] * len(ranges_ij2)), - with_results=True))) - data, i_s, j_s = map(np.hstack, zip(*results)) - diagonal_mask = i_s != j_s - data, i_s, j_s = np.hstack([data, data[diagonal_mask]]), \ - np.hstack([i_s, j_s[diagonal_mask]]), \ - np.hstack([j_s, i_s[diagonal_mask]]) - D = sparse.coo_matrix((data, (i_s, j_s)), shape=(len(x_new), len(x_new))) - - res = block_array([[self.K, B], - [B.transpose(), D]]) - - x_new_scatter_future.release() - - return res + return self.client.scatter(x, workers=self.compute_workers, broadcast=True, + direct=True, hash=False) - @staticmethod - def _harvest_result(future_result): - future, result = future_result - future.release() - return result + def _gp2Scale_covariance(self, x1, x2, hyperparameters, symmetric=False, + x1_future=None, x2_future=None): + """The single distributed kernel evaluation, shared by prior, append and posterior. + + Owns nothing but scatter lifetime: a future passed in belongs to the caller and is + left alone (this is how the persistent ``x_data_scatter_future`` survives), while + any future created here is released before returning. The scheduling and assembly + live in :py:mod:`fvgp.gp2Scale_covariance`. + """ + if self.client is None: + raise Exception("gp2Scale needs a dask client to compute covariances.") + + own1 = x1_future is None + if own1: x1_future = self._scatter(x1) + if symmetric: + # One broadcast copy, sliced on both axes -- the assembler relies on this to + # schedule only the upper triangle. + assert x2 is x1 or np.shape(x2) == np.shape(x1), "symmetric requires x1 == x2" + x2_future, own2 = x1_future, False + else: + own2 = x2_future is None + if own2: x2_future = self._scatter(x2) + + try: + return distributed_covariance( + self.client, self.kernel, hyperparameters, + x1_future=x1_future, n1=len(x1), + x2_future=x2_future, n2=len(x2), + batch_size=self.batch_size, + symmetric=symmetric, + distribution=self.gp2Scale_distribution, + k_n_params=self.k_n_params, + args=self.args) + finally: + if own1: x1_future.release() + if own2: x2_future.release() #################################################### #################################################### @@ -508,6 +454,7 @@ def __getstate__(self): m_n_params=self.m_n_params, k_n_params=self.k_n_params, batch_size=self.batch_size, + gp2Scale_distribution=self.gp2Scale_distribution, data=self.data, trainer=self.trainer, kernel=self.kernel, @@ -524,49 +471,3 @@ def __getstate__(self): def __setstate__(self, state): self.__dict__.update(state) - - -######################################################## -######################################################## -######################################################## -def kernel_function(range_ij, x1_future, x2_future, hyperparameters, kernel): - """ - Essentially, parameters other than range_ij are static across calls. range_ij defines the region of the - covariance matrix being calculated. - Rather than return a sparse array in local coordinates, we can return the COO components in global coordinates. - """ - - hps = hyperparameters - range_i, range_j = range_ij - x1 = x1_future[range_i[0]:range_i[1]] - x2 = x2_future[range_j[0]:range_j[1]] - k = kernel(x1, x2, hps) - k_sparse = sparse.coo_matrix(k) - - data, rows, cols = k_sparse.data, k_sparse.row + range_i[0], k_sparse.col + range_j[0] - - # mask lower triangular values when current chunk spans diagonal - if range_i[0] == range_j[0]: - mask = [row <= col for (row, col) in zip(rows, cols)] - return data[mask], rows[mask], cols[mask] - else: - return data, rows, cols - - -def kernel_function_update(range_ij, x1_future, x2_future, hyperparameters, kernel): - """ - Essentially, parameters other than range_ij are static across calls. range_ij defines the region of the - covariance matrix being calculated. - Rather than return a sparse array in local coordinates, we can return the COO components in global coordinates. - """ - - hps = hyperparameters - range_i, range_j = range_ij - x1 = x1_future[range_i[0]:range_i[1]] - x2 = x2_future[range_j[0]:range_j[1]] - k = kernel(x1, x2, hps) - k_sparse = sparse.coo_matrix(k) - - data, rows, cols = k_sparse.data, k_sparse.row + range_i[0], k_sparse.col + range_j[0] - - return data, rows, cols diff --git a/tests/test_fvgp.py b/tests/test_fvgp.py index aabc801..3942cc7 100755 --- a/tests/test_fvgp.py +++ b/tests/test_fvgp.py @@ -2966,3 +2966,169 @@ def smooth(t): for line in printed.splitlines() if "bo evaluation" in line] assert np.allclose(sorted(reported), sorted(info["trace f(x)"][n_design:])) assert reported != sorted(reported, reverse=True) # not a monotone best-so-far + + +########################################################################### +#################### gp2Scale distributed covariance ###################### +########################################################################### +def _reference_wendland(x1, x2, hps): + """Dense reference for the assembler tests.""" + return wendland_anisotropic_gp2Scale_cpu(x1, x2, hps) + + +def _wendland_with_args(x1, x2, hps, args): + """Four-argument kernel. The historical gp2Scale workers could not call one.""" + return args["scale"] * wendland_anisotropic_gp2Scale_cpu(x1, x2, hps) + + +def test_distributed_covariance_matches_dense(client): + """One primitive, four shapes of problem: the assembled sparse matrix must equal the + dense kernel evaluation for both distributions, symmetric and rectangular.""" + from fvgp.gp2Scale_covariance import distributed_covariance + + rng = np.random.default_rng(42) + x1 = rng.random((57, 2)) + x2 = rng.random((23, 2)) + hps = np.array([2.0, 0.4, 0.35]) + + f1 = client.scatter(x1, broadcast=True, direct=True) + f2 = client.scatter(x2, broadcast=True, direct=True) + + for distribution in ("blockwise", "rowwise"): + K = distributed_covariance(client, _reference_wendland, hps, + x1_future=f1, n1=len(x1), x2_future=f1, n2=len(x1), + batch_size=10, symmetric=True, distribution=distribution) + assert sparse.issparse(K) + assert np.allclose(K.toarray(), _reference_wendland(x1, x1, hps)) + # the mirror must actually be a mirror, not a triangle + assert np.allclose(K.toarray(), K.toarray().T) + + B = distributed_covariance(client, _reference_wendland, hps, + x1_future=f1, n1=len(x1), x2_future=f2, n2=len(x2), + batch_size=10, symmetric=False, distribution=distribution) + assert B.shape == (len(x1), len(x2)) + assert np.allclose(B.toarray(), _reference_wendland(x1, x2, hps)) + + # a single block (batch_size larger than the data) must behave identically + K1 = distributed_covariance(client, _reference_wendland, hps, + x1_future=f1, n1=len(x1), x2_future=f1, n2=len(x1), + batch_size=10000, symmetric=True, distribution=distribution) + assert np.allclose(K1.toarray(), _reference_wendland(x1, x1, hps)) + + f1.release() + f2.release() + + +def test_distributed_covariance_empty_and_args_kernel(client): + """Two things the old workers got wrong: an all-empty result, and a kernel that takes + ``args`` -- supported everywhere else in fvGP but not on the gp2Scale workers.""" + from fvgp.gp2Scale_covariance import distributed_covariance + + rng = np.random.default_rng(7) + x1 = rng.random((40, 2)) + x2 = rng.random((40, 2)) + 100.0 # far outside the Wendland support + hps = np.array([1.0, 0.1, 0.1]) + + f1 = client.scatter(x1, broadcast=True, direct=True) + f2 = client.scatter(x2, broadcast=True, direct=True) + + for distribution in ("blockwise", "rowwise"): + empty = distributed_covariance(client, _reference_wendland, hps, + x1_future=f1, n1=len(x1), x2_future=f2, n2=len(x2), + batch_size=10, distribution=distribution) + assert empty.shape == (len(x1), len(x2)) + assert empty.nnz == 0 + + with_args = distributed_covariance(client, _wendland_with_args, hps, + x1_future=f1, n1=len(x1), x2_future=f1, n2=len(x1), + batch_size=10, symmetric=True, + distribution=distribution, + k_n_params=4, args={"scale": 3.0}) + assert np.allclose(with_args.toarray(), 3.0 * _reference_wendland(x1, x1, hps)) + + f1.release() + f2.release() + + +def test_distributed_covariance_rejects_bad_distribution(client): + from fvgp.gp2Scale_covariance import distributed_covariance + + x = np.random.rand(5, 1) + f = client.scatter(x, broadcast=True, direct=True) + try: + distributed_covariance(client, _reference_wendland, np.array([1., 1.]), + x1_future=f, n1=5, x2_future=f, n2=5, + batch_size=2, distribution="columnwise") + except Exception as e: + assert "columnwise" in str(e) + else: + raise AssertionError("an unknown distribution must be rejected") + f.release() + + +def test_gp2Scale_posterior_matches_dense(client): + """The gp2Scale posterior goes through the distributed sparse cross-covariance while + the dense GP calls the kernel directly. Same kernel, same data: same answers.""" + import gc + + rng = np.random.default_rng(3) + x = rng.random((120, 1)) + y = np.sin(np.linalg.norm(x, axis=1) * 5.0) + hps = np.array([1.5, 0.3]) + x_pred_local = rng.random((17, 1)) + + dense = GP(x, y, hps, kernel_function=wendland_anisotropic_gp2Scale_cpu, + linalg_mode="Chol") + reference_mean = dense.posterior_mean(x_pred_local)["m(x)"] + reference_var = dense.posterior_covariance(x_pred_local)["v(x)"] + reference_S = dense.posterior_covariance(x_pred_local)["S"] + del dense + gc.collect() + + for distribution in ("blockwise", "rowwise"): + # batch_size below len(x) is what puts the cross-covariance on the cluster + scaled = GP(x, y, hps, gp2Scale=True, gp2Scale_batch_size=40, + gp2Scale_distribution=distribution, dask_client=client, + linalg_mode="Chol") + assert sparse.issparse(scaled.prior.compute_data_cross_covariance(x_pred_local, hps)) + assert np.allclose(scaled.posterior_mean(x_pred_local)["m(x)"], reference_mean) + post = scaled.posterior_covariance(x_pred_local) + assert np.allclose(post["v(x)"], reference_var) + assert np.allclose(post["S"], reference_S) + del scaled + gc.collect() + client.run(lambda: None) + + +def test_gp2Scale_distributions_agree(client): + """Row-wise and block-wise must produce the same prior covariance and the same + likelihood, including after an append.""" + import gc + + rng = np.random.default_rng(11) + x = rng.random((90, 2)) + y = np.sin(np.linalg.norm(x, axis=1) * 5.0) + x_add = rng.random((7, 2)) + y_add = np.sin(np.linalg.norm(x_add, axis=1) * 5.0) + hps = np.array([1.2, 0.35, 0.3]) + + results = {} + for distribution in ("blockwise", "rowwise"): + gp = GP(x, y, hps, gp2Scale=True, gp2Scale_batch_size=25, + gp2Scale_distribution=distribution, dask_client=client, + linalg_mode="sparseLU") + K = gp.prior.K.toarray().copy() + gp.update_gp_data(x_add, y_add, append=True) + results[distribution] = (K, gp.prior.K.toarray().copy(), gp.log_likelihood()) + del gp + gc.collect() + client.run(lambda: None) + + K_block, K_block_aug, ll_block = results["blockwise"] + K_row, K_row_aug, ll_row = results["rowwise"] + assert np.allclose(K_block, wendland_anisotropic_gp2Scale_cpu(x, x, hps)) + assert np.allclose(K_block, K_row) + assert np.allclose(K_block_aug, K_row_aug) + assert np.allclose(K_block_aug, wendland_anisotropic_gp2Scale_cpu( + np.vstack([x, x_add]), np.vstack([x, x_add]), hps)) + assert np.isclose(ll_block, ll_row)