Software

Run the estimator.

A focused Python package for base BNN, two-scale prediction, and regression inference, with CPU and accelerator backends.

Install from the release or source

Python 3.10+ and NumPy 1.24+. The base installation uses NumPy. The distribution is named bagged-nearest-neighbors; the Python import is bnn.

python -m pip install https://github.com/jingbowa/bnn/releases/download/v0.1.0/bagged_nearest_neighbors-0.1.0-py3-none-any.whl

For the examples and current source:

git clone https://github.com/jingbowa/bnn.git
cd bnn
python -m pip install .

The release is distributed through GitHub. A PyPI distribution has not been published.

A complete regression example

import numpy as np
from bnn import BNN, TwoScaleBNN

rng = np.random.default_rng(2026)
X = rng.uniform(0, 1, size=(500, 1))
y = np.sin(2 * np.pi * X[:, 0]) + rng.normal(0, 0.15, 500)
queries = np.linspace(0.1, 0.9, 9).reshape(-1, 1)

base = BNN(s=20).fit(X, y)
model = TwoScaleBNN(s1=20, s2=40).fit(X, y)
prediction = model.predict(queries)
bootstrap = model.bootstrap(
    queries, n_resamples=199, random_state=42
)
jackknife = model.jackknife(queries)

print(prediction)
print(bootstrap.standard_error)
print(jackknife.standard_error)

Features have shape (n, d), the response has shape (n,), and query predictions have shape (m,). Inputs must be finite. Scales are integers satisfying 1 ≤ s ≤ n or 1 ≤ s1 < s2 ≤ n. The largest jackknife scale must be smaller than n.

Bootstrap draws are paired across queries; the routines hold scales fixed. Their uncertainty estimates do not incorporate a separate scale-selection procedure. Use the paper's inference conditions when constructing confidence intervals.

Choose the computation path

Current implementation and hardware checks
BackendRoutinesVerified execution
NumPyBase/two-scale prediction, regression bootstrap and jackknife, held-out scale selectionCPU, float64
PyTorchBatched base/two-scale predictionCPU and NVIDIA CUDA, float32/float64
JAXJIT/batched base/two-scale predictionCPU, float32/float64. GPU and TPU execution has not been checked.
python -m pip install '.[torch]'
# Or, for JAX:
python -m pip install '.[jax]'
from bnn.torch_backend import TorchTwoScaleBNN

model = TorchTwoScaleBNN(
    20, 40, device="cuda", batch_size=8
).fit(X, y)
prediction = model.predict(queries)  # Tensor on the selected device
import jax
jax.config.update("jax_enable_x64", True)
from bnn.jax_backend import JAXTwoScaleBNN

model = JAXTwoScaleBNN(
    20, 40, device=jax.devices()[0], batch_size=8
).fit(X, y)
prediction = model.predict(queries)  # JAX array

Use NumPy for the package's regression resampling. Accelerator installation depends on the hardware runtime: follow the official PyTorch and JAX instructions. JAX's first call includes compilation; precision requires an explicit choice.

Make the statistical choices explicit

Distance uses the features you supply. Standardization, transformations, and feature inclusion therefore define the neighborhood. Apply preprocessing consistently to training data and queries. Adding a dimension changes both the geometry and the dimension-dependent two-scale coefficients.

Exact distance ties preserve the original training-row order. Float32 can change ordering for nearly equal distances. The two-scale coefficients use the actual integer scale ratio, and negative weights are intentional.

select_scales scores supplied pairs using global held-out MSE and reuses query orderings. This is a practical helper; the paper-specific pointwise tuning procedures remain in the replication materials.

Inspect the evidence and extend the code

Independent checks compare weights with binomial coefficients, predictions with exhaustive subsample averages, bootstrap with explicit resampling, and jackknife with delete-one refits. Backend checks examine numerical agreement and precision choices.

The verification record supplies environments, measured workloads, and limits. The replication archive contains the paper-specific extensions and saved results with provenance.

The current package is developed from the elasticity implementation. Patrick Vossler's TDNN R/C++ package and JASA replication are companion resources. Their complete R execution was not part of this verification.

Citation guidance · Report an issue · MIT license.