diff --git a/README.md b/README.md index aa23650..f6f7dea 100644 --- a/README.md +++ b/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 diff --git a/bench/README.md b/bench/README.md new file mode 100644 index 0000000..5f222c0 --- /dev/null +++ b/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")`). \ No newline at end of file diff --git a/bench/experiments.sh b/bench/experiments.sh new file mode 100755 index 0000000..76997d5 --- /dev/null +++ b/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 diff --git a/examples/heartsteps/plot_distribution_histogram.py b/examples/heartsteps/plot_distribution_histogram.py new file mode 100644 index 0000000..725c0b1 --- /dev/null +++ b/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() diff --git a/examples/heartsteps/plot_error_and_time.py b/examples/heartsteps/plot_error_and_time.py index 4a22eac..dca297c 100644 --- a/examples/heartsteps/plot_error_and_time.py +++ b/examples/heartsteps/plot_error_and_time.py @@ -32,6 +32,7 @@ files = glob(os.path.join(results_dir, "est_errors-*.csv")) df_list = [] for file in files: + print(file) df = pd.read_csv(file) df_list.append(df) df = pd.concat(df_list, ignore_index=True) @@ -43,8 +44,19 @@ .reset_index() ) # rearrange the order of the estimation methods by -# "usvt", "row-row", "col-col", "dr", "ts", "aw" -ORDER = ["usvt", "softimpute", "col-col", "row-row", "dr", "ts", "aw"] +# "usvt", "row-row", "col-col", "dr", "ts", "star" +ORDER = [ + "usvt", + "softimpute", + "col-col", + "row-row", + "dr", + "ts", + "auto", + "star", + "kernel", + "wasserstein_samples", +] df_grouped = df_grouped.sort_values( by="estimation_method", key=lambda x: x.map(lambda y: ORDER.index(y)) ) @@ -57,7 +69,7 @@ # 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 - fig = plt.figure(figsize=(2.2, 2)) + fig = plt.figure(figsize=(3.25, 2)) # Create boxplot ax = fig.add_subplot(111) box = ax.boxplot( @@ -83,7 +95,7 @@ ) ax.tick_params(axis="x", length=0) # Set tick length to 0 ax.set_ylabel(alias, fontsize=plotting_utils.LABEL_FONT_SIZE) - # ax.set_xlabel("Estimation method", fontsize=plotting_utils.LABEL_FONT_SIZE) + # ax.set_xlabel("Estimation methods",color="white", fontsize=plotting_utils.LABEL_FONT_SIZE) ax.spines["top"].set_visible(False) ax.spines["right"].set_visible(False) diff --git a/examples/heartsteps/run_distribution.py b/examples/heartsteps/run_distribution.py index 90246a1..57c5611 100644 --- a/examples/heartsteps/run_distribution.py +++ b/examples/heartsteps/run_distribution.py @@ -22,6 +22,7 @@ from nsquared.data_types import ( DistributionKernelMMD, DistributionWassersteinSamples, + Scalar, ) from nsquared import NearestNeighborImputer from nsquared.fit_methods import ( @@ -38,17 +39,11 @@ "--data_type", "-dt", type=str, - default="kernel", - choices=["kernel", "wasserstein_samples"], + default="kernel_mmd", + choices=["kernel_mmd", "wasserstein_samples"], help="Data type to use", ) args = parser.parse_args() -if args.data_type == "kernel": - data_type = DistributionKernelMMD(kernel="exponential") -elif args.data_type == "wasserstein_samples": - data_type = DistributionWassersteinSamples() -else: - raise ValueError(f"Data type {args.data_type} not supported") output_dir = args.output_dir estimation_method = args.estimation_method fit_method = args.fit_method @@ -61,29 +56,37 @@ os.makedirs(output_dir, exist_ok=True) results_dir = os.path.join(output_dir, "results") os.makedirs(results_dir, exist_ok=True) -save_path = os.path.join( - results_dir, f"est_errors-{estimation_method}-{fit_method}-{data_type}.csv" -) - -if os.path.exists(save_path) and not args.force: - logger.info(f"Results already exist at {save_path}. Use --force to overwrite.") - exit() rng = np.random.default_rng(seed=seed) # Load the heartsteps dataset # NOTE: the raw and processed data is cached in .joblib_cache start_time = time() -hs_dataloader = NNData.create("heartsteps") +hs_dataloader = NNData.create("heartsteps", freq="1min", num_measurements=60) data, mask = hs_dataloader.process_data_distribution() data = data[:, :200] # only use the first 200 timesteps mask = mask[:, :200] + elapsed_time = time() - start_time logger.info(f"Time to load and process data: {elapsed_time:.2f} seconds") logger.info("Using distribution data type") -data_type_kernel = DistributionKernelMMD(kernel="exponential") -data_type_wasserstein = DistributionWassersteinSamples() + +if args.data_type == "kernel": + data_type = DistributionKernelMMD(kernel="exponential") +elif args.data_type == "wasserstein_samples": + data_type = DistributionWassersteinSamples(num_samples=data[0, 0].shape[0]) +else: + raise ValueError(f"Data type {args.data_type} not supported") + +save_path = os.path.join( + results_dir, f"est_errors-{estimation_method}-{fit_method}-{args.data_type}.csv" +) + +if os.path.exists(save_path) and not args.force: + logger.info(f"Results already exist at {save_path}. Use --force to overwrite.") + exit() + holdout_inds = np.nonzero(mask == 1) inds_rows = holdout_inds[0] @@ -127,16 +130,16 @@ if estimation_method == "row-row": logger.info("Using row-row estimation") imputer = NearestNeighborImputer( - estimation_method=RowRowEstimator(is_percentile=False), data_type=data_type + estimation_method=RowRowEstimator(is_percentile=True), data_type=data_type ) logger.info("Using leave-block-out validation") fitter = LeaveBlockOutValidation( - block, distance_threshold_range=(0, 100), n_trials=100, data_type=data_type + block, distance_threshold_range=(0, 1), n_trials=100, data_type=data_type ) elif estimation_method == "col-col": logger.info("Using col-col estimation") imputer = NearestNeighborImputer( - estimation_method=ColColEstimator(), data_type=data_type + estimation_method=ColColEstimator(is_percentile=True), data_type=data_type ) logger.info("Using leave-block-out validation") @@ -196,15 +199,18 @@ ground_truth = data[test_inds_rows, test_inds_cols] est_errors = [] +error_datatype = Scalar() for i in range(len(imputations)): - est_errors.append(imputer.data_type.distance(imputations[i], ground_truth[i])) + est_errors.append( + error_datatype.distance(np.nanmean(imputations[i]), np.nanmean(ground_truth[i])) + ) # est_errors = np.abs(imputations - ground_truth) logger.info(f"Mean absolute error: {np.mean(est_errors)}") save_imputations = np.array(imputations, dtype=object) ground_truth = np.array(ground_truth, dtype=object) # save imputations to pkl imputations_save_path = os.path.join( - results_dir, f"imputations-{estimation_method}-{fit_method}.npy" + results_dir, f"imputations-{estimation_method}-{fit_method}-{args.data_type}.npy" ) logger.info(f"Saving imputations to {imputations_save_path}...") np.save(imputations_save_path, save_imputations) @@ -217,7 +223,7 @@ df = pd.DataFrame( data={ - "estimation_method": estimation_method, + "estimation_method": args.data_type, # "imputation": save_imputations, # "ground_truth": ground_truth, "data_type": args.data_type, diff --git a/examples/heartsteps/run_scalar.py b/examples/heartsteps/run_scalar.py index 97bc270..5f26591 100644 --- a/examples/heartsteps/run_scalar.py +++ b/examples/heartsteps/run_scalar.py @@ -21,12 +21,13 @@ # import nearest neighbor methods from nsquared.data_types import Scalar -from nsquared.estimation_methods import AWNNEstimator, TSEstimator +from nsquared.estimation_methods import AWNNEstimator, TSEstimator, AutoEstimator from nsquared import NearestNeighborImputer from nsquared.fit_methods import ( DRLeaveBlockOutValidation, TSLeaveBlockOutValidation, LeaveBlockOutValidation, + AutoDRTSLeaveBlockOutValidation, ) from nsquared.datasets.dataloader_factory import NNData from nsquared.vanilla_nn import row_row, col_col @@ -61,7 +62,7 @@ # Load the heartsteps dataset # NOTE: the raw and processed data is cached in .joblib_cache start_time = time() -hs_dataloader = NNData.create("heartsteps", agg="std") +hs_dataloader = NNData.create("heartsteps", agg="mean") data, mask = hs_dataloader.process_data_scalar() data = data[:, :200] # only use the first 200 timesteps mask = mask[:, :200] @@ -269,9 +270,25 @@ distance_threshold_range_col=(0, 1), n_trials=100, data_type=data_type, + allow_self_neighbor=True, ) # for TSNN, self neighbor is necessary (for now) allow_self_neighbor = True + elif estimation_method == "auto": + logger.info("Using AutoNN estimation") + estimator = AutoEstimator(is_percentile=True) + imputer = NearestNeighborImputer(estimator, data_type) + logger.info("Using AutoNN fit method") + # Fit the imputer using leave-block-out validation + fitter = AutoDRTSLeaveBlockOutValidation( + block, + distance_threshold_range_row=(0, 1), + distance_threshold_range_col=(0, 1), + alpha_range=(0, 1), + n_trials=200, + data_type=data_type, + allow_self_neighbor=args.allow_self_neighbor, + ) else: raise ValueError( f"Estimation method {estimation_method} and fit method {fit_method} not supported" diff --git a/examples/heartsteps/slurm_scripts/run_accuracy.sh b/examples/heartsteps/slurm_scripts/run_accuracy.sh index db7da6d..9e0d3ce 100755 --- a/examples/heartsteps/slurm_scripts/run_accuracy.sh +++ b/examples/heartsteps/slurm_scripts/run_accuracy.sh @@ -14,6 +14,7 @@ METHODS=( "dr" "ts" "aw" + "auto" ) for em in ${METHODS[@]}; do diff --git a/examples/movielens/run_scalar.py b/examples/movielens/run_scalar.py new file mode 100644 index 0000000..15e626e --- /dev/null +++ b/examples/movielens/run_scalar.py @@ -0,0 +1,327 @@ +"""Script to run NN imputers + USVT baseline on the MovieLens dataset +using 20% of the observed indices as a test block. + +Example usage (from root of repo): +```bash +python run_scalar.py -od OUTPUT_DIR -em ESTIMATION_METHOD -fm FIT_METHOD +``` +""" + +# %% +# standard imports +import numpy as np +from tqdm import tqdm +import logging +import os +from time import time +import pandas as pd +from hyperopt import Trials + +# import baseline methods +from baselines import usvt, softimpute + +# import nearest neighbor methods +from nsquared.data_types import Scalar +from nsquared.estimation_methods import ( + ColColEstimator, + RowRowEstimator, + TSEstimator, + AWNNEstimator, + AutoEstimator, +) +from nsquared import NearestNeighborImputer +from nsquared.fit_methods import ( + DRLeaveBlockOutValidation, + TSLeaveBlockOutValidation, + LeaveBlockOutValidation, + AutoDRTSLeaveBlockOutValidation, +) +from nsquared.datasets.dataloader_factory import NNData +from nsquared.vanilla_nn import row_row, col_col +from nsquared.dr_nn import dr_nn + +from nsquared.utils.experiments import get_base_parser, setup_logging + +# %% + +parser = get_base_parser() +args = parser.parse_args() +output_dir = args.output_dir +estimation_method = args.estimation_method +fit_method = args.fit_method +seed = args.seed +log_level = args.log_level + +setup_logging(log_level) +logger = logging.getLogger(__name__) + +os.makedirs(output_dir, exist_ok=True) +results_dir = os.path.join(output_dir, "results") +os.makedirs(results_dir, exist_ok=True) +save_path = os.path.join( + results_dir, f"est_errors-{estimation_method}-{fit_method}.csv" +) + +if os.path.exists(save_path) and not args.force: + logger.info(f"Results already exist at {save_path}. Use --force to overwrite.") + exit() + +rng = np.random.default_rng(seed=seed) + +# Load the movielens +# NOTE: the raw and processed data is cached in .joblib_cache +start_time = time() +sample_users = None # 1000 +sample_movies = None # 1000 +seed = 0 +ml_dataloader = NNData.create( + "movielens", sample_users=sample_users, sample_movies=sample_movies, seed=seed +) + +data, mask = ml_dataloader.process_data_scalar() +data_sparsity = 1 - np.sum(mask) / mask.size +data_shape = data.shape +logger.info(f"Data shape: {data_shape}") +logger.info(f"Data sparsity: {data_sparsity:.2%}") +elapsed_time = time() - start_time +logger.info(f"Time to load and process data: {elapsed_time:.2f} seconds") + +logger.info("Using scalar data type") +data_type = Scalar() + +holdout_inds = np.nonzero(mask == 1) +inds_rows = holdout_inds[0] +inds_cols = holdout_inds[1] + +range_inds = np.arange(len(inds_rows)) + +# randomly shuffle indices +rng.shuffle(range_inds) +# 20% of the indices will be used for testing +test_size = int(0.8 * len(range_inds)) +test_inds = range_inds[:test_size] +test_inds = test_inds[:500] +# 80% of the indices will be used for training +train_inds = range_inds[test_size:] +range_train_inds = np.arange(len(train_inds)) +rng.shuffle(range_train_inds) +# 20% of the training indices will be used for cv holdout +# cv_size = int(0.01 * len(train_inds)) +cv_size = 100 +cv_inds = range_train_inds[:cv_size] +# get the rows and columns of the train indices + +cv_inds_rows = list(inds_rows[train_inds][cv_inds]) +cv_inds_cols = list(inds_cols[train_inds][cv_inds]) +# get the rows and columns of the test indices +test_inds_rows = list(inds_rows[test_inds]) +test_inds_cols = list(inds_cols[test_inds]) + +block = list(zip(cv_inds_rows, cv_inds_cols)) + +# # Convert dense data_array to sparse, skipping NaNs using mask +valid_mask = (~np.isnan(data)) & (mask == 1) +# rows, cols = np.where(valid_mask) +# values = data[rows, cols] + +# data_sparse = coo_matrix((values, (rows, cols)), shape=data.shape).tocsc() + +test_block = list(zip(test_inds_rows, test_inds_cols)) + +mask_test = mask.copy() +mask_test[test_inds_rows, test_inds_cols] = 0 + +# %% + +num_trials = 10 +if estimation_method == "usvt": + logger.info("Using USVT estimation") + # setup usvt imputation + usvt_data = data.copy() + usvt_mask = mask.copy() + usvt_mask[test_inds_rows, test_inds_cols] = 0 + usvt_data[usvt_mask != 1] = np.nan + # impute missing values simultaneously + start_time = time() + usvt_imputed = usvt(usvt_data) + elapsed_time = time() - start_time + imputations = usvt_imputed[test_inds_rows, test_inds_cols] + # set the time to the average time per imputation + imputation_times = [elapsed_time / len(test_block)] * len(test_block) + fit_times = [0] * len(test_block) +elif estimation_method == "softimpute": + logger.info("Using SoftImpute estimation") + # setup softimpute imputation + si_data = data.copy() + si_mask = mask.copy() + si_mask[test_inds_rows, test_inds_cols] = 0 + si_data[si_mask != 1] = np.nan + # impute missing values simultaneously + start_time = time() + si_imputed = softimpute(si_data) + elapsed_time = time() - start_time + imputations = si_imputed[test_inds_rows, test_inds_cols] + # set the time to the average time per imputation + imputation_times = [elapsed_time / len(test_block)] * len(test_block) + fit_times = [0] * len(test_block) +elif estimation_method == "aw": + logger.info("Using AWNN estimation") + estimator = AWNNEstimator() + imputer = NearestNeighborImputer(estimator, data_type, distance_threshold=-1) + # Impute missing values + imputations = [] + imputation_times = [] + for row, col in tqdm(test_block, desc="Imputing missing values"): + start_time = time() + imputed_value = imputer.impute(row, col, data, mask_test) + elapsed_time = time() - start_time + imputation_times.append(elapsed_time) + imputations.append(imputed_value) + imputations = np.array(imputations) + fit_times = [0] * len(test_block) +else: + if estimation_method == "dr": + logger.info("Using doubly robust estimation") + imputer = dr_nn() + + logger.info("Using doubly robust fit method") + # Fit the imputer using leave-block-out validation + fitter = DRLeaveBlockOutValidation( + block, + distance_threshold_range_row=(0, 1), + distance_threshold_range_col=(0, 1), + n_trials=num_trials, + data_type=data_type, + ) + elif estimation_method == "row-row": + logger.info("Using row-row estimation") + imputer = row_row() + if not isinstance(imputer.estimation_method, RowRowEstimator): + raise ValueError( + f"Estimation method {imputer.estimation_method} not supported for row-row" + ) + # imputer.estimation_method._precalculate_distances( + # data, mask, np.array(cv_inds_rows) + # ) + + logger.info("Using leave-block-out validation") + fitter = LeaveBlockOutValidation( + block, + distance_threshold_range=(0, 1), + n_trials=num_trials, + data_type=data_type, + ) + elif estimation_method == "col-col": + logger.info("Using col-col estimation") + imputer = col_col() + + if not isinstance(imputer.estimation_method, ColColEstimator): + raise ValueError( + f"Estimation method {imputer.estimation_method} not supported for col-col" + ) + + # imputer.estimation_method.estimator._precalculate_distances( + # np.swapaxes(data, 0, 1), np.swapaxes(mask, 0, 1), np.array(cv_inds_cols) + # ) + + logger.info("Using leave-block-out validation") + fitter = LeaveBlockOutValidation( + block, + distance_threshold_range=(0, 1), + n_trials=num_trials, + data_type=data_type, + ) + elif estimation_method == "ts": + logger.info("Using two-sided estimation") + estimator = TSEstimator() + imputer = NearestNeighborImputer(estimator, data_type) + + logger.info("Using two-sided fit method") + # Fit the imputer using leave-block-out validation + fitter = TSLeaveBlockOutValidation( + block, + distance_threshold_range_row=(0, 1), + distance_threshold_range_col=(0, 1), + n_trials=num_trials, + data_type=data_type, + ) + elif estimation_method == "auto": + logger.info("Using AutoNN estimation") + estimator = AutoEstimator(is_percentile=True) + imputer = NearestNeighborImputer(estimator, data_type) + logger.info("Using AutoNN fit method") + # Fit the imputer using leave-block-out validation + fitter = AutoDRTSLeaveBlockOutValidation( + block, + distance_threshold_range_row=(0, 1), + distance_threshold_range_col=(0, 1), + alpha_range=(0, 1), + n_trials=num_trials, + data_type=data_type, + allow_self_neighbor=args.allow_self_neighbor, + ) + else: + raise ValueError( + f"Estimation method {estimation_method} and fit method {fit_method} not supported" + ) + + start_time = time() + trials = fitter.fit(data, valid_mask, imputer, ret_trials=True) + end_time = time() + fit_times = [end_time - start_time] * len(test_block) + + # CODE FOR EXTRACTING TRIAL METADATA + if ( + not isinstance(trials, float) + and not isinstance(trials, int) + and isinstance(trials[1], Trials) + ): + trials = trials[1] + trial_data = [] + for trial in trials.trials: + row = {} + # get param vals + params = trial["misc"]["vals"] + for param_name, param_values in params.items(): + if param_values: + row[param_name] = float(param_values[0]) + + row["loss"] = float(trial["result"]["loss"]) + trial_data.append(row) + + df_trials = pd.DataFrame(trial_data) + trials_save_path = os.path.join( + results_dir, f"cvtrials-{estimation_method}-{fit_method}.csv" + ) + logger.info(f"Saving trials data to {trials_save_path}...") + df_trials.to_csv(trials_save_path, index=False) + + # Impute missing values + imputations = [] + imputation_times = [] + for row, col in tqdm(test_block, desc="Imputing missing values"): + start_time = time() + imputed_value = imputer.impute(row, col, data, mask_test) + elapsed_time = time() - start_time + imputation_times.append(elapsed_time) + imputations.append(imputed_value) + imputations = np.array(imputations) + +ground_truth = data[test_inds_rows, test_inds_cols] +est_errors = np.abs(imputations - ground_truth) +logger.info(f"Mean absolute error: {np.mean(est_errors)}") + +df = pd.DataFrame( + data={ + "estimation_method": estimation_method, + "fit_method": fit_method, + "est_errors": est_errors, + "row": test_inds_rows, + "col": test_inds_cols, + "time_impute": imputation_times, + "time_fit": fit_times, + } +) +print(df[["est_errors", "time_impute", "time_fit"]].describe()) +logger.info(f"Saving est_errors to {save_path}...") +df.to_csv(save_path, index=False) diff --git a/examples/movielens/slurm_scripts/run_accuracy.sh b/examples/movielens/slurm_scripts/run_accuracy.sh new file mode 100755 index 0000000..0ec14c7 --- /dev/null +++ b/examples/movielens/slurm_scripts/run_accuracy.sh @@ -0,0 +1,39 @@ +# #!/bin/bash + +# # Example usage: +# # ./run_accuracy.sh OUTPUT_DIR + +# OUTPUT_DIR=$1 + +# METHODS=( +# "usvt" +# "row-row" +# "col-col" +# "dr" +# "ts" +# ) +# for em in ${METHODS[@]}; +# do +# python run_scalar.py -od $OUTPUT_DIR -em $em --force +# done +#!/bin/bash + +# Example usage: +# ./run_accuracy.sh OUTPUT_DIR + +OUTPUT_DIR=$1 + +METHODS=( + "usvt" + "softimpute" + "row-row" + "col-col" + "dr" + "ts" + "aw" + "auto" +) +for em in ${METHODS[@]}; +do + python run_scalar.py -od $OUTPUT_DIR -em $em --force +done \ No newline at end of file diff --git a/examples/prompteval/plot_dist_error.py b/examples/prompteval/plot_dist_error.py index 3c382b0..1c1658d 100644 --- a/examples/prompteval/plot_dist_error.py +++ b/examples/prompteval/plot_dist_error.py @@ -33,7 +33,8 @@ figures_dir = os.path.join(output_dir, "figures") os.makedirs(figures_dir, exist_ok=True) -files = glob(os.path.join(results_dir, f"est_errors-*-p{propensity}-*.csv")) +files = glob(os.path.join(results_dir, f"est_errors-*-p{propensity}-tp4.0.csv")) +print(files) df_list = [] for file in files: df = pd.read_csv(file) diff --git a/examples/prompteval/plot_distribution.py b/examples/prompteval/plot_distribution.py index 284f732..fa97be6 100644 --- a/examples/prompteval/plot_distribution.py +++ b/examples/prompteval/plot_distribution.py @@ -28,7 +28,7 @@ "--tuning_parameter", "-tp", type=float, - default=0.5, + default=4.0, ) args = parser.parse_args() output_dir = args.output_dir @@ -45,21 +45,63 @@ os.makedirs(figures_dir, exist_ok=True) results_dir = os.path.join(output_dir, "results") -files = glob( - os.path.join(results_dir, f"est_errors-*-p{propensity}-tp{tuning_parameter}.csv") +# print(os.path.join(results_dir, f"est_errors-*-p{0.3}-tp{tuning_parameter}.csv")) +# files01 = glob( +# os.path.join(results_dir, f"est_errors-*-p{0.3}-tp{tuning_parameter}.csv") +# ) +files05 = glob( + os.path.join(results_dir, f"est_errors-*-p{0.7}-tp{tuning_parameter}.csv") ) -df_list = [] -for file in files: +# files09 = glob( +# os.path.join(results_dir, f"est_errors-*-p{0.7}-tp{tuning_parameter}.csv") +# ) +df_list_01 = [] +# for file in files01: +# df = pd.read_csv(file) +# # Convert string representations of lists to actual lists +# df["imputation"] = df["imputation"].apply(lambda x: np.array(eval(x))) +# df["ground_truth"] = df["ground_truth"].apply(lambda x: np.array(eval(x))) +# df_list_01.append(df) +df_list_05 = [] +for file in files05: df = pd.read_csv(file) # Convert string representations of lists to actual lists df["imputation"] = df["imputation"].apply(lambda x: np.array(eval(x))) df["ground_truth"] = df["ground_truth"].apply(lambda x: np.array(eval(x))) - df_list.append(df) -df = pd.concat(df_list, ignore_index=True) + df_list_05.append(df) +# df_list_09 = [] +# for file in files09: +# df = pd.read_csv(file) +# # Convert string representations of lists to actual lists +# df["imputation"] = df["imputation"].apply(lambda x: np.array(eval(x))) +# df["ground_truth"] = df["ground_truth"].apply(lambda x: np.array(eval(x))) +# df_list_09.append(df) + +# df_01 = pd.concat(df_list_01, ignore_index=True) +df_05 = pd.concat(df_list_05, ignore_index=True) +# df_09 = pd.concat(df_list_09, ignore_index=True) # each row in the dataframe is a ground truth and imputation pair for a given estimation method, fit method, and data type at row r and column c # groupby r, c and plot the histograms of the ground truth and imputation for each group -df = ( - df.groupby(["row", "col"]) +# df_01 = ( +# df_01.groupby(["row", "col"]) +# .apply( +# lambda x: { +# "imputations": { +# (row["estimation_method"], row["fit_method"], row["data_type"]): row[ +# "imputation" +# ] +# for _, row in x.iterrows() +# }, +# "ground_truth": x["ground_truth"].iloc[ +# 0 +# ], # Take first ground truth since it should be same for all methods +# } +# ) +# .reset_index() +# .rename(columns={0: "data"}) +# ) +df_05 = ( + df_05.groupby(["row", "col"]) .apply( lambda x: { "imputations": { @@ -76,65 +118,154 @@ .reset_index() .rename(columns={0: "data"}) ) +# df_09 = ( +# df_09.groupby(["row", "col"]) +# .apply( +# lambda x: { +# "imputations": { +# (row["estimation_method"], row["fit_method"], row["data_type"]): row[ +# "imputation" +# ] +# for _, row in x.iterrows() +# }, +# "ground_truth": x["ground_truth"].iloc[ +# 0 +# ], # Take first ground truth since it should be same for all methods +# } +# ) +# .reset_index() +# .rename(columns={0: "data"}) +# ) COLORS = { ("row-row", "lbo", "kernel_mmd"): "teal", ("col-col", "lbo", "kernel_mmd"): "blue", + ("row-row", "lbo", "wasserstein_samples"): "green", + ("col-col", "lbo", "wasserstein_samples"): "orange", } -for i, row in df.head(20).iterrows(): - fig, ax = plt.subplots(figsize=(5.5 / 2, 5.5 / 2)) - # Determine common bins based on all data - all_data = [row["data"]["ground_truth"]] - for imputation in row["data"]["imputations"].values(): # type: ignore - all_data.append(imputation) +# Determine common bins based on all data + + +# Calculate the min and max across all data to create common bins +bins = list(np.linspace(0, 1, 30)) # 11 points create 10 bins - # Calculate the min and max across all data to create common bins - bins = list(np.linspace(0, 1, 40)) # 11 points create 10 bins +for i in range(0, 100): + # row_09 = df_09.iloc[i] + row_05 = df_05.iloc[i] + # row_01 = df_01.iloc[i] + fig, ax2 = plt.subplots(figsize=(3, 2.5)) + # all_data = row_09["data"]["ground_truth"] + # for imputation in row_09["data"]["imputations"].values(): # type: ignore + # all_data.append(imputation) # Plot imputation for each method - for (est_method, fit_method, data_type), imputation in row["data"][ # type: ignore + # for (est_method, fit_method, data_type), imputation in row_09["data"][ # type: ignore + # "imputations" + # ].items(): # type: ignore + # weights = np.ones_like(imputation) / len(imputation) + # ax1.hist( + # imputation, + # bins=bins, + # alpha=0.6, + # label=f"{plotting_utils.METHOD_ALIASES_SINGLE_LINE.get(est_method, est_method)}" + # f"\n({plotting_utils.DATA_TYPE_ALIASES.get(data_type, data_type)})", + # weights=weights, + # color=COLORS[(est_method, fit_method, data_type)], + # ) + # # Plot ground truth + # weights = np.ones_like(row_09["data"]["ground_truth"]) / len(row_09["data"]["ground_truth"]) + # ax1.hist( + # row_09["data"]["ground_truth"], + # weights=weights, + # bins=bins, + # alpha=0.6, + # label="Ground\nTruth", + # edgecolor="black", + # color="white", + # linestyle="--", + # ) + # ax1.set_xlabel("Score", fontsize=plotting_utils.LABEL_FONT_SIZE) + # ax1.set_ylabel("Density", fontsize=plotting_utils.LABEL_FONT_SIZE) + # ax1.set_xlim(0, 1) + + # ax1.spines["top"].set_visible(False) + # ax1.spines["right"].set_visible(False) + # # ax.spines["bottom"].set_position( + # # ("outward", plotting_utils.OUTWARD) + # # ) # Move x-axis outward + # ax1.spines["left"].set_position( + # ("outward", plotting_utils.OUTWARD) + # ) # Move y-axis outward + + # # Only add legend to the first subplot to avoid clutter + # ax1.legend(loc='upper right', fontsize=plotting_utils.LEGEND_FONT_SIZE) + + # all_data = row_05["data"]["ground_truth"] + # for imputation in row_05["data"]["imputations"].values(): # type: ignore + # all_data.append(imputation) + # Plot imputation for each method + weights = np.ones_like(row_05["data"]["ground_truth"]) / len( + row_05["data"]["ground_truth"] + ) + ax2.hist( + row_05["data"]["ground_truth"], + bins=bins, + weights=weights, + alpha=0.6, + label="Ground\nTruth", + edgecolor="black", + linestyle="--", + color="white", + ) + for (est_method, fit_method, data_type), imputation in row_05["data"][ # type: ignore "imputations" ].items(): # type: ignore - ax.hist( + if data_type == "wasserstein_samples": + # Convert to numpy array if it's not already + est_method_dt = "wasserstein_samples" + else: + est_method_dt = "kernel" + if est_method == "col-col": + est_method_lbl = "col" + else: + est_method_lbl = "row" + + weights = np.ones_like(imputation) / len(imputation) + ax2.hist( imputation, bins=bins, - alpha=0.5, - label=f"{plotting_utils.METHOD_ALIASES_SINGLE_LINE.get(est_method, est_method)}" - f"\n({plotting_utils.DATA_TYPE_ALIASES.get(data_type, data_type)})", - edgecolor="black", - density=True, + weights=weights, + alpha=0.6, + label=f"{plotting_utils.METHOD_ALIASES_SINGLE_LINE.get(est_method_dt, est_method_dt)}" + f"\n({est_method_lbl})", color=COLORS[(est_method, fit_method, data_type)], ) # Plot ground truth - ax.hist( - row["data"]["ground_truth"], - bins=bins, - alpha=0.5, - label="Ground Truth", - edgecolor="black", - density=True, - color="yellow", - ) - ax.set_xlabel("Score", fontsize=plotting_utils.LABEL_FONT_SIZE) - ax.set_ylabel("Density", fontsize=plotting_utils.LABEL_FONT_SIZE) - ax.set_xlim(0, 1) - ax.spines["top"].set_visible(False) - ax.spines["right"].set_visible(False) + ax2.set_ylabel("Proportion", fontsize=plotting_utils.LABEL_FONT_SIZE) + ax2.set_xlabel("Score", fontsize=plotting_utils.LABEL_FONT_SIZE) + ax2.legend(loc="upper right", fontsize=plotting_utils.LEGEND_FONT_SIZE) + + ax2.spines["top"].set_visible(False) + ax2.spines["right"].set_visible(False) + ax2.spines["left"].set_visible(False) + ax2.set_axisbelow(True) + ax2.grid(True, alpha=0.4) # ax.spines["bottom"].set_position( # ("outward", plotting_utils.OUTWARD) # ) # Move x-axis outward - ax.spines["left"].set_position( - ("outward", plotting_utils.OUTWARD) - ) # Move y-axis outward + # ax2.spines["left"].set_position( + # ("outward", plotting_utils.OUTWARD) + # ) # Move y-axis outward # Only add legend to the first subplot to avoid clutter - ax.legend() + ax2.legend(loc="upper left", fontsize=plotting_utils.LEGEND_FONT_SIZE) figures_path = os.path.join( figures_dir, - f"distributions-r{row.row}-c{row.col}-p{propensity}-tp{tuning_parameter}.pdf", + f"distributions-entry{i}-tp{tuning_parameter}.pdf", ) logger.info(f"Saving figure to {figures_path}") plt.savefig(figures_path, dpi=300, bbox_inches="tight") + plt.close() diff --git a/examples/prompteval/run_distribution.py b/examples/prompteval/run_distribution.py index 6281aab..c34411f 100644 --- a/examples/prompteval/run_distribution.py +++ b/examples/prompteval/run_distribution.py @@ -32,6 +32,7 @@ parser = get_base_parser() parser.add_argument( "--data_type", + "-dt", type=str, default="kernel_mmd", choices=["kernel_mmd", "wasserstein_samples"], @@ -70,24 +71,26 @@ rng = np.random.default_rng(seed=seed) +dataloader = NNData.create( + "prompteval", + # NOTE: uncomment to run on a subset of models and tasks (for debugging) + # models=['meta_llama_llama_3_8b', 'meta_llama_llama_3_8b_instruct', 'meta_llama_llama_3_70b_instruct', 'codellama_codellama_34b_instruct', ], + # tasks=['college_mathematics', 'miscellaneous', 'moral_disputes', 'jurisprudence', 'moral_scenarios', 'college_chemistry'], + propensity=propensity, + seed=seed, +) + match data_type: case "kernel_mmd": data_type = DistributionKernelMMD( kernel="exponential", tuning_parameter=tuning_parameter ) case "wasserstein_samples": - data_type = DistributionWassersteinSamples() + n = 100 + data_type = DistributionWassersteinSamples(num_samples=n) case _: raise ValueError(f"Data type {data_type} not supported") -dataloader = NNData.create( - "prompteval", - # NOTE: uncomment to run on a subset of models and tasks (for debugging) - # models=['meta_llama_llama_3_8b', 'meta_llama_llama_3_8b_instruct', 'meta_llama_llama_3_70b_instruct', 'codellama_codellama_34b_instruct', ], - # tasks=['college_mathematics', 'miscellaneous', 'moral_disputes', 'jurisprudence', 'moral_scenarios', 'college_chemistry'], - propensity=propensity, - seed=seed, -) data, mask = dataloader.process_data_distribution(data_type) holdout_inds = np.nonzero(mask == 1) @@ -132,7 +135,7 @@ raise ValueError(f"Estimation method {estimation_method} not supported") fit_method = LeaveBlockOutValidation( - block, distance_threshold_range=(0, 1.0), n_trials=10, data_type=data_type, rng=rng + block, distance_threshold_range=(0, 1.0), n_trials=50, data_type=data_type, rng=rng ) imputer = NearestNeighborImputer( estimator, diff --git a/examples/prompteval/run_ksbyp_plot.py b/examples/prompteval/run_ksbyp_plot.py new file mode 100644 index 0000000..2ac94e9 --- /dev/null +++ b/examples/prompteval/run_ksbyp_plot.py @@ -0,0 +1,243 @@ +import numpy as np +import matplotlib.pyplot as plt +import pandas as pd +import argparse +from nsquared.utils import plotting_utils + +# Define T values and corresponding error rates for each method +T_values = np.array([2**4, 2**5, 2**6, 2**7]) + + +# Define T values and corresponding error rates for each method +T_values = np.array([2**4, 2**5, 2**6, 2**7]) + + +parser = argparse.ArgumentParser(description="Plot estimation errors") +parser.add_argument("--num_sims", type=int, default=30, help="Number of simulations") +parser.add_argument("--output_dir", type=str, help="Output directory") +args = parser.parse_args() +# Define T values (sizes) +T_values = np.array([2**4, 2**5, 2**6, 2**7]) +num_sims = args.num_sims +output_dir = args.output_dir + + +# Function to process a single CSV file and extract errors by size +df_mmd_col_03 = pd.read_csv( + f"{output_dir}/results/est_errors-col-col-lbo-kernel_mmd-p0.3-tp4.0.csv" +) +df_mmd_col_05 = pd.read_csv( + f"{output_dir}/results/est_errors-col-col-lbo-kernel_mmd-p0.5-tp4.0.csv" +) +df_mmd_col_07 = pd.read_csv( + f"{output_dir}/results/est_errors-col-col-lbo-kernel_mmd-p0.7-tp4.0.csv" +) + +df_wasserstein_03 = pd.read_csv( + f"{output_dir}/results/est_errors-col-col-lbo-wasserstein_samples-p0.3-tp4.0.csv" +) +df_wasserstein_05 = pd.read_csv( + f"{output_dir}/results/est_errors-col-col-lbo-wasserstein_samples-p0.5-tp4.0.csv" +) +df_wasserstein_07 = pd.read_csv( + f"{output_dir}/results/est_errors-col-col-lbo-wasserstein_samples-p0.7-tp4.0.csv" +) + +df_mmd_row_03 = pd.read_csv( + f"{output_dir}/results/est_errors-row-row-lbo-kernel_mmd-p0.3-tp4.0.csv" +) +df_mmd_row_05 = pd.read_csv( + f"{output_dir}/results/est_errors-row-row-lbo-kernel_mmd-p0.5-tp4.0.csv" +) +df_mmd_row_07 = pd.read_csv( + f"{output_dir}/results/est_errors-row-row-lbo-kernel_mmd-p0.7-tp4.0.csv" +) + +df_wasserstein_row_03 = pd.read_csv( + f"{output_dir}/results/est_errors-row-row-lbo-wasserstein_samples-p0.3-tp4.0.csv" +) +df_wasserstein_row_05 = pd.read_csv( + f"{output_dir}/results/est_errors-row-row-lbo-wasserstein_samples-p0.5-tp4.0.csv" +) +df_wasserstein_row_07 = pd.read_csv( + f"{output_dir}/results/est_errors-row-row-lbo-wasserstein_samples-p0.7-tp4.0.csv" +) + +df_mmd_col = pd.concat([df_mmd_col_03, df_mmd_col_05, df_mmd_col_07], ignore_index=True) +df_wasserstein_col = pd.concat( + [df_wasserstein_03, df_wasserstein_05, df_wasserstein_07], ignore_index=True +) + +propensity = np.array([0.3, 0.5, 0.7]) + +mmd_col_error_03 = np.array(df_mmd_col_03["est_errors"].values) +mmd_col_error_05 = np.array(df_mmd_col_05["est_errors"].values) +mmd_col_error_07 = np.array(df_mmd_col_07["est_errors"].values) + +mmd_row_error_03 = np.array(df_mmd_row_03["est_errors"].values) +mmd_row_error_05 = np.array(df_mmd_row_05["est_errors"].values) +mmd_row_error_07 = np.array(df_mmd_row_07["est_errors"].values) + +wasserstein_col_error_03 = np.array(df_wasserstein_03["est_errors"].values) +wasserstein_col_error_05 = np.array(df_wasserstein_05["est_errors"].values) +wasserstein_col_error_07 = np.array(df_wasserstein_07["est_errors"].values) + +wasserstein_row_error_03 = np.array(df_wasserstein_row_03["est_errors"].values) +wasserstein_row_error_05 = np.array(df_wasserstein_row_05["est_errors"].values) +wasserstein_row_error_07 = np.array(df_wasserstein_row_07["est_errors"].values) + +mmd_row_errors = np.array( + [ + np.nanmean(mmd_row_error_03), + np.nanmean(mmd_row_error_05), + np.nanmean(mmd_row_error_07), + ] +) +mmd_col_errors = np.array( + [ + np.nanmean(mmd_col_error_03), + np.nanmean(mmd_col_error_05), + np.nanmean(mmd_col_error_07), + ] +) +wasserstein_row_errors = np.array( + [ + np.nanmean(wasserstein_row_error_03), + np.nanmean(wasserstein_row_error_05), + np.nanmean(wasserstein_row_error_07), + ] +) +wasserstein_col_errors = np.array( + [ + np.nanmean(wasserstein_col_error_03), + np.nanmean(wasserstein_col_error_05), + np.nanmean(wasserstein_col_error_07), + ] +) + +mmd_row_stderr = np.array( + [ + np.nanstd(mmd_row_error_03) / np.sqrt(mmd_row_error_03.shape), + np.nanstd(mmd_row_error_05) / np.sqrt(mmd_row_error_05.shape), + np.nanstd(mmd_row_error_07) / np.sqrt(mmd_row_error_07.shape), + ] +) +mmd_col_stderr = np.array( + [ + np.nanstd(mmd_col_error_03) / np.sqrt(mmd_col_error_03.shape), + np.nanstd(mmd_col_error_05) / np.sqrt(mmd_col_error_05.shape), + np.nanstd(mmd_col_error_07) / np.sqrt(mmd_col_error_07.shape), + ] +) +wasserstein_row_stderr = np.array( + [ + np.nanstd(wasserstein_row_error_03) / np.sqrt(wasserstein_row_error_03.shape), + np.nanstd(wasserstein_row_error_05) / np.sqrt(wasserstein_row_error_05.shape), + np.nanstd(wasserstein_row_error_07) / np.sqrt(wasserstein_row_error_07.shape), + ] +) +wasserstein_col_stderr = np.array( + [ + np.nanstd(wasserstein_col_error_03) / np.sqrt(wasserstein_col_error_03.shape), + np.nanstd(wasserstein_col_error_05) / np.sqrt(wasserstein_col_error_05.shape), + np.nanstd(wasserstein_col_error_07) / np.sqrt(wasserstein_col_error_07.shape), + ] +) + +# Create the plot +plt.figure(figsize=(3.25, 2.5)) +# figsize=(plotting_utils.NEURIPS_TEXTWIDTH / 2, 2.5) +# plt.plot(T_values, USVT_errors, 'r', linestyle='-', marker='D', markersize=8, linewidth=2, label=r'USVT: $T^{-0.46}$') + + +def _add_regression_line( + x: np.ndarray, y: np.ndarray, color: str, label: str, linestyle: str +) -> float: + # Fit a line to log-transformed data + slope, intercept = np.polyfit(x, y, 1) + plt.plot(x, slope * x + intercept, color=color, linestyle=linestyle, linewidth=2) + # label=f"{label} (slope: {slope:.2f})") + return slope + + +col_mmd_slope = _add_regression_line(propensity, mmd_col_errors, "blue", "", ":") +row_mmd_slope = _add_regression_line(propensity, mmd_row_errors, "magenta", "", ":") +col_wasserstein_slope = _add_regression_line( + propensity, wasserstein_col_errors, "red", "", "--" +) +row_wasserstein_slope = _add_regression_line( + propensity, wasserstein_row_errors, "green", "", "--" +) +# softimpute_slope = add_regression_line(T_values, Softimpute_errors, 'black', 'SoftImpute', ':') + + +def _format_lbl(method: str) -> str: + return str(plotting_utils.METHOD_ALIASES_SINGLE_LINE.get(method, method)) + + +plt.errorbar( + propensity, + mmd_col_errors, + fmt="Db", + yerr=mmd_col_stderr.squeeze(), + markersize=5, + linestyle="None", + label=rf"{_format_lbl('kernel')} (col)", +) +plt.errorbar( + propensity, + mmd_row_errors, + fmt="mo", + yerr=mmd_row_stderr.squeeze(), + markersize=5, + linestyle="None", + label=rf"{_format_lbl('kernel')} (row)", +) +plt.errorbar( + propensity, + wasserstein_col_errors, + fmt="vr", + yerr=wasserstein_col_stderr.squeeze(), + markersize=5, + linestyle="None", + label=rf"{_format_lbl('wasserstein_samples')} (col)", +) +plt.errorbar( + propensity, + wasserstein_row_errors, + fmt="^g", + yerr=wasserstein_row_stderr.squeeze(), + markersize=5, + linestyle="None", + label=rf"{_format_lbl('wasserstein_samples')} (row)", +) + +# Axis labels +plt.xlabel(r"Propensity (p)", fontsize=plotting_utils.LABEL_FONT_SIZE) +plt.ylabel(r"Kolmogorov-Smirnov distance", fontsize=plotting_utils.LABEL_FONT_SIZE) + +# Title +# plt.title(r'Decay of avg. error across users (N = T, 30 trials)', fontsize=16) + +# Add legend +plt.legend(fontsize=plotting_utils.LEGEND_FONT_SIZE, loc="upper right") + +# set tick font size +# print(plotting_utils.TICK_FONT_SIZE) +plt.tick_params(axis="both", which="major", labelsize=plotting_utils.LABEL_FONT_SIZE) + +# set exact x ticks +plt.xticks( + [0.3, 0.5, 0.7], + fontsize=plotting_utils.LABEL_FONT_SIZE, +) +# Grid for better readability +# plt.grid(True, which="both", linestyle='--', linewidth=0.5) +ax1 = plt.gca() +ax1.spines["top"].set_visible(False) +ax1.spines["right"].set_visible(False) +ax1.spines["left"].set_visible(False) +ax1.grid(True, alpha=0.4) +plt.tight_layout() +# Show the plot +plt.savefig(f"{output_dir}/ksbyp_prompteval_plot.pdf", bbox_inches="tight") diff --git a/examples/prompteval/run_scalar.py b/examples/prompteval/run_scalar.py index 5a9d2de..0bfc0e0 100644 --- a/examples/prompteval/run_scalar.py +++ b/examples/prompteval/run_scalar.py @@ -143,8 +143,8 @@ # Fit the imputer using leave-block-out validation fitter = DRLeaveBlockOutValidation( block, - distance_threshold_range_row=(0, 10), - distance_threshold_range_col=(0, 10), + distance_threshold_range_row=(0, 1), + distance_threshold_range_col=(0, 1), n_trials=200, data_type=data_type, ) @@ -155,7 +155,7 @@ logger.info("Using leave-block-out validation") fitter = LeaveBlockOutValidation( block, - distance_threshold_range=(0, 10), + distance_threshold_range=(0, 1), n_trials=200, data_type=data_type, ) @@ -166,7 +166,7 @@ logger.info("Using leave-block-out validation") fitter = LeaveBlockOutValidation( block, - distance_threshold_range=(0, 10), + distance_threshold_range=(0, 1), n_trials=200, data_type=data_type, ) @@ -179,8 +179,8 @@ # Fit the imputer using leave-block-out validation fitter = TSLeaveBlockOutValidation( block, - distance_threshold_range_row=(0, 10), - distance_threshold_range_col=(0, 10), + distance_threshold_range_row=(0, 1), + distance_threshold_range_col=(0, 1), n_trials=200, data_type=data_type, ) diff --git a/examples/prompteval/slurm_scripts/run_distribution.sh b/examples/prompteval/slurm_scripts/run_distribution.sh index aecabc4..6736534 100755 --- a/examples/prompteval/slurm_scripts/run_distribution.sh +++ b/examples/prompteval/slurm_scripts/run_distribution.sh @@ -9,12 +9,28 @@ METHODS=( "row-row" "col-col" ) -for em in ${METHODS[@]}; + +DATA_TYPE=( + "wasserstein_samples" + "kernel_mmd" +) + +ps=( + "0.3" + "0.5" + "0.7" +) + +for data_type in ${DATA_TYPE[@]}; do - for p in $(seq 0.1 0.1 1.0); + for em in ${METHODS[@]}; do - CMD="python run_distribution.py -od $OUTPUT_DIR -em $em -p $p -tp 4 -s 1 -f" - echo $CMD - eval $CMD + for p in ${ps[@]}; + do + echo $data_type + CMD="python run_distribution.py -em $em -p $p -tp 4.0 --data_type $data_type -s 1 --force -od $OUTPUT_DIR" + echo $CMD + eval $CMD + done done done diff --git a/examples/prop99/plot_sc_error.py b/examples/prop99/plot_sc_error.py index 45137fb..51f167f 100644 --- a/examples/prop99/plot_sc_error.py +++ b/examples/prop99/plot_sc_error.py @@ -49,7 +49,7 @@ ) # rearrange the order of the estimation methods -ORDER = ["usvt", "softimpute", "col-col", "row-row", "dr", "ts", "aw", "sc"] +ORDER = ["usvt", "softimpute", "col-col", "row-row", "dr", "ts", "auto", "aw", "sc"] df_grouped = df_grouped.sort_values( by="estimation_method", key=lambda x: x.map(lambda y: ORDER.index(y)) ) diff --git a/examples/prop99/plot_synthetic_control.py b/examples/prop99/plot_synthetic_control.py index 2a40e16..a3963f4 100644 --- a/examples/prop99/plot_synthetic_control.py +++ b/examples/prop99/plot_synthetic_control.py @@ -59,9 +59,9 @@ # NOTE: va is opposite of what you'd expect (set top to nudge down, bottom to nudge up) match estimation_method: case "dr": + va = "top" + case "col-col": va = "bottom" - # case "col-col": - # va = "bottom" case "row-row": va = "top" case _: @@ -72,13 +72,25 @@ method, # type: ignore fontsize=plotting_utils.TICK_FONT_SIZE, va=va, + color=plotting_utils.COLORS[estimation_method], ) + # # add a line segment from the text label to the last valid value + # ax.plot( + # [2000, 2010], + # [last_valid_value, last_valid_value], + # color=plotting_utils.COLORS[estimation_method], + # linestyle="-", + # ) if i == 0: ax.plot( df.index, df["obs"], label="Observed", linestyle="-", color="orange" ) ax.text( - 2001, df["obs"].iloc[-1], "Obs.", fontsize=plotting_utils.TICK_FONT_SIZE + 2001, + df["obs"].iloc[-1], + "Obs.", + fontsize=plotting_utils.TICK_FONT_SIZE, + color="orange", ) # add a vertical line at 1989 called Proposition 99 ax.axvline(x=1989, color="k", alpha=0.25, linestyle="dotted") @@ -107,7 +119,7 @@ # ax.set_ylim(40, 160) # set the y-axis label to be the number of cigarettes smoked per capita ax.set_ylabel( - "Cigarette Consumption\n(Pack Sales Per Capita)", + "Cigarette Consumption", fontsize=plotting_utils.LABEL_FONT_SIZE, ) @@ -121,7 +133,7 @@ ) # Move y-axis outward # ax.legend(fontsize=plotting_utils.LEGEND_FONT_SIZE) plt.subplots_adjust( - left=0.15, + left=0.2, right=0.85, top=0.95, bottom=0.2, diff --git a/examples/prop99/run_scalar.py b/examples/prop99/run_scalar.py index 0b6162e..ed2386d 100644 --- a/examples/prop99/run_scalar.py +++ b/examples/prop99/run_scalar.py @@ -42,12 +42,13 @@ # %% # import nearest neighbor methods from nsquared.data_types import Scalar -from nsquared.estimation_methods import AWNNEstimator, TSEstimator +from nsquared.estimation_methods import AWNNEstimator, TSEstimator, AutoEstimator from nsquared import NearestNeighborImputer from nsquared.fit_methods import ( DRLeaveBlockOutValidation, TSLeaveBlockOutValidation, LeaveBlockOutValidation, + AutoDRTSLeaveBlockOutValidation, ) from nsquared.datasets.dataloader_factory import NNData from nsquared.vanilla_nn import row_row, col_col @@ -106,7 +107,7 @@ logger.info(f"Time to load and process data: {elapsed_time:.2f} seconds") treatment_row = 0 # row corresponding to treated unit (California) -logger.info(f"Mask for row {treatment_row}:", mask[treatment_row]) +logger.info(f"Mask for row {treatment_row}: {mask[treatment_row]}") # %% logger.info("Using scalar data type") @@ -230,11 +231,11 @@ logger.info("Using leave-block-out validation") fitter = LeaveBlockOutValidation( block, - # distance_threshold_range=(0, 50), distance_threshold_range=(0, 300**2), n_trials=200, data_type=data_type, allow_self_neighbor=allow_self_neighbor, + rng=rng, ) elif estimation_method == "col-col": logger.info("Using col-col estimation") @@ -247,6 +248,7 @@ n_trials=200, data_type=data_type, allow_self_neighbor=allow_self_neighbor, + rng=rng, ) elif estimation_method == "ts": logger.info("Using two-sided estimation") @@ -264,6 +266,21 @@ allow_self_neighbor=True, ) allow_self_neighbor = True + elif estimation_method == "auto": + logger.info("Using AutoNN estimation") + estimator = AutoEstimator(is_percentile=is_percentile) + imputer = NearestNeighborImputer(estimator, data_type) + logger.info("Using AutoNN fit method") + # Fit the imputer using leave-block-out validation + fitter = AutoDRTSLeaveBlockOutValidation( + block, + distance_threshold_range_row=(0, 300**2), + distance_threshold_range_col=(0, 300**2), + alpha_range=(0, 1), + n_trials=200, + data_type=data_type, + allow_self_neighbor=allow_self_neighbor, + ) else: raise ValueError( f"Estimation method {estimation_method} and fit method {fit_method} not supported" diff --git a/examples/prop99/slurm_scripts/run_accuracy.sh b/examples/prop99/slurm_scripts/run_accuracy.sh index 3d728df..8a959a1 100755 --- a/examples/prop99/slurm_scripts/run_accuracy.sh +++ b/examples/prop99/slurm_scripts/run_accuracy.sh @@ -13,6 +13,7 @@ METHODS=( "ts" "aw" "softimpute" + "auto" ) CONTROL_STATES=( "TX" "WI" "MT" "RI" "KS" "ME" "UT" "VA" "IN" "GA" diff --git a/examples/prop99/slurm_scripts/run_california.sh b/examples/prop99/slurm_scripts/run_california.sh index 195b1fd..09dc29f 100755 --- a/examples/prop99/slurm_scripts/run_california.sh +++ b/examples/prop99/slurm_scripts/run_california.sh @@ -13,6 +13,7 @@ METHODS=( "ts" "aw" "softimpute" + "auto" ) for em in ${METHODS[@]}; do diff --git a/examples/simulations/README.md b/examples/simulations/README.md new file mode 100644 index 0000000..f454272 --- /dev/null +++ b/examples/simulations/README.md @@ -0,0 +1,19 @@ +# Replicating simulations experiments +
+ + + Figure 1: AutoNN recovers DRNN performance with High SNR +
+ + + + +To replicate Figure 1, run the following commands while in the `examples/simulations` directory. + +```bash +# Run simulated experiments with high SNR (noise std = 0.001) +./slurm_scripts run_acccuracy.sh -o OUTPUT_DIR -l ERROR -n 0.001 +# Plot the figure +python plot_size_error.py --output_dir OUTPUT_DIR +``` +To replicate the low SNR case, run the same commands with option `-n 1.0` instead. \ No newline at end of file diff --git a/examples/simulations/example_plots/sims_plot_highsnr.pdf b/examples/simulations/example_plots/sims_plot_highsnr.pdf new file mode 100644 index 0000000..dd313d5 Binary files /dev/null and b/examples/simulations/example_plots/sims_plot_highsnr.pdf differ diff --git a/examples/simulations/example_plots/sims_plot_lowsnr.pdf b/examples/simulations/example_plots/sims_plot_lowsnr.pdf new file mode 100644 index 0000000..084face Binary files /dev/null and b/examples/simulations/example_plots/sims_plot_lowsnr.pdf differ diff --git a/examples/simulations/plot_size_error.py b/examples/simulations/plot_size_error.py index b711e54..a898839 100644 --- a/examples/simulations/plot_size_error.py +++ b/examples/simulations/plot_size_error.py @@ -2,13 +2,14 @@ import matplotlib.pyplot as plt import pandas as pd import argparse +from nsquared.utils import plotting_utils # Define T values and corresponding error rates for each method -T_values = np.array([2**4, 2**5, 2**6, 2**7]) +# T_values = np.array([2**4, 2**5, 2**6, 2**7]) # Define T values and corresponding error rates for each method -T_values = np.array([2**4, 2**5, 2**6, 2**7]) +# T_values = np.array([2**4, 2**5, 2**6, 2**7]) parser = argparse.ArgumentParser(description="Plot estimation errors") @@ -41,34 +42,38 @@ def _process_csv(filepath: str) -> np.ndarray: # Process each CSV file -col_err = _process_csv(f"{output_dir}/results/est_errors-col-col-lbo.csv") -row_err = _process_csv(f"{output_dir}/results/est_errors-row-row-lbo.csv") +# col_err = _process_csv(f"{output_dir}/results/est_errors-col-col-lbo.csv") +# ow_err = _process_csv(f"{output_dir}/results/est_errors-row-row-lbo.csv") # usvt_err = process_csv(f"{output_dir}/results/est_errors-usvt-lbo.csv") drnn_err = _process_csv(f"{output_dir}/results/est_errors-dr-lbo.csv") tsnn_err = _process_csv(f"{output_dir}/results/est_errors-ts-lbo.csv") +auto_err = _process_csv(f"{output_dir}/results/est_errors-auto-lbo.csv") # softimpute_err = process_csv("cvrange_30sims/results/est_errors-softimpute-lbo.csv") # Extract errors for each method as a numpy array 4 x 30 # USVT_errors = 2**-3 * T_values**-0.10 -UserNN_errors = np.nanmean(row_err, axis=1) -TimeNN_errors = np.nanmean(col_err, axis=1) +# UserNN_errors = np.nanmean(row_err, axis=1) +# TimeNN_errors = np.nanmean(col_err, axis=1) DRNN_errors = np.nanmean(drnn_err, axis=1) # USVT_errors = np.nanmean(usvt_err, axis = 1) TSNN_errors = np.nanmean(tsnn_err, axis=1) # Softimpute_errors = np.nanmean(softimpute_err, axis = 1) +Auto_errors = np.nanmean(auto_err, axis=1) -unn_stderr = np.nanstd(row_err, axis=1) / np.sqrt(row_err.shape[1]) -tnn_stderr = np.nanstd(col_err, axis=1) / np.sqrt(col_err.shape[1]) +# unn_stderr = np.nanstd(row_err, axis=1) / np.sqrt(row_err.shape[1]) +# tnn_stderr = np.nanstd(col_err, axis=1) / np.sqrt(col_err.shape[1]) drnn_stderr = np.nanstd(drnn_err, axis=1) / np.sqrt(drnn_err.shape[1]) # usvt_stderr = np.nanstd(usvt_err, axis = 1) / np.sqrt(usvt_err.shape[1]) tsnn_stderr = np.nanstd(tsnn_err, axis=1) / np.sqrt(tsnn_err.shape[1]) +auto_stderr = np.nanstd(auto_err, axis=1) / np.sqrt(auto_err.shape[1]) # softimpute_stderr = np.nanstd(softimpute_err, axis = 1) / np.sqrt(softimpute_err.shape[1]) # Create the plot -plt.figure() +plt.figure(figsize=(plotting_utils.NEURIPS_TEXTWIDTH / 2, 2.5)) +# figsize=(plotting_utils.NEURIPS_TEXTWIDTH / 2, 2.5) # plt.plot(T_values, USVT_errors, 'r', linestyle='-', marker='D', markersize=8, linewidth=2, label=r'USVT: $T^{-0.46}$') @@ -85,52 +90,67 @@ def _add_regression_line( # usvt_slope = add_regression_line(T_values, USVT_errors, 'red', 'USVT', '--') -unn_slope = _add_regression_line(T_values, UserNN_errors, "green", "User-NN", ":") -tnn_slope = _add_regression_line(T_values, TimeNN_errors, "orange", "Time-NN", ":") -drnn_slope = _add_regression_line(T_values, DRNN_errors, "blue", "DR-NN", "--") -tsnn_slope = _add_regression_line(T_values, TSNN_errors, "purple", "Time-NN", ":") +# unn_slope = _add_regression_line(T_values, UserNN_errors, "green", "User-NN", ":") +# tnn_slope = _add_regression_line(T_values, TimeNN_errors, "orange", "Time-NN", ":") +drnn_slope = _add_regression_line(T_values, DRNN_errors, "blue", "DR-NN", ":") +tsnn_slope = _add_regression_line(T_values, TSNN_errors, "green", "Time-NN", ":") +auto_slope = _add_regression_line(T_values, Auto_errors, "red", "Auto-NN", "--") # softimpute_slope = add_regression_line(T_values, Softimpute_errors, 'black', 'SoftImpute', ':') + # Plot each method with corresponding markers, colors, and line styles # plt.errorbar(T_values, USVT_errors, yerr=usvt_stderr, fmt = 'r', marker='>', markersize=12, linestyle='None', barsabove=True, label=rf'USVT: $T^{{{usvt_slope:.2f}}}$') -plt.errorbar( - T_values, - UserNN_errors, - yerr=unn_stderr, - fmt="g", - marker="o", - markersize=12, - linestyle="None", - label=rf"Row-NN: $T^{{{unn_slope:.2f}}}$", -) -plt.errorbar( - T_values, - TimeNN_errors, - yerr=tnn_stderr, - fmt="s", - color="orange", - marker="s", - linestyle="None", - markersize=12, - label=rf"Col-NN: $T^{{{tnn_slope:.2f}}}$", -) +# plt.errorbar( +# T_values, +# UserNN_errors, +# yerr=unn_stderr, +# fmt="g", +# marker="o", +# markersize=12, +# linestyle="None", +# label=rf"Row-NN: $T^{{{unn_slope:.2f}}}$", +# ) +# plt.errorbar( +# T_values, +# TimeNN_errors, +# yerr=tnn_stderr, +# fmt="s", +# color="orange", +# marker="s", +# linestyle="None", +# markersize=12, +# label=rf"Col-NN: $T^{{{tnn_slope:.2f}}}$", +# ) +def _format_lbl(method: str) -> str: + return str(plotting_utils.METHOD_ALIASES_SINGLE_LINE.get(method, method)) + + plt.errorbar( T_values, DRNN_errors, - fmt="bD", + fmt="Db", yerr=drnn_stderr, - markersize=12, + markersize=5, linestyle="None", - label=rf"DR-NN: $T^{{{drnn_slope:.2f}}}$", + label=rf"{_format_lbl('dr')}: $T^{{{drnn_slope:.2f}}}$", ) plt.errorbar( T_values, TSNN_errors, - fmt="p", + fmt="go", yerr=tsnn_stderr, - markersize=12, + markersize=5, linestyle="None", - label=rf"TS-NN: $T^{{{tsnn_slope:.2f}}}$", + label=rf"{_format_lbl('ts')}: $T^{{{tsnn_slope:.2f}}}$", +) +plt.errorbar( + T_values, + Auto_errors, + fmt="vr", + yerr=auto_stderr, + markersize=5, + linestyle="None", + label=rf"{_format_lbl('auto')}: $T^{{{auto_slope:.2f}}}$", ) # plt.errorbar(T_values, Softimpute_errors, fmt='k^', yerr=softimpute_stderr, markersize=12, linestyle="None", label=rf'SoftImpute: $T^{{{softimpute_slope:.2f}}}$') # Logarithmic scale for both axes @@ -138,14 +158,18 @@ def _add_regression_line( plt.yscale("log", base=2) # Axis labels -plt.xlabel(r"T", fontsize=15) -plt.ylabel(r"Error for t = T, a = 1", fontsize=15) +plt.xlabel(r"\# Columns (T)", fontsize=plotting_utils.LABEL_FONT_SIZE) +plt.ylabel(r"Absolute error", color="black", fontsize=plotting_utils.LABEL_FONT_SIZE) # Title # plt.title(r'Decay of avg. error across users (N = T, 30 trials)', fontsize=16) # Add legend -plt.legend(fontsize=12) +plt.legend(fontsize=plotting_utils.LEGEND_FONT_SIZE, loc="upper right") + +# set tick font size +# print(plotting_utils.TICK_FONT_SIZE) +plt.tick_params(axis="both", which="major", labelsize=plotting_utils.LABEL_FONT_SIZE) # Grid for better readability # plt.grid(True, which="both", linestyle='--', linewidth=0.5) @@ -154,6 +178,7 @@ def _add_regression_line( ax1.spines["right"].set_visible(False) ax1.spines["left"].set_visible(False) ax1.grid(True, alpha=0.4) - -# Show the plot +plt.tight_layout() +# Show the plot as pdf (better quality) and png (for markdown) plt.savefig(f"{output_dir}/sims_plot.pdf", bbox_inches="tight") +plt.savefig(f"{output_dir}/sims_plot.png", bbox_inches="tight") diff --git a/examples/simulations/run_scalar.py b/examples/simulations/run_scalar.py index f3f33ae..de42e6f 100644 --- a/examples/simulations/run_scalar.py +++ b/examples/simulations/run_scalar.py @@ -20,13 +20,13 @@ # import nearest neighbor methods from nsquared.data_types import Scalar -from nsquared.estimation_methods import TSEstimator # , AutoEstimator +from nsquared.estimation_methods import TSEstimator, AutoEstimator from nsquared import NearestNeighborImputer from nsquared.fit_methods import ( DRLeaveBlockOutValidation, TSLeaveBlockOutValidation, LeaveBlockOutValidation, - # AutoDRTSLeaveBlockOutValidation, + AutoDRTSLeaveBlockOutValidation, ) from nsquared.datasets.dataloader_factory import NNData from nsquared.vanilla_nn import row_row, col_col @@ -35,7 +35,15 @@ from nsquared.utils.experiments import get_base_parser, setup_logging parser = get_base_parser() +parser.add_argument( + "--noise_stddev", + "-nstd", + type=float, + default=0.001, + help="Standard deviation of the noise added to the data", +) args = parser.parse_args() +noise_stddev = args.noise_stddev output_dir = args.output_dir estimation_method = args.estimation_method fit_method = args.fit_method @@ -60,7 +68,7 @@ # Load the simulated data dataset # NOTE: the raw and processed data is cached in .joblib_cache m_size = [2**4, 2**5, 2**6, 2**7] -num_trials = 15 +num_trials = 30 def random_trial() -> None: @@ -312,7 +320,7 @@ def last_col_trial() -> None: num_cols=size, seed=cantor(i, j), miss_prob=0.5, - stddev_noise=0.001, + stddev_noise=noise_stddev, latent_factor_combination_model="multiplicative", ) data, mask = sim_dataloader.process_data_scalar() @@ -454,6 +462,21 @@ def last_col_trial() -> None: allow_self_neighbor=True, ) allow_self_neighbor = True + elif estimation_method == "auto": + logger.info("Using AutoNN estimation") + estimator = AutoEstimator() + imputer = NearestNeighborImputer(estimator, data_type) + logger.info("Using AutoNN fit method") + # Fit the imputer using leave-block-out validation + fitter = AutoDRTSLeaveBlockOutValidation( + block, + distance_threshold_range_row=(0, 1), + distance_threshold_range_col=(0, 1), + alpha_range=(0, 1), + n_trials=100, + data_type=data_type, + allow_self_neighbor=args.allow_self_neighbor, + ) else: raise ValueError( f"Estimation method {estimation_method} and fit method {fit_method} not supported" diff --git a/examples/simulations/slurm_scripts/run_accuracy.sh b/examples/simulations/slurm_scripts/run_accuracy.sh index 0f12fa7..94b85e9 100755 --- a/examples/simulations/slurm_scripts/run_accuracy.sh +++ b/examples/simulations/slurm_scripts/run_accuracy.sh @@ -1,21 +1,48 @@ #!/bin/bash # Example usage: -# ./run_accuracy.sh OUTPUT_DIR +# ./run_accuracy.sh OUTPUT_DIR LOG_LEVEL NOISE_LEVEL -OUTPUT_DIR=$1 -LOG_LEVEL=$2 +OUTPUT_DIR="out" +LOG_LEVEL="WARNING" +NOISE_SIGMA=0.001 + +usage() { + echo "Usage: $0 [-o OUTPUT_DIR] [-l LOG_LEVEL] [-n NOISE_SIGMA]" + exit 1 +} + +while getopts ":o:l:n:" opt; do + case ${opt} in + o ) OUTPUT_DIR=$OPTARG ;; + l ) LOG_LEVEL=$OPTARG ;; + n ) NOISE_SIGMA=$OPTARG ;; + \? ) + echo "Invalid option: -$OPTARG" 1>&2 + usage + ;; + : ) + echo "Invalid option: -$OPTARG requires an argument" 1>&2 + usage + ;; + esac +done +shift $((OPTIND -1)) METHODS=( # "softimpute" # "usvt" - "row-row" - "col-col" + "auto" + #"row-row" + #"col-col" "dr" "ts" - "softimpute" + #"softimpute" + ) + + for em in ${METHODS[@]}; do - python run_scalar.py -od $OUTPUT_DIR -em $em --force --log_level $LOG_LEVEL + python run_scalar.py -od $OUTPUT_DIR -em $em --force --log_level $LOG_LEVEL -nstd $NOISE_SIGMA done diff --git a/pyproject.toml b/pyproject.toml index f986b73..dfef86b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta" name = "nsquared" requires-python = ">=3.10" description = "A comprehensive nearest neighbors library for matrix completion with scalar and distributional entries" -version = "1.0.0" +version = "1.1.0" authors = [ {name="Aashish Khubchandani", email="akk223@cornell.edu"}, {name="Albert Gong", email="agong@cs.cornell.edu"}, diff --git a/src/baselines/softimpute.py b/src/baselines/softimpute.py index d25f1fa..587388f 100644 --- a/src/baselines/softimpute.py +++ b/src/baselines/softimpute.py @@ -17,5 +17,5 @@ def softimpute(X: npt.NDArray) -> npt.NDArray: """ # Create a SoftImpute instance with the provided arguments - softimpute = SoftImpute(normalizer=BiScaler()) + softimpute = SoftImpute(normalizer=BiScaler(), verbose=False) return softimpute.fit_transform(X) diff --git a/src/nsquared/data_types.py b/src/nsquared/data_types.py index 15f6415..b272384 100644 --- a/src/nsquared/data_types.py +++ b/src/nsquared/data_types.py @@ -75,6 +75,11 @@ def distance(self, obj1: npt.NDArray, obj2: npt.NDArray) -> float: m = obj1.shape[0] n = obj2.shape[0] + if len(obj1.shape) == 1: + obj1 = obj1[:, np.newaxis] + if len(obj2.shape) == 1: + obj2 = obj2[:, np.newaxis] + assert obj1.shape[1] == obj2.shape[1] XX = np.matmul(obj1, np.transpose(obj1)) # m by m matrix with x_i^Tx_j @@ -114,9 +119,9 @@ def distance(self, obj1: npt.NDArray, obj2: npt.NDArray) -> float: raise ValueError(f"Unknown kernel type: {self.kernel}") val = ( - (np.nansum(kXX) - np.nansum(np.diag(kXX))) / (m * (m - 1)) - + (np.nansum(kYY) - np.nansum(np.diag(kYY))) / (n * (n - 1)) - - 2 * np.nansum(kXY) / (n * m) + (np.sum(kXX) - np.sum(np.diag(kXX))) / (m * (m - 1)) + + (np.sum(kYY) - np.sum(np.diag(kYY))) / (n * (n - 1)) + - 2 * np.sum(kXY) / (n * m) ) if val < 0: val = 0 @@ -150,6 +155,15 @@ class DistributionWassersteinSamples(DataType): where distributions are made with samples with the same number of samples. """ + def __init__(self, num_samples: int): + """Initialize the distribution data type with Wasserstein distance. + + Args: + num_samples (int): Number of samples in the distributions (n) + + """ + self.num_samples = num_samples + def distance(self, obj1: npt.NDArray, obj2: npt.NDArray) -> float: """Calculate the Wasserstein distance between two distributions with the same number of samples. @@ -183,7 +197,13 @@ def average(self, object_list: npt.NDArray[Any]) -> npt.NDArray: """ # filter out nan values # All input objects should be 1-dimensional numpy arrays - return np.mean([np.sort(obj) for obj in object_list], axis=0) + avg = np.nanmean([np.sort(obj) for obj in object_list], axis=0) + if not isinstance(avg, float): + # if avg is an array then return avg + return avg + else: + # if avg is a float then it must (?) be a nan, so return an array of nans + return np.full((self.num_samples,), np.nan) class DistributionWassersteinQuantile(DataType): diff --git a/src/nsquared/datasets/prompteval/loader.py b/src/nsquared/datasets/prompteval/loader.py index 873bf0c..0faecd3 100644 --- a/src/nsquared/datasets/prompteval/loader.py +++ b/src/nsquared/datasets/prompteval/loader.py @@ -23,6 +23,7 @@ import logging from joblib import Memory from datasets import load_dataset, Dataset +from tqdm import tqdm memory = Memory(".joblib_cache", verbose=0) @@ -164,6 +165,11 @@ def __init__( ) # instantiate random seed if provided but do it only once here self.n_examples_per_task = n_examples_per_task + # print('Loading dataset') + # for task in tqdm(self.tasks): + # # Download the dataset to cache + # load_dataset("PromptEval/PromptEval_MMLU_correctness", name=task) + @staticmethod @memory.cache def load_config_data( @@ -188,7 +194,11 @@ def load_config_data( dataset = cast( Dataset, load_dataset( - "PromptEval/PromptEval_MMLU_correctness", name=task, split=model + "PromptEval/PromptEval_MMLU_correctness", + name=task, + split=model, + num_proc=12, + keep_in_memory=True, ), ) # rows are format templates, columns are examples @@ -234,14 +244,18 @@ def process_data_scalar(self) -> tuple[np.ndarray, np.ndarray]: ds = cast( Dataset, load_dataset( - "PromptEval/PromptEval_MMLU_correctness", name=task, split=model + "PromptEval/PromptEval_MMLU_correctness", + name=task, + split=model, + keep_in_memory=True, + num_proc=8, ), ) df = cast(pd.DataFrame, ds.to_pandas()) mask = np.random.binomial(1, propensity, size=df.shape) data = df.to_numpy(dtype=float) - data[mask == 0] = np.nan + # data[mask == 0] = np.nan self.data = data self.mask = mask return data, mask @@ -275,7 +289,7 @@ def process_data_distribution( # Load the data for the multiple config and model df_list = [] - for task in tasks: + for task in tqdm(tasks): for model in models: df = self.load_config_data( task, model, self.n_examples_per_task, self.seed @@ -310,7 +324,8 @@ def process_data_distribution( # Simulate MCAR missingness mask = np.random.binomial(1, propensity, size=data.shape[:2]) - data[mask == 0] = np.nan + # TODO: change this back later + # data[mask == 0] = np.nan self.data = data self.mask = mask diff --git a/src/nsquared/estimation_methods.py b/src/nsquared/estimation_methods.py index f0a5218..543fd48 100644 --- a/src/nsquared/estimation_methods.py +++ b/src/nsquared/estimation_methods.py @@ -79,9 +79,6 @@ def impute( # Find the nearest neighbors indexes nearest_neighbors = np.where(row_dists <= eta_row)[0] - # Apply mask_array to data_array - masked_data_array = np.copy(data_array) - masked_data_array[mask_array == 0] = np.nan # NOTE: this code block will never be called since the target row # is always a nearest neighbor @@ -95,10 +92,12 @@ def impute( else: # return the average of all observed outcomes corresponding # to treatment 1 at time t. - return data_type.average(masked_data_array[:, column]) + return data_type.average(data_array[:, column]) # Calculate the average of the nearest neighbors - nearest_neighbors_data = masked_data_array[nearest_neighbors, column] + nearest_neighbors_data = data_array[nearest_neighbors, column] + nearest_neighbors_mask = mask_array[nearest_neighbors, column] + nearest_neighbors_data = nearest_neighbors_data[nearest_neighbors_mask == 1] return data_type.average(nearest_neighbors_data) def _calculate_distances( @@ -155,6 +154,12 @@ def _calculate_distances( if not overlap_columns[j]: # Skip missing values and the target column row_dists[i, j] = np.nan else: + # if data_array[row, j].dtype == np.float64: + # print (f"Row {row}, Column {j} is NaN") + # exit() + # if isinstance(data_array[i, j], np.float64): + # print (f"Row {i}, Column {j} is NaN") + # exit() row_dists[i, j] = data_type.distance( data_array[row, j], data_array[i, j] ) @@ -772,6 +777,7 @@ def impute( data_array (npt.NDArray): Data matrix containing observed and missing values. mask_array (npt.NDArray): Boolean mask matrix indicating observed values. distance_threshold (Union[float, Tuple[float, float]]): Distance threshold (unused in this method). + allow_self_neighbor (bool): Whether to allow self-neighbor. Defaults to False. (unused in this method) data_type (DataType): Data type providing methods for distance calculation and averaging. allow_self_neighbor (bool): Whether to allow self-neighbor. Defaults to False. (unused in this method) **kwargs (Any): Additional keyword arguments. @@ -878,3 +884,145 @@ def _calculate_distances( row_distances[j, i] = row_distances[i, j] self.row_distances = row_distances + + +class AutoEstimator(EstimationMethod): + """Estimate the missing entry using "Auto-NN" idea (Kyuseong Choi).""" + + def __init__(self, is_percentile: bool = True): + super().__init__(is_percentile) + self.row_distances = dict() + self.col_distances = dict() + # insert default f and g functions + # self.f = lambda x: x + # self.g = lambda x: x + self.alpha = 1.0 + self.drnn_imputer = DREstimator(is_percentile=is_percentile) + self.ts_imputer = TSEstimator(is_percentile=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 doubly robust method. + + Args: + ---- + row (int): Row index + column (int): Column index + data_array (npt.NDArray): Data matrix + mask_array (npt.NDArray): Mask matrix + distance_threshold (float or Tuple[float, float]): Distance threshold for nearest neighbors + or a tuple of (row_threshold, col_threshold) for row and column respectively. + data_type (DataType): Data type to use (e.g. scalars, distributions) + allow_self_neighbor (bool): Whether to allow self-neighbor. Defaults to False. + **kwargs (Any): Additional keyword arguments + + """ + if self.alpha is None and "alpha" not in kwargs: + raise TypeError( + "AutoEstimator.impute() missing 1 required keyword-only argument: 'alpha'" + ) + alpha = self.alpha if self.alpha else kwargs.pop("alpha") + # TODO: handle allow self neighbor (Caleb) + drnn_impute = self.drnn_imputer.impute( + row, + column, + data_array, + mask_array, + distance_threshold, + data_type, + allow_self_neighbor=False, + ) + ts_impute = self.ts_imputer.impute( + row, + column, + data_array, + mask_array, + distance_threshold, + data_type, + allow_self_neighbor=True, + ) + + # interpolate between tsnn and drnn + final_impute = (1 - alpha) * drnn_impute + alpha * ts_impute + + return final_impute + + def _calculate_distances( + self, + row: int, + col: int, + data_array: npt.NDArray, + mask_array: npt.NDArray, + data_type: DataType, + ) -> None: + """Sets the distances for the imputer. + Sets the distances as a class attribute, so returns nothing. + + Args: + row (int): Row index + col (int): Column index + data_array (npt.NDArray): Data matrix + mask_array (npt.NDArray): Mask matrix + data_type (DataType): Data type to use (e.g. scalars, distributions) + + """ + data_shape = data_array.shape + n_rows = data_shape[0] + n_cols = data_shape[1] + + if row not in self.row_distances: + # Calculate distances between rows + row_dists = np.zeros((n_rows, n_cols)) + + for i in range(n_rows): + # Get columns observed in both row i and row + overlap_columns = np.logical_and(mask_array[row], mask_array[i]) + + if not np.any(overlap_columns): + row_dists[i, :] = np.nan + continue + + # Calculate distance between rows + for j in range(n_cols): + if not overlap_columns[ + j + ]: # Skip missing values and the target column + row_dists[i, j] = np.nan + else: + row_dists[i, j] = data_type.distance( + data_array[row, j], data_array[i, j] + ) + self.row_distances[row] = row_dists + + if col not in self.col_distances: + # Calculate distances between columns + col_dists = np.zeros((n_rows, n_cols)) + + for j in range(n_cols): + # Get rows observed in both row i and row + overlap_columns = np.logical_and(mask_array[:, col], mask_array[:, j]) + + if not np.any(overlap_columns): + col_dists[:, j] = np.nan + continue + + # Calculate distance between columns + for i in range(n_rows): + if not overlap_columns[ + i + ]: # Skip missing values and the target column + col_dists[i, j] = np.nan + else: + col_dists[i, j] = data_type.distance( + data_array[i, col], data_array[i, j] + ) + self.col_distances[col] = col_dists diff --git a/src/nsquared/fit_methods.py b/src/nsquared/fit_methods.py index d231d89..dbdda86 100644 --- a/src/nsquared/fit_methods.py +++ b/src/nsquared/fit_methods.py @@ -1,5 +1,5 @@ from .nnimputer import FitMethod, DataType, NearestNeighborImputer -from .estimation_methods import DREstimator, TSEstimator # , AutoEstimator +from .estimation_methods import DREstimator, TSEstimator, AutoEstimator import numpy.typing as npt from hyperopt import hp, fmin, tpe, Trials from typing import cast, Union, Any @@ -319,3 +319,128 @@ def fit( ) imputer.estimation_method = cast(TSEstimator, imputer.estimation_method) return super().fit(data_array, mask_array, imputer) + + +class AutoDRTSLeaveBlockOutValidation(DualThresholdLeaveBlockOutValidation): + """Fit method by leaving out a block of cells using separate thresholds for rows and columns with a AutoEstimator.""" + + expected_estimator_type = AutoEstimator + + 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, + ): + """Initialize the dual threshold block fit method with AutoEstimator. + + Args: + block (list[tuple[int, int]]): List of cells as tuples of row and column indices. + distance_threshold_range_row (tuple[float, float]): Range of row distance thresholds to test. + distance_threshold_range_col (tuple[float, float]): Range of column distance thresholds to test. + alpha_range (tuple[float, float]): Range of alpha values to test. + n_trials (int): Number of trials to run. + data_type (DataType): Data type to use (e.g. scalars, distributions). + allow_self_neighbor (bool, optional): Whether to allow the entry itself as a neighbor. Defaults to False. + + """ + self.alpha_range = alpha_range + super().__init__( + block, + distance_threshold_range_row, + distance_threshold_range_col, + n_trials, + data_type, + ) + self.allow_self_neighbor = allow_self_neighbor + + 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]]: + """Find the best distance thresholds for rows and columns by leaving out a block of cells and testing imputation. + + Args: + data_array (npt.NDArray): Data matrix. + mask_array (npt.NDArray): Mask matrix. + imputer (NearestNeighborImputer): Imputer object. + ret_trials (bool): If True, return the trials object which contains metadata on hyperparameter search. + + Returns: + tuple[float, float]: Best distance thresholds for rows and columns. + + """ + if not isinstance(imputer.estimation_method, AutoEstimator): + raise ValueError( + f"The imputer must use a AutoEstimator for {self.__class__.__name__}." + ) + imputer.estimation_method = cast(AutoEstimator, imputer.estimation_method) + + def _objective(params: dict[str, float]) -> float: + """Objective function for hyperopt. + + Args: + params (dict[str, float]): Dictionary containing row and column distance thresholds + + Returns: + float: Average imputation error + + """ + row_threshold = params["distance_threshold_row"] + col_threshold = params["distance_threshold_col"] + alpha = params["alpha"] + imputer.distance_threshold = (row_threshold, col_threshold) + if not isinstance(imputer.estimation_method, AutoEstimator): + raise ValueError( + f"The imputer must use a AutoEstimator for {self.__class__.__name__}." + ) + imputer.estimation_method.alpha = alpha + return evaluate_imputation( + data_array, + mask_array, + imputer, + self.block, + self.data_type, + self.allow_self_neighbor, + ) + + lower_bound_row, upper_bound_row = self.distance_threshold_range_row + lower_bound_col, upper_bound_col = self.distance_threshold_range_col + lower_bound_alpha, upper_bound_alpha = self.alpha_range + trials = Trials() + best_params = fmin( + fn=_objective, + space={ + "distance_threshold_row": hp.uniform( + "distance_threshold_row", lower_bound_row, upper_bound_row + ), + "distance_threshold_col": hp.uniform( + "distance_threshold_col", lower_bound_col, upper_bound_col + ), + "alpha": hp.uniform("alpha", lower_bound_alpha, upper_bound_alpha), + }, + algo=tpe.suggest, + max_evals=self.n_trials, + verbose=False, + trials=trials, + ) + + if best_params is None: + return float("nan"), float("nan") + + imputer.distance_threshold = ( + best_params["distance_threshold_row"], + best_params["distance_threshold_col"], + ) + imputer.estimation_method.alpha = best_params["alpha"] + + if ret_trials: + return imputer.distance_threshold, trials + return imputer.distance_threshold diff --git a/src/nsquared/utils/plotting_utils.py b/src/nsquared/utils/plotting_utils.py index 64a2010..72cb8e4 100644 --- a/src/nsquared/utils/plotting_utils.py +++ b/src/nsquared/utils/plotting_utils.py @@ -17,6 +17,9 @@ OUTWARD = 4 METHOD_ALIASES = { + "kernel": "Kernel-\nNN", + "wasserstein_samples": "W$_2$-\nNN", + "auto": "Auto-\nNN", "row-row": "Row-\nNN", "col-col": "Col-\nNN", "dr": "DR-\nNN", @@ -25,10 +28,11 @@ "softimpute": "Soft\nImpute", "nadaraya": "NW", "aw": "AW-\nNN", - "sc": "Synth.", + "sc": "SC", } METHOD_ALIASES_SINGLE_LINE = { + "auto": "AutoNN", "row-row": "RowNN", "col-col": "ColNN", "dr": "DRNN", @@ -36,8 +40,10 @@ "usvt": "USVT", "softimpute": "SI", "nadaraya": "NW", - "aw": "aw", - "sc": "Synth.", + "aw": "AWNN", + "sc": "SC", + "kernel": "KernelNN", + "wasserstein_samples": "W$_2$NN", } METHOD_LINE_STYLES = { @@ -50,12 +56,17 @@ "nadaraya": "-", "aw": "dashdot", "sc": "-", + "auto": "-", } NEURIPS_TEXTWIDTH = 5.5 COLORS = { - "row-row": "grey", + "row-row": "red", + "auto": "grey", + "kernel": "dimgrey", + "wasserstein_samples": "dimgrey", + # "row-row": "grey", "col-col": "grey", "dr": "grey", "ts": "grey",