From ebdf2e071bb839fd71336d7c2d56be39721f8057 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 17 Jul 2026 22:29:08 +0000 Subject: [PATCH 1/2] dy_d_best: fit NNS.stack once per call (major speedup, identical results) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stack fit (CV n.best, dimension-reduction coefficients, blend weight) depends only on (x, y), not on the evaluation points, and every regressor shares the same X* = rowMeans(x) training design. Previously dy_d_best ran a full 5-fold NNS.stack CV once per (regressor x bandwidth): dy_d_best(wrt=1:5) issued 20 stack fits to predict a handful of points. Now every regressor's bandwidth blocks (and mixed-derivative corners) are gathered and predicted from a SINGLE NNS.stack fit: * _dy_d_best_scalar -> _dy_d_best_plan: builds the test chunks and returns a reducer instead of calling the engine; * dy_d_best gathers all plans' chunks, runs one _dy_d_best_reg_estimates call, and hands each regressor's slice back to its reducer. On Example 2 (n=1001, wrt=1:5): 24.2s/20 stack calls -> 1.55s/1 stack call (~15x), with byte-for-byte identical output [0.9396, 0.9476, 0.9331, 0.92, 0.9284]. Restores tractability of the large paper examples on the CV side (the remaining cost is the engine's all-pairs RPM prediction). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.8 --- src/nns/diff.py | 229 ++++++++++++++++++++++++++++++------------------ 1 file changed, 144 insertions(+), 85 deletions(-) diff --git a/src/nns/diff.py b/src/nns/diff.py index 781612f..458f982 100644 --- a/src/nns/diff.py +++ b/src/nns/diff.py @@ -321,8 +321,14 @@ def dy_d_best( raise ValueError("You have some missing values, please address.") wrt_values = _dy_d_validate_wrt(wrt, x_values.shape[1]) - outputs = [ - _dy_d_best_scalar( + + # Every regressor shares the same X* = rowMeans(x) training design, so the + # NNS.stack fit (CV n.best, dimension-reduction coefficients, blend weight) + # is identical for all of them and for every bandwidth. Gather the test + # blocks from all regressors and run a SINGLE NNS.stack fit + prediction, + # then hand each regressor's slice back to its reducer. + plans = [ + _dy_d_best_plan( x_values, y_values, int(wrt_value) - 1, @@ -333,6 +339,19 @@ def dy_d_best( ) for wrt_value in wrt_values ] + + all_chunks: list[NDArray[np.float64]] = [] + spans: list[tuple[int, int, Any]] = [] + for chunks, reduce in plans: + spans.append((len(all_chunks), len(chunks), reduce)) + all_chunks.extend(chunks) + + sizes = [chunk.shape[0] for chunk in all_chunks] + predictions = _dy_d_best_reg_estimates(x_values, y_values, np.vstack(all_chunks)) + offsets = np.concatenate(([0], np.cumsum(sizes))) + parts = [predictions[offsets[i] : offsets[i + 1]] for i in range(len(sizes))] + + outputs = [reduce(parts[start : start + count]) for start, count, reduce in spans] if len(outputs) == 1: return outputs[0] return _combine_dy_d_outputs(outputs) @@ -398,7 +417,7 @@ def _dydx_step(column: NDArray[np.float64], eval_value: float, band: float) -> f return float(upper - lower) -def _dy_d_best_scalar( +def _dy_d_best_plan( x_values: NDArray[np.float64], y_values: NDArray[np.float64], wrt_index: int, @@ -407,7 +426,16 @@ def _dy_d_best_scalar( mixed: bool, messages: bool, wrt_label: int, -) -> dict[str, NDArray[np.float64]]: +) -> tuple[ + list[NDArray[np.float64]], + Callable[[list[NDArray[np.float64]]], dict[str, NDArray[np.float64]]], +]: + """Build every test block this regressor needs, plus a reducer. + + The reducer turns the predicted chunks back into the derivative dict. All + regressors share the same X* training design, so ``dy_d_best`` gathers the + chunks from every plan and runs a single NNS.stack fit for the whole call. + """ from nns.dependence import nns_dep from nns.var import lpm_var @@ -429,9 +457,26 @@ def _dy_d_best_scalar( dependence = float(nns_dep(column, y_values)["Dependence"]) h_s = _v057_bandwidths(x_values.size, dependence) - firsts: list[NDArray[np.float64]] = [] - seconds: list[NDArray[np.float64]] = [] - mixeds: list[NDArray[np.float64]] = [] + # The NNS.stack fit (CV n.best, dimension-reduction coefficients, blend + # weight) depends only on (x, y), never on the evaluation points, so every + # bandwidth - and the mixed-derivative corners - are predicted from a single + # fit. Gather every test block, run ONE NNS.stack, then slice it back. + chunks: list[NDArray[np.float64]] = [] + sizes: list[int] = [] + + def _add(block: NDArray[np.float64]) -> int: + chunks.append(block) + sizes.append(block.shape[0]) + return len(chunks) - 1 + + def _row_nanmean(bands: list[NDArray[np.float64]]) -> NDArray[np.float64]: + matrix = np.column_stack([np.asarray(b, dtype=np.float64).reshape(-1) for b in bands]) + with np.errstate(invalid="ignore"): + return np.asarray(np.nanmean(matrix, axis=1), dtype=np.float64) + + steps_per_band: list[NDArray[np.float64]] = [] + main_chunk_idx: list[int] = [] + mixed_eval: NDArray[np.float64] | None = None if vector_branch: eval_vec = np.asarray(eval_values, dtype=np.float64).reshape(-1) @@ -446,11 +491,13 @@ def _dy_d_best_scalar( ) sample_size = grid.shape[0] k = eval_vec.size + position = np.tile( + np.repeat(np.asarray(["l", "m", "u"], dtype=object), sample_size), k + ) + ids = np.repeat(np.arange(k), 3 * sample_size) for band in h_s: steps = np.array([_dydx_step(column, ev, float(band)) for ev in eval_vec]) blocks: list[NDArray[np.float64]] = [] - position: list[str] = [] - ids: list[int] = [] for g in range(k): lower = grid.copy() middle = grid.copy() @@ -459,27 +506,10 @@ def _dy_d_best_scalar( middle[:, wrt_index] = eval_vec[g] upper[:, wrt_index] = eval_vec[g] + steps[g] blocks.extend((lower, middle, upper)) - position.extend(["l"] * sample_size + ["m"] * sample_size + ["u"] * sample_size) - ids.extend([g] * (3 * sample_size)) - estimates = _dy_d_best_reg_estimates(x_values, y_values, np.vstack(blocks)) - pos = np.asarray(position, dtype=object) - idx = np.asarray(ids) - band_first = np.empty(k) - band_second = np.empty(k) - for g in range(k): - lo = np.mean(estimates[(pos == "l") & (idx == g)]) - mid = np.mean(estimates[(pos == "m") & (idx == g)]) - up = np.mean(estimates[(pos == "u") & (idx == g)]) - h = steps[g] - if np.isfinite(h) and h != 0.0: - band_first[g] = (up - lo) / (2.0 * h) - band_second[g] = (up - 2.0 * mid + lo) / (h**2) - else: - band_first[g] = np.nan - band_second[g] = np.nan - firsts.append(band_first) - seconds.append(band_second) - mixed_eval: NDArray[np.float64] | None = None + steps_per_band.append(steps) + main_chunk_idx.append(_add(np.vstack(blocks))) + if k == 2: + mixed_eval = eval_vec.reshape(1, 2) else: eval_mat = _as_eval_matrix(eval_values, n_predictors) n_eval = eval_mat.shape[0] @@ -487,72 +517,101 @@ def _dy_d_best_scalar( steps = np.array( [_dydx_step(column, eval_mat[i, wrt_index], float(band)) for i in range(n_eval)] ) - finite = np.isfinite(steps) & (steps != 0.0) lower = eval_mat.copy() upper = eval_mat.copy() lower[:, wrt_index] = eval_mat[:, wrt_index] - steps upper[:, wrt_index] = eval_mat[:, wrt_index] + steps - if messages: - print( - "Currently generating NNS.reg finite difference estimates...bandwidth\r" - ) - estimates = _dy_d_best_reg_estimates( - x_values, y_values, np.vstack((lower, eval_mat, upper)) - ) - lo = estimates[:n_eval] - mid = estimates[n_eval : 2 * n_eval] - up = estimates[2 * n_eval :] - with np.errstate(invalid="ignore", divide="ignore"): - first = (up - lo) / (2.0 * steps) - second = (up - 2.0 * mid + lo) / (steps**2) - first[~finite] = np.nan - second[~finite] = np.nan - firsts.append(first) - seconds.append(second) + steps_per_band.append(steps) + main_chunk_idx.append(_add(np.vstack((lower, eval_mat, upper)))) mixed_eval = eval_mat - def _row_nanmean(bands: list[NDArray[np.float64]]) -> NDArray[np.float64]: - matrix = np.column_stack([np.asarray(b, dtype=np.float64).reshape(-1) for b in bands]) - with np.errstate(invalid="ignore"): - return np.asarray(np.nanmean(matrix, axis=1), dtype=np.float64) - - output = {"First": _row_nanmean(firsts), "Second": _row_nanmean(seconds)} - + # Mixed-derivative corners (also predicted from the same single fit). + mixed_meta: list[tuple[int | None, NDArray[np.float64]]] = [] if mixed: - if vector_branch: - tuple_eval = np.asarray(eval_values, dtype=np.float64).reshape(-1) - if tuple_eval.size != 2: - raise ValueError("Mixed Derivatives are only for 2 IV") - mixed_points = tuple_eval.reshape(1, 2) - else: - assert mixed_eval is not None - if mixed_eval.shape[1] != 2: - raise ValueError("Mixed Derivatives are only for 2 IV") - mixed_points = mixed_eval + if mixed_eval is None or mixed_eval.shape[1] != 2: + raise ValueError("Mixed Derivatives are only for 2 IV") for band in h_s: - band_vals: list[float] = [] - for point in mixed_points: + corner_blocks: list[NDArray[np.float64]] = [] + scales: list[float] = [] + for point in mixed_eval: s1 = _dydx_step(x_values[:, 0], point[0], float(band)) s2 = _dydx_step(x_values[:, 1], point[1], float(band)) - if not (np.isfinite(s1) and np.isfinite(s2) and s1 != 0.0 and s2 != 0.0): - band_vals.append(np.nan) - continue - corners = np.array( - [ - [point[0] + s1, point[1] + s2], - [point[0] - s1, point[1] + s2], - [point[0] + s1, point[1] - s2], - [point[0] - s1, point[1] - s2], - ] - ) - z = _dy_d_best_reg_estimates(x_values, y_values, corners) - band_vals.append((z[0] + z[3] - z[1] - z[2]) / (4.0 * s1 * s2)) - mixeds.append(np.asarray(band_vals, dtype=np.float64)) - output["Mixed"] = _row_nanmean(mixeds) + if np.isfinite(s1) and np.isfinite(s2) and s1 != 0.0 and s2 != 0.0: + corner_blocks.append( + np.array( + [ + [point[0] + s1, point[1] + s2], + [point[0] - s1, point[1] + s2], + [point[0] + s1, point[1] - s2], + [point[0] - s1, point[1] - s2], + ] + ) + ) + scales.append(4.0 * s1 * s2) + else: + scales.append(np.nan) + scale_arr = np.asarray(scales, dtype=np.float64) + if corner_blocks: + mixed_meta.append((_add(np.vstack(corner_blocks)), scale_arr)) + else: + mixed_meta.append((None, scale_arr)) + + # ---- reduction: given this plan's predicted chunks, build the result ----- + def reduce(parts: list[NDArray[np.float64]]) -> dict[str, NDArray[np.float64]]: + firsts: list[NDArray[np.float64]] = [] + seconds: list[NDArray[np.float64]] = [] + for bi, steps in enumerate(steps_per_band): + block = parts[main_chunk_idx[bi]] + if vector_branch: + k = steps.size + f = np.empty(k) + s = np.empty(k) + for g in range(k): + lo = np.mean(block[(position == "l") & (ids == g)]) + mid = np.mean(block[(position == "m") & (ids == g)]) + up = np.mean(block[(position == "u") & (ids == g)]) + h = steps[g] + if np.isfinite(h) and h != 0.0: + f[g] = (up - lo) / (2.0 * h) + s[g] = (up - 2.0 * mid + lo) / (h**2) + else: + f[g] = np.nan + s[g] = np.nan + else: + n_eval = steps.size + lo = block[:n_eval] + mid = block[n_eval : 2 * n_eval] + up = block[2 * n_eval :] + finite = np.isfinite(steps) & (steps != 0.0) + with np.errstate(invalid="ignore", divide="ignore"): + f = (up - lo) / (2.0 * steps) + s = (up - 2.0 * mid + lo) / (steps**2) + f = np.where(finite, f, np.nan) + s = np.where(finite, s, np.nan) + firsts.append(f) + seconds.append(s) + + result = {"First": _row_nanmean(firsts), "Second": _row_nanmean(seconds)} + + if mixed: + mixeds: list[NDArray[np.float64]] = [] + for chunk_idx, scale_arr in mixed_meta: + vals = np.full(scale_arr.size, np.nan) + if chunk_idx is not None: + z = parts[chunk_idx] + pos = 0 + for m in range(scale_arr.size): + if np.isfinite(scale_arr[m]): + c = z[pos : pos + 4] + vals[m] = (c[0] + c[3] - c[1] - c[2]) / scale_arr[m] + pos += 4 + mixeds.append(vals) + result["Mixed"] = _row_nanmean(mixeds) + return result + + return chunks, reduce + - if messages: - print("\r") - return output def _dy_d_scalar( From a5765d03e6d3904c5bfb53e88a67b9c34d8ac4fa Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 17 Jul 2026 23:07:14 +0000 Subject: [PATCH 2/2] engine: partial neighbor selection + streaming for multivariate RPM prediction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Universal fix below dy.d_ in the shared multivariate prediction kernels, so NNS.reg, NNS.stack, NNS.boost and dy.d_ all benefit. _mreg_predict_path previously materialised the full query-by-RPM distance matrix and ran a full stable sort of every RPM row for every query, then computed the whole k = 1..kmax path. For a Stack CV fold at n = 100k that is a ~20k x 80k (12.8 GB) matrix plus 20k full 80k-element sorts. Now: * _stable_topk(): exact partial selection of the k nearest rows - O(n + k log k) per query instead of O(n log n) - byte-identical to np.argsort(kind='stable')[:, :k], including heavy distance ties (the duplicated 1-D synthetic coordinates fed by dy.d_ cbind(X*, X*)); * _mreg_predict_path() streams queries in row-chunks (bounded memory, ~O(chunk x n) instead of O(m x n)) and uses _stable_topk; * _mreg_predict() (single selected n.best) also uses _stable_topk. Exact: 344 invariant tests pass unchanged; dy_d_best output is bit-identical. At n = 8000, kmax = 90 the CV-scale path is ~2.4x faster (grows with n) and the materialised matrix drops from 128 MB to a 32 MB chunk. New tests pin _stable_topk and the chunked path to a full-sort reference. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.8 --- src/nns/_reg_engine.py | 64 ++++++++++++++++++++++++++--------- tests/invariants/test_diff.py | 49 +++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 16 deletions(-) diff --git a/src/nns/_reg_engine.py b/src/nns/_reg_engine.py index 5ba7be8..a1f11d0 100644 --- a/src/nns/_reg_engine.py +++ b/src/nns/_reg_engine.py @@ -926,6 +926,30 @@ def _mreg_distances( return np.sqrt(np.sum(z * z, axis=2)) +def _stable_topk(d: NDArray[np.float64], k: int) -> NDArray[np.intp]: + """Indices of the k nearest RPM rows per query, in stable (distance, index) + order - identical to ``np.argsort(d, axis=1, kind="stable")[:, :k]`` but via + partial selection: O(n + k log k) per row instead of a full O(n log n) sort. + + Exactness (including heavy distance ties, e.g. the duplicated 1-D synthetic + coordinates fed by dy.d_) is preserved: candidates are every row within the + k-th smallest distance (``<=`` the threshold, so all boundary ties are + captured), then a stable sort keeps original-index order among equal + distances before truncating to k. + """ + c, n = d.shape + if k >= n: + return np.argsort(d, axis=1, kind="stable") + kth = np.partition(d, k - 1, axis=1)[:, k - 1] + out = np.empty((c, k), dtype=np.intp) + for i in range(c): + row = d[i] + cand = np.flatnonzero(row <= kth[i]) + order = np.argsort(row[cand], kind="stable") + out[i] = cand[order][:k] + return out + + def _mreg_predict( xtest: NDArray[np.float64] | None, rpm_x: NDArray[np.float64], @@ -968,7 +992,7 @@ def _mreg_predict( out[start + i] = float(nns_gravity(finite_tied)) continue - idx = np.argsort(d, axis=1, kind="stable")[:, :k] + idx = _stable_topk(d, k) dk = np.take_along_axis(d, idx, axis=1) yk = rpm_yhat[idx] w = _mreg_ensemble_weights(dk) @@ -1000,21 +1024,29 @@ def _mreg_predict_path( m = xtest.shape[0] out = np.empty((m, kmax), dtype=np.float64) - d = _mreg_distances(rpm_x, xtest, dist, mins, maxs) - idx = np.argsort(d, axis=1, kind="stable") - d_sorted = np.take_along_axis(d, idx, axis=1) - y_sorted = rpm_yhat[idx] - - # k = 1 aggregates all exact nearest-distance ties. - dmin = d.min(axis=1, keepdims=True) - for i in range(m): - tied = rpm_yhat[d[i] == dmin[i, 0]] - finite_tied = tied[np.isfinite(tied)] - out[i, 0] = float(nns_gravity(finite_tied)) - - for k in range(2, kmax + 1): - w = _mreg_ensemble_weights(d_sorted[:, :k]) - out[:, k - 1] = np.sum(y_sorted[:, :k] * w, axis=1) + # Stream the queries in row-chunks so the query-by-RPM distance matrix is + # never materialised in full, and select only the kmax nearest neighbours + # (partial selection) rather than sorting every RPM row. Results are + # byte-for-byte identical to a full stable sort of all rows. + chunk = max(1, int(4_000_000 // max(n, 1))) + for start in range(0, m, chunk): + rows = slice(start, min(start + chunk, m)) + d = _mreg_distances(rpm_x, xtest[rows], dist, mins, maxs) + c = d.shape[0] + idx = _stable_topk(d, kmax) + d_sorted = np.take_along_axis(d, idx, axis=1) + y_sorted = rpm_yhat[idx] + + # k = 1 aggregates all exact nearest-distance ties. + dmin = d.min(axis=1, keepdims=True) + for i in range(c): + tied = rpm_yhat[d[i] == dmin[i, 0]] + finite_tied = tied[np.isfinite(tied)] + out[start + i, 0] = float(nns_gravity(finite_tied)) + + for k in range(2, kmax + 1): + w = _mreg_ensemble_weights(d_sorted[:, :k]) + out[rows, k - 1] = np.sum(y_sorted[:, :k] * w, axis=1) return out diff --git a/tests/invariants/test_diff.py b/tests/invariants/test_diff.py index 81a589f..1871ff4 100644 --- a/tests/invariants/test_diff.py +++ b/tests/invariants/test_diff.py @@ -192,3 +192,52 @@ def test_dy_d_best_matches_pinned_values() -> None: rtol=0.0, atol=1e-4, ) + + +def test_stable_topk_matches_full_stable_sort() -> None: + # Partial neighbor selection must be byte-identical to a full stable argsort, + # including heavy distance ties (e.g. duplicated 1-D synthetic coordinates). + from nns._reg_engine import _stable_topk + + rng = np.random.RandomState(0) + for _ in range(300): + m = rng.randint(1, 6) + n = rng.randint(2, 40) + k = rng.randint(1, n + 1) + d = rng.randint(0, 4, size=(m, n)).astype(float) # ties on purpose + expected = np.argsort(d, axis=1, kind="stable")[:, :k] + assert np.array_equal(_stable_topk(d, k), expected) + + +def test_mreg_predict_path_chunking_is_exact() -> None: + # The chunked/partial-selection path must equal a full-sort reference. + from nns._reg_engine import _mreg_distances, _mreg_ensemble_weights, _mreg_predict_path + from nns.central_tendencies import nns_gravity + + def reference(xtest, rpm_x, rpm_yhat, kmax, dist, mins, maxs): + n = rpm_x.shape[0] + kmax = min(int(kmax), n) + m = xtest.shape[0] + out = np.empty((m, kmax)) + d = _mreg_distances(rpm_x, xtest, dist, mins, maxs) + idx = np.argsort(d, axis=1, kind="stable") + ds = np.take_along_axis(d, idx, axis=1) + ys = rpm_yhat[idx] + dmin = d.min(axis=1, keepdims=True) + for i in range(m): + t = rpm_yhat[d[i] == dmin[i, 0]] + out[i, 0] = float(nns_gravity(t[np.isfinite(t)])) + for k in range(2, kmax + 1): + w = _mreg_ensemble_weights(ds[:, :k]) + out[:, k - 1] = np.sum(ys[:, :k] * w, axis=1) + return out + + rng = np.random.RandomState(1) + rpm_x = rng.uniform(-2, 2, size=(300, 3)) + rpm_y = rng.uniform(size=300) + xt = rng.uniform(-2, 2, size=(120, 3)) + mins, maxs = rpm_x.min(0), rpm_x.max(0) + for kmax in (1, 5, 50, 300): + a = reference(xt, rpm_x, rpm_y, kmax, "NNS", mins, maxs) + b = _mreg_predict_path(xt, rpm_x, rpm_y, kmax, "NNS", mins, maxs) + np.testing.assert_allclose(a, b, rtol=0, atol=0)