Fix distributional bugs + auto nn implementation

README.md

@@ -61,7 +61,8 @@ import nsquared as nsq
 > [!NOTE]
 > If using VSCode, make sure to set the interpreter to the .venv environment using `Cmd + Shift + P` -> `Python: Select Interpreter`.
 
-
+## $N^2$ Bench and New Methods/Datasets
+To replicate the experiments in [our paper](https://arxiv.org/abs/2506.04166) 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](https://github.com/calebchin/nsquared_bench_data).
 
 ## Submitting Changes
 ### Linting

bench/README.md

@@ -0,0 +1,169 @@
+# $N^2$ Bench
+In this directory, we provide resources for replicating our experimental results, adding and testing new methods, and adding new datasets.
+## Replicating experimental results
+
+### For Windows users
+We run our experiments with shell scripts, which requires additional setup for Windows users. A recommended option is Git Bash.
+
+#### Using Git Bash
+1. Install [Git for Windows](https://gitforwindows.org/)
+2. Open Git Bash
+3. Navigate to project bench directory
+
+### For Mac / Linux Users
+No additional setup should be necessary.
+
+### Experiments:
+To replicate the experimental results in [our paper](https://arxiv.org/pdf/2506.04166), run `bench/experiments.sh` as follows:
+
+```bash
+# from the bench directory
+./experiments.sh -o OUTPUT_DIR -l LOG_LEVEL -e EXPERIMENT_NAME
+```
+If `EXPERIMENT_NAME` is `all` (default), experiments on all four datasets (heartsteps, movielens, prompteval, and prop99) will be run and stored in `OUTPUT_DIR`. You can also run the experiments for a specific dataset setting `EXPERIMENT_NAME` to `heartsteps`, `movielens`, `prompteval`, and `prop99`.
+
+The experiments for `movielens` are very memory-intensive due to the size of the dataset and could terminate unexpectedly.
+
+For plotting utilities, see the `examples` directory for files with names prefixed by `plot`.
+
+
+## Adding and testing new methods
+
+### Nearest neighbor variants:
+To add a new nearest neighbor variant:
+
+1. Define an `EstimationMethod` in `nsquared/estimation_methods.py`.
+
+   ```python
+   class MyNewEstimator(EstimationMethod):
+    """Estimate things"""
+        def __init__(self, is_percentile: bool = True):
+            super().__init__(is_percentile)
+            ...
+
+        def impute(
+            self,
+            row: int,
+            column: int,
+            data_array: npt.NDArray,
+            mask_array: npt.NDArray,
+            distance_threshold: Union[float, Tuple[float, float]],
+            data_type: DataType,
+            allow_self_neighbor: bool = False,
+            **kwargs: Any,
+        ) -> npt.NDArray:
+        """
+        Impute the missing value at
+        the given row and column using
+        XYZ method
+        """
+        ...
+   ```
+
+2. Define a `FitMethod` (if needed).
+    ```python
+    class MyNewFitMethod(FitMethod):
+        def __init__(
+            self,
+            block: list[tuple[int, int]],
+            distance_threshold_range_row: tuple[float, float],
+            distance_threshold_range_col: tuple[float, float],
+            alpha_range: tuple[float, float],
+            n_trials: int,
+            data_type: DataType,
+            allow_self_neighbor: bool = False,
+        ):
+            ...
+
+    def fit(
+            self,
+            data_array: npt.NDArray,
+            mask_array: npt.NDArray,
+            imputer: NearestNeighborImputer,
+            ret_trials: bool = False,
+        ) -> Union[tuple[float, float], tuple[tuple[float, float], Trials]]:
+            ...
+
+    ```
+
+3. To ensure compatability with the experiments and plotting functions, add an alias for your method in [`nsquared/utils/experiments.py`](https://github.com/aashish-khub/NearestNeighbors/blob/441a382efba2cf68c22ac379bd5750f91d9e03ee/src/nsquared/utils/experiments.py#L19) under the parser argument for `--estimation_method` and fill out [`nsquared/utils/plotting_utils.py`](https://github.com/aashish-khub/NearestNeighbors/blob/441a382efba2cf68c22ac379bd5750f91d9e03ee/src/nsquared/utils/plotting_utils.py#L19) with your desired plot settings.
+
+4. To test on a given dataset (for example: `heartsteps`), navigate to `examples/heartsteps/run_scalar.py` and add your method following the existing template. In general, all that is required is adding a block to [this conditional structure](https://github.com/aashish-khub/NearestNeighbors/blob/441a382efba2cf68c22ac379bd5750f91d9e03ee/examples/heartsteps/run_scalar.py#L157), but your method is not required to use the exact format of the `run_scalar.py` script. For distributional methods, do the same, but in `run_distribution.py`.
+
+5. Adjust the corresponding `slurm_scripts/*.sh` files in the `examples` directory to add your new methods alias under `METHODS`.
+
+### Non-nearest neigbor methods
+If you wish to test methods on the benchmark that do not follow the $N^2$ framework, there are two options:
+
+#### Option 1:
+Use downloaded data hosted [here](https://github.com/calebchin/nsquared_bench_data).
+
+This contains the masked (missingness included) matrix used in our experiments, the corresponding unmasked (no missingness) matrix, and the masking matrix for each dataset.
+
+#### Option 2:
+Use the dataloader in `nsquared.datasets`.
+
+For examples on how the dataloader works, check out the use of the loader [here](https://github.com/aashish-khub/NearestNeighbors/blob/main/examples/heartsteps/run_scalar.py) or [here](https://github.com/aashish-khub/NearestNeighbors/blob/main/examples/prop99/run_scalar.py).
+
+
+## Adding new datasets
+
+To add a new dataset:
+
+1. Create a new directory inside the `nsquared/datasets` directory with two files: `loader.py` and `__init__.py`.
+2. The `loader.py` file should look something like:
+   ```python
+    from nsquared.datasets.dataloader_base import NNDataLoader
+    from nsquared.datasets.dataloader_factory import register_dataset
+    # imports
+    ...
+    #
+    memory = Memory(".joblib_cache", verbose=2)
+    logger = logging.getLogger(__name__)
+    params = {
+    # specific parameters for the dataset
+    }
+
+    @register_dataset("MyNewDataset", params)
+    class MyNewDatasetLoader(NNDataLoader):
+        """
+        ...
+        """
+        def __init__(
+            self,
+            # specific params
+            agg: str = "mean",
+            **kwargs: Any,
+            ):
+            """
+            Initialize data loader
+            """
+            super().__init__(
+                agg=agg,
+                **kwargs,
+            )
+            ...
+
+        def process_data_scalar(self, agg: str = "mean") -> tuple[np.ndarray, np.ndarray]:
+            """
+            Process new dataset in the scalar setting (# of entries is 1)
+            """
+            ...
+
+        def process_data_distribution(self, data_type: DataType | None = None) -> tuple[np.ndarray, np.ndarray]:
+            """
+            Process new dataset in distribuional setting (# of entries is > 1)
+            """
+            ...
+
+        def get_full_state_as_dict(self, include_metadata: bool = False) -> dict:
+            """
+            Returns the full state as a dictionary (including data matrix, mask, and specific params).
+            """
+            ...
+   ```
+3. The `__init__.py` should contain:
+   ```python
+   from .loader import MyNewDatasetLoader  # noqa: F401
+   ```
+4. To run the existing methods on the dataset, follow the template provided in any of the example `run_scalar.py` or `run_distribution.py` files. It is possible that all you will need to change is the data loading step (e.g. `my_new_dataloader = NNData.create("MyNewDataset")`).

bench/experiments.sh

@@ -0,0 +1,94 @@
+#!/bin/bash
+
+# Example usage:
+
+OUTPUT_DIR="out"
+LOG_LEVEL="WARNING"
+EXPERIMENT="all"
+
+
+usage() {
+    echo "Usage: $0 [-o OUTPUT_DIR] [-l LOG_LEVEL] [-e EXPERIMENT]"
+    exit 1
+}
+
+while getopts ":o:l:e:" opt; do
+    case ${opt} in
+        o ) OUTPUT_DIR=$OPTARG ;;
+        l ) LOG_LEVEL=$OPTARG ;;
+        e ) EXPERIMENT=$OPTARG ;;
+        \? )
+        echo "Invalid option: -$OPTARG" 1>&2
+        usage
+        ;;
+        : )
+        echo "Invalid option: -$OPTARG requires an argument" 1>&2
+        usage
+        ;;
+    esac
+done
+shift $((OPTIND -1))
+echo $EXPERIMENT
+
+if [ "$EXPERIMENT" = "all" ]
+then
+    EXPERIMENTS=("heartsteps" "movielens" "prompteval" "prop99")
+else
+    EXPERIMENTS=($EXPERIMENT)
+fi
+
+for exper in ${EXPERIMENTS[@]};
+do
+    if [ "$exper" = "heartsteps" ]
+    then
+        HEARTSTEPS_DIR="../../bench/${OUTPUT_DIR}/heartsteps"
+        echo "Running heartsteps experiment"
+        cd ../examples/heartsteps
+        ./slurm_scripts/run_accuracy.sh $HEARTSTEPS_DIR $LOG_LEVEL
+        python run_distribution.py -od $HEARTSTEPS_DIR -dt kernel_mmd --force --log_level $LOG_LEVEL -em col-col
+        python run_distribution.py -od $HEARTSTEPS_DIR -dt wasserstein_samples --force --log_level $LOG_LEVEL -em col-col
+        cd ../../bench
+    elif [ "$exper" = "movielens" ]
+    then
+        MOVIELENS_DIR="../../bench/${OUTPUT_DIR}/movielens"
+        echo "Running movielens experiment"
+        cd ../examples/movielens
+        ./slurm_scripts/run_accuracy.sh $MOVIELENS_DIR $LOG_LEVEL
+        cd ../../bench
+    elif [ "$exper" = "prompteval" ]
+    then
+        PROMPTEVAL_DIR="../../bench/${OUTPUT_DIR}/prompteval"
+        echo "Running prompteval experiment"
+        cd ../examples/prompteval
+        ./slurm_scripts/run_accuracy.sh $PROMPTEVAL_DIR $LOG_LEVEL
+        python run_distribution.py -od $PROMPTEVAL_DIR -dt kernel_mmd --force --log_level $LOG_LEVEL -em col-col
+        python run_distribution.py -od $PROMPTEVAL_DIR -dt wasserstein_samples --force --log_level $LOG_LEVEL -em col-col
+        cd ../../bench
+    elif [ "$exper" = "prop99" ]
+    then
+        PROP99_DIR="../../bench/${OUTPUT_DIR}/prop99"
+        cd ../examples/prop99
+        echo "Running prop99 experiment"
+        ./slurm_scripts/run_accuracy.sh $PROP99_DIR $LOG_LEVEL
+        ./slurm_scripts/run_california.sh $PROP99_DIR $LOG_LEVEL
+        python proposal_99.py -od $OUTPUT_DIR --force --log_level $LOG_LEVEL
+        cd ../../bench
+    else
+        echo "$exper is not a valid experiment. Please choose from heartsteps, movielens, prompteval, or prop99."
+    fi
+done
+# SIM_METHODS=(
+#     "auto"
+#     "dr"
+#     "ts"
+# )
+# SIM_OUTPUT_DIR="${OUTPUT_DIR}/simulations"
+# for em in ${SIM_METHODS[@]};
+# do
+#     python ../examples/simulations/run_scalar.py -od "${SIM_OUTPUT_DIR}/high_snr" -em $em --force --log_level $LOG_LEVEL -nstd 0.001
+# done
+
+# for em in ${SIMMETHODS[@]};
+# do
+#     python ../examples/simulations/run_scalar.py -od ${SIM_OUTPUT_DIR}/low_snr -em $em --force --log_level $LOG_LEVEL -nstd 1.0
+# done

examples/heartsteps/plot_distribution_histogram.py

@@ -0,0 +1,99 @@
+"""Script to plot the distribution of step counts for different estimation methods
+
+NOTE: imputation time is per imputation, fit time is for the entire fitting procedure
+TODO: change fit time to be a bar plot
+
+Example usage from the examples/heartsteps directory:
+```bash
+python plot_distribution_histogram.py -od OUTPUT_DIR
+```
+"""
+
+import os
+
+import matplotlib.pyplot as plt
+import logging
+import numpy as np
+
+from nsquared.utils.experiments import get_base_parser
+from nsquared.utils import plotting_utils
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+parser = get_base_parser()
+args = parser.parse_args()
+output_dir = args.output_dir
+results_dir = os.path.join(output_dir, "results")
+figures_dir = os.path.join(output_dir, "figures")
+os.makedirs(figures_dir, exist_ok=True)
+
+wasserstein_ests = np.load(
+    os.path.join(results_dir, "imputations-col-col-lbo-wasserstein_samples.npy"),
+    allow_pickle=True,
+)
+kernel_ests = np.load(
+    os.path.join(results_dir, "imputations-col-col-lbo-kernel.npy"), allow_pickle=True
+)
+ground_truth = np.load(
+    os.path.join(results_dir, "ground_truth-col-col-lbo.npy"), allow_pickle=True
+)
+
+# NOTE: set the width to be the physical size of the figure in inches
+# The NeurIPS text is 5.5 inches wide and 9 inches long
+# If we use wrapfigure with 0.4\textwidth, then the figure needs to be 2.2 inches wide
+for i in range(0, len(ground_truth)):
+    fig = plt.figure(figsize=(2.2, 2))
+    # Create boxplot
+    ax = fig.add_subplot(111)
+    # box = ax.boxplot(
+    #     df_grouped[col_name], patch_artist=True, widths=0.6, showfliers=False
+    # )
+    bins = list(np.linspace(0, 8, 12))
+    weights_gt = np.ones_like(ground_truth[i]) / len(ground_truth[i])
+    ax.hist(
+        ground_truth[i],
+        bins=bins,
+        weights=weights_gt,
+        alpha=0.6,
+        label="Ground truth",
+        color="white",
+        edgecolor="black",
+        linestyle="--",
+    )
+    weights_kernel = np.ones_like(kernel_ests[i]) / len(kernel_ests[i])
+    ax.hist(
+        kernel_ests[i],
+        bins=bins,
+        weights=weights_kernel,
+        alpha=0.6,
+        label=str(plotting_utils.METHOD_ALIASES_SINGLE_LINE.get("kernel", "kernel")),
+        color="teal",
+    )
+    weights_wasserstein = np.ones_like(wasserstein_ests[i]) / len(wasserstein_ests[i])
+    ax.hist(
+        wasserstein_ests[i],
+        bins=bins,
+        weights=weights_wasserstein,
+        alpha=0.6,
+        label=str(
+            plotting_utils.METHOD_ALIASES_SINGLE_LINE.get("wasserstein_samples", "W2S")
+        ),
+        color="orange",
+    )
+
+    ax.set_ylim(0, None)
+
+    ax.set_ylabel("Proportion", fontsize=plotting_utils.LABEL_FONT_SIZE)
+    ax.set_xlabel("Step count", fontsize=plotting_utils.LABEL_FONT_SIZE)
+    ax.legend(loc="upper right", fontsize=plotting_utils.LEGEND_FONT_SIZE)
+
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.spines["left"].set_visible(False)
+    ax.grid(True, alpha=0.4)
+
+    save_path = os.path.join(figures_dir, f"hs_distplot_{i}.pdf")
+    logger.info(f"Saving plot to {save_path}...")
+    plt.savefig(save_path, bbox_inches="tight")
+    plt.close()