Add ordered_complement helper, vectorize diff branch, and tweak vignette headers

examples/vignettes/03_correlation_and_dependence.py

@@ -1,8 +1,8 @@
 # %% [markdown]
 # # 03. Getting Started with NNS: Correlation and Dependence
 #
-# Section-for-section port of the R vignette.  Figures are saved beside the
-# script under `output/03_correlation_and_dependence/`.
+# Section-for-section port of `NNSvignette_03_Correlation_and_Dependence.Rmd`.
+# Figures are saved beside the script under `output/03_correlation_and_dependence/`.
 
 from __future__ import annotations
 

examples/vignettes/04_normalization_and_rescaling.py

@@ -1,4 +1,5 @@
-"""04. Getting Started with NNS: Normalization and Rescaling.
+"""# %% [markdown]
+04. Getting Started with NNS: Normalization and Rescaling.
 
 This is an instructional, section-for-section Python port of
 ``NNSvignette_04_Normalization_and_Rescaling.Rmd``.  It preserves the R

src/nns/_indices.py

@@ -0,0 +1,24 @@
+"""Index helpers shared by resampling-heavy estimators."""
+
+from __future__ import annotations
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def ordered_complement(
+    all_index: NDArray[np.int64],
+    excluded: NDArray[np.int64],
+) -> NDArray[np.int64]:
+    """Return ``all_index`` values not present in ``excluded`` without sorting.
+
+    The stack/boost split builders already create unique validation indices
+    against a dense ``0..n-1`` index vector.  A boolean complement avoids
+    ``np.setdiff1d``'s sort/unique work while preserving the input row order.
+    """
+    if all_index.size == 0 or excluded.size == 0:
+        return all_index.copy()
+    mask_size = int(max(np.max(all_index), np.max(excluded))) + 1
+    excluded_mask = np.zeros(mask_size, dtype=bool)
+    excluded_mask[excluded] = True
+    return all_index[~excluded_mask[all_index]]

src/nns/boost.py

@@ -25,6 +25,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from nns._indices import ordered_complement
 from nns._reg_engine import _validate_dist, nns_reg_engine
 from nns._rrng import RRNG
 from nns.central_tendencies import nns_gravity
@@ -422,7 +423,7 @@ def fit_subset(
             trial_validation_index = random_validation_index()
         else:
             trial_validation_index = validation_index
-        trial_train_index = np.setdiff1d(all_index, trial_validation_index)
+        trial_train_index = ordered_complement(all_index, trial_validation_index)
         predicted = fit_subset(subset, trial_train_index, trial_validation_index)
         learner_results[i] = score(predicted, y[trial_validation_index])
 
@@ -516,7 +517,7 @@ def best_trial_subset() -> list[tuple[int, ...]]:
 
             if ts_test_value is None:
                 epoch_validation_index = random_validation_index()
-                epoch_train_index = np.setdiff1d(all_index, epoch_validation_index)
+                epoch_train_index = ordered_complement(all_index, epoch_validation_index)
             else:
                 assert epoch_split_id is not None
                 epoch_train_index, epoch_validation_index = chronological_splits[

src/nns/diff.py

@@ -500,10 +500,6 @@ def _row_nanmean(bands: list[NDArray[np.float64]]) -> NDArray[np.float64]:
         )
         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]] = []
@@ -573,19 +569,16 @@ def reduce(parts: list[NDArray[np.float64]]) -> dict[str, NDArray[np.float64]]:
             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
+                means = np.mean(block.reshape(k, 3, sample_size), axis=2)
+                lo = means[:, 0]
+                mid = means[:, 1]
+                up = means[:, 2]
+                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)
             else:
                 n_eval = steps.size
                 lo = block[:n_eval]

src/nns/stack.py

@@ -25,6 +25,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
+from nns._indices import ordered_complement
 from nns._reg_engine import (
     _mreg_predict_path,
     _mreg_prepare,
@@ -540,7 +541,7 @@ def make_splits() -> list[dict[str, NDArray[np.int64]]]:
                     validation = np.sort(
                         rng.sample_int(n_obs, size, replace=False) - 1
                     ).astype(np.int64)
-                training = np.setdiff1d(all_index, validation)
+                training = ordered_complement(all_index, validation)
                 if training.size < 3 or not has_all_classes(y[training]):
                     raise ValueError(
                         "Unable to create a repeated holdout retaining enough "
@@ -578,7 +579,7 @@ def make_splits() -> list[dict[str, NDArray[np.int64]]]:
         out = []
         for b in range(1, use_folds + 1):
             validation = np.flatnonzero(fold_id == b).astype(np.int64)
-            training = np.setdiff1d(all_index, validation)
+            training = ordered_complement(all_index, validation)
             if validation.size > 0 and training.size >= 3 and has_all_classes(y[training]):
                 out.append({"train": training, "validation": validation})
         if len(out) < 2:

tests/invariants/test_indices.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import numpy as np
+
+from nns._indices import ordered_complement
+
+
+def test_ordered_complement_preserves_source_index_order() -> None:
+    all_index = np.array([4, 2, 0, 3, 1], dtype=np.int64)
+    excluded = np.array([0, 3], dtype=np.int64)
+
+    result = ordered_complement(all_index, excluded)
+
+    np.testing.assert_array_equal(result, np.array([4, 2, 1], dtype=np.int64))