$N^2$

This package requires Python 3.10 or later. Please verify your Python version by running python --version in your terminal. If you're not running Python 3.10+, please adjust your environment accordingly (for example, if you use pyenv: pyenv local 3.10 or any later version like pyenv local 3.11).

Setup for production usage

To install the package in production, run the following command:

pip install nsquared

This will install the package and all its dependencies. You can then use the package with:

import nsquared as nsq

Setup for Development

If using last release

If you want to use the last release with development functionality enabled, you can install the package with the following command:

pip install nsquared[dev]

This will install the package and all its dependencies, including the development dependencies. You can then use the package with:

import nsquared as nsq

If using latest code (not yet tagged and released)

Dependencies are managed in pyproject.toml. To install the dependencies, run the following commands, based on your Operating System:

POSIX Systems (MacOS/Linux):

python --version   # Ensure this outputs Python 3.10 or later
python -m venv .venv
source .venv/bin/activate
pip install -U pip
pip install -e . # Install in editable mode

Windows Systems:

python --version   # Ensure this outputs Python 3.10 or later
python -m venv .venv
.venv\Scripts\activate
pip install -U pip
pip install -e . # Install in editable mode

Once installed, you can use the package with:

import nsquared as nsq

$N^2$ Bench and New Methods/Datasets

To replicate the experiments in our paper and test out new methods or datasets, check out the bench directory. For direct access to the data used in the $N^2$ bench, we host the data for download on this repo.

Submitting Changes

Linting

Before submitting changes, please run pre-commit hooks to ensure that the code is formatted correctly. To do so, run the following command:

pre-commit run --a

The linter should run without any errors and autofix any issues that it can. If there are any issues that the linter cannot fix, please fix them manually before committing your changes.

Tests

Please ensure that all tests pass before submitting your changes. To run the tests, run the following command:

pytest

Once all tests pass, you may submit your changes.

Tagging and Releasing a New Version

Detailed instructions for releasing a new version of the package can be found in the Release Instructions file. This includes instructions for updating the version number, creating a new release branch, and tagging the release.

Contributing

We welcome contributions to this package! If you have any ideas for new features, bug fixes, or improvements, please feel free to open an issue or a pull request. We appreciate your help in making this package better!

License

This package is licensed under the MIT License. See the LICENSE file for more details.

Examples

This repository implements the methods from the following papers:

Vanilla Nearest Neighbors

@article{dwivedi2022counterfactual,
  title={Counterfactual inference for sequential experiments},
  author={Dwivedi, Raaz and Tian, Katherine and Tomkins, Sabina and Klasnja, Predrag and Murphy, Susan and Shah, Devavrat},
  journal={arXiv preprint arXiv:2202.06891},
  year={2022}
}

Two Sided Nearest Neighbors (TS-NN)

@article{sadhukhan2024adaptivity,
  title={On adaptivity and minimax optimality of two-sided nearest neighbors},
  author={Sadhukhan, Tathagata and Paul, Manit and Dwivedi, Raaz},
  journal={arXiv preprint arXiv:2411.12965},
  year={2024}
}

Doubly Robust Nearest Neighbors (DRNN)

@article{dwivedi2022doubly,
  title={Doubly robust nearest neighbors in factor models},
  author={Dwivedi, Raaz and Tian, Katherine and Tomkins, Sabina and Klasnja, Predrag and Murphy, Susan and Shah, Devavrat},
  journal={arXiv preprint arXiv:2211.14297},
  year={2022}
}

Distributional (Wasserstein) Nearest Neighbors (WassersteinNN)

@article{feitelberg2024distributional,
  title={Distributional Matrix Completion via Nearest Neighbors in the Wasserstein Space},
  author={Feitelberg, Jacob and Choi, Kyuseong and Agarwal, Anish and Dwivedi, Raaz},
  journal={arXiv preprint arXiv:2410.13112},
  year={2024}
}

Distributional (Kernel) Nearest Neighbors

@article{choi2024learning,
  title={Learning Counterfactual Distributions via Kernel Nearest Neighbors},
  author={Choi, Kyuseong and Feitelberg, Jacob and Chin, Caleb and Agarwal, Anish and Dwivedi, Raaz},
  journal={arXiv preprint arXiv:2410.13381},
  year={2024}
}

Adaptively-weighted Nearest Neighbors (AWNN)

@misc{sadhukhan2025adaptivelyweightednearestneighborsmatrix,
      title={Adaptively-weighted Nearest Neighbors for Matrix Completion},
      author={Tathagata Sadhukhan and Manit Paul and Raaz Dwivedi},
      year={2025},
      eprint={2505.09612},
      archivePrefix={arXiv},
      primaryClass={stat.ML},
      url={https://arxiv.org/abs/2505.09612},
}