dy_d_best: fit NNS.stack once per call (major speedup, identical resu…

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
 
 

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,100 +506,112 @@ 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]
         for band in h_s:
             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(