diff --git a/.gitignore b/.gitignore index 69a8754b..dde72bbc 100644 --- a/.gitignore +++ b/.gitignore @@ -29,7 +29,7 @@ wheels/ # Synthcity backups **/workspace/*.bkp -# Dataset files +# Data files examples/**/data/ # Trained metaclassifiers diff --git a/examples/ensemble_attack/config.yaml b/examples/ensemble_attack/config.yaml index fcd35ffb..d080e614 100644 --- a/examples/ensemble_attack/config.yaml +++ b/examples/ensemble_attack/config.yaml @@ -11,14 +11,13 @@ data_paths: attack_results_path: ${base_example_dir}/attack_results # Path where the attack results will be stored model_paths: - shadow_models_path: ${base_example_dir}/shadow_models # Path where the shadow models are stored metaclassifier_model_path: ${base_example_dir}/trained_models # Path where the trained metaclassifier model will be saved # Pipeline control pipeline: run_data_processing: false # Set this to false if you have already saved the processed data - run_shadow_model_training: true - run_metaclassifier_training: false + run_shadow_model_training: false # Set this to false if shadow models are already trained and saved + run_metaclassifier_training: true # Dataset specific information used for processing in this example @@ -54,7 +53,18 @@ shadow_training: tabddpm_training_config_path: ${base_example_dir}/data_configs/trans.json # Model training artifacts are saved under shadow_models_data_path/workspace_name/exp_name # Also, training configs for each shadow model are created under shadow_models_data_path. - shadow_models_output_path: ${base_data_dir}/shadow_models_data + shadow_models_output_path: ${base_data_dir}/shadow_models_and_data + target_model_output_path: ${base_data_dir}/target_model_and_data + final_shadow_models_path: [ + "${shadow_training.shadow_models_output_path}/initial_model_rmia_1/shadow_workspace/pre_trained_model/rmia_shadows.pkl", + "${shadow_training.shadow_models_output_path}/initial_model_rmia_2/shadow_workspace/pre_trained_model/rmia_shadows.pkl", + "${shadow_training.shadow_models_output_path}/shadow_model_rmia_third_set/shadow_workspace/trained_model/rmia_shadows_third_set.pkl", + ] # Paths to final shadow models used for metaclassifier training (relative to shadow_models_output_path) + # These paths are a result of running the shadow model training pipeline, specifically the + # train_three_sets_of_shadow_models in shadow_model_training.py + # Each .pkl file contains the training data, trained model and training results for all shadow models in a list. + final_target_model_path: ${shadow_training.target_model_output_path}/target_model/shadow_workspace/trained_target_model/target_model.pkl + # Path to final target model (relative to target_model_output_path) fine_tuning_config: fine_tune_diffusion_iterations: 2 fine_tune_classifier_iterations: 2 @@ -66,10 +76,13 @@ metaclassifier: # Data types json file is used for xgboost model training. data_types_file_path: ${base_example_dir}/data_configs/data_types.json model_type: "xgb" - use_gpu: true + # Model training parameters + num_optuna_trials: 10 # Original code: 100 + num_kfolds: 5 + use_gpu: false # Temporary. Might remove having an epoch parameter. epochs: 1 # General settings -random_seed: 42 +random_seed: 42 # Set to null for no seed, or an integer for a fixed seed diff --git a/examples/ensemble_attack/data_configs/data_types.json b/examples/ensemble_attack/data_configs/data_types.json index a7a93569..1851838e 100644 --- a/examples/ensemble_attack/data_configs/data_types.json +++ b/examples/ensemble_attack/data_configs/data_types.json @@ -1,5 +1,6 @@ { "numerical": ["trans_date", "amount", "balance", "account"], "categorical": ["trans_type", "operation", "k_symbol", "bank"], - "variable_to_predict": "trans_type" + "variable_to_predict": "trans_type", + "id_column_name": "trans_id" } diff --git a/examples/ensemble_attack/run_attack.py b/examples/ensemble_attack/run_attack.py index c77b2e6b..e206f5fc 100644 --- a/examples/ensemble_attack/run_attack.py +++ b/examples/ensemble_attack/run_attack.py @@ -59,10 +59,25 @@ def main(config: DictConfig) -> None: # TODO: Investigate the source of error. if config.pipeline.run_shadow_model_training: shadow_pipeline = importlib.import_module("examples.ensemble_attack.run_shadow_model_training") - shadow_pipeline.run_shadow_model_training(config) + attack_data_paths = shadow_pipeline.run_shadow_model_training(config) + attack_data_paths = [Path(path) for path in attack_data_paths] + + target_data_path = shadow_pipeline.run_target_model_training(config) + target_data_path = Path(target_data_path) + if config.pipeline.run_metaclassifier_training: + if not config.pipeline.run_shadow_model_training: + # If shadow model training is skipped, we need to provide the previous shadow model and target model paths. + + shadow_data_paths = [Path(path) for path in config.shadow_training.final_shadow_models_path] + + target_data_path = Path(config.shadow_training.final_target_model_path) + + assert len(shadow_data_paths) == 3, "The attack_data_paths list must contain exactly three elements." + assert target_data_path is not None, "The target_data_path must be provided for metaclassifier training." + meta_pipeline = importlib.import_module("examples.ensemble_attack.run_metaclassifier_training") - meta_pipeline.run_metaclassifier_training(config) + meta_pipeline.run_metaclassifier_training(config, shadow_data_paths, target_data_path) if __name__ == "__main__": diff --git a/examples/ensemble_attack/run_metaclassifier_training.py b/examples/ensemble_attack/run_metaclassifier_training.py index 711ea643..51e48563 100644 --- a/examples/ensemble_attack/run_metaclassifier_training.py +++ b/examples/ensemble_attack/run_metaclassifier_training.py @@ -11,19 +11,30 @@ from midst_toolkit.common.logger import log -def run_metaclassifier_training(config: DictConfig) -> None: +def run_metaclassifier_training( + config: DictConfig, + shadow_data_paths: list[Path], + target_data_path: Path, +) -> None: """ Fuction to run the metaclassifier training and evaluation. Args: config: Configuration object set in config.yaml. + shadow_data_paths: List of paths to the trained shadow models and all their attributes and synthetic data. + The list should contain three paths, one for each set of shadow models. + target_data_path: Path to the target model and all its attributes and synthetic data. """ log(INFO, "Running metaclassifier training...") + # Load the processed data splits. df_meta_train = load_dataframe( Path(config.data_paths.processed_attack_data_path), "master_challenge_train.csv", ) + + # y_meta_train consists of binary labels (0s and 1s) indicating whether each row in df_meta_train + # belongs to the target model's training set. y_meta_train = np.load( Path(config.data_paths.processed_attack_data_path) / "master_challenge_train_labels.npy", ) @@ -35,23 +46,48 @@ def run_metaclassifier_training(config: DictConfig) -> None: Path(config.data_paths.processed_attack_data_path) / "master_challenge_test_labels.npy", ) - # Synthetic data borrowed from the attack implementation repository. - # From (https://github.com/CRCHUM-CITADEL/ensemble-mia/tree/main/input/tabddpm_black_box/meta_classifier) - # TODO: Change this file path to the path where the synthetic data is stored. - df_synthetic = load_dataframe( - Path(config.data_paths.processed_attack_data_path), - "synth.csv", + # Three sets of shadow models are trained separately and their paths are provided here. + + assert len(shadow_data_paths) == 3, ( + "At this point of development, the shadow_data_paths list must contain exactly three elements." ) + shadow_data_collection = [] + + for model_path in shadow_data_paths: + assert model_path.exists(), ( + f"No file found at {model_path}. Make sure the path is correct, or run shadow model training first." + ) + + with open(model_path, "rb") as f: + shadow_data_and_result = pickle.load(f) + shadow_data_collection.append(shadow_data_and_result) + + assert target_data_path.exists(), ( + f"No file found at {target_data_path}. Make sure the path is correct and that you have trained the target model." + ) + + with open(target_data_path, "rb") as f: + target_data_and_result = pickle.load(f) + + target_synthetic = target_data_and_result["trained_results"][0].synthetic_data + assert target_synthetic is not None, "Target model pickle missing synthetic_data." + target_synthetic = target_synthetic.copy() + df_reference = load_dataframe( Path(config.data_paths.population_path), "population_all_with_challenge_no_id.csv", ) - # We should drop the id column from master metaclassifier train data. - if "trans_id" in df_meta_train.columns: - df_meta_train = df_meta_train.drop(columns=["trans_id", "account_id"]) - if "trans_id" in df_meta_test.columns: - df_meta_test = df_meta_test.drop(columns=["trans_id", "account_id"]) + + # Extract trans_id from both train and test dataframes + assert "trans_id" in df_meta_train.columns, "Meta train data must have trans_id column" + train_trans_ids = df_meta_train["trans_id"] + + assert "trans_id" in df_meta_test.columns, "Meta test data must have trans_id column" + test_trans_ids = df_meta_test["trans_id"] + + df_meta_train = df_meta_train.drop(columns=["trans_id", "account_id"]) + df_meta_test = df_meta_test.drop(columns=["trans_id", "account_id"]) # Fit the metaclassifier. meta_classifier_enum = MetaClassifierType(config.metaclassifier.model_type) @@ -59,45 +95,52 @@ def run_metaclassifier_training(config: DictConfig) -> None: # 1. Initialize the attacker blending_attacker = BlendingPlusPlus( config=config, + shadow_data_collection=shadow_data_collection, + target_data=target_data_and_result, meta_classifier_type=meta_classifier_enum, random_seed=config.random_seed, ) - log( - INFO, - f"{meta_classifier_enum} created with random seed {config.random_seed}, starting training...", - ) + + log(INFO, f"{meta_classifier_enum} created with random seed {config.random_seed}.") # 2. Train the attacker on the meta-train set blending_attacker.fit( df_train=df_meta_train, y_train=y_meta_train, - df_synthetic=df_synthetic, + df_target_synthetic=target_synthetic, df_reference=df_reference, + id_column_data=train_trans_ids, use_gpu=config.metaclassifier.use_gpu, epochs=config.metaclassifier.epochs, ) - log(INFO, "Metaclassifier training finished.") - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - # TODO: Create the directory folder if it does not exist. model_filename = f"{timestamp}_{config.metaclassifier.model_type}_trained_metaclassifier.pkl" with open(Path(config.model_paths.metaclassifier_model_path) / model_filename, "wb") as f: pickle.dump(blending_attacker.trained_model, f) log(INFO, "Metaclassifier model saved, starting evaluation...") + # Get the synthetic data provided by the challenge for evaluation + # TODO: Check if the file is the correct one. + df_synthetic_original = load_dataframe( + Path(config.data_paths.processed_attack_data_path), + "synth.csv", + ) + # 3. Get predictions on the test set probabilities, pred_score = blending_attacker.predict( df_test=df_meta_test, - df_synthetic=df_synthetic, + df_original_synthetic=df_synthetic_original, df_reference=df_reference, + id_column_data=test_trans_ids, y_test=y_meta_test, ) # Save the prediction probabilities - # TODO: Create the attack results directory folder if it does not exist. + attack_results_path = Path(config.data_paths.attack_results_path) + attack_results_path.mkdir(parents=True, exist_ok=True) np.save( Path(config.data_paths.attack_results_path) / f"{timestamp}_{config.metaclassifier.model_type}_test_pred_proba.npy", diff --git a/examples/ensemble_attack/run_shadow_model_training.py b/examples/ensemble_attack/run_shadow_model_training.py index 62cf076c..66659373 100644 --- a/examples/ensemble_attack/run_shadow_model_training.py +++ b/examples/ensemble_attack/run_shadow_model_training.py @@ -1,5 +1,8 @@ +import pickle +import shutil from logging import INFO from pathlib import Path +from typing import Any from omegaconf import DictConfig @@ -7,15 +10,94 @@ from midst_toolkit.attacks.ensemble.rmia.shadow_model_training import ( train_three_sets_of_shadow_models, ) +from midst_toolkit.attacks.ensemble.shadow_model_utils import ( + save_additional_tabddpm_config, + train_tabddpm_and_synthesize, +) from midst_toolkit.common.logger import log -def run_shadow_model_training(config: DictConfig) -> None: +def run_target_model_training(config: DictConfig) -> Path: + """ + Function to run the target model training for RMIA attack. + + Args: + config: Configuration object set in config.yaml. + + Returns: + Path to the saved target model results. + """ + log(INFO, "Running target model training...") + + # Load the required dataframe for target model training. + df_real_data = load_dataframe( + Path(config.data_paths.processed_attack_data_path), + "real_train.csv", + ) + + # TODO: Test when pipeline is complete to make sure real_data is correct. + + target_model_output_path = Path(config.shadow_training.target_model_output_path) + target_training_json_config_paths = config.shadow_training.training_json_config_paths + + # TODO: Add this to config or .json files + table_name = "trans" + id_column_name = "trans_id" + + target_folder = target_model_output_path / "target_model" + + target_folder.mkdir(parents=True, exist_ok=True) + shutil.copyfile( + target_training_json_config_paths.table_domain_file_path, + target_folder / f"{table_name}_domain.json", + ) + shutil.copyfile( + target_training_json_config_paths.dataset_meta_file_path, + target_folder / "dataset_meta.json", + ) + configs, save_dir = save_additional_tabddpm_config( + data_dir=target_folder, + training_config_json_path=Path(target_training_json_config_paths.tabddpm_training_config_path), + final_config_json_path=target_folder / f"{table_name}.json", # Path to the new json + experiment_name="trained_target_model", + ) + + train_result = train_tabddpm_and_synthesize( + train_set=df_real_data, + configs=configs, + save_dir=save_dir, + synthesize=True, + ) + + # TODO: Check: Selected_id_lists should be of form [[]] + selected_id_lists = [df_real_data[id_column_name].tolist()] + + attack_data: dict[str, Any] = { + "selected_sets": selected_id_lists, + "trained_results": [], + } + + attack_data["trained_results"].append(train_result) + + # Pickle dump the results + result_path = Path(save_dir, "target_model.pkl") + with open(result_path, "wb") as file: + pickle.dump(attack_data, file) + + return result_path + + +def run_shadow_model_training(config: DictConfig) -> list[Path]: """ Function to run the shadow model training for RMIA attack. Args: config: Configuration object set in config.yaml. + + Returns: + Paths to the saved shadow model results for the three sets of shadow models. For more details, + see the documentation and return value of `train_three_sets_of_shadow_models` + at src/midst_toolkit/attacks/ensemble/rmia/shadow_model_training.py. """ log(INFO, "Running shadow model training...") # Load the required dataframes for shadow model training. @@ -55,5 +137,7 @@ def run_shadow_model_training(config: DictConfig) -> None: ) log( INFO, - f"Shadow model training finished and saved at 1) {first_set_result_path}, 2) {second_set_result_path}, 3) {third_set_result_path}", + f"Shadow model training finished and saved at \n1) {first_set_result_path} \n2) {second_set_result_path} \n3) {third_set_result_path}", ) + + return [first_set_result_path, second_set_result_path, third_set_result_path] diff --git a/src/midst_toolkit/attacks/ensemble/blending.py b/src/midst_toolkit/attacks/ensemble/blending.py index 1f156828..83c31804 100644 --- a/src/midst_toolkit/attacks/ensemble/blending.py +++ b/src/midst_toolkit/attacks/ensemble/blending.py @@ -2,6 +2,8 @@ import json from enum import Enum +from logging import INFO +from typing import Any import numpy as np import pandas as pd @@ -9,8 +11,10 @@ from sklearn.linear_model import LogisticRegression from midst_toolkit.attacks.ensemble.distance_features import calculate_domias_score, calculate_gower_features +from midst_toolkit.attacks.ensemble.rmia.rmia_calculation import calculate_rmia_signals from midst_toolkit.attacks.ensemble.train_utils import get_tpr_at_fpr from midst_toolkit.attacks.ensemble.xgboost_tuner import XgBoostHyperparameterTuner +from midst_toolkit.common.logger import log class MetaClassifierType(Enum): @@ -22,6 +26,8 @@ class BlendingPlusPlus: def __init__( self, config: DictConfig, + shadow_data_collection: list[dict[str, list[Any]]], + target_data: dict[str, list[Any]], meta_classifier_type: MetaClassifierType = MetaClassifierType.XGB, random_seed: int | None = None, ) -> None: @@ -30,12 +36,27 @@ def __init__( This class encapsulates the entire workflow: 1. Generates features from Gower distance and DOMIAS. - 2. Assembles a meta-feature set. - 3. Trains a meta-classifier on these features. - 4. Predicts membership probability on new data. + 2. Calculates RMIA signals using the provided attack data, which contains training/fine-tuning data + and the synthetic data generated. + 3. Assembles a meta-feature set (original numerical features + Gower + DOMIAS + RMIA). + 4. Trains a meta-classifier on these features. + 5. Predicts membership probability on new data. Args: config: Dictionary storing data configuration paths and parameters, used to load data properties. + shadow_data_collection: List of training data of the shadow models and their generated synthetic data. + Each list element is a dict with keys "fine_tuning_sets" and "fine_tuned_results". + Fine_tuning_sets is a list of dataframes used to fine-tune the shadow models, and fine_tuned_results + is a list of type TrainingResult containing model training information and generated synthetic data. + For more details, see the documentation of `train_three_sets_of_shadow_models` at + attacks/ensemble/rmia/shadow_model_training.py. + target_data: Dictionary containing the training data of the target model and its generated synthetic data. + The dictionary contains the keys "selected_sets" and "trained_results". + Selected_sets is a list of dataframes used to train the target model, and trained_results + is a list of type TrainingResult containing model training information and generated synthetic data. + For more details, see the documentation of `train_three_sets_of_shadow_models` at + attacks/ensemble/rmia/shadow_model_training.py. + meta_classifier_type: Type of meta classifier model. Defaults to MetaClassifierType.XGB. random_seed: Random seed for reproducibility. Defaults to None. @@ -43,18 +64,23 @@ def __init__( # TODO: We can directly pass the `data_types_file_path` as a parameter to this class. with open(config.metaclassifier.data_types_file_path, "r") as f: self.column_types = json.load(f) + + self.shadow_data_collection = shadow_data_collection + self.target_data = target_data self.meta_classifier_type = meta_classifier_type self.trained_model = None self.random_seed = random_seed + self.training_config = config.metaclassifier - # TODO: Add RMIA function def _prepare_meta_features( self, df_input: pd.DataFrame, df_synthetic: pd.DataFrame, df_reference: pd.DataFrame, + id_column_data: pd.Series, categorical_cols: list[str], numerical_cols: list[str], + id_column_name: str, ) -> pd.DataFrame: """ Prepares meta-classifier features by combining original continuous features, @@ -64,8 +90,10 @@ def _prepare_meta_features( df_input: Input dataframe (e.g., meta-classifier train or test set). df_synthetic: Synthetic dataframe. df_reference: Real population dataframe, used as a reference for calculating the DOMIAS score. + id_column_data: The data in the ID column, used to ensure correct alignment of results. categorical_cols: Categorical column names. numerical_cols: Numerical column names. + id_column_name: Name of the ID column. Returns: A dataframe with the meta-classifier features. @@ -74,23 +102,36 @@ def _prepare_meta_features( """ df_synthetic = df_synthetic.reset_index(drop=True)[df_input.columns] - # 1. Get Gower distance features + # 1. Get RMIA signals + + log(INFO, "Calculating RMIA signals...") + + rmia_signals = calculate_rmia_signals( + df_input=df_input, + shadow_data_collection=self.shadow_data_collection, + target_data=self.target_data, + categorical_column_names=categorical_cols, + id_column_name=id_column_name, + id_column_data=id_column_data, + random_seed=self.random_seed, + ) + + # 2. Get Gower distance features + + log(INFO, "Calculating Gower features...") + gower_features = calculate_gower_features( df_input=df_input, df_synthetic=df_synthetic, categorical_column_names=categorical_cols ) - # 2. Get DOMIAS predictions + # 3. Get DOMIAS predictions + + log(INFO, "Calculating DOMIAS scores...") + domias_features = calculate_domias_score( df_input=df_input, df_synthetic=df_synthetic, df_reference=df_reference ) - # 3. Get RMIA signals (borrowed from the attack implementation repository, - # at https://github.com/CRCHUM-CITADEL/ensemble-mia/tree/main/input/tabddpm_black_box/meta_classifier) - # Will be removed after our own implementation is ready. - rmia_signals = pd.read_csv( - "examples/ensemble_attack/data/attack_data/og_rmia_train_meta_pred.csv" - ) # Placeholder for RMIA features - original_numerical_features = df_input[numerical_cols] # Numerical features from original data return pd.concat( @@ -103,14 +144,16 @@ def _prepare_meta_features( axis=1, ) + # TODO: Handle epochs parameter def fit( self, df_train: pd.DataFrame, y_train: np.ndarray, - df_synthetic: pd.DataFrame, + df_target_synthetic: pd.DataFrame, df_reference: pd.DataFrame, + id_column_data: pd.Series, use_gpu: bool = True, - epochs: int = 1, + epochs: int | None = None, ) -> None: """ Trains the Blending++ meta-classifier. @@ -119,21 +162,27 @@ def fit( df_train: Dataframe for training the meta-classifier. This training set is derived from the population dataset which is all the data the attacker has access to (all the other attacks' training data, holdout data, and the challenge dataset). - The meta training set is a combination of the "real train" data and "real control val", which is - the data used to validate the diffusion model to generate synthetic data. - y_train: Labels for the meta-classifier training data. - df_synthetic: Synthetic dataframe, generated by the diffusion model. + The meta training set is a combination of the "real train" data and "real control val". + y_train: Labels for the meta-classifier training data, indicating whether rows in df_train are a + member of the target model's train data or not. + df_target_synthetic: Synthetic dataframe, generated by the simulated target diffusion model. df_reference: Reference (real) population dataframe. + id_column_data: The data in the ID column, used to ensure correct alignment of results. use_gpu: Whether to use GPU acceleration. Defaults to True. - epochs: Number of training iterations. Defaults to 1. + epochs: Number of training iterations. Defaults to None, in which case self.training_config.epochs is used. """ + if epochs is None: + epochs = self.training_config.epochs + meta_features = self._prepare_meta_features( df_input=df_train, - df_synthetic=df_synthetic, + df_synthetic=df_target_synthetic, df_reference=df_reference, + id_column_data=id_column_data, categorical_cols=self.column_types["categorical"], numerical_cols=self.column_types["numerical"], + id_column_name=self.column_types["id_column_name"], ) if self.meta_classifier_type == MetaClassifierType.XGB: @@ -146,8 +195,8 @@ def fit( # Run the tuning process self.trained_model = tuner.tune_hyperparameters( - num_optuna_trials=100, - num_kfolds=5, + num_optuna_trials=self.training_config.num_optuna_trials, + num_kfolds=self.training_config.num_kfolds, ) elif self.meta_classifier_type == MetaClassifierType.LR: @@ -160,8 +209,9 @@ def fit( def predict( self, df_test: pd.DataFrame, - df_synthetic: pd.DataFrame, + df_original_synthetic: pd.DataFrame, df_reference: pd.DataFrame, + id_column_data: pd.Series, y_test: np.ndarray, ) -> tuple[np.ndarray, float | None]: """ @@ -174,8 +224,9 @@ def predict( attacks' training data, holdout data, and the challenge dataset). The meta-test set includes "real train" data and "real control test" data used to evaluate the diffusion model's synthetic data generation. - df_synthetic: DataFrame containing synthetic data generated by the diffusion model. + df_original_synthetic: DataFrame containing synthetic data generated by the diffusion model. df_reference: DataFrame of the real population data, used as a reference for calculating the DOMIAS score. + id_column_data: The data in the ID column, used to ensure correct alignment of results. y_test: Optional array of test labels for evaluation. A label of "1" indicates membership in the diffusion model's training set, while "0" indicates non-membership. @@ -191,10 +242,12 @@ def predict( df_test_features = self._prepare_meta_features( df_input=df_test, - df_synthetic=df_synthetic, + df_synthetic=df_original_synthetic, df_reference=df_reference, + id_column_data=id_column_data, categorical_cols=self.column_types["categorical"], numerical_cols=self.column_types["numerical"], + id_column_name=self.column_types["id_column_name"], ) probabilities = self.trained_model.predict_proba(df_test_features)[:, 1] diff --git a/src/midst_toolkit/attacks/ensemble/rmia/rmia_calculation.py b/src/midst_toolkit/attacks/ensemble/rmia/rmia_calculation.py new file mode 100644 index 00000000..7b9cb4f4 --- /dev/null +++ b/src/midst_toolkit/attacks/ensemble/rmia/rmia_calculation.py @@ -0,0 +1,369 @@ +"""Computes the Gower distance and membership signals between the challenge points + and the synthetic data generated by the shadow models. + Overall workflow and numerical decisions are taken with from CITADEL & UQAM team's attack implementation at +https://github.com/CRCHUM-CITADEL/ensemble-mia. +""" + +from enum import Enum +from logging import INFO +from typing import Any + +import gower +import numpy as np +import pandas as pd + +from midst_toolkit.common.logger import log + + +class Key(Enum): + TRAINED_RESULTS = "trained_results" + FINE_TUNED_RESULTS = "fine_tuned_results" + + +def get_rmia_gower( + df_input: pd.DataFrame, + model_data: dict[str, list[Any]], + min_length: int, + key: Key, + categorical_column_names: list[str], + id_column_name: str, + random_seed: int | None = None, +) -> list[np.ndarray]: + """ + Computes the Gower distance between the challenge points and the synthetic data generated by the shadow models. + + Args: + df_input: The dataframe to generate features for (e.g., meta classifier train or test set), derived + from the challenge dataset and processed in process_split_data.py. + model_data: A dictionary with keys "fine_tuning_sets" and "fine_tuned_results". + Fine_tuning_sets is a list of dataframes used to fine-tune the shadow models, and fine_tuned_results + is a list of type TrainingResult containing model training information and generated synthetic data. + min_length: Minimum length across all training data, fine tuning data, and synthetic data sizes. This length + will be used for downsampling to ensure consistent Gower distance calculations. + key: An instance of the Key Enum, either Key.TRAINED_RESULTS or Key.FINE_TUNED_RESULTS, + depending on which set of shadow models to use. + categorical_column_names: A list of categorical column names. We assume that all other columns are numerical. + id_column_name: Name of the ID column. + random_seed: Random seed for reproducibility. + + Returns: + A list of numpy arrays, each representing the Gower distance matrix between the input dataframe and the + synthetic data from a model. Each entry in the list corresponds to one of the synthetically generated + datasets provided within ``model_data``. + + """ + # Check if any specified categorical columns are missing from the dataframe + missing_categorical_columns = set(categorical_column_names) - set(df_input.columns) + if missing_categorical_columns: + log( + INFO, + f"Warning: The following categorical columns are missing from the dataframe, but have been passed as categorical columns: {missing_categorical_columns}. " + "These columns will be ignored. Ensure that the specified categorical columns in your dataset's data_types.json file are in sync with the actual dataframe columns, unless it is intentional.", + ) + + categorical_features = [column in categorical_column_names for column in df_input.columns] + + gower_matrices = [] + + numerical_columns = [col for col in df_input.columns if col not in categorical_column_names] + + df_input[numerical_columns] = df_input[numerical_columns].astype(float) + + for i in range(len(model_data[key.value])): + df_synthetic = model_data[key.value][i].synthetic_data.copy() + + # Convert numerical columns to float (otherwise error in the numpy divide) + df_synthetic[numerical_columns] = df_synthetic[numerical_columns].astype(float) + + # Sample synthetic data points if there's too many + if len(df_synthetic) > min_length: + df_synthetic = df_synthetic.sample(n=min_length, random_state=random_seed) + + if id_column_name in df_synthetic.columns: + df_synthetic = df_synthetic.drop(columns=[id_column_name]) + + gower_matrix = gower.gower_matrix(data_x=df_input, data_y=df_synthetic, cat_features=categorical_features) + + assert np.all((gower_matrix >= 0) & (gower_matrix <= 1)), "Distances are falling outside of range [0, 1]." + + gower_matrices.append(gower_matrix) + + return gower_matrices + + +def conditional_average(values: np.ndarray, condition_mask: np.ndarray) -> np.ndarray: + """ + Calculate the conditional average of values based on a condition mask. + This function computes the mean of the ``values`` array where the corresponding + entries in the ``condition_mask`` are ``True``. If no entries satisfy the condition + for a particular record, the result will be ``NaN`` for that record. + + Args: + values: A NumPy array containing the values to average. + condition_mask (numpy.ndarray): A boolean NumPy array of the same shape as + ``values``, where ``True`` indicates the values to include in the average. + + Returns: + An array of the same shape as the input, containing the conditional averages + where the condition is met, or ``NaN`` where no values satisfy the condition. + + """ + assert values.shape == condition_mask.shape, "condition_mask must have the same shape as values" + + mask_sum = condition_mask.sum(axis=0) + return np.where(mask_sum > 0, np.sum(values * condition_mask, axis=0) / mask_sum, np.nan) + + +def calculate_rmia_signals( + df_input: pd.DataFrame, + shadow_data_collection: list[dict[str, list[Any]]], + target_data: dict[str, list[Any]], + categorical_column_names: list[str], + id_column_name: str, + id_column_data: pd.Series, + k: int = 5, + random_seed: int | None = None, +) -> pd.DataFrame: + """ + Main orchestration function to compute Robust Membership Inference Attack (RMIA) signals + with respect to shadow and target models. + + This function implements a data-based RMIA to determine the likelihood that a + given data record was part of a target model's training set. The attack is + data-based because it only uses synthetic data generated by the models, not the + models themselves. + + The core idea is to compare a data record's similarity to two groups of synthetic data: + 1. Data generated by the 'target model' (the model we are simulating to attack). + 2. Data generated by several 'shadow models' (models trained on similar, but + distinct, datasets). + + The intuition is that if a record was used to train the target model, the synthetic + data it generates should be "closer" or more similar to that record than the + synthetic data generated by the shadow models. This similarity is measured using + Gower distance, which is suitable for datasets with mixed data types (both + categorical and numerical). + + The final RMIA score is a ratio: (Similarity to Target Synthetic Data)/(Mean Similarity to Shadows Synthetic Data). + A high RMIA score suggests that the record is disproportionately more similar to the + target model's output, indicating a higher probability of being a training member. + + This implementation is based on how they used RMIA in this paper by Meeus et al.: + https://arxiv.org/abs/2502.14921 + + The original RMIA paper by Zarifzadeh et al. can be found at: + https://arxiv.org/abs/2312.03262 + + Args: + df_input: The dataframe to generate features for (e.g., meta classifier train or test set), derived + from the challenge dataset and processed in ``process_split_data.py``. + shadow_data_collection: A list containing three dictionaries, each representing a collection of shadow + models with their training data and generated synthetic outputs. Each collection can contain multiple + shadow models. The first two dictionaries have keys ``fine_tuning_sets`` and ``fine_tuned_results``, while + the third has keys ``selected_sets`` and ``trained_results``. The ``fine_tuning_sets`` and + ``selected_sets`` keys map to lists of DataFrames containing the data used to fine-tune or train the shadow + models. The ``fine_tuned_results`` and ``trained_results`` keys map to lists of ``TrainingResult`` objects, + which store model training metadata and the corresponding generated synthetic data. + See ``train_three_sets_of_shadow_models`` in attacks/ensemble/rmia/shadow_model_training.py + for additional details. + target_data: A dictionary containing information about the target model. It includes: + - ``selected_sets``: A list of DataFrames used to train the target model. + - ``trained_results``: A list of ``TrainingResult`` objects, each containing details about the model's + training process and the synthetic data generated during training. + categorical_column_names: A list of categorical column names. + id_column_name: Name of the ID column. + id_column_data: The data in the ID column extracted from df_input, ensuring that output signals are + correctly aligned with their corresponding input records. + k: The number of nearest neighbors to consider when computing the final membership inference signals. + A higher ``k`` provides a more stable but less sensitive signal. Default is 5. + random_seed: Random seed for reproducibility. + + Returns: + A pandas DataFrame of shape (num_samples, 13) containing the following signals: + - ``{id_column_name}``: The ID column for result alignment. + - ``signal_shadow_k_1``: The smallest Gower distance from the input data to any shadow model's data + (averaged across all shadow models). Lower is more similar. + - ``signal_shadow_k_{k}``: The mean of the ``k`` smallest Gower distances from the + input data to shadow models' data. + - ``signal_shadows_in_k_1``: Smallest distance from the input data for records known to be IN the + shadow models' training sets (a baseline for member behavior). + - ``signal_shadows_in_k_{k}``: Mean of ``k`` smallest distances from the input data for IN records. + - ``signal_shadows_out_k_1``: Smallest distance from the input data for records known to be OUT of + the shadow models' training sets (a baseline for non-member behavior). + - ``signal_shadows_out_k_{k}``: Mean of ``k`` smallest distances from the input data for OUT records. + - ``signal_target_k_1``: The smallest Gower distance from the input data to the target model's data. + - ``signal_target_k_{k}``: The mean of the ``k`` smallest Gower distances from the input data to the + target model's data. + - ``rmia_k_1``: The final RMIA score using the single nearest neighbor. A higher + value indicates a stronger membership signal. + - ``rmia_k_{k}``: The final RMIA score using ``k`` nearest neighbors. + - ``rmia_out_k_1``: RMIA score calibrated against the 'OUT' shadow signals, + providing a more robust measure of membership. + - ``rmia_out_k_{k}``: RMIA score using ``k`` neighbors, calibrated against the + 'OUT' shadow signals. + """ + # Extract shadow data collections. The first two elements are fine-tuned shadow models, + # and the third element is the trained shadow models. + + fine_tuned_shadow_data_0 = shadow_data_collection[0] + fine_tuned_shadow_data_1 = shadow_data_collection[1] + trained_shadow_data = shadow_data_collection[2] + + all_lengths = [ + [len(data.synthetic_data) for data in fine_tuned_shadow_data_0["fine_tuned_results"]], + [len(data.synthetic_data) for data in fine_tuned_shadow_data_1["fine_tuned_results"]], + [len(data.synthetic_data) for data in trained_shadow_data["trained_results"]], + [len(data) for data in fine_tuned_shadow_data_0["fine_tuning_sets"]], + [len(data) for data in fine_tuned_shadow_data_1["fine_tuning_sets"]], + [len(data) for data in trained_shadow_data["selected_sets"]], + ] + + # Validate lengths + if any(len(group) == 0 for group in all_lengths): + raise ValueError("shadow_data_collection/target_data contain empty sets; cannot compute RMIA.") + + min_length = min(min(group) for group in all_lengths) + if not (1 <= k <= min_length): + raise ValueError(f"k={k} must be within [1, {min_length}]") + + shadow_model_gower_0 = get_rmia_gower( + df_input=df_input, + model_data=fine_tuned_shadow_data_0, + min_length=min_length, + key=Key.FINE_TUNED_RESULTS, + categorical_column_names=categorical_column_names, + id_column_name=id_column_name, + random_seed=random_seed, + ) + + shadow_model_gower_1 = get_rmia_gower( + df_input=df_input, + model_data=fine_tuned_shadow_data_1, + min_length=min_length, + key=Key.FINE_TUNED_RESULTS, + categorical_column_names=categorical_column_names, + id_column_name=id_column_name, + random_seed=random_seed, + ) + + shadow_model_gower_2 = get_rmia_gower( + df_input=df_input, + model_data=trained_shadow_data, + min_length=min_length, + key=Key.TRAINED_RESULTS, + categorical_column_names=categorical_column_names, + id_column_name=id_column_name, + random_seed=random_seed, + ) + + gower_shadows = np.vstack( + [np.array(shadow_model_gower_0), np.array(shadow_model_gower_1), np.array(shadow_model_gower_2)] + ) + + # TODO: ideally remove hard-copied keys + shadow_training_data = ( + fine_tuned_shadow_data_0["fine_tuning_sets"] + + fine_tuned_shadow_data_1["fine_tuning_sets"] + + trained_shadow_data["selected_sets"] + ) + + # TODO: check key after we have the official target model + target_model_gower = get_rmia_gower( + df_input=df_input, + model_data=target_data, + min_length=min_length, + key=Key.TRAINED_RESULTS, + categorical_column_names=categorical_column_names, + id_column_name=id_column_name, + ) + + # gower_target is a 2D NumPy array that stores the Gower distances between + # the input data (df_input) and the synthetic data generated by the target model. + # so its shape is (len(df_input), len(target_synthetic_data)). + # gower_target[i, j] = Gower distance between i-th input record and j-th target synthetic record + + gower_target = target_model_gower[0] # There is only one target model + + # Compute the signals based on the gower matrix wrt to the target synthetic data. + # Here, for each input record, we sort the distances to all target synthetic records, + # then keep only the first k columns (the k smallest distances). + k_nearest_target_distances = np.sort(gower_target, axis=1)[:, :k] + signal_target_k_mean = np.mean(k_nearest_target_distances, axis=1) + signal_target_k_1 = gower_target[:, 0] # First element is the minimum + + # Process shadow model distances. Gower_shadows is a 3D matrix of shape: + # ((total number of shadow models), len(df_input), len(shadow_synthetic)) + sorted_shadow_gower = np.sort(gower_shadows, axis=2) + + # For k nearest neighbors + # Similar to target, we sort the distances in each shadow model's gower matrix, + # and keep only the first k columns (the k smallest distances). + k_nearest_shadow_distances = sorted_shadow_gower[:, :, :k] + signal_shadow_k_mean = np.mean(k_nearest_shadow_distances, axis=2) + + # Taking the mean across ALL shadow models (axis=0) to compute signal_shadows + signal_shadows = np.mean(signal_shadow_k_mean, axis=0) + + # For single nearest neighbor (k=1) + nearest_shadow_distances = sorted_shadow_gower[:, :, 0] + # Taking the mean across all shadow models (axis=0) to compute k=1 signal for each input sample + signal_shadows_k_1 = np.mean(nearest_shadow_distances, axis=0) + + # Create a dataframe for the computed scores + results_df = pd.DataFrame( + { + id_column_name: id_column_data.values, + "signal_shadow_k_1": signal_shadows_k_1, + f"signal_shadow_k_{k}": signal_shadows, + } + ) + + # Create masks for records in/out of training sets. We're creating masks for all the samples in train_df, + # as opposed to the original implementation which only creates masks a sample of 200 records. We've also + # changed the way the masks are created to improve efficiency. + + shadow_training_id_data = [set(id_list) for id_list in shadow_training_data] + + mask_in_training = np.array( + [results_df[id_column_name].isin(id_set).to_numpy() for id_set in shadow_training_id_data] + ) + + mask_not_in_training = ~mask_in_training + + # Calculate signals based on membership status + results_df["signal_shadows_in_k_1"] = conditional_average(nearest_shadow_distances, mask_in_training) + results_df[f"signal_shadows_in_k_{k}"] = conditional_average(signal_shadow_k_mean, mask_in_training) + results_df["signal_shadows_out_k_1"] = conditional_average(nearest_shadow_distances, mask_not_in_training) + results_df[f"signal_shadows_out_k_{k}"] = conditional_average(signal_shadow_k_mean, mask_not_in_training) + + # Add target signals to results + results_df["signal_target_k_1"] = signal_target_k_1 + results_df[f"signal_target_k_{k}"] = signal_target_k_mean + + # Calculate RMIA scores (ratios of target to shadow signals) + results_df["rmia_k_1"] = np.divide( + results_df["signal_target_k_1"], + results_df["signal_shadow_k_1"], + out=np.full_like(results_df["signal_target_k_1"], np.nan, dtype=float), + where=results_df["signal_shadow_k_1"] > 0, + ) + results_df[f"rmia_k_{k}"] = np.divide( + results_df[f"signal_target_k_{k}"], + results_df[f"signal_shadow_k_{k}"], + out=np.full_like(results_df[f"signal_target_k_{k}"], np.nan, dtype=float), + where=results_df[f"signal_shadow_k_{k}"] > 0, + ) + results_df["rmia_out_k_1"] = np.divide( + results_df["signal_target_k_1"], + results_df["signal_shadows_out_k_1"], + out=np.full_like(results_df["signal_target_k_1"], np.nan, dtype=float), + where=results_df["signal_shadows_out_k_1"] > 0, + ) + results_df[f"rmia_out_k_{k}"] = np.divide( + results_df[f"signal_target_k_{k}"], + results_df[f"signal_shadows_out_k_{k}"], + out=np.full_like(results_df[f"signal_target_k_{k}"], np.nan, dtype=float), + where=results_df[f"signal_shadows_out_k_{k}"] > 0, + ) + + return results_df diff --git a/src/midst_toolkit/attacks/ensemble/rmia/shadow_model_training.py b/src/midst_toolkit/attacks/ensemble/rmia/shadow_model_training.py index 6fe66e80..ea2e1e8a 100644 --- a/src/midst_toolkit/attacks/ensemble/rmia/shadow_model_training.py +++ b/src/midst_toolkit/attacks/ensemble/rmia/shadow_model_training.py @@ -309,7 +309,7 @@ def train_three_sets_of_shadow_models( in each set. Each observation in the challenge points is repeated ``n_reps`` times in the training set of each shadow model. - This attack by default trains 16 shadow models in total, 8 of which are fine-tuned from a pre-traine model and + This attack by default trains 16 shadow models in total, 8 of which are fine-tuned from a pre-trained model and 8 of which are trained from scratch on the challenge points. Pre-training is done on a number of samples from the population data as stated in the ``fine_tuning_config`` or `60K` samples by default. We keep track of and save the challenge id's used to train each shadow model to be used by RMIA. diff --git a/tests/unit/attacks/ensemble/test_meta_classifier.py b/tests/unit/attacks/ensemble/test_meta_classifier.py index ab60c313..ecc16c1a 100644 --- a/tests/unit/attacks/ensemble/test_meta_classifier.py +++ b/tests/unit/attacks/ensemble/test_meta_classifier.py @@ -10,7 +10,16 @@ from midst_toolkit.attacks.ensemble.blending import BlendingPlusPlus, MetaClassifierType -MOCK_COLUMN_TYPES_CONTENT = {"numerical": ["numerical_col1", "numerical_col2"], "categorical": ["cat_col1"]} +MOCK_COLUMN_TYPES_CONTENT = { + "numerical": ["numerical_col1", "numerical_col2"], + "categorical": ["cat_col1"], + "id_column_name": "id_col", +} + +MOCK_TARGET_DATA = { + "selected_sets": [pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})], + "trained_results": [{"model_info": "mock_model", "synthetic_data": [5, 6]}], +} @pytest.fixture(scope="module") @@ -21,14 +30,25 @@ def cfg() -> DictConfig: @pytest.fixture def mock_config_with_json_path(): - """Provides a mock DictConfig object with the structure required by BlendingPlusPlus.__init__.""" - return DictConfig({"metaclassifier": {"data_types_file_path": "/mock/path/to/data_types.json"}}) + """Provides a mock DictConfig object with the structure required by BlendingPlusPlus.""" + return DictConfig( + { + "metaclassifier": { + "data_types_file_path": "/mock/path/to/data_types.json", + "num_optuna_trials": 100, + "num_kfolds": 5, + "epochs": 1, + } + } + ) @pytest.fixture def sample_dataframes(): + """Provides sample dataframes, now including an ID column.""" df = pd.DataFrame( { + "id_col": [10, 20, 30, 40], "cat_col1": ["A", "B", "A", "C"], "numerical_col1": [1.0, 2.0, 3.0, 4.0], "numerical_col2": [0.1, 0.2, 0.3, 0.4], @@ -37,6 +57,7 @@ def sample_dataframes(): df_synth = pd.DataFrame( { + "id_col": [11, 22, 33, 44], "cat_col1": ["A", "B", "C", "C"], "numerical_col1": [1.5, 2.5, 3.5, 4.5], "numerical_col2": [0.15, 0.25, 0.35, 0.45], @@ -45,6 +66,7 @@ def sample_dataframes(): df_ref = pd.DataFrame( { + "id_col": [1, 2, 3, 4, 5, 6], "cat_col1": ["A", "B", "C", "A", "B", "C"], "numerical_col1": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], "numerical_col2": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6], @@ -70,7 +92,12 @@ def test_init_success(self, mock_file, mock_config_with_json_path): json_content_str = json.dumps(MOCK_COLUMN_TYPES_CONTENT) mock_file.return_value.read.return_value = json_content_str - bpp_xgb = BlendingPlusPlus(config=mock_config_with_json_path, meta_classifier_type=MetaClassifierType("xgb")) + bpp_xgb = BlendingPlusPlus( + config=mock_config_with_json_path, + shadow_data_collection=[], + target_data=MOCK_TARGET_DATA, + meta_classifier_type=MetaClassifierType("xgb"), + ) file_path = mock_config_with_json_path.metaclassifier.data_types_file_path mock_file.assert_called_once_with(file_path, "r") @@ -81,7 +108,12 @@ def test_init_success(self, mock_file, mock_config_with_json_path): mock_file.reset_mock() - bpp_lr = BlendingPlusPlus(config=mock_config_with_json_path, meta_classifier_type=MetaClassifierType("lr")) + bpp_lr = BlendingPlusPlus( + config=mock_config_with_json_path, + shadow_data_collection=[], + target_data=MOCK_TARGET_DATA, + meta_classifier_type=MetaClassifierType("lr"), + ) assert bpp_lr.meta_classifier_type == MetaClassifierType.LR assert bpp_lr.trained_model is None @@ -90,19 +122,23 @@ def test_init_success(self, mock_file, mock_config_with_json_path): @patch("builtins.open", new_callable=mock_open) def test_init_invalid_type_raises_error(self, mock_file, mock_config_with_json_path): """Tests that initialization with an invalid type raises a ValueError.""" - # Configure the mock file json_content_str = json.dumps(MOCK_COLUMN_TYPES_CONTENT) mock_file.return_value.read.return_value = json_content_str with pytest.raises(ValueError): - BlendingPlusPlus(config=mock_config_with_json_path, meta_classifier_type=MetaClassifierType("svm")) + BlendingPlusPlus( + config=mock_config_with_json_path, + shadow_data_collection=[], + target_data=MOCK_TARGET_DATA, + meta_classifier_type=MetaClassifierType("svm"), + ) @patch("builtins.open", new_callable=mock_open) @patch("midst_toolkit.attacks.ensemble.blending.calculate_gower_features") @patch("midst_toolkit.attacks.ensemble.blending.calculate_domias_score") - @patch("pandas.read_csv") + @patch("midst_toolkit.attacks.ensemble.blending.calculate_rmia_signals") def test_prepare_meta_features( - self, mock_read_csv, mock_domias, mock_gower, mock_file, mock_config_with_json_path, sample_dataframes + self, mock_rmia, mock_domias, mock_gower, mock_file, mock_config_with_json_path, sample_dataframes ): """Tests the _prepare_meta_features method for correct feature assembly.""" json_content_str = json.dumps(MOCK_COLUMN_TYPES_CONTENT) @@ -110,24 +146,32 @@ def test_prepare_meta_features( mock_gower.return_value = pd.DataFrame({"gower_1": [0.1] * 4, "gower_2": [0.2] * 4}) mock_domias.return_value = pd.DataFrame({"domias": [0.9, 0.8, 0.7, 0.6]}) - mock_read_csv.return_value = pd.DataFrame({"rmia": [1, 0, 1, 0]}) + mock_rmia.return_value = pd.DataFrame({"rmia": [1, 0, 1, 0]}) - bpp = BlendingPlusPlus(config=mock_config_with_json_path) + bpp = BlendingPlusPlus( + config=mock_config_with_json_path, + shadow_data_collection=[], + target_data=MOCK_TARGET_DATA, + ) categorical_cols = MOCK_COLUMN_TYPES_CONTENT["categorical"] numerical_cols = MOCK_COLUMN_TYPES_CONTENT["numerical"] + id_col_name = MOCK_COLUMN_TYPES_CONTENT["id_column_name"] + id_col_data = sample_dataframes["df_train"][id_col_name] meta_features = bpp._prepare_meta_features( df_input=sample_dataframes["df_train"], df_synthetic=sample_dataframes["df_synth"], df_reference=sample_dataframes["df_ref"], + id_column_data=id_col_data, categorical_cols=categorical_cols, numerical_cols=numerical_cols, + id_column_name=id_col_name, ) mock_gower.assert_called_once() mock_domias.assert_called_once() - mock_read_csv.assert_called_once() + mock_rmia.assert_called_once() expected_columns = ["numerical_col1", "numerical_col2", "gower_1", "gower_2", "domias", "rmia"] assert meta_features.shape == (4, 6) @@ -136,6 +180,54 @@ def test_prepare_meta_features( meta_features["numerical_col1"], sample_dataframes["df_train"]["numerical_col1"], check_names=False ) + @patch("builtins.open", new_callable=mock_open) + @patch("midst_toolkit.attacks.ensemble.blending.calculate_gower_features") + @patch("midst_toolkit.attacks.ensemble.blending.calculate_domias_score") + @patch("midst_toolkit.attacks.ensemble.blending.calculate_rmia_signals") + def test_prepare_meta_features_rmia_calculation( + self, mock_rmia, mock_domias, mock_gower, mock_file, mock_config_with_json_path, sample_dataframes + ): + """Tests that calculate_rmia_signals is called with the correct arguments.""" + json_content_str = json.dumps(MOCK_COLUMN_TYPES_CONTENT) + mock_file.return_value.read.return_value = json_content_str + + # Mock the return values of other feature calculators + mock_gower.return_value = pd.DataFrame({"gower_1": [0.1] * 4}) + mock_domias.return_value = pd.DataFrame({"domias": [0.9] * 4}) + mock_rmia.return_value = pd.DataFrame({"rmia": [1] * 4}) + + attack_collection = [{"name": "attack_model_1"}] + + bpp = BlendingPlusPlus( + config=mock_config_with_json_path, + shadow_data_collection=attack_collection, + target_data=MOCK_TARGET_DATA, + ) + + df_train = sample_dataframes["df_train"] + id_col_name = MOCK_COLUMN_TYPES_CONTENT["id_column_name"] + id_col_data = df_train[id_col_name] + + bpp._prepare_meta_features( + df_input=df_train, + df_synthetic=sample_dataframes["df_synth"], + df_reference=sample_dataframes["df_ref"], + id_column_data=id_col_data, + categorical_cols=MOCK_COLUMN_TYPES_CONTENT["categorical"], + numerical_cols=MOCK_COLUMN_TYPES_CONTENT["numerical"], + id_column_name=id_col_name, + ) + + mock_rmia.assert_called_once() + _, call_kwargs = mock_rmia.call_args + + # Verify the arguments + pd.testing.assert_frame_equal(call_kwargs["df_input"], df_train) + assert call_kwargs["shadow_data_collection"] == attack_collection + assert call_kwargs["categorical_column_names"] == MOCK_COLUMN_TYPES_CONTENT["categorical"] + assert call_kwargs["id_column_name"] == id_col_name + pd.testing.assert_series_equal(call_kwargs["id_column_data"], id_col_data) + @patch("builtins.open", new_callable=mock_open) @patch("midst_toolkit.attacks.ensemble.blending.BlendingPlusPlus._prepare_meta_features") @patch("midst_toolkit.attacks.ensemble.blending.LogisticRegression") @@ -150,12 +242,18 @@ def test_fit_logistic_regression( mock_lr.return_value = mock_lr_instance mock_lr_instance.fit.return_value = mock_lr_instance - bpp = BlendingPlusPlus(config=mock_config_with_json_path, meta_classifier_type=MetaClassifierType("lr")) + bpp = BlendingPlusPlus( + config=mock_config_with_json_path, + shadow_data_collection=[], + target_data=MOCK_TARGET_DATA, + meta_classifier_type=MetaClassifierType("lr"), + ) bpp.fit( df_train=sample_dataframes["df_train"], y_train=sample_dataframes["y_train"], - df_synthetic=sample_dataframes["df_synth"], + df_target_synthetic=sample_dataframes["df_synth"], df_reference=sample_dataframes["df_ref"], + id_column_data=sample_dataframes["df_train"]["id_col"], ) mock_prepare_features.assert_called_once() @@ -178,31 +276,43 @@ def test_fit_xgboost( mock_tuner_instance.tune_hyperparameters.return_value = mock_fitted_xgb mock_tuner_class.return_value = mock_tuner_instance - bpp = BlendingPlusPlus(config=mock_config_with_json_path, meta_classifier_type=MetaClassifierType("xgb")) + bpp = BlendingPlusPlus( + config=mock_config_with_json_path, + shadow_data_collection=[], + target_data=MOCK_TARGET_DATA, + meta_classifier_type=MetaClassifierType("xgb"), + ) bpp.fit( df_train=sample_dataframes["df_train"], y_train=sample_dataframes["y_train"], - df_synthetic=sample_dataframes["df_synth"], + df_target_synthetic=sample_dataframes["df_synth"], df_reference=sample_dataframes["df_ref"], + id_column_data=sample_dataframes["df_train"]["id_col"], ) mock_prepare_features.assert_called_once() mock_tuner_class.assert_called_once() - mock_tuner_instance.tune_hyperparameters.assert_called_once_with(num_optuna_trials=100, num_kfolds=5) + # Assert that hyperparameters are taken from the config + mock_tuner_instance.tune_hyperparameters.assert_called_once_with( + num_optuna_trials=mock_config_with_json_path.metaclassifier.num_optuna_trials, + num_kfolds=mock_config_with_json_path.metaclassifier.num_kfolds, + ) assert bpp.trained_model is mock_fitted_xgb @patch("builtins.open", new_callable=mock_open) def test_predict_raises_error_if_not_fit(self, mock_file, mock_config_with_json_path, sample_dataframes): """Tests that calling .predict() before .fit() raises a RuntimeError.""" - # Configure the mock file for __init__ mock_file.return_value.read.return_value = json.dumps(MOCK_COLUMN_TYPES_CONTENT) - bpp = BlendingPlusPlus(config=mock_config_with_json_path) + bpp = BlendingPlusPlus( + config=mock_config_with_json_path, shadow_data_collection=[], target_data=MOCK_TARGET_DATA + ) with pytest.raises(AssertionError): bpp.predict( df_test=sample_dataframes["df_test"], - df_synthetic=sample_dataframes["df_synth"], + df_original_synthetic=sample_dataframes["df_synth"], df_reference=sample_dataframes["df_ref"], + id_column_data=sample_dataframes["df_test"]["id_col"], y_test=sample_dataframes["y_test"], ) @@ -220,14 +330,18 @@ def test_predict_flow( mock_classifier.predict_proba.return_value = np.array([[0.9, 0.1], [0.2, 0.8], [0.6, 0.4], [0.05, 0.95]]) mock_get_tpr.return_value = 0.99 - bpp = BlendingPlusPlus(config=mock_config_with_json_path) - + bpp = BlendingPlusPlus( + config=mock_config_with_json_path, + shadow_data_collection=[], + target_data=MOCK_TARGET_DATA, + ) bpp.trained_model = mock_classifier probabilities, score = bpp.predict( df_test=sample_dataframes["df_test"], - df_synthetic=sample_dataframes["df_synth"], + df_original_synthetic=sample_dataframes["df_synth"], df_reference=sample_dataframes["df_ref"], + id_column_data=sample_dataframes["df_test"]["id_col"], y_test=sample_dataframes["y_test"], ) diff --git a/tests/unit/attacks/ensemble/test_rmia.py b/tests/unit/attacks/ensemble/test_rmia.py new file mode 100644 index 00000000..b349fb84 --- /dev/null +++ b/tests/unit/attacks/ensemble/test_rmia.py @@ -0,0 +1,346 @@ +from collections import namedtuple +from typing import Any + +import numpy as np +import numpy.testing as npt +import pandas as pd +import pandas.testing as pdt +import pytest + +from midst_toolkit.attacks.ensemble.rmia.rmia_calculation import ( + Key, + calculate_rmia_signals, + conditional_average, + get_rmia_gower, +) + + +MockTrainingResult = namedtuple("TrainingResult", ["synthetic_data"]) + + +@pytest.fixture +def base_data() -> dict[str, Any]: + """Provides base data for testing.""" + df_input = pd.DataFrame( + { + "age": [30, 40, 50], + "city": ["A", "B", "C"], + "score": [100.0, 200.0, 300.0], + } + ) + + df_syn1 = pd.DataFrame( + { + "id": [101, 102], + "age": [30, 55], + "city": ["A", "C"], + "score": [100.0, 350.0], + } + ) + df_syn2 = pd.DataFrame( + { + "id": [201, 202, 203], + "age": [42, 45, 48], + "city": ["B", "B", "C"], + "score": [210.0, 220.0, 280.0], + } + ) + + model_data = { + "trained_results": [ + MockTrainingResult(synthetic_data=df_syn1), + MockTrainingResult(synthetic_data=df_syn2), + ], + "fine_tuned_results": [ + MockTrainingResult(synthetic_data=df_syn1), + ], + } + + return { + "df_input": df_input, + "model_data": model_data, + "categorical_column_names": ["city"], + "id_column_name": "id", + "random_seed": 42, + } + + +@pytest.fixture +def rmia_signal_data() -> dict[str, Any]: + """Provides complex mock data for the main calculate_rmia_signals function.""" + k = 2 + df_input = pd.DataFrame({"age": [30, 40, 50], "city": ["A", "B", "C"], "score": [100, 200, 300]}) + id_column_data = pd.Series([1, 2, 3], name="id") + + train_set_0 = pd.DataFrame({"id": [1, 101], "age": [30, 31], "city": ["A", "A"], "score": [100, 110]}) + train_set_1 = pd.DataFrame({"id": [2, 202], "age": [40, 41], "city": ["B", "B"], "score": [200, 210]}) + train_set_2 = pd.DataFrame({"id": [1, 303], "age": [30, 32], "city": ["A", "A"], "score": [100, 120]}) + + syn_data_5 = pd.DataFrame(np.random.rand(5, 3), columns=["age", "city", "score"]) + + shadow_data_collection = [ + { + "fine_tuning_sets": [train_set_0["id"].tolist()], + "fine_tuned_results": [MockTrainingResult(syn_data_5.copy())], + }, + { + "fine_tuning_sets": [train_set_1["id"].tolist()], + "fine_tuned_results": [MockTrainingResult(syn_data_5.copy())], + }, + { + "selected_sets": [train_set_2["id"].tolist()], + "trained_results": [MockTrainingResult(syn_data_5.copy())], + }, + ] + + target_data = { + "selected_sets": [pd.DataFrame(np.random.rand(5, 2))], + "trained_results": [MockTrainingResult(syn_data_5.copy())], + } + + return { + "df_input": df_input, + "id_column_data": id_column_data, + "shadow_data_collection": shadow_data_collection, + "target_data": target_data, + "categorical_column_names": ["city"], + "id_column_name": "id", + "k": k, + "random_seed": 42, + } + + +class TestConditionalAverage: + def test_conditional_average_basic(self): + """Tests standard column-wise conditional averaging.""" + values = np.array([[10, 20, 30], [2, 4, 6], [100, 200, 300]]) + mask = np.array([[True, False, True], [True, True, False], [False, True, True]]) + expected = np.array([6.0, 102.0, 165.0]) + result = conditional_average(values, mask) + npt.assert_allclose(result, expected) + + def test_conditional_average_all_false_column(self): + """Tests that a column with no True in mask results in NaN.""" + values = np.array([[10, 20], [2, 4]]) + mask = np.array([[True, False], [True, False]]) + expected = np.array([6.0, np.nan]) + result = conditional_average(values, mask) + npt.assert_allclose(result, expected, equal_nan=True) + + def test_conditional_average_all_false_mask(self): + """Tests that an all-False mask results in all NaNs.""" + values = np.array([[10, 20], [2, 4]]) + mask = np.array([[False, False], [False, False]]) + expected = np.array([np.nan, np.nan]) + result = conditional_average(values, mask) + npt.assert_allclose(result, expected, equal_nan=True) + + def test_conditional_average_shape_mismatch(self): + """Tests that mismatched shapes raise an AssertionError.""" + values = np.array([[1, 2], [3, 4]]) + mask = np.array([True, False]) + with pytest.raises(AssertionError, match="condition_mask must have the same shape as values"): + conditional_average(values, mask) + + +class TestGetRmiaGower: + def test_get_rmia_gower_basic_run(self, base_data, mocker): + """Tests a basic run without sampling, checking mock calls.""" + mock_gower_matrix = mocker.patch( + "midst_toolkit.attacks.ensemble.rmia.rmia_calculation.gower.gower_matrix", + side_effect=[ + np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), + np.array([[0.7, 0.8, 0.9], [0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]), + ], + ) + + min_length = 3 + results = get_rmia_gower( + df_input=base_data["df_input"], + model_data=base_data["model_data"], + min_length=min_length, + key=Key.TRAINED_RESULTS, + categorical_column_names=base_data["categorical_column_names"], + id_column_name=base_data["id_column_name"], + random_seed=base_data["random_seed"], + ) + + assert len(results) == 2 + npt.assert_array_equal(results[0], np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])) + npt.assert_array_equal( + results[1], + np.array([[0.7, 0.8, 0.9], [0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]), + ) + + assert mock_gower_matrix.call_count == 2 + + call_args_1 = mock_gower_matrix.call_args_list[0].kwargs + pdt.assert_frame_equal(call_args_1["data_x"], base_data["df_input"], check_dtype=False) + syn_data_1_dropped = base_data["model_data"]["trained_results"][0].synthetic_data.drop(columns=["id"]) + pdt.assert_frame_equal(call_args_1["data_y"], syn_data_1_dropped, check_dtype=False) + assert call_args_1["cat_features"] == [False, True, False] + + call_args_2 = mock_gower_matrix.call_args_list[1].kwargs + syn_data_2_dropped = base_data["model_data"]["trained_results"][1].synthetic_data.drop(columns=["id"]) + pdt.assert_frame_equal(call_args_2["data_y"], syn_data_2_dropped, check_dtype=False) + + def test_get_rmia_gower_with_sampling(self, base_data, mocker): + """Tests that sampling is triggered and random_state is used.""" + mock_gower_matrix = mocker.patch( + "midst_toolkit.attacks.ensemble.rmia.rmia_calculation.gower.gower_matrix", + return_value=np.array([[0.1], [0.2], [0.3]]), + ) + + original_syn_data = base_data["model_data"]["trained_results"][1].synthetic_data + + mock_sample = mocker.patch("pandas.DataFrame.sample", wraps=original_syn_data.sample) + + min_length = 2 + get_rmia_gower( + df_input=base_data["df_input"], + model_data=base_data["model_data"], + min_length=min_length, + key=Key.TRAINED_RESULTS, + categorical_column_names=base_data["categorical_column_names"], + id_column_name=base_data["id_column_name"], + random_seed=base_data["random_seed"], + ) + + assert mock_gower_matrix.call_count == 2 + mock_sample.assert_called_once_with(n=min_length, random_state=base_data["random_seed"]) + + call_args_2 = mock_gower_matrix.call_args_list[1].kwargs + expected_sampled_data = original_syn_data.sample(n=min_length, random_state=base_data["random_seed"]).drop( + columns=[base_data["id_column_name"]] + ) + pdt.assert_frame_equal(call_args_2["data_y"], expected_sampled_data, check_dtype=False) + + def test_get_rmia_gower_missing_categorical_column(self, base_data, mocker, caplog): + """Tests that a warning is logged for missing categorical columns.""" + mocker.patch( + "midst_toolkit.attacks.ensemble.rmia.rmia_calculation.gower.gower_matrix", return_value=np.array([[0.1]]) + ) + + missing_cat_cols = ["city", "non_existent_column"] + + with caplog.at_level("INFO"): + get_rmia_gower( + df_input=base_data["df_input"], + model_data=base_data["model_data"], + min_length=1, + key=Key.FINE_TUNED_RESULTS, + categorical_column_names=missing_cat_cols, + id_column_name=base_data["id_column_name"], + ) + + assert "Warning: The following categorical columns are missing" in caplog.text + assert "{'non_existent_column'}" in caplog.text + + +class TestCalculateRmiaSignals: + @pytest.fixture + def mock_dependencies(self, mocker, rmia_signal_data): + """Mocks dependencies for calculate_rmia_signals.""" + gower_shadow_0 = [np.array([[0.1, 0.2, 0.3, 0.4, 0.5], [0.6, 0.5, 0.4, 0.3, 0.2], [0.9, 0.9, 0.8, 0.7, 0.6]])] + gower_shadow_1 = [np.array([[0.2, 0.3, 0.4, 0.5, 0.6], [0.5, 0.4, 0.3, 0.2, 0.1], [0.8, 0.7, 0.6, 0.9, 1.0]])] + gower_shadow_2 = [np.array([[0.3, 0.4, 0.5, 0.6, 0.7], [0.4, 0.3, 0.2, 0.1, 0.0], [0.7, 0.6, 0.8, 0.9, 0.5]])] + gower_target = [ + np.array([[0.05, 0.1, 0.15, 0.2, 0.25], [0.01, 0.02, 0.03, 0.04, 0.05], [0.8, 0.82, 0.84, 0.86, 0.88]]) + ] + + mock_get_gower = mocker.patch( + "midst_toolkit.attacks.ensemble.rmia.rmia_calculation.get_rmia_gower", + side_effect=[gower_shadow_0, gower_shadow_1, gower_shadow_2, gower_target], + autospec=True, + ) + + return mock_get_gower, gower_shadow_0, gower_shadow_1, gower_shadow_2, gower_target + + def test_calculate_rmia_signals_main_logic(self, rmia_signal_data, mock_dependencies): + """Tests the main orchestration and calculation logic of the function.""" + k = rmia_signal_data["k"] + + result_df = calculate_rmia_signals(**rmia_signal_data) + + signal_target_k_2 = np.array([0.075, 0.015, 0.81]) + signal_target_k_1 = np.array([0.05, 0.01, 0.8]) + + signal_shadows_k_1 = np.array([0.2, 0.1, 0.56666667]) + signal_shadows_k_2 = np.array([0.25, 0.15, 0.61666667]) + + signal_shadows_in_k_1 = np.array([0.2, 0.1, np.nan]) + signal_shadows_in_k_2 = np.array([0.25, 0.15, np.nan]) + + signal_shadows_out_k_1 = np.array([0.2, 0.1, 0.56666667]) + signal_shadows_out_k_2 = np.array([0.25, 0.15, 0.61666667]) + + rmia_k_1 = np.array([0.25, 0.1, 1.4117647]) + rmia_k_2 = np.array([0.3, 0.1, 1.3135593]) + rmia_out_k_1 = np.array([0.25, 0.1, 1.4117647]) + rmia_out_k_2 = np.array([0.3, 0.1, 1.3135593]) + + expected_df = pd.DataFrame( + { + "id": [1, 2, 3], + "signal_shadow_k_1": signal_shadows_k_1, + f"signal_shadow_k_{k}": signal_shadows_k_2, + "signal_shadows_in_k_1": signal_shadows_in_k_1, + f"signal_shadows_in_k_{k}": signal_shadows_in_k_2, + "signal_shadows_out_k_1": signal_shadows_out_k_1, + f"signal_shadows_out_k_{k}": signal_shadows_out_k_2, + "signal_target_k_1": signal_target_k_1, + f"signal_target_k_{k}": signal_target_k_2, + "rmia_k_1": rmia_k_1, + f"rmia_k_{k}": rmia_k_2, + "rmia_out_k_1": rmia_out_k_1, + f"rmia_out_k_{k}": rmia_out_k_2, + } + ) + + pdt.assert_frame_equal(result_df, expected_df, check_dtype=False, atol=0.005) + + def test_calculate_rmia_signals_value_errors(self, rmia_signal_data): + """Tests that ValueErrors are raised for invalid k or empty data.""" + data_k0 = rmia_signal_data.copy() + data_k0["k"] = 0 + with pytest.raises(ValueError): + calculate_rmia_signals(**data_k0) + + data_empty = rmia_signal_data.copy() + data_empty["shadow_data_collection"][0]["fine_tuning_sets"] = [] + with pytest.raises(ValueError, match="contain empty sets"): + calculate_rmia_signals(**data_empty) + + def test_calculate_rmia_signals_division_by_zero(self, rmia_signal_data, mocker): + """Tests that division by zero in RMIA score results in NaN.""" + k = rmia_signal_data["k"] + + gower_zeros = [np.array([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])] + gower_target = [np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])] + + mocker.patch( + "midst_toolkit.attacks.ensemble.rmia.rmia_calculation.get_rmia_gower", + side_effect=[gower_zeros, gower_zeros, gower_zeros, gower_target], + autospec=True, + ) + + mocker.patch( + "midst_toolkit.attacks.ensemble.rmia.rmia_calculation.conditional_average", + return_value=np.array([0.0, 0.0, 0.0]), + autospec=True, + ) + + result_df = calculate_rmia_signals(**rmia_signal_data) + + assert (result_df["signal_shadow_k_1"] == 0.0).all() + assert (result_df[f"signal_shadow_k_{k}"] == 0.0).all() + assert (result_df["signal_shadows_out_k_1"] == 0.0).all() + assert (result_df[f"signal_shadows_out_k_{k}"] == 0.0).all() + + assert (result_df["signal_target_k_1"] > 0.0).all() + assert (result_df[f"signal_target_k_{k}"] > 0.0).all() + + assert result_df["rmia_k_1"].isna().all() + assert result_df[f"rmia_k_{k}"].isna().all() + assert result_df["rmia_out_k_1"].isna().all() + assert result_df[f"rmia_out_k_{k}"].isna().all()