Unify the gp2Scale covariance computation behind one primitive

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.

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)

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]####