diff --git a/.gitignore b/.gitignore index 6a17814..fc94012 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ .idea/ *.ini __pycache__/ +**/*.pyc .venv*/ data/ src/__pycache__/ @@ -8,3 +9,7 @@ weights/ wandb/ *.png *.pdf +*.out +interactive.sh +*.pth +./datasets/ \ No newline at end of file diff --git a/README.md b/README.md index a970a7a..f8f12f7 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,3 @@ -# Near Field Localization via AI-Aided Subspace Methods - -## Related Publications -- [1] [DCD-MUSIC: Deep-Learning-Aided Cascaded Differentiable MUSIC Algorithm for Near-Field Localization of Multiple Sources](https://ieeexplore.ieee.org/abstract/document/10888295) (ICASSP 2025) -- [2] [Near Field Localization via AI-Aided Subspace Methods (preprint)](http://arxiv.org/abs/2504.00599) - -## Introduction -This repository contain the implementation of the AI-aided subspace methods for near-field localizaiton: -- DCD-MUSIC -- NF-SSN - ## Getting Started ### Prerequisites - Python 3.10+ @@ -17,25 +6,10 @@ This repository contain the implementation of the AI-aided subspace methods for - For conda, `conda env create -f full_environment.yml` and active it by `conda activate ai_subspace_env` - For venv, `py -m venv ai_subspace_env` - Activate the virtual environment - - Go to [Pytorch official website](https://pytorch.org/) to install the correct version of Pytorch - Install the required packages by running `pip install -r requirements.txt` -- See example usage in 'main.py' -## Citation -If you find this work useful, please cite: +- For both cases, you'll need to install seperately a torch package (>=2.7.0) that matches your CUDA version, + visit [PyTorch's official website](https://pytorch.org/get-started/locally/) for more details. -```bibtex -@inproceedings{gast2025dcdmusic, - author = {Arad Gast, Luc Le Magoarou, Nir Shlezinger}, - title = {DCD-MUSIC: Deep-Learning-Aided Cascaded Differentiable MUSIC Algorithm for Near-Field Localization of Multiple Sources}, - booktitle = {ICASSP}, - year = {2025}, - publisher = {IEEE} -} +- See main.py for configuring parameters and running the simulation. -@article{gast2025aisubspacenear, - author = {Arad Gast, Luc Le Magoarou, Nir Shlezinger}, - title = {Near Field Localization via AI-Aided Subspace Methods}, - journal = {arXiv preprint arXiv:2504.00599}, - year = {2025} -} \ No newline at end of file diff --git a/archive/main.py b/archive/main.py new file mode 100644 index 0000000..bdae120 --- /dev/null +++ b/archive/main.py @@ -0,0 +1,290 @@ +""" +This script is used to run the simulation with the given parameters. The parameters can be set in the script or +by using the command line arguments. The script will run the simulation with the given parameters and save the +results to the results folder. The results will include the learning curves, RMSE results, and the accuracy results +of the evaluation. The results will be saved in the results folder in the project directory. + +The script can be run with the following command line arguments: + --snr: SNR value + --N: Number of antennas + --M: Number of sources + --field_type: Field type + --signal_nature: Signal nature + --model_type: Model type + --train: Train model + --train_criteria: Training criteria + --eval: Evaluate model + --eval_criteria: Evaluation criteria + --samples_size: Samples size + --train_test_ratio: Train test ratio + +""" +# Imports +import os +import warnings +import time +import matplotlib.pyplot as plt +from run_simulation import run_simulation +import argparse + +#TODO: add appropriate model parameters for diffmusic after being built. +#ORI: here we set the parameters manually, but we can also argparse them which allows command line +#execution. + +# Initialization +os.system("cls||clear") +plt.close("all") + +scenario_dict = { + # "SNR": [-10, -5, 0, 5, 10], + # "T": [10, 20, 30, 50, 70, 100], + # "eta": [0.0, 0.01, 0.02, 0.03, 0.04], + # "M": [2, 3, 4, 5, 6, 7], +} + +simulation_commands = { + "SAVE_TO_FILE": False, + "CREATE_DATA": False, + "SAVE_DATASET": True, + "LOAD_MODEL": False, + "TRAIN_MODEL": True, + "SAVE_MODEL": True, + "EVALUATE_MODE": True, + "PLOT_RESULTS": True, # if True, the learning curves will be plotted + "PLOT_LOSS_RESULTS": True, # if True, the RMSE results of evaluation will be plotted + "PLOT_ACC_RESULTS": True, # if True, the accuracy results of evaluation will be plotted + "SAVE_PLOTS": False, # if True, the plots will be saved to the results folder +} + +system_model_params = { + "N": 15, # number of antennas + "M": 2, # number of sources + "T": 100, # number of snapshots + "snr": 10, # if defined, values in scenario_dict will be ignored + "field_type": "Far", # Near, Far + "signal_type": "Narrowband", # Narrowband, broadband + "signal_nature": "non-coherent", # if defined, values in scenario_dict will be ignored + "eta": 0.0, # steering vector uniform error variance with respect to the wavelength. + "bias": 0, # steering vector bias error + "sv_noise_var": 0.0, # steering vector addative gaussian error noise variance + "doa_range": 60, # The range of the DOA values [-doa_range, doa_range] + "doa_resolution": .5, # The resolution of the DOA values in degrees + "max_range_ratio_to_limit": 0.5, # The ratio of the maximum range in respect to the Fraunhofer distance + "range_resolution": 1, # The resolution of the range values in meters + "wavelength": 1, # The carrier wavelength of the signal in meters +} +model_config = { + "model_type": "SubspaceNet", # SubspaceNet, DCD-MUSIC, DeepCNN, TransMUSIC, DR_MUSIC + "model_params": {} +} +if model_config.get("model_type") == "SubspaceNet": + model_config["model_params"]["diff_method"] = "music_2D" # esprit, music_1D, music_2D, beamformer + model_config["model_params"]["train_loss_type"] = "music_spectrum" # music_spectrum, rmspe, beamformerloss + model_config["model_params"]["tau"] = 8 + model_config["model_params"]["field_type"] = "Near" # Far, Near + model_config["model_params"]["regularization"] = None # aic, mdl, threshold, None + model_config["model_params"]["variant"] = "small" # big, small + model_config["model_params"]["norm_layer"] = True + model_config["model_params"]["batch_norm"] = False + +elif model_config.get("model_type") == "DCD-MUSIC": + model_config["model_params"]["tau"] = 8 + model_config["model_params"]["diff_method"] = ("esprit", "music_1D") # ("esprit", "music_1D") + model_config["model_params"]["train_loss_type"] = ("rmspe", "rmspe") # ("rmspe", "rmspe"), ("rmspe", + # "music_spectrum"), ("music_spectrum", "rmspe") + model_config["model_params"]["regularization"] = None # aic, mdl, threshold, None + model_config["model_params"]["variant"] = "small" # big, small + model_config["model_params"]["norm_layer"] = True + +elif model_config.get("model_type") == "DeepCNN": + model_config["model_params"]["grid_size"] = 361 + +training_params = { + "samples_size": 4096, + "train_test_ratio": 0.1, + "training_objective": "angle, range", # angle, range, source_estimation + "batch_size": 128, + "epochs": 50, + "optimizer": "Adam", # Adam, SGD + "scheduler": "ReduceLROnPlateau", # StepLR, ReduceLROnPlateau + "learning_rate": 0.001, + "weight_decay": 1e-9, + "step_size": 50, + "gamma": 0.5, + "true_doa_train": None, # if set, this doa will be set to all samples in the train dataset + "true_range_train": None, # if set, this range will be set to all samples in the train dataset + "true_doa_test": None, # if set, this doa will be set to all samples in the test dataset + "true_range_test": None, # if set, this range will be set to all samples in the train dataset + "use_wandb": False, + "simulation_name": None, +} +evaluation_params = { + "models": { + # "TransMUSIC": { + # "model_name": "TransMUSIC", + # }, + # "DCD-MUSIC": { + # "model_name": "DCD-MUSIC", + # "tau": 8, + # "diff_method": ("esprit", "music_1d"), + # "regularization": None, + # }, + # "DCD-MUSIC_V2": { + # "model_name": "DCD-MUSIC", + # "tau": 8, + # "diff_method": ("esprit", "music_1d"), + # "regularization": "aic", + # "variant": "big" + # }, + # "NFSubspaceNet": { + # "model_name": "SubspaceNet", + # "tau": 8, + # "diff_method": "music_2D", + # "train_loss_type": "music_spectrum", + # "field_type": "near", + # "regularization": None, + # }, + # "NFSubspaceNet_V2": { + # "model_name": "SubspaceNet", + # "tau": 8, + # "diff_method": "music_2D", + # "train_loss_type": "music_spectrum", + # "field_type": "near", + # "regularization": "aic", + # "variant": "big", + # }, + }, + "augmented_methods": [ + # ("SubspaceNet", "beamformer", {"tau": 8, "diff_method": "music_2D", "train_loss_type": "music_spectrum", "field_type": "near"}), + # ("SubspaceNet", "beamformer", {"tau": 8, "diff_method": "esprit", "train_loss_type": "rmspe", "field_type": "far"}), + # ("SubspaceNet", "esprit", {"tau": 8, "diff_method": "esprit", "train_loss_type": "rmspe", "field_type": "far"}), + ], + "subspace_methods": [ + # "CCRB", + "2D-MUSIC", + "Beamformer", + # "CS_Estimator", + # "ESPRIT", + # "1D-MUSIC", + # "Root-MUSIC", + # "TOPS", + ] +} + + + +def parse_arguments(): + parser = argparse.ArgumentParser(description="Run simulation with optional parameters.") + parser.add_argument('-s', "--snr" ,type=int, help='SNR value', default=None) + parser.add_argument('-n', "--number_of_sensors",type=int, help='Number of antennas', default=None) + parser.add_argument('-m', "--number_of_sources", help='Number of sources, could be int or tuple for random case', default=None) + parser.add_argument('-snap', "--number_of_snapshots", type=int, help='Number of snapshots', default=None) + parser.add_argument('-eta', "--sv_error_var", type=float, help='Steering vector uniform error variance', default=None) + parser.add_argument('-ft', '--field_type', type=str, help='Field type, far or near field.', default=None) + parser.add_argument('-sn', '--signal_nature', type=str, help='Signal nature; non-coherent or coherent', default=None) + parser.add_argument('-wav', '--wavelength', type=float, help='Wavelength of the signal in meters', default=None) + + parser.add_argument('-mt', '--model_type', type=str, help='Model type; SubspaceNet, DCD-MUSIC, DeepCNN, TransMUSIC, DR_MUSIC', default=None) + parser.add_argument('-reg', '--regularization', type=str, help='Regularization method for SubspaceNet of DCD', default=model_config["model_params"]["regularization"]) + parser.add_argument('-tau', '--tau', type=int, help='Tau value for SubspaceNet or DCD-MUSIC', default=model_config["model_params"].get("tau")) + parser.add_argument("-v", "--variant", type=str, help="Variant of the SubspaceNet model; big, small", default=model_config["model_params"].get("variant")) + + parser.add_argument('-ss', '--samples_size', type=int, help='Samples size', default=None) + parser.add_argument('-ttr', '--train_test_ratio', type=float, help='Train test ratio', default=None) + parser.add_argument('-to', '--training_objective', type=str, help='Training objective; angle, range or angle, range.', default=None) + parser.add_argument('-bs', '--batch_size', type=int, help='Batch size', default=None) + parser.add_argument('-ep', '--epochs', type=int, help='Number of epochs', default=None) + parser.add_argument('-op', '--optimizer', type=str, help='Optimizer; Adam, SGD', default=None) + parser.add_argument('-sch', '--scheduler', type=str, help='Scheduler; StepLR, ReduceLROnPlateau', default=None) + parser.add_argument('-lr', '--learning_rate', type=float, help='Learning rate', default=None) + parser.add_argument('-wd', '--weight_decay', type=float, help='Weight decay', default=None) + parser.add_argument('-step', '--step_size', type=int, help='Step size', default=None) + parser.add_argument('-g', '--gamma', type=float, help='Gamma', default=None) + parser.add_argument('-w', '--wandb', action="store_true", help='Use wandb', default=training_params["use_wandb"]) + + parser.add_argument('-t', '--train', action="store_true", help='Train model', default=simulation_commands["TRAIN_MODEL"]) + parser.add_argument('-no_t', "--no_train", action="store_false", help='Do not train model', dest='train') + parser.add_argument('-e', '--eval', action="store_true", help='Evaluate model', default=simulation_commands["EVALUATE_MODE"]) + + parser.add_argument('-c', '--create', action="store_true", help='create a new dataset', default=simulation_commands["CREATE_DATA"]) + parser.add_argument('-sv', '--save', action='store_true', help="save dataset", default=simulation_commands["SAVE_DATASET"]) + + return parser.parse_args() + + +if __name__ == "__main__": + # torch.set_printoptions(precision=12) + + args = parse_arguments() + if args.snr is not None: + system_model_params["snr"] = args.snr + if args.number_of_sensors is not None: + system_model_params["N"] = args.number_of_sensors + if args.number_of_sources is not None: + # catch a case of random number of sources between two possible values + str_m = args.number_of_sources + if str_m.isnumeric(): + system_model_params["M"] = int(str_m) + else: + system_model_params["M"] = tuple(map(int, str_m.split(','))) + if args.number_of_snapshots is not None: + system_model_params["T"] = args.number_of_snapshots + if args.sv_error_var is not None: + system_model_params["eta"] = args.sv_error_var + if args.field_type is not None: + system_model_params["field_type"] = args.field_type + if args.signal_nature is not None: + system_model_params["signal_nature"] = args.signal_nature + if args.wavelength is not None: + system_model_params["wavelength"] = args.wavelength + + if args.model_type is not None: + warnings.warn("Please make sure to configure the model parameters in the script.") + model_config["model_type"] = args.model_type + if model_config["model_type"] == "SubspaceNet": + model_config["model_params"]["regularization"] = None if args.regularization == "None" else args.regularization + model_config["model_params"]["tau"] = args.tau + model_config["model_params"]["variant"] = args.variant + + if args.samples_size is not None: + training_params["samples_size"] = args.samples_size + if args.train_test_ratio is not None: + training_params["train_test_ratio"] = args.train_test_ratio + if args.training_objective is not None: + if args.training_objective.startswith("angle,range"): + training_params["training_objective"] = "angle, range" + else: + training_params["training_objective"] = args.training_objective + if args.batch_size is not None: + training_params["batch_size"] = args.batch_size + if args.epochs is not None: + training_params["epochs"] = args.epochs + if args.optimizer is not None: + training_params["optimizer"] = args.optimizer + if args.scheduler is not None: + training_params["scheduler"] = args.scheduler + if args.learning_rate is not None: + training_params["learning_rate"] = args.learning_rate + if args.weight_decay is not None: + training_params["weight_decay"] = args.weight_decay + if args.step_size is not None: + training_params["step_size"] = args.step_size + if args.gamma is not None: + training_params["gamma"] = args.gamma + if args.wandb is not None: + training_params["use_wandb"] = args.wandb + + simulation_commands["TRAIN_MODEL"] = args.train + simulation_commands["EVALUATE_MODE"] = args.eval + + simulation_commands["CREATE_DATA"] = args.create + simulation_commands["SAVE_DATASET"] = args.save + + start = time.time() + loss = run_simulation(simulation_commands=simulation_commands, + system_model_params=system_model_params, + model_config=model_config, + training_params=training_params, + evaluation_params=evaluation_params, + scenario_dict=scenario_dict) + print("Total time: ", time.time() - start) diff --git a/archive/run_simulation.py b/archive/run_simulation.py new file mode 100644 index 0000000..35bc83e --- /dev/null +++ b/archive/run_simulation.py @@ -0,0 +1,286 @@ +""" +This script is used to run an end to end simulation, including creating or loading data, training or loading +a NN model and do an evaluation for the algorithms. +The type of the run is based on the scenrio_dict. the avialble scrorios are: +- SNR: a list of SNR values to be tested +- T: a list of number of snapshots to be tested +- eta: a list of steering vector error values to be tested +- M: a list of number of sources to be tested + + +""" +# Imports +import sys +from src.data_handler import * +from src.training import * +from src.plotting import * +from src.evaluation import evaluate +from pathlib import Path +from src.models import ModelGenerator +from src.system_model import SystemModel, SystemModelParams +from src.utils import set_unified_seed, initialize_data_paths, print_loss_results_from_simulation + + +def __run_simulation(**kwargs): + SIMULATION_COMMANDS = kwargs["simulation_commands"] + SYSTEM_MODEL_PARAMS = kwargs["system_model_params"] + MODEL_CONFIG = kwargs["model_config"] + TRAINING_PARAMS = kwargs["training_params"] + EVALUATION_PARAMS = kwargs["evaluation_params"] + save_to_file = SIMULATION_COMMANDS["SAVE_TO_FILE"] # Saving results to file or present them over CMD + create_data = SIMULATION_COMMANDS["CREATE_DATA"] # Creating new dataset + load_model = SIMULATION_COMMANDS["LOAD_MODEL"] # Load specific model for training + train_model = SIMULATION_COMMANDS["TRAIN_MODEL"] # Applying training operation + save_model = SIMULATION_COMMANDS["SAVE_MODEL"] # Saving tuned model + evaluate_mode = SIMULATION_COMMANDS["EVALUATE_MODE"] # Evaluating desired algorithms + plot_mode = SIMULATION_COMMANDS["PLOT_RESULTS"] # Plotting results + save_plots = SIMULATION_COMMANDS["SAVE_PLOTS"] # Saving plots + load_data = not create_data # Loading data from exist dataset + print("Running simulation...") + if train_model: + print("Training model - ", MODEL_CONFIG.get('model_type')) + print("Training objective - ", TRAINING_PARAMS.get('training_objective')) + + now = datetime.now() + plot_path = Path(__file__).parent / "plots" + plot_path.mkdir(parents=True, exist_ok=True) + dt_string_for_save = now.strftime("%d_%m_%Y_%H_%M") + # torch.set_printoptions(precision=12) + + # Initialize seed + set_unified_seed() + + # Initialize paths + datasets_path, simulations_path, saving_path = initialize_data_paths(Path(__file__).parent / "data") + + # Saving simulation scores to external file + suffix = "" + if train_model: + suffix += f"_train_{MODEL_CONFIG.get('model_type')}_{TRAINING_PARAMS.get('training_objective')}" + suffix += (f"_{SYSTEM_MODEL_PARAMS['signal_nature']}_SNR_{SYSTEM_MODEL_PARAMS['snr']}_T_{SYSTEM_MODEL_PARAMS['T']}" + f"_eta{SYSTEM_MODEL_PARAMS['eta']}.txt") + + if save_to_file: + orig_stdout = sys.stdout + file_path = ( + simulations_path / "results" / "scores" / Path(dt_string_for_save + suffix) + ) + sys.stdout = open(file_path, "w") + # Define system model parameters + system_model_params = ( + SystemModelParams() + .set_parameter("N", SYSTEM_MODEL_PARAMS["N"]) + .set_parameter("M", SYSTEM_MODEL_PARAMS["M"]) + .set_parameter("T", SYSTEM_MODEL_PARAMS["T"]) + .set_parameter("snr", SYSTEM_MODEL_PARAMS["snr"]) + .set_parameter("field_type", SYSTEM_MODEL_PARAMS["field_type"]) + .set_parameter("signal_nature", SYSTEM_MODEL_PARAMS["signal_nature"]) + .set_parameter("signal_type", SYSTEM_MODEL_PARAMS["signal_type"]) + .set_parameter("eta", SYSTEM_MODEL_PARAMS["eta"]) + .set_parameter("bias", SYSTEM_MODEL_PARAMS["bias"]) + .set_parameter("sv_noise_var", SYSTEM_MODEL_PARAMS["sv_noise_var"]) + .set_parameter("doa_range", SYSTEM_MODEL_PARAMS["doa_range"]) + .set_parameter("doa_resolution", SYSTEM_MODEL_PARAMS["doa_resolution"]) + .set_parameter("max_range_ratio_to_limit", SYSTEM_MODEL_PARAMS["max_range_ratio_to_limit"]) + .set_parameter("range_resolution", SYSTEM_MODEL_PARAMS["range_resolution"]) + .set_parameter("wavelength", SYSTEM_MODEL_PARAMS["wavelength"]) + ) + + # Define samples size + samples_size = TRAINING_PARAMS["samples_size"] # Overall dateset size + train_test_ratio = TRAINING_PARAMS["train_test_ratio"] # training and testing datasets ratio + # Sets simulation filename + # simulation_filename = get_simulation_filename(system_model_params=system_model_params, + # model_config=model_config) + # Print new simulation intro + print("------------------------------------") + print("---------- New Simulation ----------") + print("------------------------------------") + + if load_data: + if train_model: + try: + start = time.time() + train_dataset = load_datasets( + system_model_params=system_model_params, + samples_size=samples_size, + datasets_path=datasets_path, + is_training=True, + ) + print(f"Load the data took {time.time() - start} sec") + except Exception as e: + print(e) + print("#############################################") + print("load_datasets: Error loading train dataset") + print("#############################################") + create_data = True + load_data = False + if evaluate_mode: + try: + generic_test_dataset = load_datasets( + system_model_params=system_model_params, + samples_size=samples_size * train_test_ratio, + datasets_path=datasets_path, + is_training=False, + ) + except Exception as e: + print(e) + print("#############################################") + print("load_datasets: Error loading test dataset") + print("#############################################") + create_data = True + load_data = False + if create_data and not load_data: + # Define which datasets to generate + print("Creating Data...") + # init sample model + samples_model = Samples(system_model_params) + if train_model: + # Generate training dataset + start = time.time() + train_dataset, _ = create_dataset( + samples_model=samples_model, + samples_size=samples_size, + save_datasets=SIMULATION_COMMANDS["SAVE_DATASET"], + datasets_path=datasets_path, + true_doa=TRAINING_PARAMS["true_doa_train"], + true_range=TRAINING_PARAMS["true_range_train"], + phase="train", + ) + print(f"Create the data took {time.time() - start} sec") + if evaluate_mode: + # Generate test dataset + generic_test_dataset, _ = create_dataset( + samples_model=samples_model, + samples_size=int(train_test_ratio * samples_size), + save_datasets=SIMULATION_COMMANDS["SAVE_DATASET"], + datasets_path=datasets_path, + true_doa=TRAINING_PARAMS["true_doa_test"], + true_range=TRAINING_PARAMS["true_range_test"], + phase="test", + ) + + if train_model: + # Generate model configuration + model_config = ( + ModelGenerator() + .set_model_type(MODEL_CONFIG.get("model_type")) + .set_system_model(system_model_params) + .set_model_params(MODEL_CONFIG.get("model_params")) + .set_model() + ) + + trainingparams = TrainingParamsNew(learning_rate=TRAINING_PARAMS["learning_rate"], + weight_decay=TRAINING_PARAMS["weight_decay"], + epochs=TRAINING_PARAMS["epochs"], + optimizer=TRAINING_PARAMS["optimizer"], + step_size=TRAINING_PARAMS["step_size"], + gamma=TRAINING_PARAMS["gamma"], + training_objective=TRAINING_PARAMS["training_objective"], + scheduler=TRAINING_PARAMS["scheduler"], + batch_size=TRAINING_PARAMS["batch_size"], + simulation_name=TRAINING_PARAMS["simulation_name"], + ) + train_dataloader, valid_dataloader = train_dataset.get_dataloaders(batch_size=TRAINING_PARAMS["batch_size"]) + trainer = Trainer(model=model_config.model, training_params=trainingparams, show_plots=True) + model = trainer.train(train_dataloader, valid_dataloader, + use_wandb=TRAINING_PARAMS["use_wandb"], + save_final=save_model, load_model=load_model) + + # Evaluation stage + if evaluate_mode: + if not train_model: + model = None + # Define loss measure for evaluation + if isinstance(system_model_params.M, int): + generic_test_dataset = torch.utils.data.DataLoader(generic_test_dataset, + batch_size=100, + shuffle=False) + else: + batch_sampler_test = SameLengthBatchSampler(generic_test_dataset, batch_size=100) + generic_test_dataset = torch.utils.data.DataLoader(generic_test_dataset, + collate_fn=collate_fn, + batch_sampler=batch_sampler_test, + shuffle=False) + + # Evaluate DNN models, augmented and subspace methods + loss = evaluate( + generic_test_dataset=generic_test_dataset, + system_model_params=system_model_params, + models=EVALUATION_PARAMS["models"], + augmented_methods=EVALUATION_PARAMS["augmented_methods"], + subspace_methods=EVALUATION_PARAMS["subspace_methods"], + model_tmp=model + ) + # plt.show() + print("END OF EVALUATION") + if save_to_file: + sys.stdout.close() + sys.stdout = orig_stdout + return loss + return + + +def run_simulation(**kwargs): + """ + This function is used to run an end to end simulation, including creating or loading data, training or loading + a NN model and do an evaluation for the algorithms. + The type of the run is based on the scenrio_dict. the avialble scrorios are: + - SNR: a list of SNR values to be tested + - T: a list of number of snapshots to be tested + - eta: a list of steering vector error values to be tested + - M: a list of number of sources to be tested + """ + if kwargs["scenario_dict"] == {}: + loss = __run_simulation(**kwargs) + return loss + loss_dict = {} + default_snr = kwargs["system_model_params"]["snr"] + default_T = kwargs["system_model_params"]["T"] + default_eta = kwargs["system_model_params"]["eta"] + default_m = kwargs["system_model_params"]["M"] + for key, value in kwargs["scenario_dict"].items(): + if key == "SNR": + loss_dict["SNR"] = {snr: None for snr in value} + print(f"Testing SNR values: {value}") + for snr in value: + kwargs["system_model_params"]["snr"] = snr + loss = __run_simulation(**kwargs) + loss_dict["SNR"][snr] = loss + kwargs["system_model_params"]["snr"] = default_snr + if key == "T": + loss_dict["T"] = {T: None for T in value} + print(f"Testing T values: {value}") + for T in value: + kwargs["system_model_params"]["T"] = T + loss = __run_simulation(**kwargs) + loss_dict["T"][T] = loss + kwargs["system_model_params"]["T"] = default_T + if key == "eta": + loss_dict["eta"] = {eta: None for eta in value} + print(f"Testing eta values: {value}") + for eta in value: + kwargs["system_model_params"]["eta"] = eta + loss = __run_simulation(**kwargs) + loss_dict["eta"][eta] = loss + kwargs["system_model_params"]["eta"] = default_eta + if key == "M": + loss_dict["M"] = {m: None for m in value} + print(f"Testing M values: {value}") + for m in value: + kwargs["system_model_params"]["M"] = m + loss = __run_simulation(**kwargs) + loss_dict["M"][m] = loss + kwargs["system_model_params"]["M"] = default_m + if None not in list(next(iter(loss_dict.values())).values()): + print_loss_results_from_simulation(loss_dict) + if kwargs["simulation_commands"]["PLOT_LOSS_RESULTS"]: + plot_results(loss_dict, kwargs["system_model_params"]["field_type"], + plot_acc=kwargs["simulation_commands"]["PLOT_ACC_RESULTS"], + save_to_file=kwargs["simulation_commands"]["SAVE_PLOTS"]) + + return loss_dict + + +if __name__ == "__main__": + now = datetime.now() diff --git a/src/__init__.py b/archive/src/__init__.py similarity index 100% rename from src/__init__.py rename to archive/src/__init__.py diff --git a/src/config.py b/archive/src/config.py similarity index 100% rename from src/config.py rename to archive/src/config.py diff --git a/src/data_handler.py b/archive/src/data_handler.py similarity index 100% rename from src/data_handler.py rename to archive/src/data_handler.py diff --git a/src/evaluation.py b/archive/src/evaluation.py similarity index 100% rename from src/evaluation.py rename to archive/src/evaluation.py diff --git a/src/methods_pack/.gitignore b/archive/src/methods_pack/.gitignore similarity index 100% rename from src/methods_pack/.gitignore rename to archive/src/methods_pack/.gitignore diff --git a/src/methods_pack/__init__.py b/archive/src/methods_pack/__init__.py similarity index 100% rename from src/methods_pack/__init__.py rename to archive/src/methods_pack/__init__.py diff --git a/src/methods_pack/beamformer.py b/archive/src/methods_pack/beamformer.py similarity index 100% rename from src/methods_pack/beamformer.py rename to archive/src/methods_pack/beamformer.py diff --git a/src/methods_pack/csestimator.py b/archive/src/methods_pack/csestimator.py similarity index 100% rename from src/methods_pack/csestimator.py rename to archive/src/methods_pack/csestimator.py diff --git a/src/methods_pack/esprit.py b/archive/src/methods_pack/esprit.py similarity index 100% rename from src/methods_pack/esprit.py rename to archive/src/methods_pack/esprit.py diff --git a/src/methods_pack/mle.py b/archive/src/methods_pack/mle.py similarity index 100% rename from src/methods_pack/mle.py rename to archive/src/methods_pack/mle.py diff --git a/src/methods_pack/music.py b/archive/src/methods_pack/music.py similarity index 100% rename from src/methods_pack/music.py rename to archive/src/methods_pack/music.py diff --git a/src/methods_pack/root_music.py b/archive/src/methods_pack/root_music.py similarity index 100% rename from src/methods_pack/root_music.py rename to archive/src/methods_pack/root_music.py diff --git a/src/methods_pack/subspace_method.py b/archive/src/methods_pack/subspace_method.py similarity index 99% rename from src/methods_pack/subspace_method.py rename to archive/src/methods_pack/subspace_method.py index dc39054..058ab35 100644 --- a/src/methods_pack/subspace_method.py +++ b/archive/src/methods_pack/subspace_method.py @@ -11,10 +11,9 @@ from src.config import device - class SubspaceMethod(nn.Module): """ - + Basic methods for all subspace methods. """ def __init__(self, system_model: SystemModel, model_order_estimation: str = None): diff --git a/src/metrics/__init__.py b/archive/src/metrics/__init__.py similarity index 100% rename from src/metrics/__init__.py rename to archive/src/metrics/__init__.py diff --git a/src/metrics/beamformer_loss.py b/archive/src/metrics/beamformer_loss.py similarity index 100% rename from src/metrics/beamformer_loss.py rename to archive/src/metrics/beamformer_loss.py diff --git a/src/metrics/cartesian_loss.py b/archive/src/metrics/cartesian_loss.py similarity index 100% rename from src/metrics/cartesian_loss.py rename to archive/src/metrics/cartesian_loss.py diff --git a/src/metrics/music_spec_loss.py b/archive/src/metrics/music_spec_loss.py similarity index 100% rename from src/metrics/music_spec_loss.py rename to archive/src/metrics/music_spec_loss.py diff --git a/src/metrics/rmse_loss.py b/archive/src/metrics/rmse_loss.py similarity index 100% rename from src/metrics/rmse_loss.py rename to archive/src/metrics/rmse_loss.py diff --git a/src/metrics/rmspe_loss.py b/archive/src/metrics/rmspe_loss.py similarity index 100% rename from src/metrics/rmspe_loss.py rename to archive/src/metrics/rmspe_loss.py diff --git a/src/models.py b/archive/src/models.py similarity index 100% rename from src/models.py rename to archive/src/models.py diff --git a/src/models_pack/.gitignore b/archive/src/models_pack/.gitignore similarity index 100% rename from src/models_pack/.gitignore rename to archive/src/models_pack/.gitignore diff --git a/src/models_pack/__init__.py b/archive/src/models_pack/__init__.py similarity index 100% rename from src/models_pack/__init__.py rename to archive/src/models_pack/__init__.py diff --git a/src/models_pack/dcd_music.py b/archive/src/models_pack/dcd_music.py similarity index 100% rename from src/models_pack/dcd_music.py rename to archive/src/models_pack/dcd_music.py diff --git a/src/models_pack/deep_augmented_music.py b/archive/src/models_pack/deep_augmented_music.py similarity index 100% rename from src/models_pack/deep_augmented_music.py rename to archive/src/models_pack/deep_augmented_music.py diff --git a/src/models_pack/deep_cnn.py b/archive/src/models_pack/deep_cnn.py similarity index 100% rename from src/models_pack/deep_cnn.py rename to archive/src/models_pack/deep_cnn.py diff --git a/src/models_pack/deep_root_music.py b/archive/src/models_pack/deep_root_music.py similarity index 100% rename from src/models_pack/deep_root_music.py rename to archive/src/models_pack/deep_root_music.py diff --git a/src/models_pack/parent_model.py b/archive/src/models_pack/parent_model.py similarity index 100% rename from src/models_pack/parent_model.py rename to archive/src/models_pack/parent_model.py diff --git a/src/models_pack/subspacenet.py b/archive/src/models_pack/subspacenet.py similarity index 100% rename from src/models_pack/subspacenet.py rename to archive/src/models_pack/subspacenet.py diff --git a/src/models_pack/trans_music.py b/archive/src/models_pack/trans_music.py similarity index 100% rename from src/models_pack/trans_music.py rename to archive/src/models_pack/trans_music.py diff --git a/src/plotting.py b/archive/src/plotting.py similarity index 100% rename from src/plotting.py rename to archive/src/plotting.py diff --git a/archive/src/signal_creation.py b/archive/src/signal_creation.py new file mode 100644 index 0000000..381fefe --- /dev/null +++ b/archive/src/signal_creation.py @@ -0,0 +1,315 @@ +"""Subspace-Net +Details +---------- +Name: signal_creation.py +Authors: D. H. Shmuel +Created: 01/10/21 +Edited: 02/06/23 + +Purpose: +-------- +This script defines the Samples class, which inherits from SystemModel class. +This class is used for defining the samples model. +""" + +# Imports +from random import sample +from src.system_model import SystemModel, SystemModelParams +# from src.utils import * +import numpy as np +import torch + +class Samples(SystemModel): + """ + Class used for defining and creating signals and observations. + Inherits from SystemModel class. + + ... + + Attributes: + ----------- + doa (np.ndarray): Array of angels (directions) of arrival. + + Methods: + -------- + set_doa(doa): Sets the direction of arrival (DOA) for the signals. + samples_creation(noise_mean: float = 0, noise_variance: float = 1, signal_mean: float = 0, + signal_variance: float = 1): Creates samples based on the specified mode and parameters. + noise_creation(noise_mean, noise_variance): Creates noise based on the specified mean and variance. + signal_creation(signal_mean=0, signal_variance=1, SNR=10): Creates signals based on the specified mode and parameters. + """ + + def __init__(self, system_model_params: SystemModelParams): + """Initializes a Samples object. + + Args: + ----- + system_model_params (SystemModelParams): an instance of SystemModelParams, + containing all relevant system model parameters. + + """ + super().__init__(system_model_params) + self.angles = None + self.distances = None + + def set_labels(self, number_of_sources: int, angles: list, distances: list): + if self.params.field_type.lower() == "far": + self.set_angles(angles, number_of_sources) + elif self.params.field_type.lower() in {"near", "full"}: + self.set_angles(angles, number_of_sources) + self.set_distances(distances, number_of_sources) + else: + raise ValueError(f"Samples.set_labels: Field type {self.params.field_type} is not defined") + + def get_labels(self): + if self.params.field_type.lower() == "far": + return torch.tensor(self.angles, dtype=torch.float32) + elif self.params.field_type.lower() in {"near", "full"}: + labels = torch.cat((torch.tensor(self.angles, dtype=torch.float32), torch.tensor(self.distances, dtype=torch.float32)), dim=0) + return labels + else: + raise ValueError(f"Samples.get_labels: Field type {self.params.field_type} is not defined") + + + def set_angles(self, doa: list, M: int): + """ + Sets the direction of arrival (DOA) for the signals. + + Args: + ----- + doa (np.ndarray): Array containing the DOA values. + + """ + + def create_doa_with_gap(gap: float, M: int): + """Create angles with a value gap. + + Args: + ----- + gap (float): Minimal gap value. + + Returns: + -------- + np.ndarray: DOA array. + + """ + # LEGACY CODE + # while True: + # # DOA = np.round(np.random.rand(M) * 180, decimals=2) - 90 + # DOA = np.random.randint(-55, 55, M) + # DOA.sort() + # diff_angles = np.array( + # [np.abs(DOA[i + 1] - DOA[i]) for i in range(M - 1)] + # ) + # if (np.sum(diff_angles > gap) == M - 1) and ( + # np.sum(diff_angles < (180 - gap)) == M - 1 + # ): + # break + + # based on https://stackoverflow.com/questions/51918580/python-random-list-of-numbers-in-a-range-keeping-with-a-minimum-distance + doa_range = self.params.doa_range + doa_resolution = self.params.doa_resolution + if doa_resolution <= 0: + raise ValueError("DOA resolution must be positive.") + if M <= 0: + raise ValueError("M (number of elements) must be positive.") + if gap <= 0: + raise ValueError("Gap must be positive.") + + # Compute the range of possible DOA values + # Ensure the sampled DOAs do not exceed [-doa_range, +doa_range] + max_offset = (gap - 1) * (M - 1) + effective_range = 2 * doa_range - max_offset + if effective_range <= 0: + raise ValueError(f"Invalid effective range: {effective_range}. Check your parameters.") + + # Define the valid range for sampling + if doa_resolution >= 1: + valid_range = range(0, effective_range, doa_resolution) + sampled_values = sorted(sample(valid_range, M)) + else: + step_count = int(effective_range // doa_resolution) + valid_range = range(step_count) + sampled_values = sorted(sample(valid_range, M)) + sampled_values = [x * doa_resolution for x in sampled_values] + + # Compute DOAs + DOA = [(gap - 1) * i + x - doa_range for i, x in enumerate(sampled_values)] + + # Ensure all DOAs fall naturally within the valid range + if any(d < -doa_range or d > doa_range for d in DOA): + raise ValueError("Computed DOAs exceed the valid range. Check your logic.") + + # Round results to 3 decimal places + DOA = np.round(DOA, 3) + + return DOA + + if doa == None: + # Generate angels with gap greater than 0.2 rad (nominal case) + self.angles = np.deg2rad(np.array(create_doa_with_gap(gap=10, M=M))) + else: + # Generate + self.angles = np.deg2rad(doa) + + def set_distances(self, distance: list | np.ndarray, M: int) -> np.ndarray: + """ + + Args: + distance: + + Returns: + + """ + + def choose_distances(M, min_val: float, max_val: int, distance_resolution: float = 1.0) -> np.ndarray: + """ + Choose distances for the sources. + + Args: + M (int): Number of sources. + min_val (float): Minimal value of the distances. + max_val (int): Maximal value of the distances. + distance_resolution (float, optional): Resolution of the distances. Defaults to 1.0. + + """ + distances_options = np.arange(min_val, max_val, distance_resolution) + distances = np.random.choice(distances_options, M, replace=True) + return np.round(distances, 3) + + if distance is None: + self.distances = choose_distances(M, min_val=np.ceil(self.fresnel) + self.params.range_resolution, + max_val=np.floor(self.fraunhofer * self.params.max_range_ratio_to_limit), + distance_resolution=self.params.range_resolution) + else: + self.distances = np.array(distance) + + def samples_creation( + self, + noise_mean: float = 0, + noise_variance: float = 1, + signal_mean: float = 0, + signal_variance: float = 1, + source_number: int = None, + ): + """Creates samples based on the specified mode and parameters. + + Args: + ----- + noise_mean (float, optional): Mean of the noise. Defaults to 0. + noise_variance (float, optional): Variance of the noise. Defaults to 1. + signal_mean (float, optional): Mean of the signal. Defaults to 0. + signal_variance (float, optional): Variance of the signal. Defaults to 1. + + Returns: + -------- + tuple: Tuple containing the created samples, signal, steering vectors, and noise. + + Raises: + ------- + Exception: If the signal_type is not defined. + + """ + # Generate signal matrix + signal = self.signal_creation(signal_mean, signal_variance, source_number=source_number) + signal = torch.from_numpy(signal) + # Generate noise matrix + noise = self.noise_creation(noise_mean, noise_variance) + noise = torch.from_numpy(noise) + if self.params.signal_type.startswith("broadband"): + raise Exception("Samples.samples_creation: Broadband signal type is not defined for far field") + if self.params.field_type.startswith("far"): + A = self.steering_vec(self.angles, f_c=self.f_rng[self.params.signal_type]) + samples = (A @ signal) + noise + elif self.params.field_type.startswith("near"): + A = self.steering_vec(angles=self.angles, ranges=self.distances, nominal=False, generate_search_grid=False, + f_c=self.f_rng[self.params.signal_type]) + samples = (A @ signal) + noise + elif self.params.field_type.startswith("full"): + A = self.steering_vec_full_model(angles=self.angles, + ranges=self.distances) + samples = (A @ signal) + noise + else: + raise Exception(f"Samples.params.field_type: Field type {self.params.field_type} is not defined") + return samples, signal, A, noise + + def noise_creation(self, noise_mean, noise_variance): + """Creates noise based on the specified mean and variance. + + Args: + ----- + noise_mean (float): Mean of the noise. + noise_variance (float): Variance of the noise. + + Returns: + -------- + np.ndarray: Generated noise. + + """ + # for NarrowBand signal_type Noise represented in the time domain + noise = ( + np.sqrt(noise_variance) + * (np.sqrt(2) / 2) + * ( + np.random.randn(self.params.N, self.params.T) + + 1j * np.random.randn(self.params.N, self.params.T) + ) + + noise_mean + ) + return noise + + def signal_creation(self, signal_mean: float = 0, signal_variance: float = 1, source_number: int = None): + """ + Creates signals based on the specified signal nature and parameters. + + Args: + ----- + signal_mean (float, optional): Mean of the signal. Defaults to 0. + signal_variance (float, optional): Variance of the signal. Defaults to 1. + + Returns: + -------- + np.ndarray: Created signals. + + Raises: + ------- + Exception: If the signal type is not defined. + Exception: If the signal nature is not defined. + """ + M = source_number + if self.params.snr is None: + snr = np.random.uniform(-5, 5) + else: + snr = self.params.snr + amplitude = 10 ** (snr / 10) + # NarrowBand signal creation + if self.params.signal_type == "narrowband": + if self.params.signal_nature == "non-coherent": + # create M non-coherent signals + return ( + amplitude + * (np.sqrt(2) / 2) + * np.sqrt(signal_variance) + * ( + np.random.randn(M, self.params.T) + + 1j * np.random.randn(M, self.params.T) + ) + + signal_mean + ) + + elif self.params.signal_nature == "coherent": + # Coherent signals: same amplitude and phase for all signals + sig = ( + amplitude + * (np.sqrt(2) / 2) + * np.sqrt(signal_variance) + * ( + np.random.randn(1, self.params.T) + + 1j * np.random.randn(1, self.params.T) + ) + + signal_mean + ) + return np.repeat(sig, M, axis=0) + + else: + raise Exception(f"signal type {self.params.signal_type} is not defined") diff --git a/src/system_model.py b/archive/src/system_model.py similarity index 100% rename from src/system_model.py rename to archive/src/system_model.py diff --git a/src/training.py b/archive/src/training.py similarity index 100% rename from src/training.py rename to archive/src/training.py diff --git a/archive/src/utils.py b/archive/src/utils.py new file mode 100644 index 0000000..79293ea --- /dev/null +++ b/archive/src/utils.py @@ -0,0 +1,566 @@ +"""Subspace-Net +Details +---------- +Name: utils.py +Authors: D. H. Shmuel +Created: 01/10/21 +Edited: 17/03/23 + +Purpose: +-------- +This script defines some helpful functions: + * sum_of_diag: returns the some of each diagonal in a given matrix. + * sum_of_diag_torch: returns the some of each diagonal in a given matrix, Pytorch oriented. + * find_roots: solves polynomial equation defines by polynomial coefficients. + * find_roots_torch: solves polynomial equation defines by polynomial coefficients, Pytorch oriented.. + * set_unified_seed: Sets unified seed for all random attributed in the simulation. + * get_k_angles: Retrieves the top-k angles from a prediction tensor. + * get_k_peaks: Retrieves the top-k peaks (angles) from a prediction tensor using peak finding. + * gram_diagonal_overload(self, Kx: torch.Tensor, eps: float): generates Hermitian and PSD (Positive Semi-Definite) matrix, + using gram operation and diagonal loading. +""" + +# Imports +import numpy as np +import torch + +torch.cuda.empty_cache() +import random +import scipy +import warnings + +from pathlib import Path +from src.config import device +import matplotlib.pyplot as plt +import torch.nn as nn + +# Constants +R2D = 180 / np.pi +D2R = 1 / R2D +plot_styles = { + 'CCRB': {'color': 'r', 'linestyle': '-', 'marker': 'o', "markersize": 8}, + 'Beamformer': {'color': 'r', 'linestyle': '--', 'marker': 's', "markersize": 8}, + 'DCD-MUSIC': {'color': 'g', 'linestyle': '-', 'marker': 'D', "markersize": 8}, + 'DCD-MUSIC_V2': {'color': 'g', 'linestyle': '--', 'marker': 'd', "markersize": 8}, + 'TransMUSIC': {'color': 'm', 'linestyle': '-.', 'marker': 'P', "markersize": 8}, + '2D-MUSIC': {'color': 'c', 'linestyle': ':', 'marker': '^', "markersize": 8}, + '2D-MUSIC(SPS)': {'color': 'c', 'linestyle': '--', 'marker': 'v', "markersize": 8}, + 'SubspaceNet': {'color': 'k', 'linestyle': '-', 'marker': 'X', "markersize": 8}, + 'NFSubspaceNet': {'color': 'k', 'linestyle': '--', 'marker': 'p', "markersize": 8}, + 'NFSubspaceNet_V2': {'color': 'b', 'linestyle': '-.', 'marker': 'h', "markersize": 8}, + 'ESPRIT': {'color': 'r', 'linestyle': '-', 'marker': 'v', "markersize": 8}, + 'esprit(SPS)': {'color': 'r', 'linestyle': '--', 'marker': 'v', "markersize": 8}, + '1D-MUSIC': {'color': 'y', 'linestyle': '-.', 'marker': 's', "markersize": 8}, + 'music(SPS)': {'color': 'y', 'linestyle': ':', 'marker': 's', "markersize": 8}, +} + +def validate_constant_sources_number(number_of_sources: torch.tensor): + """ + Validate that the number of sources in the batch is equal for all samples. + Args: + number_of_sources: The number of sources in the batch. + + Returns: + None + + Raises: + ValueError: If the number of sources in the batch is not equal for all samples + + """ + if (number_of_sources != number_of_sources[0]).any(): + raise ValueError(f"validate_constant_sources_number: " + f"Number of sources in the batch is not equal for all samples.") + +def initialize_data_paths(path: Path): + datasets_path = path / "datasets" + simulations_path = path / "simulations" + saving_path = path / "weights" + + # create folders if not exists + datasets_path.mkdir(parents=True, exist_ok=True) + (datasets_path / "train").mkdir(parents=True, exist_ok=True) + (datasets_path / "test").mkdir(parents=True, exist_ok=True) + simulations_path.mkdir(parents=True, exist_ok=True) + saving_path.mkdir(parents=True, exist_ok=True) + (saving_path / "final_models").mkdir(parents=True, exist_ok=True) + + return datasets_path, simulations_path, saving_path + +def sample_covariance(x: torch.Tensor) -> torch.Tensor: + """ + Calculates the sample covariance matrix for each element in the batch. + + Args: + ----- + X (np.ndarray): Input samples matrix. + + Returns: + -------- + covariance_mat (np.ndarray): Covariance matrix. + """ + if x.dim() == 2: + x = x[None, :, :] + batch_size, sensor_number, samples_number = x.shape + Rx = torch.einsum("bmt, btl -> bml", x, torch.conj(x).transpose(1, 2)) / samples_number + return Rx + +def spatial_smoothing_covariance(x: torch.Tensor): + """ + Calculates the covariance matrix using spatial smoothing technique for each element in the batch. + + Args: + ----- + X (np.ndarray): Input samples matrix. + + Returns: + -------- + covariance_mat (np.ndarray): Covariance matrix. + """ + + if x.dim() == 2: + x = x[None, :, :] + batch_size, sensor_number, samples_number = x.shape + # Define the sub-arrays size + sub_array_size = sensor_number // 2 + 1 + # Define the number of sub-arrays + number_of_sub_arrays = sensor_number - sub_array_size + 1 + # Initialize covariance matrix + Rx_smoothed = torch.zeros(batch_size, sub_array_size, sub_array_size, dtype=torch.complex128, device=device) + Rx = sample_covariance(x) + for j in range(number_of_sub_arrays): + Rx_smoothed += Rx[:, j:j + sub_array_size, j:j + sub_array_size] / number_of_sub_arrays + # Divide overall matrix by the number of sources + return Rx_smoothed + +def tops_covariance(x: torch.Tensor, number_of_bins: int=1): + """ + Tops algorithm uses K bins to calculate the covariance by using STFT. + Args: + x: + number_of_bins: + + + Returns: + + """ + Rx = torch.zeros(x.shape[0], number_of_bins,x.shape[1], x.shape[1], dtype=torch.complex128, device=device) + bin_size = x.shape[2] // number_of_bins + for i in range(number_of_bins): + x_bin = x[:, :, i*bin_size:(i+1)*bin_size] + Rx[:, i, :, :] = sample_covariance(x_bin) + return Rx + + +def keep_far_enough_points(tensor, M, D): + # # Calculate pairwise distances between columns + # distances = cdist(tensor.T, tensor.T, metric="euclidean") + # + # # Keep the first M columns as far enough points + # selected_cols = [] + # for i in range(tensor.shape[1]): + # if len(selected_cols) >= M: + # break + # if all(distances[i, col] >= D for col in selected_cols): + # selected_cols.append(i) + # + # # Remove columns that are less than distance D from each other + # filtered_tensor = tensor[:, selected_cols] + # retrun filtered_tensor + ############################################## + # Extract x_coords (first dimension) + x_coords = tensor[0, :] + + # Keep the first M columns that are far enough apart in x_coords + selected_cols = [] + for i in range(tensor.shape[1]): + if len(selected_cols) >= M: + break + if i == 0: + selected_cols.append(i) + continue + if all(abs(x_coords[i] - x_coords[col]) >= D for col in selected_cols): + selected_cols.append(i) + + # Select the columns that meet the distance criterion + filtered_tensor = tensor[:, selected_cols] + + return filtered_tensor + +# Functions +# def sum_of_diag(matrix: np.ndarray) -> list: +def sum_of_diag(matrix: np.ndarray): + """Calculates the sum of diagonals in a square matrix. + + Args: + matrix (np.ndarray): Square matrix for which diagonals need to be summed. + + Returns: + list: A list containing the sums of all diagonals in the matrix, from left to right. + + Raises: + None + + Examples: + >>> matrix = np.array([[1, 2, 3], + [4, 5, 6], + [7, 8, 9]]) + >>> sum_of_diag(matrix) + [7, 12, 15, 8, 3] + + """ + diag_sum = [] + diag_index = np.linspace( + -matrix.shape[0] + 1, + matrix.shape[0] + 1, + 2 * matrix.shape[0] - 1, + endpoint=False, + dtype=int, + ) + for idx in diag_index: + diag_sum.append(np.sum(matrix.diagonal(idx))) + return diag_sum + + +def sum_of_diags_torch(matrix: torch.Tensor): + """Calculates the sum of diagonals in a square matrix. + equivalent sum_of_diag, but support Pytorch. + + Args: + matrix (torch.Tensor): Square matrix for which diagonals need to be summed. + + Returns: + torch.Tensor: A list containing the sums of all diagonals in the matrix, from left to right. + + Raises: + None + + Examples: + >>> matrix = torch.tensor([[1, 2, 3], + [4, 5, 6], + [7, 8, 9]]) + >>> sum_of_diag(matrix) + torch.tensor([7, 12, 15, 8, 3]) + """ + diag_sum = [] + diag_index = torch.linspace( + -matrix.shape[0] + 1, matrix.shape[0] - 1, 2 * matrix.shape[0] - 1, dtype=int + ) + for idx in diag_index: + diag_sum.append(torch.sum(torch.diagonal(matrix, idx))) + return torch.stack(diag_sum, dim=0) + + +# def find_roots(coefficients: list) -> np.ndarray: +def find_roots(coefficients: list): + """Finds the roots of a polynomial defined by its coefficients. + + Args: + coefficients (list): List of polynomial coefficients in descending order of powers. + + Returns: + np.ndarray: An array containing the roots of the polynomial. + + Raises: + None + + Examples: + >>> coefficients = [1, -5, 6] # x^2 - 5x + 6 + >>> find_roots(coefficients) + array([3., 2.]) + + """ + coefficients = np.array(coefficients) + A = np.diag(np.ones((len(coefficients) - 2,), coefficients.dtype), -1) + if np.abs(coefficients[0]) == 0: + A[0, :] = -coefficients[1:] / (coefficients[0] + 1e-9) + else: + A[0, :] = -coefficients[1:] / coefficients[0] + roots = np.array(np.linalg.eigvals(A)) + return roots + + +def find_roots_torch(coefficients: torch.Tensor): + """Finds the roots of a polynomial defined by its coefficients. + equivalent to src.utils.find_roots, but support Pytorch. + + Args: + coefficients (torch.Tensor): List of polynomial coefficients in descending order of powers. + + Returns: + torch.Tensor: An array containing the roots of the polynomial. + + Raises: + None + + Examples: + >>> coefficients = torch.tensor([1, -5, 6]) # x^2 - 5x + 6 + >>> find_roots(coefficients) + tensor([3., 2.]) + + """ + A = torch.diag(torch.ones(len(coefficients) - 2, dtype=coefficients.dtype), -1) + A[0, :] = -coefficients[1:] / coefficients[0] + roots = torch.linalg.eigvals(A) + return roots + + +def set_unified_seed(seed: int = 42): + """ + Sets the seed value for random number generators in Python libraries. + + Args: + seed (int): The seed value to set for the random number generators. Defaults to 42. + + Returns: + None + + Examples: + >>> set_unified_seed(42) + + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + if torch.cuda.is_available(): + torch.use_deterministic_algorithms(False) + else: + torch.use_deterministic_algorithms(True) + + +# def get_k_angles(grid_size: float, k: int, prediction: torch.Tensor) -> torch.Tensor: +def get_k_angles(grid_size: float, k: int, prediction: torch.Tensor): + """ + Retrieves the top-k angles from a prediction tensor. + + Args: + grid_size (float): The size of the angle grid (range) in degrees. + k (int): The number of top angles to retrieve. + prediction (torch.Tensor): The prediction tensor containing angle probabilities, sizeof equal to grid_size . + + Returns: + torch.Tensor: A tensor containing the top-k angles in degrees. + + Raises: + None + + Examples: + >>> grid_size = 6 + >>> k = 3 + >>> prediction = torch.tensor([0.1, 0.3, 0.5, 0.2, 0.4, 0.6]) + >>> get_k_angles(grid_size, k, prediction) + tensor([ 90., -18., 54.]) + + """ + angles_grid = torch.linspace(-90, 90, grid_size) + doa_prediction = angles_grid[torch.topk(prediction.flatten(), k).indices] + return doa_prediction + + +# def get_k_peaks(grid_size, k: int, prediction) -> torch.Tensor: +def get_k_peaks(grid_size: int, k: int, prediction: torch.Tensor): + """ + Retrieves the top-k peaks (angles) from a prediction tensor using peak finding. + + Args: + grid_size (int): The size of the angle grid (range) in degrees. + k (int): The number of top peaks (angles) to retrieve. + prediction (torch.Tensor): The prediction tensor containing the peak values. + + Returns: + torch.Tensor: A tensor containing the top-k angles in degrees. + + Raises: + None + + Examples: + >>> grid_size = 6 + >>> k = 3 + >>> prediction = torch.tensor([0.1, 0.3, 0.5, 0.2, 0.4, 0.6]) + >>> get_k_angles(grid_size, k, prediction) + tensor([ 90., -18., 54.]) + + """ + angels_grid = torch.linspace(-90, 90, grid_size) + peaks, peaks_data = scipy.signal.find_peaks( + prediction.detach().numpy().flatten(), prominence=0.05, height=0.01 + ) + peaks = peaks[np.argsort(peaks_data["peak_heights"])[::-1]] + doa_prediction = angels_grid[peaks] + while doa_prediction.shape[0] < k: + doa_prediction = torch.cat( + ( + doa_prediction, + torch.Tensor(np.round(np.random.rand(1) * 180, decimals=2) - 90.00), + ), + 0, + ) + + return doa_prediction[:k] + + +# def gram_diagonal_overload(Kx: torch.Tensor, eps: float) -> torch.Tensor: +def gram_diagonal_overload(Kx: torch.Tensor, eps: float): + """Multiply a matrix Kx with its Hermitian conjecture (gram matrix), + and adds eps to the diagonal values of the matrix, + ensuring a Hermitian and PSD (Positive Semi-Definite) matrix. + + Args: + ----- + Kx (torch.Tensor): Complex matrix with shape [BS, N, N], + where BS is the batch size and N is the matrix size. + eps (float): Constant added to each diagonal element. + + Returns: + -------- + torch.Tensor: Hermitian and PSD matrix with shape [BS, N, N]. + + """ + # Insuring Tensor input + if not isinstance(Kx, torch.Tensor): + Kx = torch.tensor(Kx) + Kx = Kx.to(device) + + # Kx_garm = torch.matmul(torch.transpose(Kx.conj(), 1, 2).to("cpu"), Kx.to("cpu")).to(device) + Kx_garm = torch.bmm(Kx.conj().transpose(1, 2), Kx) + eps_addition = (eps * torch.diag(torch.ones(Kx_garm.shape[-1]))).to(device) + Kx_Out = Kx_garm + eps_addition + + # check if the matrix is Hermitian - A^H = A + mask = (torch.abs(Kx_Out - Kx_Out.conj().transpose(1, 2)) > 1e-6) + if mask.any(): + batch_mask = mask.any(dim=(1,2)) + warnings.warn(f"gram_diagonal_overload: {batch_mask.sum()} matrices in the batch aren't hermitian, taking the average of R and R^H.") + Kx_Out[batch_mask] = 0.5 * (Kx_Out[batch_mask] + Kx_Out[batch_mask].conj().transpose(1, 2)) + + return Kx_Out + + +# def _spatial_smoothing_covariance(sampels: torch.Tensor): +# """ +# Calculates the covariance matrix using spatial smoothing technique. +# +# Args: +# ----- +# X (np.ndarray): Input samples matrix. +# +# Returns: +# -------- +# covariance_mat (np.ndarray): Covariance matrix. +# """ +# +# X = sampels.squeeze() +# N = X.shape[0] +# # Define the sub-arrays size +# sub_array_size = int(N / 2) + 1 +# # Define the number of sub-arrays +# number_of_sub_arrays = N - sub_array_size + 1 +# # Initialize covariance matrix +# covariance_mat = torch.zeros((sub_array_size, sub_array_size), dtype=torch.complex128) +# +# for j in range(number_of_sub_arrays): +# # Run over all sub-arrays +# x_sub = X[j: j + sub_array_size, :] +# # Calculate sample covariance matrix for each sub-array +# sub_covariance = torch.cov(x_sub) +# # Aggregate sub-arrays covariances +# covariance_mat += sub_covariance / number_of_sub_arrays +# # Divide overall matrix by the number of sources +# return covariance_mat + + +def parse_loss_results_for_plotting(loss_results: dict, tested_param: str): + plt_res = {} + plt_acc = False + for test, results in loss_results.items(): + for method, loss_ in results.items(): + if plt_res.get(method) is None: + plt_res[method] = {tested_param: []} + try: + plt_res[method][tested_param].append(loss_[tested_param]) + except KeyError: + plt_res[method][tested_param].append(loss_["Overall"]) + if loss_.get("Accuracy") is not None: + if "Accuracy" not in plt_res[method].keys(): + plt_res[method]["Accuracy"] = [] + plt_acc = True + plt_res[method]["Accuracy"].append(loss_["Accuracy"]) + return plt_res, plt_acc + + +def print_loss_results_from_simulation(loss_results: dict): + """ + Print the loss results from the simulation. + """ + for test, value_dict in loss_results.items(): + print("#" * 10 + f"{test} TEST RESULTS" + "#" * 10) + for test_value, results in value_dict.items(): + if test == "SNR": + print(f"{test} = {test_value} [dB]: ") + else: + print(f"{test} = {test_value}: ") + for method, loss in results.items(): + txt = f"\t{method.upper(): <30}: " + for key, value in loss.items(): + if value is not None: + if key == "Accuracy": + txt += f"{key}: {value * 100:.2f} %|" + else: + txt += f"{key}: {value:.6e} |" + print(txt) + print("\n") + print("\n") + +class AntiRectifier(nn.Module): + def __init__(self, relu_inplace=False): + super(AntiRectifier, self).__init__() + self.relu = nn.ReLU(inplace=relu_inplace) + + def forward(self, x): + return torch.cat((self.relu(x), self.relu(-x)), 1) + +class L2NormLayer(nn.Module): + def __init__(self, dim=(1, 2), eps=1e-6): + super(L2NormLayer, self).__init__() + self.dim = dim + self.eps = eps + + def forward(self, x): + return torch.nn.functional.normalize(x, p=2, dim=self.dim, eps=self.eps) + self.eps * torch.diag(torch.ones(x.shape[-1], device=x.device)) + +class TraceNorm(nn.Module): + def __init__(self, eps=1e-8): + super().__init__() + self.eps = eps + + def forward(self, Rz): + trace = torch.real(Rz.diagonal(dim1=-2, dim2=-1).sum(-1)).clamp(min=self.eps) # shape [B] + trace = trace.view(-1, 1, 1) + return Rz / trace + + +if __name__ == "__main__": + # sum_of_diag example + matrix = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + sum_of_diag(matrix) + + matrix = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + sum_of_diags_torch(matrix) + + # find_roots example + coefficients = [1, -5, 6] + find_roots(coefficients) + + # get_k_angles example + grid_size = 6 + k = 3 + prediction = torch.tensor([0.1, 0.3, 0.5, 0.2, 0.4, 0.6]) + get_k_angles(grid_size, k, prediction) + + # get_k_peaks example + grid_size = 6 + k = 3 + prediction = torch.tensor([0.1, 0.3, 0.5, 0.2, 0.4, 0.6]) + get_k_peaks(grid_size, k, prediction) diff --git a/train_dcd.py b/archive/train_dcd.py similarity index 100% rename from train_dcd.py rename to archive/train_dcd.py diff --git a/check.py b/check.py new file mode 100644 index 0000000..5bd6f41 --- /dev/null +++ b/check.py @@ -0,0 +1,30 @@ +import torch + + +dict1= { + "a": 1, + "b": 2, + "c": 3 +} + +dict2 = {"d": 4, "e": 5} + +dict3 = {"f": 6, "g": 7} + +def check_some(**kwargs): + kwargs = {**kwargs} + print("Received arguments:") + for key, value in kwargs.items(): + print(f"{key}: {value}") + +tuple1 = (1, 2, 3) +tuple2 = (4, 5, 6) + + +def check_2(*args): + for arg in args: + print(f"Argument: {arg}") + + +if __name__ == "__main__": + print(("d","e") in dict2.items()) \ No newline at end of file diff --git a/doa_runner.py b/doa_runner.py new file mode 100644 index 0000000..04342a6 --- /dev/null +++ b/doa_runner.py @@ -0,0 +1,1166 @@ +""" +DoA Algorithm Runner - Handles training and evaluation of DoA estimation algorithms +""" + +import torch +import torch.nn as nn +import torch.optim as optim +import torch.optim.lr_scheduler as lr_scheduler +import numpy as np +import warnings +from pathlib import Path +from typing import Dict, Any, Optional, Tuple +from tqdm import tqdm +from datetime import datetime +import time +from typing import Dict, Any, Optional, Tuple, List + +from src.utils import sample_covariance + + +class AlgorithmFactory: + """ + Factory class for creating DoA algorithms and training components + Encapsulated within DoARunner to avoid circular imports and organize related functionality + """ + + def __init__(self, device): + self.device = device + + def create_algorithm(self, model_type: str, system_model_params, **kwargs): + """ + Create DoA estimation algorithm based on model type + + Args: + model_type: Type of algorithm ("diffMUSIC", "MUSIC", etc.) + system_model_params: System model parameters + **kwargs: Additional algorithm-specific parameters + + Returns: + Algorithm instance + """ + from src.diffmusic import DiffMUSIC + + if model_type.lower() == "diffmusic": + algorithm = DiffMUSIC( + system_model_params=system_model_params, + N=system_model_params.N, + model_order_estimation=kwargs.get('model_order_estimation', None), + physical_array=kwargs.get('physical_array', None), + physical_gains=kwargs.get('physical_gains', None) + ) + algorithm.train() + elif model_type.lower() == "music": + # For now, use DiffMUSIC in eval mode for classical MUSIC + algorithm = DiffMUSIC( + system_model_params=system_model_params, + N=system_model_params.N, + model_order_estimation=kwargs.get('model_order_estimation', None), + physical_array=kwargs.get('physical_array', None), + physical_gains=kwargs.get('physical_gains', None) + ) + algorithm.eval() # Set to evaluation mode for classical MUSIC behavior + else: + raise ValueError(f"Unknown model type: {model_type}") + + return algorithm.to(self.device) + + def create_loss_function(self, loss_type: str, **kwargs): + """ + Create loss function based on type + + Args: + loss_type: Type of loss ("rmspe", "spectrum", "unsupervised") + **kwargs: Additional loss-specific parameters + + Returns: + Loss function instance + """ + from src.diffmusic import DiffMUSICLoss + + return DiffMUSICLoss(loss_type=loss_type, **kwargs) + + def create_optimizer(self, algorithm, optimizer_type: str, learning_rate: float, **kwargs): + """ + Create optimizer for algorithm parameters + + Args: + algorithm: Algorithm instance with learnable parameters + optimizer_type: Type of optimizer ("Adam", "SGD") + learning_rate: Learning rate + **kwargs: Additional optimizer parameters + + Returns: + Optimizer instance + """ + if optimizer_type.lower() == "adam": + return optim.Adam( + algorithm.parameters(), + lr=learning_rate, + weight_decay=kwargs.get('weight_decay', 0) + ) + elif optimizer_type.lower() == "sgd": + return optim.SGD( + algorithm.parameters(), + lr=learning_rate, + momentum=kwargs.get('momentum', 0), + weight_decay=kwargs.get('weight_decay', 0) + ) + else: + raise ValueError(f"Unknown optimizer type: {optimizer_type}") + + def create_scheduler(self, optimizer, scheduler_type: str, **kwargs): + """ + Create learning rate scheduler + + Args: + optimizer: Optimizer instance + scheduler_type: Type of scheduler ("StepLR", "ReduceLROnPlateau") + **kwargs: Additional scheduler parameters + + Returns: + Scheduler instance or None + """ + if scheduler_type is None: + return None + + if scheduler_type.lower() == "steplr": + return lr_scheduler.StepLR( + optimizer, + step_size=kwargs.get('step_size', 50), + gamma=kwargs.get('gamma', 0.1) + ) + elif scheduler_type.lower() == "reducelronplateau": + return lr_scheduler.ReduceLROnPlateau( + optimizer, + mode='min', + factor=kwargs.get('factor', 0.5), + patience=kwargs.get('patience', 10) + ) + else: + raise ValueError(f"Unknown scheduler type: {scheduler_type}") + + +class DoARunner: + """ + Main runner class for DoA estimation algorithms + Handles both training and evaluation phases + """ + + def __init__(self, system_model_params, data_dict: Dict[str, Any]): + """ + Initialize DoA Runner + + Args: + system_model_params: SystemModelParams object containing all configuration + data_dict: Dictionary containing: + - 'measurements': Measurement data (N, T) + - 'true_angles': True DoA angles (M,) + - 'steering_matrix': Steering matrix (N, M) + - 'noise': Noise data (N, T) + - 'physical_array': Physical array positions (N,) + - 'physical_antennas_gains': Physical antenna gains (N,) + """ + self.system_params = system_model_params + self.data_dict = data_dict + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Precompute covariance matrix once (reused across all trials) + measurements = self.data_dict['measurements'].to(self.device).unsqueeze(0) + self.cov_matrix = sample_covariance(measurements) + + # Create factory for algorithm and training components + self.factory = AlgorithmFactory(self.device) + + # Create algorithm (may be recreated during window size optimization) + self.algorithm = self.factory.create_algorithm( + model_type=system_model_params.model_type, + system_model_params=system_model_params + ) + + # Initialize training components if needed + self.loss_fn = None + self.optimizer = None + self.scheduler = None + self.results = None + self.training_history = { + 'train_loss': [], + 'val_loss': [], + 'learning_rates': [] + } + + if self._needs_training(): + self._setup_training() + + def _needs_training(self) -> bool: + """Check if algorithm needs training""" + return self.algorithm.training + + + def _setup_training(self): + """Setup training components (loss, optimizer, scheduler)""" + # Determine loss type based on system parameters + if self.system_params.loss_type is None: + warnings.warn("Loss type not specified, defaulting to 'rmspe'.") + self.system_params.loss_type = 'rmspe' + loss_type = getattr(self.system_params, 'loss_type', 'rmspe') + + + # Create loss function + self.loss_fn = self.factory.create_loss_function(loss_type) + + # Create optimizer + self.optimizer = self.factory.create_optimizer( + algorithm=self.algorithm, + optimizer_type=self.system_params.optimizer, + learning_rate=self.system_params.learning_rate, + weight_decay=getattr(self.system_params, 'weight_decay', 0) + ) + + # Create scheduler + self.scheduler = self.factory.create_scheduler( + optimizer=self.optimizer, + scheduler_type=self.system_params.scheduler, + step_size=getattr(self.system_params, 'step_size', 50), + patience=getattr(self.system_params, 'patience', 10) + ) + + def _setup_wandb(self, window_size, trial_idx=None): + """ + Setup wandb logging if enabled + + Args: + window_size: Current window size being used + trial_idx: Trial index for optimization mode (None for single mode) + """ + if not getattr(self.system_params, 'use_wandb', False): + return + + try: + import wandb + + # Create project name with model type and timestamp + dt_string = getattr(self.system_params, 'dt_string_for_save', + datetime.now().strftime("%d_%m_%Y_%H_%M")) + project_name = f"{self.system_params.model_type}_{dt_string}" + + # Create run name + if trial_idx is not None: + run_name = f"trial_{trial_idx+1}_window_size_{window_size}_{self.system_params.loss_type}" + else: + run_name = f"winsow_size_{window_size}_{self.system_params.loss_type}" + + # Initialize wandb + wandb.init( + project=project_name, + name=run_name, + config={ + "model_type": self.system_params.model_type, + "loss_type": self.system_params.loss_type, + "window_size": window_size, + "learning_rate": self.system_params.learning_rate, + "epochs": self.system_params.epochs, + "N": self.system_params.N, + "M": self.system_params.M, + "T": self.system_params.T, + "snr": self.system_params.snr, + "optimizer": self.system_params.optimizer, + "scheduler": self.system_params.scheduler, + }, + reinit=True # Allow multiple runs in same script + ) + + except ImportError: + warnings.warn("wandb not installed. Install with 'pip install wandb' to enable logging.") + except Exception as e: + warnings.warn(f"Failed to initialize wandb: {e}") + + def _close_wandb(self): + """Close wandb run if active""" + if getattr(self.system_params, 'use_wandb', False): + try: + import wandb + wandb.finish() + except: + pass + + def run(self) -> Dict[str, Any]: + """ + Main run method - handles both training and evaluation + Supports single window size or window size optimization + + Returns: + Dictionary with results including estimated DoAs and metrics + """ + # Normalize window_size to always be a list/array for unified handling + window_size = getattr(self.system_params, 'softmax_window_size', 21) + if not hasattr(window_size, '__len__'): + window_sizes = [window_size] # Single case + optimization_mode = False + else: + window_sizes = window_size # Multiple cases + optimization_mode = len(window_sizes) > 1 + + if optimization_mode: + print(f"Starting window size optimization over {len(window_sizes)} configurations...") + else: + print("Running single configuration...") + + results = self._run_with_window_optimization(window_sizes, optimization_mode) + self.results = self.data_dict.copy() # Copy to avoid modifying original + self.results.update(results) # Update with results + + + return self.results.copy() # Return a copy to avoid external modifications + + def _create_fresh_algorithm(self, window_size): + """ + Create a new algorithm instance with fresh parameters + + Args: + window_size: Window size (int, float, or relative float) + + Returns: + Fresh algorithm instance with specified window size + """ + algorithm = self.factory.create_algorithm( + model_type=self.system_params.model_type, + system_model_params=self.system_params + ) + + # Set the specific window size + if isinstance(window_size, float) and 0 < window_size < 1: + # Case 2: relative to grid size + algorithm.window_size = int(window_size * len(algorithm.angles_grid)) + else: + # Case 1: absolute size + algorithm.window_size = int(window_size) + + return algorithm.to(self.device) + + def _run_with_window_optimization(self, window_sizes, optimization_mode=True) -> Dict[str, Any]: + """ + Unified method for both single and multiple window size configurations + + Args: + window_sizes: List of window sizes to try + optimization_mode: Whether to print optimization-specific messages + + Returns: + Results from best configuration (or single configuration) plus stats + """ + + + best_rmspe = float('inf') + best_results = None + best_window_size = None + + + trial_times = [] + + for i, window_size in enumerate(window_sizes): + trial_start_time = time.time() + + if optimization_mode: + print(f"Trial {i+1}/{len(window_sizes)}: window_size={window_size}") + else: + print(f"Running with window_size={window_size}") + + # Create fresh algorithm instance with reset parameters + self.algorithm = self._create_fresh_algorithm(window_size) + + # Setup fresh training components and wandb + if self._needs_training(): + self._setup_training() + self._setup_wandb(window_size, trial_idx=i if optimization_mode else None) + + # Run complete training + evaluation cycle + results = {} + if self._needs_training(): + train_results = self._run_training() + results.update(train_results) + + eval_results = self._run_evaluation() + results.update(eval_results) + + # Close wandb for this trial + self._close_wandb() + + trial_end_time = time.time() + trial_duration = trial_end_time - trial_start_time + trial_times.append(trial_duration) + + # Check if this is the best configuration + if results['rmspe'] < best_rmspe: + best_rmspe = results['rmspe'] + best_results = results.copy() + best_window_size = window_size + status_msg = f"New best RMSPE: {best_rmspe:.6f}" if optimization_mode else f"RMSPE: {best_rmspe:.6f}" + print(f" {status_msg} (Time: {trial_duration:.2f}s)") + else: + print(f" RMSPE: {results['rmspe']:.6f} (Time: {trial_duration:.2f}s)") + + total_time = np.sum(np.array(trial_times)) + + if optimization_mode: + average_duration = np.mean(trial_times) + print(f"\nOptimization complete!") + print(f"Best window_size: {best_window_size}, RMSPE: {best_rmspe:.6f}") + print(f"Total time: {total_time:.2f}s") + print(f"Average time per trial: {average_duration:.2f}s") + + # Add optimization stats + best_results['optimal_window_size'] = best_window_size + best_results['optimization_stats'] = { + 'total_time': total_time, + 'average_time_per_trial': average_duration, + 'trial_times': trial_times, + 'num_trials': len(window_sizes) + } + else: + print(f"Completed in {total_time:.2f}s") + best_results['execution_time'] = total_time + + return best_results + + def _run_training(self) -> Dict[str, Any]: + """ + Run training phase using precomputed covariance matrix + + Returns: + Training results + """ + + # Extract data + true_angles = self.data_dict['true_angles'].to(self.device).unsqueeze(0) # (1, M) + M = len(self.data_dict['true_angles']) + + # Reset training history for fresh instance + self.training_history = { + 'train_loss': [], + 'val_loss': [], + 'learning_rates': [], + 'steering_mse': [] + } + + print(f"Training for {self.system_params.epochs} epochs...") + + for epoch in tqdm(range(self.system_params.epochs), desc="Training"): + epoch_loss = self._train_epoch(self.cov_matrix, true_angles, M) + + # Compute steering MSE for this epoch + with torch.no_grad(): + steering_mse = self._compute_steering_matrix_mse(true_angles.squeeze(0)) + + # Store training history + self.training_history['train_loss'].append(epoch_loss) + self.training_history['steering_mse'].append(steering_mse) + if self.optimizer: + current_lr = self.optimizer.param_groups[0]['lr'] + self.training_history['learning_rates'].append(current_lr) + + # Log to wandb if enabled + if getattr(self.system_params, 'use_wandb', False): + try: + import wandb + wandb.log({ + "epoch": epoch, + "train_loss": epoch_loss, + "steering_mse": steering_mse, + "learning_rate": current_lr if self.optimizer else 0 + }) + except: + pass + + # Step scheduler + if self.scheduler: + if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + self.scheduler.step(epoch_loss) + else: + self.scheduler.step() + + # Print progress + if (epoch + 1) % 10 == 0 or epoch == 0: + print(f"Epoch {epoch+1}/{self.system_params.epochs}, Loss: {epoch_loss:.6f}, Steering MSE: {steering_mse:.6f}") + + return { + 'training_history': self.training_history, + 'final_train_loss': epoch_loss + } + + def _train_epoch(self, cov_matrix: torch.Tensor, true_angles: torch.Tensor, M: int) -> float: + """ + Train for one epoch + + Args: + cov_matrix: Covariance matrix (1, N, N) + true_angles: True angles (1, M) + M: Number of sources + + Returns: + Epoch loss + """ + self.optimizer.zero_grad() + + # Forward pass + algorithm_result = self.algorithm(cov_matrix, M) + if self.system_params.model_type.lower() == "diffmusic": + estimated_angles, peaks_masks, _ , _ = algorithm_result + if peaks_masks is None: + warnings.warn("No peaks masks returned from algorithm, something is weird and you need to check that.") + + # Compute loss based on loss type + loss_kwargs = { + 'predictions': estimated_angles, + 'targets': true_angles + } + + # Add additional arguments for spectrum-based losses + if hasattr(self.loss_fn, 'loss_type') and self.loss_fn.loss_type in ['spectrum', 'unsupervised']: + loss_kwargs['spectrum'] = self.algorithm.music_spectrum + loss_kwargs['angles_grid'] = self.algorithm.angles_grid + + if self.loss_fn.loss_type == 'unsupervised': + + loss_kwargs['peak_masks'] = peaks_masks + + # Compute loss + loss = self.loss_fn(**loss_kwargs) + + # Handle different loss return types + if isinstance(loss, torch.Tensor) and loss.dim() > 0: + loss = loss.mean() + + # Backward pass + loss.backward() + self.optimizer.step() + + return loss.item() + + def _run_evaluation(self) -> Dict[str, Any]: + """ + Run evaluation phase using precomputed covariance matrix + + Returns: + Evaluation results + """ + + with torch.no_grad(): + # Extract data + true_angles = self.data_dict['true_angles'].to(self.device).unsqueeze(0) # (1, M) + M = len(self.data_dict['true_angles']) + + # Forward pass with precomputed covariance + algorithm_result = self.algorithm(self.cov_matrix, M) + if "music" in self.system_params.model_type.lower(): # Handle MUSIC and DiffMUSIC + estimated_angles, _, _, _ = algorithm_result + + # Remove batch dimension for output + estimated_angles = estimated_angles.squeeze(0) # (M,) + true_angles = true_angles.squeeze(0) # (M,) + + learned_antenna_positions, learned_coplex_gains = self.algorithm.get_array_learnable_parameters(learnable=False) + + # Compute MSE between real and estimated steering matrices + steering_matrix_mse = self._compute_steering_matrix_mse(true_angles) + + # Compute RMSPE for evaluation + from src.metrics import RMSPELoss + rmspe_fn = RMSPELoss() + rmspe = rmspe_fn(estimated_angles.unsqueeze(0), true_angles.unsqueeze(0)) + + + + return { + 'estimated_angles': np.sort(estimated_angles.cpu().numpy()), + 'true_angles': true_angles.cpu().numpy(), + 'rmspe': rmspe.item(), + "learned_antenna_positions": learned_antenna_positions, + "learned_antennas_gains": learned_coplex_gains, + 'music_spectrum': self.algorithm.music_spectrum.cpu().numpy() if self.algorithm.music_spectrum is not None else None, + 'angles_grid': self.algorithm.angles_grid.cpu().numpy() if hasattr(self.algorithm, 'angles_grid') else None, + 'steering_matrix_mse': steering_matrix_mse + } + + def save_model(self, path: Path): + """Save trained model""" + if self._needs_training(): + torch.save({ + 'model_state_dict': self.algorithm.state_dict(), + 'optimizer_state_dict': self.optimizer.state_dict(), + 'training_history': self.training_history, + 'system_params': self.system_params + }, path) + + def load_model(self, path: Path): + """Load trained model""" + checkpoint = torch.load(path, map_location=self.device) + self.algorithm.load_state_dict(checkpoint['model_state_dict']) + if self.optimizer and 'optimizer_state_dict' in checkpoint: + self.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + if 'training_history' in checkpoint: + self.training_history = checkpoint['training_history'] + + def plot_graphs(self, path: Path): + """ + Plot specified graphs by case + """ + + if self.results is None: + raise ValueError("No results available. Run the algorithm first.") + true_angles = torch.from_numpy(self.results['true_angles']).to(self.device) + + if "music" in self.system_params.model_type.lower() and self.system_params.plot_results: + if self.algorithm.music_spectrum is None: + raise ValueError("Music spectrum is None the algorithm hasn't runned yet.") + self.algorithm.plot_spectrum( + true_angles, + path + ) + # Plot learned parameters (Figure 3 style) for diffMUSIC + if self.system_params.model_type.lower() == "diffmusic" and self.system_params.plot_results: # TODO: check with Ori + # Get physical parameters for comparison + physical_array = self.data_dict.get('physical_array', None) + physical_gains = self.data_dict.get('physical_antennas_gains', None) + + # Convert to torch tensors if they're numpy arrays + if physical_array is not None and not isinstance(physical_array, torch.Tensor): + physical_array = torch.from_numpy(physical_array) + if physical_gains is not None and not isinstance(physical_gains, torch.Tensor): + physical_gains = torch.from_numpy(physical_gains) + + # Create title based on loss type and performance + rmspe = self.results.get('rmspe', 0) + loss_type = getattr(self.system_params, 'loss_type', 'unknown') + title = f"Learned Parameters ({loss_type.upper()}) - RMSPE: {rmspe:.3f}°" + + self.algorithm.plot_learned_parameters( + true_positions=physical_array, + true_gains=physical_gains, + path_to_save=path, + title=title + ) + return None + + + def run_multi_loss_spectrum_comparison(self, loss_functions: List[str]) -> Dict[str, Any]: + """ + Run the same experiment with multiple loss functions and collect spectra for comparison + + Args: + loss_functions: List of loss functions to compare (e.g., ['rmspe', 'spectrum', 'unsupervised']) + + Returns: + Dictionary containing spectra and results for each loss function + """ + print(f"Starting multi-loss spectrum comparison with {len(loss_functions)} loss functions...") + print(f"Loss functions: {loss_functions}") + + # Initialize results storage + spectra_results = {} + full_results = {} + execution_times = {} + + # Get window size for consistent comparison + window_size = getattr(self.system_params, 'softmax_window_size', 21) + if hasattr(window_size, '__len__'): + # If multiple window sizes, use the first one for comparison + window_size = window_size[0] if len(window_size) > 0 else 21 + + # Run experiment for each loss function + for i, loss_function in enumerate(loss_functions): + print(f"\n{'='*50}") + print(f"Running with Loss Function: {loss_function} ({i+1}/{len(loss_functions)})") + print(f"{'='*50}") + + start_time = time.time() + + try: + # Create fresh algorithm instance with reset parameters + self.algorithm = self._create_fresh_algorithm(window_size) + + # Update system params for this loss function + original_loss_type = getattr(self.system_params, 'loss_type', None) + self.system_params.loss_type = loss_function + + # Setup fresh training components + if self._needs_training(): + self._setup_training() + self._setup_wandb(window_size, trial_idx=i) + + # Run complete training + evaluation cycle + results = {} + if self._needs_training(): + train_results = self._run_training() + results.update(train_results) + + eval_results = self._run_evaluation() + results.update(eval_results) + + # Store spectrum with proper dimension handling + spectrum = results.get('music_spectrum', None) + if spectrum is not None: + # Ensure consistent storage format (remove batch dimension) + if isinstance(spectrum, torch.Tensor): + spectrum = spectrum.cpu().numpy() + if spectrum.ndim == 2 and spectrum.shape[0] == 1: + spectrum = spectrum.squeeze(0) # Remove batch dimension + elif spectrum.ndim == 2 and spectrum.shape[0] > 1: + spectrum = spectrum[0] # Take first batch element + + spectra_results[loss_function] = spectrum + else: + spectra_results[loss_function] = None + + full_results[loss_function] = results + execution_times[loss_function] = time.time() - start_time + + print(f" RMSPE: {results['rmspe']:.6f} (Time: {execution_times[loss_function]:.2f}s)") + + # Close wandb for this trial + self._close_wandb() + + # Restore original loss type + if original_loss_type is not None: + self.system_params.loss_type = original_loss_type + + except Exception as e: + print(f" Error with {loss_function}: {e}") + import traceback + traceback.print_exc() # Print full traceback for debugging + + # Store error results + spectra_results[loss_function] = None + full_results[loss_function] = {'error': str(e), 'rmspe': float('inf')} + execution_times[loss_function] = time.time() - start_time + + # Compile consolidated results + total_time = sum(execution_times.values()) + consolidated_results = { + 'comparison_type': 'multi_loss_spectrum', + 'loss_functions': loss_functions, + 'window_size_used': window_size, + 'spectra': spectra_results, + 'results': full_results, + 'execution_times': execution_times, + 'total_execution_time': total_time, + 'shared_data': { + 'angles_grid': self.algorithm.angles_grid.cpu().numpy() if hasattr(self.algorithm, 'angles_grid') else None, + 'true_angles': self.data_dict['true_angles'].cpu().numpy() if isinstance(self.data_dict['true_angles'], torch.Tensor) else self.data_dict['true_angles'], + 'system_params': self.system_params + } + } + + print(f"\nMulti-loss spectrum comparison completed in {total_time:.2f}s") + + # Store results for potential later use + self.multi_loss_results = consolidated_results + + return consolidated_results + + def plot_multi_loss_spectrum_comparison(self, multi_loss_results: Dict[str, Any], path: Path): + """ + Plot spectra from multiple loss functions on the same figure + + Args: + multi_loss_results: Results from run_multi_loss_spectrum_comparison() + path: Path to save the plot + """ + import matplotlib.pyplot as plt + import numpy as np + + if multi_loss_results is None: + raise ValueError("No multi-loss results available. Run run_multi_loss_spectrum_comparison() first.") + + # Extract data + loss_functions = multi_loss_results['loss_functions'] + spectra = multi_loss_results['spectra'] + results = multi_loss_results['results'] + shared_data = multi_loss_results['shared_data'] + + angles_grid = shared_data['angles_grid'] + true_angles = shared_data['true_angles'] + + if angles_grid is None: + print("Warning: No angles grid available for plotting") + return + + # Convert angles to degrees for plotting + angles_deg = np.rad2deg(angles_grid) + true_angles_deg = np.rad2deg(true_angles) + + # Create figure + plt.figure(figsize=(14, 8)) + + # Define colors and line styles for different loss functions FIRST + colors = {'rmspe': 'blue', 'spectrum': 'red', 'unsupervised': 'green', 'physical': 'black'} + line_styles = {'rmspe': '-', 'spectrum': '--', 'unsupervised': '-.', 'physical': ':'} + default_colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown'] + default_styles = ['-', '--', '-.', ':', '-', '--'] + + # Compute and plot physical parameters spectrum if available + physical_spectrum = self._compute_physical_spectrum() + if physical_spectrum is not None: + plt.plot(angles_deg, physical_spectrum, + color=colors['physical'], linestyle=line_styles['physical'], linewidth=3, + label='Physical Parameters', alpha=0.8) + + # Plot spectrum for each loss function + for i, loss_function in enumerate(loss_functions): + spectrum = spectra.get(loss_function, None) + result = results.get(loss_function, {}) + + if spectrum is not None and 'error' not in result: + # Handle spectrum dimensions - remove batch dimension if present + if isinstance(spectrum, np.ndarray): + if spectrum.ndim == 2 and spectrum.shape[0] == 1: + spectrum = spectrum.squeeze(0) # Remove batch dimension (1, 481) -> (481,) + elif spectrum.ndim == 2 and spectrum.shape[0] > 1: + spectrum = spectrum[0] # Take first batch element + elif isinstance(spectrum, torch.Tensor): + spectrum = spectrum.cpu().numpy() + if spectrum.ndim == 2 and spectrum.shape[0] == 1: + spectrum = spectrum.squeeze(0) + elif spectrum.ndim == 2 and spectrum.shape[0] > 1: + spectrum = spectrum[0] + + # Ensure spectrum is 1D + if spectrum.ndim != 1: + print(f"Warning: Unexpected spectrum shape for {loss_function}: {spectrum.shape}") + continue + + # Ensure angles_deg and spectrum have the same length + if len(angles_deg) != len(spectrum): + print(f"Warning: Length mismatch for {loss_function}: angles={len(angles_deg)}, spectrum={len(spectrum)}") + continue + + # Get color and style + color = colors.get(loss_function, default_colors[i % len(default_colors)]) + style = line_styles.get(loss_function, default_styles[i % len(default_styles)]) + + # Get RMSPE for label + rmspe = result.get('rmspe', float('inf')) + + # Plot spectrum + plt.plot(angles_deg, spectrum, + color=color, linestyle=style, linewidth=2, + label=f'{loss_function.upper()} (RMSPE: {rmspe:.3f}°)') + else: + print(f"Warning: No valid spectrum for {loss_function}") + + # Add true angle markers + for i, true_angle in enumerate(true_angles_deg): + plt.axvline(x=true_angle, color='black', linestyle=':', alpha=0.7, + label='True DoA' if i == 0 else "") + + # Customize plot + plt.xlabel('Angle [degrees]', fontsize=12, fontweight='bold') + plt.ylabel('Spectrum Power', fontsize=12, fontweight='bold') + plt.title('Multi-Loss Function Spectrum Comparison', fontsize=14, fontweight='bold') + plt.grid(True, alpha=0.3) + plt.legend(fontsize=10) + + # Add system info as text + steering_mse_text = ", ".join([f"{lf}: {results[lf].get('steering_matrix_mse', float('nan')):.3f}" + for lf in loss_functions if 'error' not in results.get(lf, {})]) + + info_text = (f"N={shared_data['system_params'].N}, " + f"M={shared_data['system_params'].M}, " + f"T={shared_data['system_params'].T}, " + f"SNR={shared_data['system_params'].snr}dB\n" + f"Steering MSE - {steering_mse_text}") + plt.text(0.02, 0.98, info_text, transform=plt.gca().transAxes, + verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + plt.tight_layout() + + # Save plot + save_path = path / "multi_loss_spectrum_comparison.png" + plt.savefig(save_path, dpi=300, bbox_inches='tight') + plt.savefig(path / "multi_loss_spectrum_comparison.pdf", bbox_inches='tight') + + print(f"Multi-loss spectrum comparison plot saved to: {save_path}") + plt.show() + + + def plot_multi_loss_learned_parameters(self, multi_loss_results: Dict[str, Any], path: Path): + """ + Plot learned parameters from multiple loss functions on the same figure + + Args: + multi_loss_results: Results from run_multi_loss_spectrum_comparison() + path: Path to save the plot + """ + import matplotlib.pyplot as plt + import matplotlib.patches as patches + import numpy as np + + if multi_loss_results is None: + raise ValueError("No multi-loss results available. Run run_multi_loss_spectrum_comparison() first.") + + # Extract data + loss_functions = multi_loss_results['loss_functions'] + results = multi_loss_results['results'] + shared_data = multi_loss_results['shared_data'] + + # Get physical parameters for comparison + true_positions = self.data_dict.get('physical_array', None) + true_gains = self.data_dict.get('physical_antennas_gains', None) + + # Convert to numpy if they're torch tensors + if true_positions is not None and isinstance(true_positions, torch.Tensor): + true_positions = true_positions.cpu().numpy() + if true_gains is not None and isinstance(true_gains, torch.Tensor): + true_gains = true_gains.cpu().numpy() + + # Create figure with appropriate height for all rows + num_loss_functions = len([lf for lf in loss_functions if 'error' not in results.get(lf, {})]) + total_rows = num_loss_functions + (1 if true_positions is not None else 0) + fig, ax = plt.subplots(1, 1, figsize=(16, 3 + total_rows * 1.5)) + + # Define colors for different loss functions + colors = {'rmspe': 'blue', 'spectrum': 'red', 'unsupervised': 'green'} + default_colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown'] + + # Y positions for different rows + row_spacing = 1.5 + current_y = (total_rows - 1) * row_spacing / 2 + + # Plot learned parameters for each loss function + plotted_loss_functions = [] + for i, loss_function in enumerate(loss_functions): + result = results.get(loss_function, {}) + + if 'error' in result: + print(f"Skipping {loss_function} due to error: {result['error']}") + continue + + learned_positions = result.get('learned_antenna_positions', None) + learned_gains = result.get('learned_antennas_gains', None) + rmspe = result.get('rmspe', float('inf')) + + if learned_positions is None or learned_gains is None: + print(f"Warning: No learned parameters for {loss_function}") + continue + + # Convert positions to wavelength units for display + wavelength = shared_data['system_params'].wavelength + learned_positions_wl = learned_positions / (wavelength / 2) + + # Get color for this loss function + color = colors.get(loss_function, default_colors[i % len(default_colors)]) + + # Plot learned parameters + for j, (pos, gain) in enumerate(zip(learned_positions_wl, learned_gains)): + # Circle radius represents gain magnitude + radius = abs(gain) * 0.2 # Smaller radius for multiple rows + + # Circle color and segment angle represent gain phase + phase = np.angle(gain) + + # Draw circle with color corresponding to loss function + circle = patches.Circle((pos, current_y), radius, + facecolor=color, alpha=0.3, + edgecolor=color, + linewidth=2, + label=f'{loss_function.upper()} (RMSPE: {rmspe:.3f}°)' if j == 0 else "") + ax.add_patch(circle) + + # Draw phase segment (line from center to edge) + segment_x = pos + radius * np.cos(phase) + segment_y = current_y + radius * np.sin(phase) + ax.plot([pos, segment_x], [current_y, segment_y], color=color, linewidth=2) + + # Add antenna number + if j == 0: # Only add for first antenna to avoid clutter + ax.text(pos - 0.3, current_y, f'{loss_function.upper()}', + ha='right', va='center', fontsize=10, fontweight='bold', color=color) + + # Add horizontal reference line + ax.axhline(y=current_y, color=color, linestyle=':', alpha=0.3, linewidth=1) + + plotted_loss_functions.append(loss_function) + current_y -= row_spacing + + # Plot physical parameters if available (bottom row) + if true_positions is not None and true_gains is not None: + true_positions_wl = true_positions / (wavelength / 2) + + for j, (pos, gain) in enumerate(zip(true_positions_wl, true_gains)): + radius = abs(gain) * 0.2 + phase = np.angle(gain) + + # Draw circle with different style for physical parameters + circle = patches.Circle((pos, current_y), radius, + facecolor='lightgray', + edgecolor='black', + linewidth=2, + alpha=0.7, + label='Physical' if j == 0 else "") + ax.add_patch(circle) + + # Draw phase segment + segment_x = pos + radius * np.cos(phase) + segment_y = current_y + radius * np.sin(phase) + ax.plot([pos, segment_x], [current_y, segment_y], 'k-', linewidth=2) + + # Add label for physical parameters + ax.text(true_positions_wl[0] - 0.3, current_y, 'PHYSICAL', + ha='right', va='center', fontsize=10, fontweight='bold', color='black') + + # Add horizontal reference line + ax.axhline(y=current_y, color='black', linestyle=':', alpha=0.3, linewidth=1) + + # Set axis limits and labels + all_positions = [] + for loss_function in plotted_loss_functions: + result = results.get(loss_function, {}) + if 'learned_antenna_positions' in result: + positions_wl = result['learned_antenna_positions'] / (wavelength / 2) + all_positions.extend(positions_wl) + + if true_positions is not None: + all_positions.extend(true_positions_wl) + + if all_positions: + ax.set_xlim(min(all_positions) - 0.5, max(all_positions) + 0.5) + + y_range = total_rows * row_spacing / 2 + 0.8 + ax.set_ylim(-y_range, y_range) + ax.set_xlabel('x [λ/2]', fontsize=12, fontweight='bold') + ax.set_ylabel('') + + # Create title with summary + rmspe_summary = ', '.join([f"{lf.upper()}: {results[lf]['rmspe']:.3f}°" + for lf in plotted_loss_functions]) + title = f'Multi-Loss Learned Parameters Comparison\n{rmspe_summary}' + ax.set_title(title, fontsize=14, fontweight='bold') + + ax.grid(True, alpha=0.3) + + # Create custom legend + legend_elements = [] + for i, loss_function in enumerate(plotted_loss_functions): + color = colors.get(loss_function, default_colors[i % len(default_colors)]) + rmspe = results[loss_function]['rmspe'] + legend_elements.append( + plt.Line2D([0], [0], marker='o', color='w', markerfacecolor=color, + markersize=10, markeredgecolor=color, markeredgewidth=2, + alpha=0.7, label=f'{loss_function.upper()} (RMSPE: {rmspe:.3f}°)') + ) + + if true_positions is not None: + legend_elements.append( + plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='lightgray', + markersize=10, markeredgecolor='black', markeredgewidth=2, + alpha=0.7, label='Physical') + ) + + ax.legend(handles=legend_elements, loc='upper right', fontsize=10) + + # Remove y-axis ticks for cleaner look + ax.set_yticks([]) + + plt.tight_layout() + + # Save plot + save_path = path / "multi_loss_learned_parameters.png" + plt.savefig(save_path, dpi=300, bbox_inches='tight') + plt.savefig(path / "multi_loss_learned_parameters.pdf", bbox_inches='tight') + + print(f"Multi-loss learned parameters plot saved to: {save_path}") + plt.show() + + + def _compute_steering_matrix_mse(self, true_angles: torch.Tensor) -> float: + """ + Compute MSE between real and estimated steering matrices + + Args: + true_angles: True DoA angles in radians + + Returns: + MSE value + """ + # Get real steering matrix from data_dict (computed during data creation) + real_steering_matrix = self.data_dict.get('steering_matrix', None) + + if real_steering_matrix is None: + return float('nan') # Cannot compute if real steering matrix not available + + # Convert to torch tensor if needed and move to device + if not isinstance(real_steering_matrix, torch.Tensor): + real_steering_matrix = torch.from_numpy(real_steering_matrix) + real_steering_matrix = real_steering_matrix.to(self.device) + + true_angles = true_angles.to(self.device) + + # Compute estimated steering matrix using learned parameters + estimated_steering_matrix = self.algorithm.compute_steering_matrix(true_angles) + + # Compute MSE + mse = torch.mean(torch.abs(real_steering_matrix - estimated_steering_matrix) ** 2) + + return mse.item() + + + def _compute_physical_spectrum(self) -> Optional[np.ndarray]: + """ + Compute MUSIC spectrum using physical parameters for comparison + + Returns: + Physical spectrum as numpy array, or None if physical parameters not available + """ + # Check if physical parameters are available + physical_array = self.data_dict.get('physical_array', None) + physical_gains = self.data_dict.get('physical_antennas_gains', None) + + if physical_array is None or physical_gains is None: + print("Warning: Physical parameters not available for spectrum computation") + return None + + try: + # Convert to torch tensors if needed + if not isinstance(physical_array, torch.Tensor): + physical_array = torch.from_numpy(physical_array) + if not isinstance(physical_gains, torch.Tensor): + physical_gains = torch.from_numpy(physical_gains) + + # Move to device + physical_array = physical_array.to(self.device).to(torch.float64) + physical_gains = physical_gains.to(self.device).to(torch.complex64) + + # Check if noise subspace is available from the algorithm + if not hasattr(self.algorithm, 'noise_subspace') or self.algorithm.noise_subspace is None: + print("Warning: No noise subspace available for physical spectrum computation") + return None + + # Save current learned parameters + with torch.no_grad(): + current_positions = self.algorithm.antenna_positions.clone() + current_gains = self.algorithm.complex_gain.clone() + + # Temporarily set to physical parameters + self.algorithm.antenna_positions.copy_(physical_array) + self.algorithm.complex_gain.copy_(physical_gains) + + # Compute spectrum using physical parameters + inverse_spectrum = self.algorithm._compute_inverse_spectrum(self.algorithm.noise_subspace) + physical_spectrum = 1 / (inverse_spectrum + 1e-10) + + # Remove batch dimension and convert to numpy + if physical_spectrum.dim() == 2 and physical_spectrum.shape[0] == 1: + physical_spectrum = physical_spectrum.squeeze(0) + physical_spectrum = physical_spectrum.cpu().detach().numpy() + + # Restore learned parameters + self.algorithm.antenna_positions.copy_(current_positions) + self.algorithm.complex_gain.copy_(current_gains) + + return physical_spectrum + + except Exception as e: + print(f"Error computing physical spectrum: {e}") + import traceback + traceback.print_exc() # This will help debug any remaining issues + return None \ No newline at end of file diff --git a/full_environment.yml b/full_environment.yml index c51ba88..f9311e4 100644 --- a/full_environment.yml +++ b/full_environment.yml @@ -1,4 +1,4 @@ -name: ai_subspace_env +name: calibration_through_doa channels: - conda-forge - defaults @@ -74,3 +74,4 @@ dependencies: - typing-inspection==0.4.0 - urllib3==2.3.0 - wandb==0.19.8 + - tqdm==4.67.1 diff --git a/main.py b/main.py index 079933b..1384a40 100644 --- a/main.py +++ b/main.py @@ -1,8 +1,7 @@ """ This script is used to run the simulation with the given parameters. The parameters can be set in the script or by using the command line arguments. The script will run the simulation with the given parameters and save the -results to the results folder. The results will include the learning curves, RMSE results, and the accuracy results -of the evaluation. The results will be saved in the results folder in the project directory. +results to the results folder. The script can be run with the following command line arguments: --snr: SNR value @@ -11,160 +10,154 @@ --field_type: Field type --signal_nature: Signal nature --model_type: Model type - --train: Train model - --train_criteria: Training criteria - --eval: Evaluate model - --eval_criteria: Evaluation criteria - --samples_size: Samples size - --train_test_ratio: Train test ratio + --loss_type: Training loss type (rmspe, spectrum, unsupervised) + --epochs: Number of training epochs + --batch_size: Batch size + --learning_rate: Learning rate + --optimizer: Optimizer type + --scheduler: Scheduler type + --create: Create new dataset + --wandb: Use wandb """ # Imports import os import warnings import time +import numpy as np import matplotlib.pyplot as plt from run_simulation import run_simulation import argparse +import torch + +#TODO: +# 1. Right now, each experiment is saved alone (while optimizing the sotfmax window size). +# we need to create also a multi-case scenario where we can test different cross products of parameters. +# Something similiar to the old scenario_dict, but not necessirilly it. for the clearness of the code +# we might want to create a new class for this, including the plotting capabilities (next point) + +# 2. based on the multi-case results, create functions plotting relevanmt graphs from the paper to +# test our recreation. not such a big deal, please arrange it under one class. + +#NOTE: +# points for the future: + +# 1. The unsupervised loss suffers from problems at the endfire due to smaller number of samples in the window. +# We need to think about maybe mirroring at the edges or renormalizing the Jain's index somehow. + +# 2. The losses minimize thr rmspe, leading to higher DoA accuracy, but the actual configuration paraameters +# diverge heavily. It'll be interesting to compare the learned steering matrix itself to the physical one although we will get +# probably the same conclusion. +# SUPER IMPORTANT IMPLICATION: If we don't really learn the rifht parameters, tracking with this garbage observations will not lead us far. # Initialization os.system("cls||clear") plt.close("all") -scenario_dict = { - # "SNR": [-10, -5, 0, 5, 10], - # "T": [10, 20, 30, 50, 70, 100], - # "eta": [0.0, 0.01, 0.02, 0.03, 0.04], - # "M": [2, 3, 4, 5, 6, 7], +scenario_dict = { # Fill in to do multiple experiments + # # Example 1: Single loss function sweep (original behavior) + # "SNR_sweep_rmspe": { + # "parameter": "snr", + # "values": [-10, -5, 0, 5, 10, 15, 20, 25, 30], + # "fixed_params": {"T": 100}, + # "plot_config": { + # "title": "RMSPE vs SNR (T=100, RMSPE Loss)", + # "x_label": "SNR (dB)", + # "y_label": "RMSPE (degrees)", + # "save_name": "rmspe_vs_snr_rmspe_loss" + # } + # }, + + # # Example 2: Multi-loss function sweep + # "SNR_sweep_multi_loss": { + # "parameter": "snr", + # "values": [-5, 0, 5, 10, 15, 20, 25, 30], + # "loss_functions": ["rmspe", "spectrum", "unsupervised"], # This enables multi-loss mode + # "fixed_params": {"T": 100}, + # "plot_config": { + # "title": "RMSPE vs SNR (T=100) - Loss Function Comparison", + # "x_label": "SNR (dB)", + # "y_label": "RMSPE (degrees)", + # "save_name": "rmspe_vs_snr_multi_loss" + # } + # }, + + # # Example 3: T sweep with multi-loss + # "T_sweep_multi_loss": { + # "parameter": "T", + # "values": [20, 30, 50, 70, 100, 150, 200], + # "loss_functions": ["rmspe", "spectrum", "unsupervised"], + # "fixed_params": {"snr": 30}, + # "plot_config": { + # "title": "RMSPE vs T (SNR=30dB) - Loss Function Comparison", + # "x_label": "T (snapshots)", + # "y_label": "RMSPE (degrees)", + # "save_name": "rmspe_vs_T_multi_loss" + # } + # } } simulation_commands = { - "SAVE_TO_FILE": False, - "CREATE_DATA": False, - "SAVE_DATASET": True, - "LOAD_MODEL": False, - "TRAIN_MODEL": True, - "SAVE_MODEL": True, - "EVALUATE_MODE": True, - "PLOT_RESULTS": True, # if True, the learning curves will be plotted - "PLOT_LOSS_RESULTS": True, # if True, the RMSE results of evaluation will be plotted - "PLOT_ACC_RESULTS": True, # if True, the accuracy results of evaluation will be plotted - "SAVE_PLOTS": False, # if True, the plots will be saved to the results folder + "create_data": True, + "save_data": False, # Save data after creation + "plot_results": True, # Plot data after creation + "multi_loss_comparison": True, # Enable multi-loss spectrum and learned parameters comparison, if multiple experiments then need to be False + "spectrum_loss_functions": ["rmspe", "spectrum", "unsupervised"], # Loss functions to compare + "data_loading_path": "datasets/N:16_M:5_T:100_snr:10_location_pert_boundary:0.25_gain_perturbation_var:0.36_seed:42/03_06_2025_15_06/data.pkl" + # This is the path to the data file, ONLY USED if CREATE_DATA is False! + # By now, this gets set manually. } system_model_params = { - "N": 15, # number of antennas - "M": 2, # number of sources + "N": 16, # number of antennas + "M": 5, # number of sources "T": 100, # number of snapshots - "snr": 10, # if defined, values in scenario_dict will be ignored - "field_type": "near", # Near, Far - "signal_type": "Narrowband", # Narrowband, broadband - "signal_nature": "coherent", # if defined, values in scenario_dict will be ignored - "eta": 0.0, # steering vector uniform error variance with respect to the wavelength. + "snr": 10, # if defined, values in scenario_dict will be ignored "bias": 0, # steering vector bias error "sv_noise_var": 0.0, # steering vector addative gaussian error noise variance "doa_range": 60, # The range of the DOA values [-doa_range, doa_range] "doa_resolution": .5, # The resolution of the DOA values in degrees - "max_range_ratio_to_limit": 0.5, # The ratio of the maximum range in respect to the Fraunhofer distance - "range_resolution": 1, # The resolution of the range values in meters - "wavelength": 1, # The carrier wavelength of the signal in meters -} -model_config = { - "model_type": "SubspaceNet", # SubspaceNet, DCD-MUSIC, DeepCNN, TransMUSIC, DR_MUSIC - "model_params": {} + "wavelength": 1, # The carrier wavelength of the signal in meters, 1 can be fine for research, + # 0.06 is for wifi 5 GHz for example. + + "location_perturbation": "wavelength/4", # The boundaries of the location perturbation in meters, + # insert any valid float between 0 and wavelength/4 or "wavelength/n" to use with reference to the wavelength + + "gain_perturbation_var": 0.36, # The variance of the gain perturbation + "seed": 42, # Seed for reproducibility + ############################### Fixed for now ################################## + "field_type": "Far", # Near, Far + "signal_type": "Narrowband", # Narrowband, broadband + "signal_nature": "non-coherent" # if defined, values in scenario_dict will be ignored } -if model_config.get("model_type") == "SubspaceNet": - model_config["model_params"]["diff_method"] = "music_2D" # esprit, music_1D, music_2D, beamformer - model_config["model_params"]["train_loss_type"] = "music_spectrum" # music_spectrum, rmspe, beamformerloss - model_config["model_params"]["tau"] = 8 - model_config["model_params"]["field_type"] = "Near" # Far, Near - model_config["model_params"]["regularization"] = None # aic, mdl, threshold, None - model_config["model_params"]["variant"] = "small" # big, small - model_config["model_params"]["norm_layer"] = True - model_config["model_params"]["batch_norm"] = False -elif model_config.get("model_type") == "DCD-MUSIC": - model_config["model_params"]["tau"] = 8 - model_config["model_params"]["diff_method"] = ("esprit", "music_1D") # ("esprit", "music_1D") - model_config["model_params"]["train_loss_type"] = ("rmspe", "rmspe") # ("rmspe", "rmspe"), ("rmspe", - # "music_spectrum"), ("music_spectrum", "rmspe") - model_config["model_params"]["regularization"] = None # aic, mdl, threshold, None - model_config["model_params"]["variant"] = "small" # big, small - model_config["model_params"]["norm_layer"] = True +model_config = \ +{ + "model_type": "diffMUSIC", # or "MUSIC" + + # Case 1: Fixed integer window size (original behavior) + # "softmax_window_size": 21, + + # Case 2: Relative window size (float between 0-1) + "softmax_window_size": 0.1, # % of angle grid length + + # Case 3: Window size optimization (array/list of values) + # "softmax_window_size": np.arange(0.01, 0.4, 0.01), # Relative sizes: [0.2, 0.25, 0.3, 0.35] + # "softmax_window_size": [15, 21, 27, 33], # Absolute sizes + # "softmax_window_size": [0.25, 21, 0.35, 27], # Mixed relative and absolute +} -elif model_config.get("model_type") == "DeepCNN": - model_config["model_params"]["grid_size"] = 361 training_params = { - "samples_size": 4096, - "train_test_ratio": 0.1, - "training_objective": "angle, range", # angle, range, source_estimation - "batch_size": 128, - "epochs": 50, + # "batch_size": 128, # Note: This is legacy parameter, actual batch size is handled by snapshots + "epochs": 100, + "loss_type": "spectrum", # rmspe, spectrum, unsupervised "optimizer": "Adam", # Adam, SGD "scheduler": "ReduceLROnPlateau", # StepLR, ReduceLROnPlateau "learning_rate": 0.001, - "weight_decay": 1e-9, "step_size": 50, - "gamma": 0.5, - "true_doa_train": None, # if set, this doa will be set to all samples in the train dataset - "true_range_train": None, # if set, this range will be set to all samples in the train dataset - "true_doa_test": None, # if set, this doa will be set to all samples in the test dataset - "true_range_test": None, # if set, this range will be set to all samples in the train dataset - "use_wandb": False, - "simulation_name": None, -} -evaluation_params = { - "models": { - # "TransMUSIC": { - # "model_name": "TransMUSIC", - # }, - # "DCD-MUSIC": { - # "model_name": "DCD-MUSIC", - # "tau": 8, - # "diff_method": ("esprit", "music_1d"), - # "regularization": None, - # }, - # "DCD-MUSIC_V2": { - # "model_name": "DCD-MUSIC", - # "tau": 8, - # "diff_method": ("esprit", "music_1d"), - # "regularization": "aic", - # "variant": "big" - # }, - # "NFSubspaceNet": { - # "model_name": "SubspaceNet", - # "tau": 8, - # "diff_method": "music_2D", - # "train_loss_type": "music_spectrum", - # "field_type": "near", - # "regularization": None, - # }, - # "NFSubspaceNet_V2": { - # "model_name": "SubspaceNet", - # "tau": 8, - # "diff_method": "music_2D", - # "train_loss_type": "music_spectrum", - # "field_type": "near", - # "regularization": "aic", - # "variant": "big", - # }, - }, - "augmented_methods": [ - # ("SubspaceNet", "beamformer", {"tau": 8, "diff_method": "music_2D", "train_loss_type": "music_spectrum", "field_type": "near"}), - # ("SubspaceNet", "beamformer", {"tau": 8, "diff_method": "esprit", "train_loss_type": "rmspe", "field_type": "far"}), - # ("SubspaceNet", "esprit", {"tau": 8, "diff_method": "esprit", "train_loss_type": "rmspe", "field_type": "far"}), - ], - "subspace_methods": [ - # "CCRB", - "2D-MUSIC", - "Beamformer", - # "CS_Estimator", - # "ESPRIT", - # "1D-MUSIC", - # "Root-MUSIC", - # "TOPS", - ] + "weight_decay": 0.0, + "use_wandb": False } @@ -179,39 +172,31 @@ def parse_arguments(): parser.add_argument('-ft', '--field_type', type=str, help='Field type, far or near field.', default=None) parser.add_argument('-sn', '--signal_nature', type=str, help='Signal nature; non-coherent or coherent', default=None) parser.add_argument('-wav', '--wavelength', type=float, help='Wavelength of the signal in meters', default=None) + parser.add_argument("--location_perturbation", type=float, help="Location perturbation variance", default=None) + parser.add_argument("--gain_perturbation_var", type=float, help="Gain perturbation variance", default=None) + parser.add_argument("--seed", type=int, help="Seed for reproducibility", default=None) - parser.add_argument('-mt', '--model_type', type=str, help='Model type; SubspaceNet, DCD-MUSIC, DeepCNN, TransMUSIC, DR_MUSIC', default=None) - parser.add_argument('-reg', '--regularization', type=str, help='Regularization method for SubspaceNet of DCD', default=model_config["model_params"]["regularization"]) - parser.add_argument('-tau', '--tau', type=int, help='Tau value for SubspaceNet or DCD-MUSIC', default=model_config["model_params"].get("tau")) - parser.add_argument("-v", "--variant", type=str, help="Variant of the SubspaceNet model; big, small", default=model_config["model_params"].get("variant")) + parser.add_argument('-mt', '--model_type', type=str, help='Model type; diffMUSIC, MUSIC', default=None) + parser.add_argument('-lt', '--loss_type', type=str, help='Loss type; rmspe, spectrum, unsupervised', default=None) - parser.add_argument('-ss', '--samples_size', type=int, help='Samples size', default=None) - parser.add_argument('-ttr', '--train_test_ratio', type=float, help='Train test ratio', default=None) - parser.add_argument('-to', '--training_objective', type=str, help='Training objective; angle, range or angle, range.', default=None) - parser.add_argument('-bs', '--batch_size', type=int, help='Batch size', default=None) + parser.add_argument('-bs', '--batch_size', type=int, help='Batch size (legacy parameter)', default=None) parser.add_argument('-ep', '--epochs', type=int, help='Number of epochs', default=None) parser.add_argument('-op', '--optimizer', type=str, help='Optimizer; Adam, SGD', default=None) parser.add_argument('-sch', '--scheduler', type=str, help='Scheduler; StepLR, ReduceLROnPlateau', default=None) parser.add_argument('-lr', '--learning_rate', type=float, help='Learning rate', default=None) parser.add_argument('-wd', '--weight_decay', type=float, help='Weight decay', default=None) parser.add_argument('-step', '--step_size', type=int, help='Step size', default=None) - parser.add_argument('-g', '--gamma', type=float, help='Gamma', default=None) - parser.add_argument('-w', '--wandb', action="store_true", help='Use wandb', default=training_params["use_wandb"]) - - parser.add_argument('-t', '--train', action="store_true", help='Train model', default=simulation_commands["TRAIN_MODEL"]) - parser.add_argument('-no_t', "--no_train", action="store_false", help='Do not train model', dest='train') - parser.add_argument('-e', '--eval', action="store_true", help='Evaluate model', default=simulation_commands["EVALUATE_MODE"]) - parser.add_argument('-c', '--create', action="store_true", help='create a new dataset', default=simulation_commands["CREATE_DATA"]) - parser.add_argument('-sv', '--save', action='store_true', help="save dataset", default=simulation_commands["SAVE_DATASET"]) + parser.add_argument('-w', '--wandb', action="store_true", help='Use wandb') + parser.add_argument('-c', '--create', action="store_true", help='create a new dataset') return parser.parse_args() if __name__ == "__main__": - # torch.set_printoptions(precision=12) - args = parse_arguments() + + # Update system model parameters from command line arguments if args.snr is not None: system_model_params["snr"] = args.snr if args.number_of_sensors is not None: @@ -234,23 +219,25 @@ def parse_arguments(): if args.wavelength is not None: system_model_params["wavelength"] = args.wavelength + if args.location_perturbation is not None: + system_model_params["location_perturbation"] = args.location_perturbation + elif (isinstance(system_model_params["location_perturbation"], str)) and ("wavelength" in system_model_params["location_perturbation"]): + int_idx = system_model_params["location_perturbation"].find("/") + 1 + system_model_params["location_perturbation"] = (system_model_params["wavelength"] / + float(system_model_params["location_perturbation"][int_idx:int_idx + 1])) + + if args.gain_perturbation_var is not None: + system_model_params["gain_perturbation_var"] = args.gain_perturbation_var + if args.seed is not None: + system_model_params["seed"] = args.seed + + # Update model configuration if args.model_type is not None: - warnings.warn("Please make sure to configure the model parameters in the script.") model_config["model_type"] = args.model_type - if model_config["model_type"] == "SubspaceNet": - model_config["model_params"]["regularization"] = None if args.regularization == "None" else args.regularization - model_config["model_params"]["tau"] = args.tau - model_config["model_params"]["variant"] = args.variant - if args.samples_size is not None: - training_params["samples_size"] = args.samples_size - if args.train_test_ratio is not None: - training_params["train_test_ratio"] = args.train_test_ratio - if args.training_objective is not None: - if args.training_objective.startswith("angle,range"): - training_params["training_objective"] = "angle, range" - else: - training_params["training_objective"] = args.training_objective + # Update training parameters + if args.loss_type is not None: + training_params["loss_type"] = args.loss_type if args.batch_size is not None: training_params["batch_size"] = args.batch_size if args.epochs is not None: @@ -265,22 +252,50 @@ def parse_arguments(): training_params["weight_decay"] = args.weight_decay if args.step_size is not None: training_params["step_size"] = args.step_size - if args.gamma is not None: - training_params["gamma"] = args.gamma - if args.wandb is not None: - training_params["use_wandb"] = args.wandb - simulation_commands["TRAIN_MODEL"] = args.train - simulation_commands["EVALUATE_MODE"] = args.eval + if args.wandb: + training_params["use_wandb"] = args.wandb + if args.create: + simulation_commands["create_data"] = args.create + + simulation_commands["load_data"] = not simulation_commands["create_data"] - simulation_commands["CREATE_DATA"] = args.create - simulation_commands["SAVE_DATASET"] = args.save + # Validate location perturbation + if system_model_params["location_perturbation"] > system_model_params["wavelength"] / 4: + raise ValueError("Location perturbation should be less than wavelength/4, " + "This may result in order switching between neighboring array sensors.") start = time.time() - loss = run_simulation(simulation_commands=simulation_commands, - system_model_params=system_model_params, - model_config=model_config, - training_params=training_params, - evaluation_params=evaluation_params, - scenario_dict=scenario_dict) + results = run_simulation(simulation_commands=simulation_commands, + system_model_params=system_model_params, + model_config=model_config, + training_params=training_params, + scenario_dict=scenario_dict) print("Total time: ", time.time() - start) + + # Print final results summary + if simulation_commands.get("multi_loss_spectrum_comparison", False): + print(f"\nMulti-Loss Spectrum Comparison Completed!") + print(f"Loss functions compared: {simulation_commands['spectrum_loss_functions']}") + print(f"Results saved and plots generated.") + elif isinstance(results, dict) and 'rmspe' in results: + print(f"\nFinal Results:") + print(f"RMSPE: {results['rmspe']:.6f} degrees") + if 'steering_matrix_mse' in results: + print(f"Steering Matrix MSE: {results['steering_matrix_mse']:.6f}") # Add this line + if 'estimated_angles' in results and 'true_angles' in results: + print(f"True angles: {np.rad2deg(results['true_angles'])}") + print(f"Estimated angles: {np.rad2deg(results['estimated_angles'])}") + if 'final_train_loss' in results: + print(f"Final training loss: {results['final_train_loss']:.6f}") + if "learned_antenna_positions" in results.keys() and "learned_antennas_gains" in results.keys(): + if model_config["model_type"].lower() == "music": + print("These should be the tandard one:") + print(f"Learned antennas positions: {np.round(results['learned_antenna_positions'], 4)}") + print("Compared to the physical array positions:") + print(f"Physical antennas positions: {results['physical_array']}") + print(f"Learned antennas gains: {results['learned_antennas_gains']}") + print(f"Learned antennas gains phase: {np.round(np.angle(results['learned_antennas_gains']), 4)}") + print("Compared to the physical antennas gains:") + print(f"Physical antennas gains: {results['physical_antennas_gains']}") + print(f"Physical antennas gains phase: {np.round(np.angle(results['physical_antennas_gains']), 4)}") \ No newline at end of file diff --git a/run_simulation.py b/run_simulation.py index 24d5c81..2a19026 100644 --- a/run_simulation.py +++ b/run_simulation.py @@ -11,214 +11,281 @@ """ # Imports import sys -from src.data_handler import * -from src.training import * -from src.plotting import * -from src.evaluation import evaluate from pathlib import Path -from src.models import ModelGenerator -from src.system_model import SystemModel, SystemModelParams -from src.utils import set_unified_seed, initialize_data_paths, print_loss_results_from_simulation +from src.signal_creation import Samples, SystemModelParams +import src.utils as utils +from datetime import datetime +import torch +import numpy as np +from doa_runner import DoARunner -def __run_simulation(**kwargs): +def run_single_simulation(**kwargs): SIMULATION_COMMANDS = kwargs["simulation_commands"] SYSTEM_MODEL_PARAMS = kwargs["system_model_params"] MODEL_CONFIG = kwargs["model_config"] TRAINING_PARAMS = kwargs["training_params"] - EVALUATION_PARAMS = kwargs["evaluation_params"] - save_to_file = SIMULATION_COMMANDS["SAVE_TO_FILE"] # Saving results to file or present them over CMD - create_data = SIMULATION_COMMANDS["CREATE_DATA"] # Creating new dataset - load_model = SIMULATION_COMMANDS["LOAD_MODEL"] # Load specific model for training - train_model = SIMULATION_COMMANDS["TRAIN_MODEL"] # Applying training operation - save_model = SIMULATION_COMMANDS["SAVE_MODEL"] # Saving tuned model - evaluate_mode = SIMULATION_COMMANDS["EVALUATE_MODE"] # Evaluating desired algorithms - plot_mode = SIMULATION_COMMANDS["PLOT_RESULTS"] # Plotting results - save_plots = SIMULATION_COMMANDS["SAVE_PLOTS"] # Saving plots - load_data = not create_data # Loading data from exist dataset + plot_results = SIMULATION_COMMANDS["plot_results"] + create_data = SIMULATION_COMMANDS["create_data"] # Creating new dataset + save_data = SIMULATION_COMMANDS["save_data"] # Save created data to file + load_data = SIMULATION_COMMANDS["load_data"] # Load specific model for training + print("Running simulation...") - if train_model: - print("Training model - ", MODEL_CONFIG.get('model_type')) - print("Training objective - ", TRAINING_PARAMS.get('training_objective')) now = datetime.now() plot_path = Path(__file__).parent / "plots" plot_path.mkdir(parents=True, exist_ok=True) dt_string_for_save = now.strftime("%d_%m_%Y_%H_%M") - # torch.set_printoptions(precision=12) - + # Initialize seed - set_unified_seed() + utils.set_unified_seed(SYSTEM_MODEL_PARAMS["seed"]) + # Define system model parameters - unify all parameters + system_model_params = SystemModelParams(**SYSTEM_MODEL_PARAMS, **MODEL_CONFIG, **TRAINING_PARAMS, **SIMULATION_COMMANDS) + + # Add dt_string for wandb project naming + system_model_params.dt_string_for_save = dt_string_for_save + # Initialize paths - datasets_path, simulations_path, saving_path = initialize_data_paths(Path(__file__).parent / "data") + #TODO: change the paths to be per model_connfig["model_type"] + data_saving_path, results_path = utils.initialize_paths(Path(__file__).parent, system_model_params) + data_loading_path = SIMULATION_COMMANDS["data_loading_path"] #ONLY USED if CREATE_DATA is False! + + # Prepare data dictionary + data_dict = {} - # Saving simulation scores to external file - suffix = "" - if train_model: - suffix += f"_train_{MODEL_CONFIG.get('model_type')}_{TRAINING_PARAMS.get('training_objective')}" - suffix += (f"_{SYSTEM_MODEL_PARAMS['signal_nature']}_SNR_{SYSTEM_MODEL_PARAMS['snr']}_T_{SYSTEM_MODEL_PARAMS['T']}" - f"_eta{SYSTEM_MODEL_PARAMS['eta']}.txt") + if create_data: + signals_creator = Samples(system_model_params) + signals_creator.set_labels(None) # creates random angles + measurements, signals, steering_mat, noise = signals_creator.samples_creation() + true_angles = signals_creator.get_labels() + physical_array = signals_creator.get_array() + physical_antennas_gains = signals_creator.get_antenna_gains() + + # Pack data into dictionary + data_dict = { + 'measurements': measurements, + 'signals': signals, + 'steering_matrix': steering_mat, + 'noise': noise, + 'true_angles': true_angles, + 'physical_array': physical_array, + 'physical_antennas_gains': physical_antennas_gains + } + + # Save the created data under the data_path + if save_data: + utils.save_data_to_file(data_saving_path, measurements, signals, steering_mat, + noise, true_angles, physical_array, physical_antennas_gains, system_model_params) + else: + # Load data from file + loaded_data = utils.load_data_from_file(data_loading_path) + measurements, signals, steering_mat, noise, true_angles, physical_array, physical_antennas_gains, system_model_params = loaded_data + + # Add dt_string for wandb project naming (in case of loaded data) + system_model_params.dt_string_for_save = dt_string_for_save + + # Pack loaded data into dictionary + data_dict = { + 'measurements': measurements, + 'signals': signals, + 'steering_matrix': steering_mat, + 'noise': noise, + 'true_angles': true_angles, + 'physical_array': physical_array, + 'physical_antennas_gains': physical_antennas_gains + } - if save_to_file: - orig_stdout = sys.stdout - file_path = ( - simulations_path / "results" / "scores" / Path(dt_string_for_save + suffix) - ) - sys.stdout = open(file_path, "w") - # Define system model parameters - system_model_params = ( - SystemModelParams() - .set_parameter("N", SYSTEM_MODEL_PARAMS["N"]) - .set_parameter("M", SYSTEM_MODEL_PARAMS["M"]) - .set_parameter("T", SYSTEM_MODEL_PARAMS["T"]) - .set_parameter("snr", SYSTEM_MODEL_PARAMS["snr"]) - .set_parameter("field_type", SYSTEM_MODEL_PARAMS["field_type"]) - .set_parameter("signal_nature", SYSTEM_MODEL_PARAMS["signal_nature"]) - .set_parameter("signal_type", SYSTEM_MODEL_PARAMS["signal_type"]) - .set_parameter("eta", SYSTEM_MODEL_PARAMS["eta"]) - .set_parameter("bias", SYSTEM_MODEL_PARAMS["bias"]) - .set_parameter("sv_noise_var", SYSTEM_MODEL_PARAMS["sv_noise_var"]) - .set_parameter("doa_range", SYSTEM_MODEL_PARAMS["doa_range"]) - .set_parameter("doa_resolution", SYSTEM_MODEL_PARAMS["doa_resolution"]) - .set_parameter("max_range_ratio_to_limit", SYSTEM_MODEL_PARAMS["max_range_ratio_to_limit"]) - .set_parameter("range_resolution", SYSTEM_MODEL_PARAMS["range_resolution"]) - .set_parameter("wavelength", SYSTEM_MODEL_PARAMS["wavelength"]) - ) + # Create and run DoA algorithm + print(f"Running {system_model_params.model_type} algorithm...") + doa_runner = DoARunner(system_model_params, data_dict) + results = doa_runner.run() + if plot_results: + doa_runner.plot_graphs(results_path) + + # Save results + results_file = results_path / "results.pkl" + utils.save_data_to_file(results_path, results, system_model_params) + + # # Save trained model if training was performed + # if doa_runner._needs_training(): + # model_file = results_path / "trained_model.pth" + # doa_runner.save_model(model_file) + + print(f"Results saved to: {results_path}") + print(f"RMSPE: {results.get('rmspe', 'N/A'):.6f}") + return results + +def __run_parameter_sweeps(**kwargs): + """ + Run parameter sweeps based on scenario_dict configuration + + Args: + **kwargs: Contains scenario_dict and other configuration + + Returns: + Dictionary with results from all sweeps + """ + from src.multi_experiment_runner import MultiExperimentRunner + + scenario_dict = kwargs["scenario_dict"] + all_sweep_results = {} + + print(f"Starting parameter sweeps for {len(scenario_dict)} configurations...") + + for sweep_name, sweep_config in scenario_dict.items(): + print(f"\n{'='*60}") + print(f"Running sweep: {sweep_name}") + print(f"{'='*60}") + + try: + # Create and run sweep + runner = MultiExperimentRunner(sweep_config, kwargs) + sweep_results = runner.run_sweep() + all_sweep_results[sweep_name] = sweep_results + + print(f"Sweep '{sweep_name}' completed successfully") + print(f"Mean RMSPE: {sweep_results['aggregated_metrics']['mean_rmspe']:.6f}") + + except Exception as e: + print(f"Error in sweep '{sweep_name}': {e}") + all_sweep_results[sweep_name] = {'error': str(e)} + + print(f"\nAll parameter sweeps completed!") + return all_sweep_results + +def run_multi_loss_comparison(**kwargs): + """ + Run multi-loss spectrum and learned parameters comparison for the same data with different loss functions + + Args: + **kwargs: Contains simulation_commands, system_model_params, model_config, training_params + + Returns: + Dictionary with multi-loss spectrum comparison results + """ + SIMULATION_COMMANDS = kwargs["simulation_commands"] + SYSTEM_MODEL_PARAMS = kwargs["system_model_params"] + MODEL_CONFIG = kwargs["model_config"] + TRAINING_PARAMS = kwargs["training_params"] + + create_data = SIMULATION_COMMANDS["create_data"] + save_data = SIMULATION_COMMANDS["save_data"] + plot_results = SIMULATION_COMMANDS["plot_results"] + loss_functions = SIMULATION_COMMANDS["spectrum_loss_functions"] + + print("=== MULTI-LOSS SPECTRUM COMPARISON MODE ===") + print(f"Comparing loss functions: {loss_functions}") - # Define samples size - samples_size = TRAINING_PARAMS["samples_size"] # Overall dateset size - train_test_ratio = TRAINING_PARAMS["train_test_ratio"] # training and testing datasets ratio - # Sets simulation filename - # simulation_filename = get_simulation_filename(system_model_params=system_model_params, - # model_config=model_config) - # Print new simulation intro - print("------------------------------------") - print("---------- New Simulation ----------") - print("------------------------------------") + now = datetime.now() + plot_path = Path(__file__).parent / "plots" + plot_path.mkdir(parents=True, exist_ok=True) + dt_string_for_save = now.strftime("%d_%m_%Y_%H_%M") + + # Initialize seed + utils.set_unified_seed(SYSTEM_MODEL_PARAMS["seed"]) - if load_data: - if train_model: - try: - start = time.time() - train_dataset = load_datasets( - system_model_params=system_model_params, - samples_size=samples_size, - datasets_path=datasets_path, - is_training=True, - ) - print(f"Load the data took {time.time() - start} sec") - except Exception as e: - print(e) - print("#############################################") - print("load_datasets: Error loading train dataset") - print("#############################################") - create_data = True - load_data = False - if evaluate_mode: - try: - generic_test_dataset = load_datasets( - system_model_params=system_model_params, - samples_size=samples_size * train_test_ratio, - datasets_path=datasets_path, - is_training=False, - ) - except Exception as e: - print(e) - print("#############################################") - print("load_datasets: Error loading test dataset") - print("#############################################") - create_data = True - load_data = False - if create_data and not load_data: - # Define which datasets to generate - print("Creating Data...") - # init sample model - samples_model = Samples(system_model_params) - if train_model: - # Generate training dataset - start = time.time() - train_dataset, _ = create_dataset( - samples_model=samples_model, - samples_size=samples_size, - save_datasets=SIMULATION_COMMANDS["SAVE_DATASET"], - datasets_path=datasets_path, - true_doa=TRAINING_PARAMS["true_doa_train"], - true_range=TRAINING_PARAMS["true_range_train"], - phase="train", - ) - print(f"Create the data took {time.time() - start} sec") - if evaluate_mode: - # Generate test dataset - generic_test_dataset, _ = create_dataset( - samples_model=samples_model, - samples_size=int(train_test_ratio * samples_size), - save_datasets=SIMULATION_COMMANDS["SAVE_DATASET"], - datasets_path=datasets_path, - true_doa=TRAINING_PARAMS["true_doa_test"], - true_range=TRAINING_PARAMS["true_range_test"], - phase="test", - ) + # Define system model parameters - unify all parameters + system_model_params = SystemModelParams(**SYSTEM_MODEL_PARAMS, **MODEL_CONFIG, **TRAINING_PARAMS, **SIMULATION_COMMANDS) + + # Add dt_string for naming + system_model_params.dt_string_for_save = dt_string_for_save + + # Initialize paths - modify to indicate multi-loss comparison + data_saving_path, results_path = utils.initialize_paths(Path(__file__).parent, system_model_params) + # Update results path name to indicate multi-loss comparison + results_path = results_path.parent / (results_path.name + "_multi_loss_spectrum") + results_path.mkdir(parents=True, exist_ok=True) + + data_loading_path = SIMULATION_COMMANDS["data_loading_path"] - if train_model: - # Generate model configuration - model_config = ( - ModelGenerator() - .set_model_type(MODEL_CONFIG.get("model_type")) - .set_system_model(system_model_params) - .set_model_params(MODEL_CONFIG.get("model_params")) - .set_model() - ) + # Prepare data dictionary + data_dict = {} - trainingparams = TrainingParamsNew(learning_rate=TRAINING_PARAMS["learning_rate"], - weight_decay=TRAINING_PARAMS["weight_decay"], - epochs=TRAINING_PARAMS["epochs"], - optimizer=TRAINING_PARAMS["optimizer"], - step_size=TRAINING_PARAMS["step_size"], - gamma=TRAINING_PARAMS["gamma"], - training_objective=TRAINING_PARAMS["training_objective"], - scheduler=TRAINING_PARAMS["scheduler"], - batch_size=TRAINING_PARAMS["batch_size"], - simulation_name=TRAINING_PARAMS["simulation_name"], - ) - train_dataloader, valid_dataloader = train_dataset.get_dataloaders(batch_size=TRAINING_PARAMS["batch_size"]) - trainer = Trainer(model=model_config.model, training_params=trainingparams, show_plots=True) - model = trainer.train(train_dataloader, valid_dataloader, - use_wandb=TRAINING_PARAMS["use_wandb"], - save_final=save_model, load_model=load_model) + if create_data: + signals_creator = Samples(system_model_params) + signals_creator.set_labels(None) # creates random angles + measurements, signals, steering_mat, noise = signals_creator.samples_creation() + true_angles = signals_creator.get_labels() + physical_array = signals_creator.get_array() + physical_antennas_gains = signals_creator.get_antenna_gains() + + # Pack data into dictionary + data_dict = { + 'measurements': measurements, + 'signals': signals, + 'steering_matrix': steering_mat, + 'noise': noise, + 'true_angles': true_angles, + 'physical_array': physical_array, + 'physical_antennas_gains': physical_antennas_gains + } + + # Save the created data if requested + if save_data: + utils.save_data_to_file(data_saving_path, measurements, signals, steering_mat, + noise, true_angles, physical_array, physical_antennas_gains, system_model_params) + else: + # Load data from file + loaded_data = utils.load_data_from_file(data_loading_path) + measurements, signals, steering_mat, noise, true_angles, physical_array, physical_antennas_gains, system_model_params = loaded_data + + # Add dt_string for naming (in case of loaded data) + system_model_params.dt_string_for_save = dt_string_for_save + + # Pack loaded data into dictionary + data_dict = { + 'measurements': measurements, + 'signals': signals, + 'steering_matrix': steering_mat, + 'noise': noise, + 'true_angles': true_angles, + 'physical_array': physical_array, + 'physical_antennas_gains': physical_antennas_gains + } - # Evaluation stage - if evaluate_mode: - if not train_model: - model = None - # Define loss measure for evaluation - if isinstance(system_model_params.M, int): - generic_test_dataset = torch.utils.data.DataLoader(generic_test_dataset, - batch_size=100, - shuffle=False) + # Create DoA runner + print(f"Creating DoA runner for multi-loss comparison...") + doa_runner = DoARunner(system_model_params, data_dict) + + # Run multi-loss spectrum comparison + multi_loss_results = doa_runner.run_multi_loss_spectrum_comparison(loss_functions) + + # Plot comparisons if requested + if plot_results: + print("Generating plots...") + + # Plot spectrum comparison + doa_runner.plot_multi_loss_spectrum_comparison(multi_loss_results, results_path) + + # Plot learned parameters comparison (only for diffMUSIC) + if system_model_params.model_type.lower() == "diffmusic": + print("Generating learned parameters comparison plot...") + doa_runner.plot_multi_loss_learned_parameters(multi_loss_results, results_path) + else: + print("Learned parameters plotting skipped (only available for diffMUSIC)") + + # Save results + results_file = results_path / "multi_loss_results.pkl" + utils.save_data_to_file(results_path, multi_loss_results, system_model_params) + + print(f"Multi-loss comparison results saved to: {results_path}") + + # Print summary + print(f"\n=== MULTI-LOSS COMPARISON SUMMARY ===") + for loss_func in loss_functions: + result = multi_loss_results['results'].get(loss_func, {}) + if 'error' not in result: + rmspe = result.get('rmspe', float('inf')) + steering_mse = result.get('steering_matrix_mse', float('nan')) + print(f"{loss_func.upper()}: RMSPE = {rmspe:.6f}°, Steering MSE = {steering_mse:.6f}") else: - batch_sampler_test = SameLengthBatchSampler(generic_test_dataset, batch_size=100) - generic_test_dataset = torch.utils.data.DataLoader(generic_test_dataset, - collate_fn=collate_fn, - batch_sampler=batch_sampler_test, - shuffle=False) + print(f"{loss_func.upper()}: ERROR - {result.get('error', 'Unknown error')}") + + print(f"Total execution time: {multi_loss_results['total_execution_time']:.2f}s") + + return multi_loss_results - # Evaluate DNN models, augmented and subspace methods - loss = evaluate( - generic_test_dataset=generic_test_dataset, - system_model_params=system_model_params, - models=EVALUATION_PARAMS["models"], - augmented_methods=EVALUATION_PARAMS["augmented_methods"], - subspace_methods=EVALUATION_PARAMS["subspace_methods"], - model_tmp=model - ) - # plt.show() - print("END OF EVALUATION") - if save_to_file: - sys.stdout.close() - sys.stdout = orig_stdout - return loss - return def run_simulation(**kwargs): @@ -229,57 +296,22 @@ def run_simulation(**kwargs): - SNR: a list of SNR values to be tested - T: a list of number of snapshots to be tested - eta: a list of steering vector error values to be tested + - M: a list of number of sources to be tested """ - if kwargs["scenario_dict"] == {}: - loss = __run_simulation(**kwargs) - return loss - loss_dict = {} - default_snr = kwargs["system_model_params"]["snr"] - default_T = kwargs["system_model_params"]["T"] - default_eta = kwargs["system_model_params"]["eta"] - default_m = kwargs["system_model_params"]["M"] - for key, value in kwargs["scenario_dict"].items(): - if key == "SNR": - loss_dict["SNR"] = {snr: None for snr in value} - print(f"Testing SNR values: {value}") - for snr in value: - kwargs["system_model_params"]["snr"] = snr - loss = __run_simulation(**kwargs) - loss_dict["SNR"][snr] = loss - kwargs["system_model_params"]["snr"] = default_snr - if key == "T": - loss_dict["T"] = {T: None for T in value} - print(f"Testing T values: {value}") - for T in value: - kwargs["system_model_params"]["T"] = T - loss = __run_simulation(**kwargs) - loss_dict["T"][T] = loss - kwargs["system_model_params"]["T"] = default_T - if key == "eta": - loss_dict["eta"] = {eta: None for eta in value} - print(f"Testing eta values: {value}") - for eta in value: - kwargs["system_model_params"]["eta"] = eta - loss = __run_simulation(**kwargs) - loss_dict["eta"][eta] = loss - kwargs["system_model_params"]["eta"] = default_eta - if key == "M": - loss_dict["M"] = {m: None for m in value} - print(f"Testing M values: {value}") - for m in value: - kwargs["system_model_params"]["M"] = m - loss = __run_simulation(**kwargs) - loss_dict["M"][m] = loss - kwargs["system_model_params"]["M"] = default_m - if None not in list(next(iter(loss_dict.values())).values()): - print_loss_results_from_simulation(loss_dict) - if kwargs["simulation_commands"]["PLOT_LOSS_RESULTS"]: - plot_results(loss_dict, kwargs["system_model_params"]["field_type"], - plot_acc=kwargs["simulation_commands"]["PLOT_ACC_RESULTS"], - save_to_file=kwargs["simulation_commands"]["SAVE_PLOTS"]) + #TODO: check if anything is missing for the scenario_dict option once needed. + scenario_dict = kwargs["scenario_dict"] + simulation_commands = kwargs["simulation_commands"] + + + if simulation_commands.get("multi_loss_comparison", False): # multi-loss spectrum comparison + return run_multi_loss_comparison(**kwargs) + elif not scenario_dict: # Single experiment (current behavior) + return run_single_simulation(**kwargs) + else: # Multi-experiment sweep + return __run_parameter_sweeps(**kwargs) + - return loss_dict if __name__ == "__main__": - now = datetime.now() + now = datetime.now() \ No newline at end of file diff --git a/src/diffmusic.py b/src/diffmusic.py new file mode 100644 index 0000000..f6b92db --- /dev/null +++ b/src/diffmusic.py @@ -0,0 +1,573 @@ +""" +diffMUSIC: Differentiable MUSIC for DoA Estimation with Hardware Impairment Learning + +This module implements the diffMUSIC algorithm as described in the paper: +"Physically Parameterized Differentiable MUSIC for DoA Estimation with Uncalibrated Arrays" + +Key Features: +- Learnable antenna positions and complex gains +- Differentiable steering matrix computation +- Softmax-based peak finding for end-to-end learning +- Support for both supervised and unsupervised learning +""" + +import warnings +import numpy as np +import torch +import torch.nn as nn +import matplotlib.pyplot as plt +import scipy as sc + +from src.subspace_method import SubspaceMethod +from src.signal_creation import SystemModel +from src.utils import * +from src.metrics import SpectrumLoss, UnsupervisedSpectrumLoss, RMSPELoss +# from src.metrics import RMSPELoss + + +class DiffMUSIC(SubspaceMethod): + """ + Differentiable MUSIC (diffMUSIC) and classical MUSIC implementation for DoA estimation with hardware impairment learning. + When model.training = True it runs diffMUSIC, otherwise, it runs classical MUSIC. + + This implementation is focused on: + - Far-field scenarios only + - Non-coherent sources + - Narrowband signals + + Key features: + - Learnable antenna positions (nn.Parameter) + - Learnable complex gains (nn.Parameter) + - Differentiable steering matrix computation + - Softmax-based peak finding for differentiability + """ + + def __init__(self, system_model_params, N: int, model_order_estimation: str = None, + physical_array: torch.Tensor = None, physical_gains: torch.Tensor = None): + """ + Initialize diffMUSIC + + Args: + system_model: System model object (kept for compatibility) + N: Number of antennas + wavelength: Signal wavelength + model_order_estimation: Model order estimation method + """ + system_model = SystemModel(system_model_params) + super().__init__(system_model, model_order_estimation=model_order_estimation, + physical_array=physical_array, physical_gains=physical_gains) + + self.params = system_model.params + self.N = N + self.wavelength = self.params.wavelength + self.window_size = self.params.softmax_window_size + + # Initialize learnable parameters + self._init_learnable_parameters() + + # Initialize DoA grid for far-field + self._init_angle_grid() + + # Precompute steering matrix on grid + self._precompute_steering_grid() + + self.music_spectrum = None + self.noise_subspace = None + + def _init_learnable_parameters(self): + """Initialize learnable antenna positions and complex gains""" + + # Initialize antenna positions - nominal ULA with half-wavelength spacing + nominal_positions = torch.arange(self.N, dtype=torch.float64) * (self.wavelength / 2) + self.antenna_positions = nn.Parameter(nominal_positions) + + # Initialize complex gains - start with unit gains (real=1, imag=0) + gains_real = torch.ones(self.N, dtype=torch.float64) + gains_imag = torch.zeros(self.N, dtype=torch.float64) + complex_gain = torch.complex(gains_real, gains_imag).to(torch.complex64) + self.complex_gain = nn.Parameter(complex_gain) + + def _init_angle_grid(self): + """Initialize angle grid for DOA estimation""" + angle_range = np.deg2rad(self.params.doa_range) + angle_resolution = np.deg2rad(self.params.doa_resolution / 2) # Higher resolution by 2 than the original grid. + angle_decimals = int(np.ceil(np.log10(1 / angle_resolution))) # Formula to determine floating point accuracy based on the resolution. + + self.angles_grid = torch.arange(-angle_range, angle_range + angle_resolution, + angle_resolution, dtype=torch.float64) + self.angles_grid = torch.round(self.angles_grid, decimals=angle_decimals) + + def get_angles_grid(self) -> torch.Tensor: + """ + Get the angle grid for DoA estimation + + Returns: + Tensor of angles in radians, shape (num_angles,) + """ + return self.angles_grid.to(self.device) + + def _precompute_steering_grid(self): + """Precompute steering vectors for the angular grid""" + # This will be computed dynamically during forward pass since parameters are learnable + pass + + def get_array_learnable_parameters(self, learnable = True) -> tuple: + """ + Get the learnable antenna positions and complex gains - if learnable return nn.Parameters, else return numpy arrays detached from the graph. + """ + if learnable: + return self.antenna_positions, self.complex_gain + else: + return self.antenna_positions.detach().cpu().numpy(), self.complex_gain.detach().cpu().numpy() + + def set_window_size(self, window_size): + """ + Dynamically update window size for diffMUSIC + + Args: + window_size: New window size (int for absolute, float 0-1 for relative) + """ + if isinstance(window_size, float) and 0 < window_size < 1: + # Relative to grid size + self.window_size = int(window_size * len(self.angles_grid)) + else: + # Absolute size + self.window_size = int(window_size) + + + def compute_steering_matrix(self, angles: torch.Tensor) -> torch.Tensor: + """ + Compute differentiable steering matrix for given angles + + Args: + angles: Tensor of angles in radians, shape (num_angles,) + + Returns: + Complex steering matrix of shape (N, num_angles) + """ + if angles.dim() == 0: + angles = angles.unsqueeze(0) + + # Ensure angles have the same dtype as antenna_positions (float64) + angles = angles.to(torch.float64) + + # Get complex gains + antenna_positions, complex_gains = self.get_array_learnable_parameters(learnable=True) + complex_gains = complex_gains.to(torch.complex128) # Ensure complex gains are in complex128 format + + # Compute steering vectors: a(θ) = g ⊙ exp(-j * 2π * p * sin(θ) / λ) + # where ⊙ is element-wise multiplication + + # Phase computation: (N, 1) * (1, num_angles) -> (N, num_angles) + phase_delays = antenna_positions.unsqueeze(1) @ torch.sin(angles).unsqueeze(0) + + # Steering matrix without gains + steering_base = torch.exp(-2j * torch.pi * phase_delays / self.wavelength) + + # Apply complex gains: (N, N) * (N, num_angles) -> (N, num_angles) + steering_matrix = torch.diag(complex_gains) @ steering_base + + # Normalization (as in paper equation 2) + norm_factor = 1 / torch.linalg.norm(complex_gains, ord=2) + steering_matrix = norm_factor * steering_matrix + + return steering_matrix.to(torch.complex64) + +#NOTE: from now on, the first dimension is always batch size, in our case it's dummy for compatibility. It'll be always 1. + + def forward(self, cov: torch.Tensor, number_of_sources: int, known_angles=None): + """ + Forward pass of diffMUSIC + + Args: + cov: Covariance matrix of shape (BATCH_SIZE, N, N). We leave the batch dimension although for now it's going to be dummy just for the sake of compatibility. + number_of_sources: Number of sources to estimate, if None, will estimate the number of sources + known_angles: Not used in far-field case (kept for compatibility) + known_distances: Not used in far-field case (kept for compatibility) + + Returns: + tuple: (estimated_angles, source_estimation, eigen_regularization) + """ + # Ensure covariance is complex + cov = cov.to(torch.complex128) + + # Subspace decomposition + _, noise_subspace, source_estimation, eigen_regularization = self.subspace_separation(cov, number_of_sources) + self.noise_subspace = noise_subspace.to(self.device) + + # Compute MUSIC spectrum + inverse_spectrum = self._compute_inverse_spectrum(self.noise_subspace) + self.music_spectrum = 1 / (inverse_spectrum + 1e-10) + + # Peak finding (differentiable during training, hard during inference) + estimated_angles, peaks_masks = self._peak_finder(number_of_sources) + + return estimated_angles, peaks_masks, source_estimation, eigen_regularization #eigen_regularization is not used in this implementation + + def _compute_inverse_spectrum(self, noise_subspace: torch.Tensor) -> torch.Tensor: + """ + Compute the inverse MUSIC spectrum using current learnable parameters + + Args: + noise_subspace: Noise subspace of shape (batch_size, N, N-M) + + Returns: + Inverse spectrum of shape (batch_size, num_angles) + """ + # Compute steering matrix for all angles in the grid + steering_grid = self.compute_steering_matrix(self.angles_grid.to(self.device)) # (N, num_angles) + + # Compute inverse spectrum: ||U_N^H * A(θ)||² + # steering_grid: (N, num_angles), noise_subspace: (batch_size, N, N-M) + var1 = torch.einsum("na, bnm -> bam", + steering_grid.conj(), + noise_subspace.to(torch.complex64)) + # Computes the matrix multiplication using Einstein notation, just a fancier way to write it. + + inverse_spectrum = torch.norm(var1, dim=2) ** 2 # (batch_size, num_angles) + + return inverse_spectrum + + def _peak_finder(self, number_of_sources: int) -> torch.Tensor: + """ + Peak finding - differentiable during training, hard during inference + + Args: + number_of_sources: Number of peaks to find + + Returns: + Estimated angles in radians + """ + if self.training: #This is built-in since it is a Module, just use model.train() or model.eval() + return self._differentiable_peak_finder(number_of_sources) + else: + return self._hard_peak_finder(number_of_sources, return_angles=True) + + def _hard_peak_finder(self, number_of_sources: int, return_angles = True) -> torch.Tensor: + """ + Non-differentiable peak finding for inference + + Args: + number_of_sources: Number of peaks to find + return_angles: Wheter to return angles or peaks_indices for DiffMUSIC's internal use + + Returns: + Estimated angles in radians, shape (batch_size, number_of_sources) + """ + batch_size = self.music_spectrum.shape[0] + peaks = torch.zeros(batch_size, number_of_sources, dtype=torch.int64, device=self.device) + + for batch in range(batch_size): + spectrum = self.music_spectrum[batch].cpu().detach().numpy() + + # Find peaks using scipy + peaks_indices = sc.signal.find_peaks(spectrum, threshold=0.0)[0] + + if len(peaks_indices) < number_of_sources: + warnings.warn(f"diffMUSIC: Not enough peaks found, trying to find another {number_of_sources - len(peaks_indices)} peaks by top amplitude.") + # Use highest values instead + additional_peaks = torch.topk(torch.from_numpy(spectrum), + number_of_sources - len(peaks_indices), + largest=True).indices.numpy() + peaks_indices = np.concatenate([peaks_indices, additional_peaks]) + + # Sort by amplitude and take top peaks + sorted_peaks = peaks_indices[np.argsort(spectrum[peaks_indices])[::-1]] + peaks[batch] = torch.from_numpy(sorted_peaks[:number_of_sources]).to(self.device) + + if not return_angles: + # Return peak indices directly + return peaks + else: + # Convert indices to angles + estimated_angles = torch.gather( + self.angles_grid.unsqueeze(0).repeat(batch_size, 1).to(self.device), + 1, peaks + ) + + return estimated_angles, None # Here for compatibility with the differentiable peak finder, we return None for peaks_masks. + + def _differentiable_peak_finder(self, number_of_sources: int) -> torch.Tensor: + """ + Differentiable peak finding using softmax (Algorithm 2 from paper) + + Args: + number_of_sources: Number of sources to estimate + + Returns: + Estimated angles in radians, shape (batch_size, number_of_sources) + """ + batch_size = self.music_spectrum.shape[0] + estimated_angles = torch.zeros(batch_size, number_of_sources, + dtype=torch.float64, device=self.device) + + peaks_indices = self._hard_peak_finder(number_of_sources, return_angles=False) # torch.Tensor of (batch_size, number_of_sources) + peaks_masks = [] # a list of lists, each containing indices of peaks for each batch, needed for the unsupervised loss. + + for batch in range(batch_size): + batch_masks = [] + spectrum = self.music_spectrum[batch] + + # Find initial peaks (non-differentiable, but gradients will flow through softmax) + peaks_indices_per_batch = peaks_indices[batch] + + # For each peak, apply differentiable refinement + for source_idx in range(number_of_sources): + peak_idx = peaks_indices_per_batch[source_idx] + + # Create angular mask around peak + mask_indices = self._create_angular_mask(peak_idx, spectrum.shape[0]) + batch_masks.append(mask_indices) # Store for unsupervised loss + + # Extract spectrum values in the mask + masked_spectrum = spectrum[mask_indices] + + # Apply softmax to get weights + weights = torch.softmax(masked_spectrum, dim=0) + + # Compute weighted average of angles (Equation 13 from paper) + masked_angles = self.angles_grid[mask_indices].to(self.device) + estimated_angles[batch, source_idx] = torch.sum(weights * masked_angles) + + peaks_masks.append(batch_masks) # Store masks for unsupervised loss + + return estimated_angles, peaks_masks + + + def _create_angular_mask(self, center_idx: int, spectrum_length: int) -> torch.Tensor: + """ + Create angular mask around peak center (ΠL operation from paper) + + Args: + center_idx: Center index of the peak + spectrum_length: Length of the spectrum + + Returns: + Indices for the angular mask + """ + half_window = self.window_size // 2 + + # Create mask indices around center + start_idx = max(0, center_idx - half_window) + end_idx = min(spectrum_length, center_idx + half_window + 1) + + mask_indices = torch.arange(start_idx, end_idx, device=self.device) + + return mask_indices + + + def set_nominal_parameters(self, positions: torch.Tensor = None, gains: torch.Tensor = None): + """ + Set parameters to nominal values (useful for initialization or comparison) + + Args: + positions: Nominal antenna positions (if None, use half-wavelength spacing) + gains: Nominal complex gains (if None, use unit gains) + """ + with torch.no_grad(): + if positions is None: + # Half-wavelength spacing + positions = torch.arange(self.N, dtype=torch.float64) * (self.wavelength / 2) + + if gains is None: + # Unit gains + gains = torch.ones(self.N, dtype=torch.complex64) + + self.antenna_positions.copy_(positions) + self.gains_real.copy_(gains.real.to(torch.float64)) + self.gains_imag.copy_(gains.imag.to(torch.float64)) + + def plot_spectrum(self, highlight_angles: torch.Tensor = None, + path_to_save: str | None = False, batch_idx: int = 0): + """ + Plot the MUSIC spectrum + + Args: + batch_idx: Which batch element to plot + highlight_angles: True angles to highlight on the plot + save: Whether to save the plot + title: Plot title + """ + if self.music_spectrum is None: + warnings.warn("No spectrum computed yet. Run forward pass first.") + return + + angles_deg = torch.rad2deg(self.angles_grid).cpu().numpy() + spectrum = self.music_spectrum[batch_idx].cpu().detach().numpy() + + plt.figure(figsize=(10, 6)) + plt.plot(angles_deg, spectrum, 'b-', linewidth=2, label='diffMUSIC Spectrum') + + if highlight_angles is not None: + highlight_deg = torch.rad2deg(highlight_angles).cpu().numpy() + for i, angle in enumerate(highlight_deg): + plt.axvline(x=angle, color='r', linestyle='--', alpha=0.7, + label='True DoA' if i == 0 else "") + + plt.xlabel('Angle [degrees]') + plt.ylabel('Spectrum Power') + plt.title(("diffMUSIC" if self.training else "MUSIC") + " Spectrum") + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + + if path_to_save is not None: + plt.savefig(path_to_save / f"Spectrum.png", dpi=300) + plt.show() + + def plot_learned_parameters(self, true_positions: torch.Tensor = None, true_gains: torch.Tensor = None, + path_to_save: str = None, title: str = "Learned Parameters"): + """ + Plot learned antenna parameters similar to Figure 3 in the paper + + Args: + true_positions: True antenna positions for comparison + true_gains: True antenna gains for comparison + path_to_save: Path to save the plot + title: Plot title + """ + import matplotlib.pyplot as plt + import matplotlib.patches as patches + + # Get learned parameters + learned_positions = self.antenna_positions.detach().cpu().numpy() + learned_gains = self.complex_gain.detach().cpu().numpy() + + # Convert positions to wavelength units for display + learned_positions_wl = learned_positions / (self.wavelength / 2) + + fig, ax = plt.subplots(1, 1, figsize=(16, 6)) + + # Vertical separation between learned and physical parameters + learned_y = 0.5 + physical_y = -0.5 + + # Plot learned parameters (blue, top row) + for i, (pos, gain) in enumerate(zip(learned_positions_wl, learned_gains)): + # Circle radius represents gain magnitude + radius = abs(gain) * 0.25 # Slightly smaller for better visibility + + # Circle color and segment angle represent gain phase + phase = np.angle(gain) + + # Draw circle + circle = patches.Circle((pos, learned_y), radius, + facecolor='lightblue', + edgecolor='blue', + linewidth=2, + alpha=0.7, + label='Learned' if i == 0 else "") + ax.add_patch(circle) + + # Draw phase segment (line from center to edge) + segment_x = pos + radius * np.cos(phase) + segment_y = learned_y + radius * np.sin(phase) + ax.plot([pos, segment_x], [learned_y, segment_y], 'b-', linewidth=2) + + # Add antenna number above the circle + ax.text(pos, learned_y + 0.4, f'{i}', ha='center', va='center', fontsize=10, fontweight='bold') + + # Plot true parameters if provided (red, bottom row) + if true_positions is not None and true_gains is not None: + true_positions_np = true_positions.cpu().numpy() if isinstance(true_positions, torch.Tensor) else true_positions + true_gains_np = true_gains.cpu().numpy() if isinstance(true_gains, torch.Tensor) else true_gains + + true_positions_wl = true_positions_np / (self.wavelength / 2) + + for i, (pos, gain) in enumerate(zip(true_positions_wl, true_gains_np)): + radius = abs(gain) * 0.25 # Same scaling as learned + phase = np.angle(gain) + + # Draw circle with different style + circle = patches.Circle((pos, physical_y), radius, + facecolor='lightcoral', + edgecolor='red', + linewidth=2, + alpha=0.7, + label='Physical' if i == 0 else "") + ax.add_patch(circle) + + # Draw phase segment + segment_x = pos + radius * np.cos(phase) + segment_y = physical_y + radius * np.sin(phase) + ax.plot([pos, segment_x], [physical_y, segment_y], 'r-', linewidth=2) + + # Add antenna number below the circle + ax.text(pos, physical_y - 0.4, f'{i}', ha='center', va='center', fontsize=10, fontweight='bold') + + # Add horizontal reference lines + ax.axhline(y=learned_y, color='blue', linestyle=':', alpha=0.3, linewidth=1) + if true_positions is not None and true_gains is not None: + ax.axhline(y=physical_y, color='red', linestyle=':', alpha=0.3, linewidth=1) + + # Set axis limits and labels + ax.set_xlim(-0.5, max(learned_positions_wl) + 0.5) + ax.set_ylim(-1.2, 1.2) + ax.set_xlabel('x [λ/2]', fontsize=12, fontweight='bold') + ax.set_ylabel('') + ax.set_title(title, fontsize=14, fontweight='bold') + ax.grid(True, alpha=0.3) + + # Create custom legend + legend_elements = [ + plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='lightblue', + markersize=10, markeredgecolor='blue', markeredgewidth=2, label='Learned'), + ] + if true_positions is not None and true_gains is not None: + legend_elements.append( + plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='lightcoral', + markersize=10, markeredgecolor='red', markeredgewidth=2, label='Physical') + ) + + ax.legend(handles=legend_elements, loc='upper right', fontsize=12) + + # Remove y-axis ticks for cleaner look + ax.set_yticks([]) + + plt.tight_layout() + + if path_to_save is not None: + plt.savefig(path_to_save / "learned_parameters.png", dpi=300, bbox_inches='tight') + + plt.show() + + +class DiffMUSICLoss(nn.Module): + """ + Wrapper loss class for diffMUSIC training + Supports different loss strategies: RMSPE, spectrum, and unsupervised + """ + + def __init__(self, loss_type: str = "rmspe", **kwargs): + """ + Args: + loss_type: "rmspe" for LSL,θ, "spectrum" for LSL,P, or "unsupervised" for LUL + **kwargs: Additional arguments for specific loss functions + """ + super(DiffMUSICLoss, self).__init__() + self.loss_type = loss_type + + # Import loss functions from metrics + from src.metrics import RMSPELoss, SpectrumLoss, UnsupervisedSpectrumLoss + + if loss_type == "rmspe": + self.loss_fn = RMSPELoss(**kwargs) + elif loss_type == "spectrum": + self.loss_fn = SpectrumLoss() + elif loss_type == "unsupervised": + self.loss_fn = UnsupervisedSpectrumLoss() + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + def forward(self, **kwargs): + """ + Forward pass - delegates to appropriate loss function + """ + if self.loss_type == "rmspe": + return self.loss_fn(kwargs['predictions'], kwargs['targets']) + elif self.loss_type == "spectrum": + return self.loss_fn(kwargs['spectrum'], kwargs['targets'], kwargs['angles_grid']) + elif self.loss_type == "unsupervised": + return self.loss_fn(kwargs['spectrum'], kwargs['peak_masks']) + else: + raise ValueError(f"Unknown loss type: {self.loss_type}") \ No newline at end of file diff --git a/src/metrics.py b/src/metrics.py new file mode 100644 index 0000000..ebafd80 --- /dev/null +++ b/src/metrics.py @@ -0,0 +1,179 @@ +import torch +import torch.nn as nn +import numpy as np +from itertools import permutations + + +class RMSPELoss(nn.Module): + """ + Root Mean Squared Periodic Error (RMSPE) loss for angle estimation + Handles the periodic nature of angles and permutation invariance + """ + + def __init__(self, balance_factor: float = None): + """ + Args: + balance_factor: Weighting factor for the loss (1.0 for getting the loss as is, + sqrt(M) is the deafult value if not speccified - by the definition in the paper) + """ + super(RMSPELoss, self).__init__() + self.balance_factor = balance_factor + + def forward(self, predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + """ + Compute RMSPE loss between predicted and target angles + + Args: + predictions: Predicted angles in radians, shape (batch_size, M) + targets: Target angles in radians, shape (batch_size, M) + + Returns: + Loss tensor, shape (batch_size,) + """ + batch_size, M = predictions.shape + + if M == 1: + # Single source case - no permutation needed + return self._periodic_mse(predictions, targets) + + # Multi-source case - find best permutation + min_loss = torch.full((batch_size,), float('inf'), device=predictions.device) + + for perm in permutations(range(M)): + perm_tensor = torch.tensor(perm, device=predictions.device) + perm_predictions = predictions[:, perm_tensor] + loss = self._periodic_mse(perm_predictions, targets) + min_loss = torch.min(min_loss, loss) + + if self.balance_factor is None: + # Default balance factor is sqrt(M) as per the paper + self.balance_factor = 1/torch.sqrt(torch.tensor(M, dtype=torch.float32, device=predictions.device)) + + return min_loss * self.balance_factor + + def _periodic_mse(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """ + Compute periodic MSE for angles + + Args: + pred: Predicted angles, shape (batch_size, M) + target: Target angles, shape (batch_size, M) + + Returns: + MSE loss per batch, shape (batch_size,) + """ + # Compute angular difference considering periodicity + diff = pred - target + # Wrap to [-π, π] (not critical for angles in [-π/2, π/2] but here for generality) + diff = torch.atan2(torch.sin(diff), torch.cos(diff)) + + # Compute MSE + mse = torch.mean(diff ** 2, dim=1) + + return torch.sqrt(mse) + + +class SpectrumLoss(nn.Module): + """ + Spectrum-based loss (LSL,P from paper) + Maximizes spectrum amplitude at true DoA locations + """ + + def __init__(self): + super(SpectrumLoss, self).__init__() + + def forward(self, spectrum: torch.Tensor, true_angles: torch.Tensor, + angles_grid: torch.Tensor) -> torch.Tensor: + """ + Compute spectrum loss by maximizing amplitude at true DoA locations + + Args: + spectrum: MUSIC spectrum, shape (batch_size, num_angles) + true_angles: True DoAs, shape (batch_size, M) + angles_grid: Angular grid, shape (num_angles,) + + Returns: + Loss value (scalar) + """ + batch_size = spectrum.shape[0] + total_loss = 0 + + for batch in range(batch_size): + for angle in true_angles[batch]: + # Find closest angle in grid + angle_idx = torch.argmin(torch.abs(angles_grid - angle)) + # Negative spectrum value (to maximize) + total_loss -= spectrum[batch, angle_idx] + + # Uncomment the following row to normalize also by the number of true angles, this is not done + # here for consistency with the paper loss definition. Doesn't influence the optimization. + # total_loss /= true_angles.shape[1] + return total_loss / batch_size + + +class JainsIndexLoss(nn.Module): + """ + Jain's Index loss for unsupervised learning + Encourages sharp, concentrated peaks in the spectrum + """ + + def __init__(self): + super(JainsIndexLoss, self).__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Compute Jain's Index: J(x) = (sum(x))^2 / (n * sum(x^2)) + + Args: + x: Input tensor (spectrum values) + + Returns: + Jain's index value (higher = more concentrated) + """ + n = x.shape[0] + sum_x = torch.sum(x) + sum_x_squared = torch.sum(x ** 2) + + jains_index = (sum_x ** 2) / (n * sum_x_squared + 1e-8) + + # Return negative to minimize (we want to maximize Jain's index) + return -jains_index + + +class UnsupervisedSpectrumLoss(nn.Module): + """ + Unsupervised loss using Jain's Index on spectrum peaks (LUL from paper) + """ + + def __init__(self): + super(UnsupervisedSpectrumLoss, self).__init__() + self.jains_index_loss = JainsIndexLoss() + + def forward(self, spectrum: torch.Tensor, peak_masks: list) -> torch.Tensor: + """ + Compute unsupervised loss using Jain's Index on spectrum peaks + + Args: + spectrum: MUSIC spectrum, shape (batch_size, num_angles) + peak_masks: List of masks for each peak region + + Returns: + Loss value encouraging sharp peaks + """ + total_loss = 0 + batch_size = spectrum.shape[0] + + for batch in range(batch_size): + spectrum_batch = spectrum[batch] + + # Apply Jain's index to each peak region + for mask_indices in peak_masks[batch]: + masked_spectrum = spectrum_batch[mask_indices] + + # Jain's index encourages sharp peaks + jains_loss = self.jains_index_loss(masked_spectrum) + #Uncomment the following line to normalize by the number of peaks + # jains_loss /= len(peak_masks[batch]) # Normalize by number of peaks/ + total_loss += jains_loss + + return total_loss / batch_size \ No newline at end of file diff --git a/src/multi_experiment_runner.py b/src/multi_experiment_runner.py new file mode 100644 index 0000000..0885b55 --- /dev/null +++ b/src/multi_experiment_runner.py @@ -0,0 +1,614 @@ +""" +Multi-Experiment Runner - Handles parameter sweeps and multi-experiment orchestration +""" + +import torch +import numpy as np +import time +from pathlib import Path +from typing import Dict, Any, List, Union +from tqdm import tqdm +import copy + +from doa_runner import DoARunner +from src.signal_creation import SystemModelParams +import src.utils as utils + + +class MultiExperimentRunner: + """ + Orchestrates parameter sweeps across multiple experiments + """ + + def __init__(self, sweep_config: Dict[str, Any], base_kwargs: Dict[str, Any]): + """ + Initialize Multi-Experiment Runner + + Args: + sweep_config: Configuration for the parameter sweep + base_kwargs: Base configuration (system_model_params, model_config, etc.) + """ + self.sweep_config = sweep_config + self.base_kwargs = base_kwargs + self.sweep_parameter = sweep_config["parameter"] + self.sweep_values = sweep_config["values"] + self.fixed_params = sweep_config.get("fixed_params", {}) + self.plot_config = sweep_config.get("plot_config", {}) + + # Multi-loss configuration + self.loss_functions = sweep_config.get("loss_functions", None) + self.multi_loss_mode = self.loss_functions is not None + + # Initialize results manager + self.results_manager = SweepResultsManager(sweep_config) + + # Setup paths + self.base_path = Path(__file__).parent.parent + self.sweep_results_path = None + + def run_sweep(self) -> Dict[str, Any]: + """ + Execute the complete parameter sweep + Supports both single loss and multi-loss modes + + Returns: + Aggregated sweep results + """ + if self.multi_loss_mode: + return self._run_multi_loss_sweep() + else: + return self._run_single_loss_sweep() + + def _run_single_loss_sweep(self) -> Dict[str, Any]: + """Execute parameter sweep with single loss function""" + print(f"Starting parameter sweep: {self.sweep_parameter}") + print(f"Values: {self.sweep_values}") + print(f"Fixed parameters: {self.fixed_params}") + + # Setup results directory + self._setup_results_directory() + + # Initialize progress tracking + sweep_start_time = time.time() + individual_times = [] + + # Run experiments for each parameter value + for i, param_value in enumerate(tqdm(self.sweep_values, desc=f"Sweeping {self.sweep_parameter}")): + exp_start_time = time.time() + + print(f"\nExperiment {i+1}/{len(self.sweep_values)}: {self.sweep_parameter}={param_value}") + + # Run single experiment + try: + experiment_kwargs = self._prepare_experiment_kwargs(param_value) + results = self._run_single_experiment(experiment_kwargs) + + # Store results + self.results_manager.add_experiment_result(param_value, results) + + exp_duration = time.time() - exp_start_time + individual_times.append(exp_duration) + + print(f" RMSPE: {results.get('rmspe', 'N/A'):.6f} (Time: {exp_duration:.2f}s)") + + except Exception as e: + print(f" Error in experiment {param_value}: {e}") + # Store error result + self.results_manager.add_experiment_result(param_value, { + 'error': str(e), + 'rmspe': float('inf') + }) + individual_times.append(0) + + # Finalize results + total_time = time.time() - sweep_start_time + sweep_results = self.results_manager.finalize_results(total_time, individual_times) + + # Save results and generate plots + self._save_sweep_results(sweep_results) + self._generate_plots(sweep_results) + + print(f"\nSweep completed in {total_time:.2f}s") + print(f"Results saved to: {self.sweep_results_path}") + + return sweep_results + + def _run_multi_loss_sweep(self) -> Dict[str, Any]: + """Execute parameter sweep with multiple loss functions""" + print(f"Starting MULTI-LOSS parameter sweep: {self.sweep_parameter}") + print(f"Parameter values: {self.sweep_values}") + print(f"Loss functions: {self.loss_functions}") + print(f"Fixed parameters: {self.fixed_params}") + + # Setup results directory + self._setup_results_directory() + + # Initialize multi-loss results manager + multi_loss_results = MultiLossResultsManager(self.sweep_config) + + total_experiments = len(self.sweep_values) * len(self.loss_functions) + experiment_count = 0 + + sweep_start_time = time.time() + + # Run experiments for each loss function + for loss_idx, loss_function in enumerate(self.loss_functions): + print(f"\n{'='*50}") + print(f"Running with Loss Function: {loss_function} ({loss_idx+1}/{len(self.loss_functions)})") + print(f"{'='*50}") + + loss_results = {} + loss_times = [] + + # Run parameter sweep for this loss function + for param_idx, param_value in enumerate(self.sweep_values): + experiment_count += 1 + exp_start_time = time.time() + + print(f"Experiment {experiment_count}/{total_experiments}: " + f"{self.sweep_parameter}={param_value}, loss={loss_function}") + + try: + # Prepare experiment with specific loss function + experiment_kwargs = self._prepare_experiment_kwargs(param_value, loss_function) + results = self._run_single_experiment(experiment_kwargs) + + # Store results + loss_results[param_value] = results + + exp_duration = time.time() - exp_start_time + loss_times.append(exp_duration) + + print(f" RMSPE: {results.get('rmspe', 'N/A'):.6f} (Time: {exp_duration:.2f}s)") + + except Exception as e: + print(f" Error: {e}") + loss_results[param_value] = {'error': str(e), 'rmspe': float('inf')} + loss_times.append(0) + + # Add results for this loss function + multi_loss_results.add_loss_function_results(loss_function, loss_results, loss_times) + + # Finalize multi-loss results + total_time = time.time() - sweep_start_time + final_results = multi_loss_results.finalize_results(total_time) + + # Save results and generate plots + self._save_sweep_results(final_results) + self._generate_multi_loss_plots(final_results) + + print(f"\nMulti-loss sweep completed in {total_time:.2f}s") + print(f"Results saved to: {self.sweep_results_path}") + + return final_results + + def _prepare_experiment_kwargs(self, param_value, loss_function=None) -> Dict[str, Any]: + """ + Prepare kwargs for a single experiment with the specified parameter value + + Args: + param_value: Value for the sweep parameter + loss_function: Optional loss function override + + Returns: + Modified kwargs for the experiment + """ + # Deep copy to avoid modifying original + experiment_kwargs = copy.deepcopy(self.base_kwargs) + + # Apply fixed parameters + for param_name, param_val in self.fixed_params.items(): + if param_name in experiment_kwargs["system_model_params"]: + experiment_kwargs["system_model_params"][param_name] = param_val + + # Apply sweep parameter + if self.sweep_parameter in experiment_kwargs["system_model_params"]: + experiment_kwargs["system_model_params"][self.sweep_parameter] = param_value + + # Apply loss function if specified + if loss_function is not None: + experiment_kwargs["training_params"]["loss_type"] = loss_function + + # Force data creation for sweeps (don't load from file) + experiment_kwargs["simulation_commands"]["create_data"] = True + experiment_kwargs["simulation_commands"]["load_data"] = False + experiment_kwargs["simulation_commands"]["plot_results"] = False # Handle plotting at sweep level + + return experiment_kwargs + + def _setup_results_directory(self): + """Setup directory structure for sweep results""" + from datetime import datetime + + # Create timestamp + dt_string = datetime.now().strftime("%d_%m_%Y_%H_%M") + + # Create sweep directory name + sweep_name = f"{self.sweep_parameter}_sweep" + if self.fixed_params: + fixed_str = "_".join([f"{k}{v}" for k, v in self.fixed_params.items()]) + sweep_name += f"_{fixed_str}" + + # Add multi-loss indicator + if self.multi_loss_mode: + loss_str = "_".join(self.loss_functions) + sweep_name += f"_losses_{loss_str}" + + sweep_name += f"_{dt_string}" + + # Setup paths + self.sweep_results_path = self.base_path / "results" / "parameter_sweeps" / sweep_name + self.sweep_results_path.mkdir(parents=True, exist_ok=True) + + # Create subdirectories + (self.sweep_results_path / "plots").mkdir(exist_ok=True) + (self.sweep_results_path / "individual_experiments").mkdir(exist_ok=True) + + def _run_single_experiment(self, experiment_kwargs: Dict[str, Any]) -> Dict[str, Any]: + """ + Run a single experiment with the given parameters + + Args: + experiment_kwargs: Experiment configuration + + Returns: + Experiment results + """ + # Import here to avoid circular imports + from run_simulation import run_single_simulation + + return run_single_simulation(**experiment_kwargs) + + def _save_sweep_results(self, sweep_results: Dict[str, Any]): + """Save sweep results to file""" + results_file = self.sweep_results_path / "sweep_results.pkl" + utils.save_pickle(sweep_results, results_file) + + # Also save a summary text file + summary_file = self.sweep_results_path / "summary.txt" + self._save_summary_text(sweep_results, summary_file) + + def _save_summary_text(self, sweep_results: Dict[str, Any], summary_file: Path): + """Save human-readable summary""" + with open(summary_file, 'w') as f: + f.write(f"Parameter Sweep Summary\n") + f.write(f"======================\n\n") + f.write(f"Sweep Parameter: {sweep_results['sweep_parameter']}\n") + f.write(f"Sweep Values: {sweep_results['sweep_values']}\n") + f.write(f"Fixed Parameters: {sweep_results['fixed_params']}\n") + + # Handle multi-loss results + if 'loss_functions' in sweep_results: + f.write(f"Loss Functions: {sweep_results['loss_functions']}\n\n") + + # Write results for each loss function + for loss_func in sweep_results['loss_functions']: + f.write(f"Results for {loss_func}:\n") + f.write(f"{''.join(['-'] * (len(loss_func) + 12))}\n") + + loss_metrics = sweep_results['aggregated_metrics'][loss_func] + for param_val, rmspe in zip(sweep_results['sweep_values'], + loss_metrics['rmspe_values']): + f.write(f"{sweep_results['sweep_parameter']}={param_val}: RMSPE={rmspe:.6f}\n") + + f.write(f"Mean RMSPE: {loss_metrics['mean_rmspe']:.6f}\n") + f.write(f"Std RMSPE: {loss_metrics['std_rmspe']:.6f}\n\n") + else: + # Single loss function + f.write(f"\nResults:\n") + f.write(f"--------\n") + for param_val, rmspe in zip(sweep_results['sweep_values'], + sweep_results['aggregated_metrics']['rmspe_values']): + f.write(f"{sweep_results['sweep_parameter']}={param_val}: RMSPE={rmspe:.6f}\n") + + f.write(f"\nSummary Statistics:\n") + f.write(f"Mean RMSPE: {sweep_results['aggregated_metrics']['mean_rmspe']:.6f}\n") + f.write(f"Std RMSPE: {sweep_results['aggregated_metrics']['std_rmspe']:.6f}\n") + + f.write(f"Total Time: {sweep_results['aggregated_metrics']['total_execution_time']:.2f}s\n") + + def _generate_plots(self, sweep_results: Dict[str, Any]): + """Generate plots for single loss function sweep results""" + plotter = SweepPlotter(sweep_results, self.sweep_results_path / "plots") + plotter.plot_rmspe_vs_parameter() + + def _generate_multi_loss_plots(self, sweep_results: Dict[str, Any]): + """Generate plots for multi-loss function sweep results""" + plotter = MultiLossSweepPlotter(sweep_results, self.sweep_results_path / "plots") + plotter.plot_multi_loss_comparison() + + +class MultiLossResultsManager: + """ + Manages aggregation and organization of multi-loss sweep results + """ + + def __init__(self, sweep_config: Dict[str, Any]): + self.sweep_config = sweep_config + self.loss_functions = sweep_config["loss_functions"] + self.results_by_loss = {} + self.times_by_loss = {} + + def add_loss_function_results(self, loss_function: str, results: Dict[str, Any], times: List[float]): + """ + Add results for a specific loss function + + Args: + loss_function: Name of the loss function + results: Results dictionary for this loss function + times: List of execution times for each experiment + """ + self.results_by_loss[loss_function] = results + self.times_by_loss[loss_function] = times + + def finalize_results(self, total_time: float) -> Dict[str, Any]: + """ + Finalize and aggregate all multi-loss results + + Args: + total_time: Total execution time for all experiments + + Returns: + Complete aggregated multi-loss results dictionary + """ + # Aggregate metrics for each loss function + aggregated_metrics = {} + + for loss_function in self.loss_functions: + results = self.results_by_loss[loss_function] + times = self.times_by_loss[loss_function] + + # Extract RMSPE values in order + rmspe_values = [] + valid_results = [] + + for param_val in self.sweep_config["values"]: + if param_val in results: + result = results[param_val] + if 'error' not in result: + rmspe_values.append(result.get('rmspe', float('inf'))) + valid_results.append(result) + else: + rmspe_values.append(float('inf')) + + # Calculate aggregated metrics for this loss function + valid_rmspe = [r for r in rmspe_values if r != float('inf')] + aggregated_metrics[loss_function] = { + 'rmspe_values': rmspe_values, + 'mean_rmspe': np.mean(valid_rmspe) if valid_rmspe else float('inf'), + 'std_rmspe': np.std(valid_rmspe) if valid_rmspe else 0, + 'min_rmspe': np.min(valid_rmspe) if valid_rmspe else float('inf'), + 'max_rmspe': np.max(valid_rmspe) if valid_rmspe else float('inf'), + 'execution_times': times, + 'average_time_per_experiment': np.mean(times) if times else 0, + 'num_successful_experiments': len(valid_results), + 'num_failed_experiments': len(results) - len(valid_results) + } + + # Overall statistics + aggregated_metrics['total_execution_time'] = total_time + aggregated_metrics['total_experiments'] = len(self.sweep_config["values"]) * len(self.loss_functions) + + # Compile final results + final_results = { + 'sweep_type': 'multi_loss_parameter_sweep', + 'sweep_parameter': self.sweep_config['parameter'], + 'sweep_values': self.sweep_config['values'], + 'loss_functions': self.loss_functions, + 'fixed_params': self.sweep_config.get('fixed_params', {}), + 'plot_config': self.sweep_config.get('plot_config', {}), + 'results_by_loss': self.results_by_loss, + 'aggregated_metrics': aggregated_metrics, + 'timestamp': time.strftime("%d_%m_%Y_%H_%M") + } + + return final_results + + +# Keep the original SweepResultsManager for single loss function sweeps +class SweepResultsManager: + """ + Manages aggregation and organization of sweep results + """ + + def __init__(self, sweep_config: Dict[str, Any]): + self.sweep_config = sweep_config + self.results = {} + + def add_experiment_result(self, param_value: Union[int, float], result: Dict[str, Any]): + """ + Add result from a single experiment + + Args: + param_value: Parameter value for this experiment + result: Results dictionary from the experiment + """ + self.results[param_value] = result + + def finalize_results(self, total_time: float, individual_times: List[float]) -> Dict[str, Any]: + """ + Finalize and aggregate all results + + Args: + total_time: Total execution time for the sweep + individual_times: List of individual experiment times + + Returns: + Complete aggregated results dictionary + """ + # Extract RMSPE values in order + rmspe_values = [] + valid_results = [] + + for param_val in self.sweep_config["values"]: + if param_val in self.results: + result = self.results[param_val] + if 'error' not in result: + rmspe_values.append(result.get('rmspe', float('inf'))) + valid_results.append(result) + else: + rmspe_values.append(float('inf')) + + # Calculate aggregated metrics + valid_rmspe = [r for r in rmspe_values if r != float('inf')] + aggregated_metrics = { + 'rmspe_values': rmspe_values, + 'mean_rmspe': np.mean(valid_rmspe) if valid_rmspe else float('inf'), + 'std_rmspe': np.std(valid_rmspe) if valid_rmspe else 0, + 'min_rmspe': np.min(valid_rmspe) if valid_rmspe else float('inf'), + 'max_rmspe': np.max(valid_rmspe) if valid_rmspe else float('inf'), + 'total_execution_time': total_time, + 'individual_times': individual_times, + 'average_time_per_experiment': np.mean(individual_times) if individual_times else 0, + 'num_successful_experiments': len(valid_results), + 'num_failed_experiments': len(self.results) - len(valid_results) + } + + # Compile final results + final_results = { + 'sweep_type': self.sweep_config.get('sweep_type', 'parameter_sweep'), + 'sweep_parameter': self.sweep_config['parameter'], + 'sweep_values': self.sweep_config['values'], + 'fixed_params': self.sweep_config.get('fixed_params', {}), + 'plot_config': self.sweep_config.get('plot_config', {}), + 'results': self.results, + 'aggregated_metrics': aggregated_metrics, + 'timestamp': time.strftime("%d_%m_%Y_%H_%M") + } + + return final_results + + +class MultiLossSweepPlotter: + """ + Handles plotting of multi-loss sweep results + """ + + def __init__(self, sweep_results: Dict[str, Any], plots_path: Path): + self.sweep_results = sweep_results + self.plots_path = plots_path + self.plots_path.mkdir(parents=True, exist_ok=True) + + def plot_multi_loss_comparison(self): + """Plot comparison of all loss functions on the same graph""" + import matplotlib.pyplot as plt + + # Extract data + param_values = self.sweep_results['sweep_values'] + loss_functions = self.sweep_results['loss_functions'] + param_name = self.sweep_results['sweep_parameter'] + plot_config = self.sweep_results.get('plot_config', {}) + + plt.figure(figsize=(12, 8)) + + # Plot each loss function + colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown'] + markers = ['o', 's', '^', 'D', 'v', '<'] + + for i, loss_function in enumerate(loss_functions): + rmspe_values = self.sweep_results['aggregated_metrics'][loss_function]['rmspe_values'] + + # Filter out infinite values for plotting + valid_indices = [j for j, r in enumerate(rmspe_values) if r != float('inf')] + valid_param_values = [param_values[j] for j in valid_indices] + valid_rmspe_values = [rmspe_values[j] for j in valid_indices] + + if valid_param_values: + color = colors[i % len(colors)] + marker = markers[i % len(markers)] + + plt.plot(valid_param_values, valid_rmspe_values, + color=color, marker=marker, linewidth=2, markersize=6, + label=f'{loss_function} (mean: {self.sweep_results["aggregated_metrics"][loss_function]["mean_rmspe"]:.4f}°)') + + plt.grid(True, alpha=0.3) + plt.xlabel(plot_config.get('x_label', param_name)) + plt.ylabel(plot_config.get('y_label', 'RMSPE (degrees)')) + plt.title(plot_config.get('title', f'RMSPE vs {param_name} - Loss Function Comparison')) + plt.legend() + + # Save plot + save_name = f'multi_loss_comparison_{param_name}' + plt.tight_layout() + plt.savefig(self.plots_path / f'{save_name}.png', dpi=300, bbox_inches='tight') + plt.savefig(self.plots_path / f'{save_name}.pdf', bbox_inches='tight') + plt.show() + + print(f"Multi-loss comparison plot saved to: {self.plots_path / f'{save_name}.png'}") + + def plot_individual_loss_functions(self): + """Plot individual graphs for each loss function""" + for loss_function in self.sweep_results['loss_functions']: + # Create individual plot for this loss function + single_loss_results = { + 'sweep_parameter': self.sweep_results['sweep_parameter'], + 'sweep_values': self.sweep_results['sweep_values'], + 'plot_config': self.sweep_results['plot_config'], + 'aggregated_metrics': { + 'rmspe_values': self.sweep_results['aggregated_metrics'][loss_function]['rmspe_values'], + 'mean_rmspe': self.sweep_results['aggregated_metrics'][loss_function]['mean_rmspe'], + 'std_rmspe': self.sweep_results['aggregated_metrics'][loss_function]['std_rmspe'] + } + } + + # Use the single loss plotter + individual_plotter = SweepPlotter(single_loss_results, self.plots_path) + individual_plotter.plot_rmspe_vs_parameter(suffix=f'_{loss_function}') + + +class SweepPlotter: + """ + Handles plotting of sweep results + """ + + def __init__(self, sweep_results: Dict[str, Any], plots_path: Path): + self.sweep_results = sweep_results + self.plots_path = plots_path + self.plots_path.mkdir(parents=True, exist_ok=True) + + def plot_rmspe_vs_parameter(self, suffix=''): + """Plot RMSPE vs the sweep parameter""" + import matplotlib.pyplot as plt + + # Extract data + param_values = self.sweep_results['sweep_values'] + rmspe_values = self.sweep_results['aggregated_metrics']['rmspe_values'] + param_name = self.sweep_results['sweep_parameter'] + plot_config = self.sweep_results.get('plot_config', {}) + + # Filter out infinite values for plotting + valid_indices = [i for i, r in enumerate(rmspe_values) if r != float('inf')] + valid_param_values = [param_values[i] for i in valid_indices] + valid_rmspe_values = [rmspe_values[i] for i in valid_indices] + + if not valid_param_values: + print("No valid results to plot") + return + + # Create plot + plt.figure(figsize=(10, 6)) + plt.plot(valid_param_values, valid_rmspe_values, 'b-o', linewidth=2, markersize=6) + plt.grid(True, alpha=0.3) + + # Customize plot + plt.xlabel(plot_config.get('x_label', param_name)) + plt.ylabel(plot_config.get('y_label', 'RMSPE (degrees)')) + plt.title(plot_config.get('title', f'RMSPE vs {param_name}')) + + # Add statistics text + mean_rmspe = self.sweep_results['aggregated_metrics']['mean_rmspe'] + std_rmspe = self.sweep_results['aggregated_metrics']['std_rmspe'] + plt.text(0.02, 0.98, f'Mean RMSPE: {mean_rmspe:.4f}°\nStd RMSPE: {std_rmspe:.4f}°', + transform=plt.gca().transAxes, verticalalignment='top', + bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) + + # Save plot + save_name = plot_config.get('save_name', f'rmspe_vs_{param_name}') + suffix + plt.tight_layout() + plt.savefig(self.plots_path / f'{save_name}.png', dpi=300, bbox_inches='tight') + plt.savefig(self.plots_path / f'{save_name}.pdf', bbox_inches='tight') + + # Show plot + plt.show() + + print(f"Plot saved to: {self.plots_path / f'{save_name}.png'}") \ No newline at end of file diff --git a/src/signal_creation.py b/src/signal_creation.py index 381fefe..a1beb75 100644 --- a/src/signal_creation.py +++ b/src/signal_creation.py @@ -1,315 +1,332 @@ -"""Subspace-Net -Details ----------- -Name: signal_creation.py -Authors: D. H. Shmuel -Created: 01/10/21 -Edited: 02/06/23 - -Purpose: --------- -This script defines the Samples class, which inherits from SystemModel class. -This class is used for defining the samples model. -""" - -# Imports -from random import sample -from src.system_model import SystemModel, SystemModelParams -# from src.utils import * import numpy as np import torch +from random import sample -class Samples(SystemModel): +class SystemModelParams: + """ + Parameters for the for the simulation at one claass. """ - Class used for defining and creating signals and observations. - Inherits from SystemModel class. - - ... - - Attributes: - ----------- - doa (np.ndarray): Array of angels (directions) of arrival. - Methods: - -------- - set_doa(doa): Sets the direction of arrival (DOA) for the signals. - samples_creation(noise_mean: float = 0, noise_variance: float = 1, signal_mean: float = 0, - signal_variance: float = 1): Creates samples based on the specified mode and parameters. - noise_creation(noise_mean, noise_variance): Creates noise based on the specified mean and variance. - signal_creation(signal_mean=0, signal_variance=1, SNR=10): Creates signals based on the specified mode and parameters. + def __init__(self, **kwargs): + kwargs = {**kwargs} + for key, value in kwargs.items(): + if isinstance(value, str): + value = value.lower() + setattr(self, key, value) + + def __repr__(self): + attrs = [f"{k}={v}" for k, v in self.__dict__.items()] + return f"SystemModelParams({', '.join(attrs)})" + + +class SystemModel: + """ + Simplified SystemModel class for + far-field + non-coherent + narrow-band + DoA estimation """ def __init__(self, system_model_params: SystemModelParams): - """Initializes a Samples object. + self.params = system_model_params + self.array = None + self.dist_array_elems = None + self.gain_petrurbation = np.ones(self.params.N, dtype=np.complex64) # Default gain perturbation + + # Set inter-element spacing + self.define_scenario_params() + # Initialize array geometry + self.create_array() + + + def define_scenario_params(self): + """Simplified parameter definition for narrowband far-field""" + # Distance between array elements (half wavelength spacing) + self.dist_array_elems = self.params.wavelength / 2 + + def create_array(self): + """ + Create the antenna array geometry + Location perturbation are applied here""" + N = self.params.N + # Linear array: [0, 1, 2, ..., N-1] * (λ/2) + base_array = np.arange(N, dtype=float) * self.dist_array_elems + # add a perturbation sampled uniformly between + # -self.params.location_perturbation and self.params.location_perturbation for middle elements, + # for the first element sampled uniformly between 0 and self.params.location_perturbation, + # and for the last element sampled uniformly between -self.params.location_perturbation and 0. + # This makes sure that the over all size is no more than (N-1)*\lambda/2 + if self.params.location_perturbation is not None: + first_element_perturbation = np.random.uniform(0, self.params.location_perturbation) + mid_elements_perturbation = np.random.uniform(-self.params.location_perturbation, + self.params.location_perturbation, N-2) + last_element_perturbation = np.random.uniform(-self.params.location_perturbation, 0) + perturbation = np.concatenate(([first_element_perturbation], mid_elements_perturbation, + [last_element_perturbation]), axis=0) + base_array += perturbation + + self.array = torch.from_numpy(base_array).to(torch.float64) # Convert to torch tensor + + def steering_vec(self, angles: np.ndarray) -> torch.Tensor: + """ + Compute steering vector for far-field sources + Args: - ----- - system_model_params (SystemModelParams): an instance of SystemModelParams, - containing all relevant system model parameters. + angles: Array of source angles in radians + + Returns: + Complex steering matrix of shape (N, M) + """ + return self.steering_vec_far_field(angles) + + #TODO: Check that the steering matrix is implemented correctly for the far-field case + + def steering_vec_far_field(self, angles: np.ndarray | torch.Tensor) -> torch.Tensor: + """ + Compute far-field steering vectors + + Args: + angles: Array of source angles in RADIANS!!!! (M,) + + Returns: + Complex steering matrix (N, M) including gain and location perturbations if specified + """ + if isinstance(angles, np.ndarray): + angles = torch.from_numpy(angles) + + # Convert to float64 for precision + angles = angles.to(torch.float64) + array = self.array.view(-1,1) # (N, 1) + + # Ensure angles is 2D: (1, M) for broadcasting + if angles.dim() == 1: + angles = angles.unsqueeze(0) # (1, M) + + + # Phase delays: (N, 1) * sin(angles) -> (N, M) + time_delay = array @ torch.sin(angles) # Broadcasting: (N,1) * (1,M) -> (N,M) + + # Steering matrix: complex exponentials + steering_matrix = torch.exp(-2j * torch.pi * time_delay / self.params.wavelength).to(torch.complex64) # (N, M) + + if self.params.gain_perturbation_var > 0.0: + gain_impairments = (torch.ones(self.params.N, dtype=torch.complex64) + + torch.randn(self.params.N, dtype=torch.complex64) * torch.sqrt(torch.tensor(self.params.gain_perturbation_var))) + + self.gain_petrurbation = gain_impairments + + #make a diagonal matrix of the gain impairments + diag_gain_impairments = torch.diag(gain_impairments) + normalization_factor = 1/torch.linalg.norm(gain_impairments, 2, dtype=torch.complex64) + # Apply gain perturbations + steering_matrix = normalization_factor * diag_gain_impairments @ steering_matrix + + return steering_matrix + + def get_array(self): + """ + Get the current array geometry + + Returns: + Numpy array of antenna positions + """ + return self.array + def get_antenna_gains(self): + """ + Get the gain perturbation applied to the steering vector + + Returns: + Numpy array of gain perturbations """ + return self.gain_petrurbation if self.gain_petrurbation is not None else np.ones(self.params.N, dtype=torch.complex64) + +class Samples(SystemModel): + """ + Simplified Samples class for signal and noise generation + Inherits from SystemModel for steering vector computation + + Removed: + - Distance handling (near-field) + - Coherent signal support + - Broadband signal support + - Variable M support within single sample + """ + + def __init__(self, system_model_params: SystemModelParams): super().__init__(system_model_params) self.angles = None - self.distances = None - - def set_labels(self, number_of_sources: int, angles: list, distances: list): - if self.params.field_type.lower() == "far": - self.set_angles(angles, number_of_sources) - elif self.params.field_type.lower() in {"near", "full"}: - self.set_angles(angles, number_of_sources) - self.set_distances(distances, number_of_sources) - else: - raise ValueError(f"Samples.set_labels: Field type {self.params.field_type} is not defined") - + + def set_labels(self, angles: list | None = None): + """ + Set the angles for the sources + + Args: + angles: List of angles in degrees, or None for random generation + """ + self.set_angles(angles) + def get_labels(self): - if self.params.field_type.lower() == "far": - return torch.tensor(self.angles, dtype=torch.float32) - elif self.params.field_type.lower() in {"near", "full"}: - labels = torch.cat((torch.tensor(self.angles, dtype=torch.float32), torch.tensor(self.distances, dtype=torch.float32)), dim=0) - return labels - else: - raise ValueError(f"Samples.get_labels: Field type {self.params.field_type} is not defined") - - - def set_angles(self, doa: list, M: int): + """Get the current source angles as tensor""" + return torch.tensor(self.angles, dtype=torch.float32) + + def set_angles(self, doa: list | None = None): """ - Sets the direction of arrival (DOA) for the signals. - + Set direction of arrival angles + Args: - ----- - doa (np.ndarray): Array containing the DOA values. - + doa: List of angles in degrees, or None for random generation """ - - def create_doa_with_gap(gap: float, M: int): - """Create angles with a value gap. - - Args: - ----- - gap (float): Minimal gap value. - - Returns: - -------- - np.ndarray: DOA array. - - """ - # LEGACY CODE - # while True: - # # DOA = np.round(np.random.rand(M) * 180, decimals=2) - 90 - # DOA = np.random.randint(-55, 55, M) - # DOA.sort() - # diff_angles = np.array( - # [np.abs(DOA[i + 1] - DOA[i]) for i in range(M - 1)] - # ) - # if (np.sum(diff_angles > gap) == M - 1) and ( - # np.sum(diff_angles < (180 - gap)) == M - 1 - # ): - # break - - # based on https://stackoverflow.com/questions/51918580/python-random-list-of-numbers-in-a-range-keeping-with-a-minimum-distance - doa_range = self.params.doa_range - doa_resolution = self.params.doa_resolution - if doa_resolution <= 0: - raise ValueError("DOA resolution must be positive.") - if M <= 0: - raise ValueError("M (number of elements) must be positive.") - if gap <= 0: - raise ValueError("Gap must be positive.") - - # Compute the range of possible DOA values - # Ensure the sampled DOAs do not exceed [-doa_range, +doa_range] - max_offset = (gap - 1) * (M - 1) - effective_range = 2 * doa_range - max_offset - if effective_range <= 0: - raise ValueError(f"Invalid effective range: {effective_range}. Check your parameters.") - - # Define the valid range for sampling - if doa_resolution >= 1: - valid_range = range(0, effective_range, doa_resolution) - sampled_values = sorted(sample(valid_range, M)) - else: - step_count = int(effective_range // doa_resolution) - valid_range = range(step_count) - sampled_values = sorted(sample(valid_range, M)) - sampled_values = [x * doa_resolution for x in sampled_values] - - # Compute DOAs - DOA = [(gap - 1) * i + x - doa_range for i, x in enumerate(sampled_values)] - - # Ensure all DOAs fall naturally within the valid range - if any(d < -doa_range or d > doa_range for d in DOA): - raise ValueError("Computed DOAs exceed the valid range. Check your logic.") - - # Round results to 3 decimal places - DOA = np.round(DOA, 3) - - return DOA - - if doa == None: - # Generate angels with gap greater than 0.2 rad (nominal case) - self.angles = np.deg2rad(np.array(create_doa_with_gap(gap=10, M=M))) + if doa is None: + # Generate random angles with minimum separation + self.angles = np.deg2rad(self._create_doa_with_gap(gap=10)) else: - # Generate - self.angles = np.deg2rad(doa) - - def set_distances(self, distance: list | np.ndarray, M: int) -> np.ndarray: + # Use provided angles + self.angles = np.deg2rad(np.array(doa)) + + def _create_doa_with_gap(self, gap: float = 10): """ - + Create angles with minimum separation (same logic as original) + Args: - distance: - + gap: Minimum separation in degrees + Returns: - + List of angles in degrees """ - - def choose_distances(M, min_val: float, max_val: int, distance_resolution: float = 1.0) -> np.ndarray: - """ - Choose distances for the sources. - - Args: - M (int): Number of sources. - min_val (float): Minimal value of the distances. - max_val (int): Maximal value of the distances. - distance_resolution (float, optional): Resolution of the distances. Defaults to 1.0. - - """ - distances_options = np.arange(min_val, max_val, distance_resolution) - distances = np.random.choice(distances_options, M, replace=True) - return np.round(distances, 3) - - if distance is None: - self.distances = choose_distances(M, min_val=np.ceil(self.fresnel) + self.params.range_resolution, - max_val=np.floor(self.fraunhofer * self.params.max_range_ratio_to_limit), - distance_resolution=self.params.range_resolution) - else: - self.distances = np.array(distance) - - def samples_creation( - self, - noise_mean: float = 0, - noise_variance: float = 1, - signal_mean: float = 0, - signal_variance: float = 1, - source_number: int = None, - ): - """Creates samples based on the specified mode and parameters. - + M = self.params.M + doa_range = self.params.doa_range # 0-90 degrees + doa_resolution = self.params.doa_resolution + + # Compute effective range - reserving slots for the M-1 gaps between the M sources for allocation + max_offset = (gap - 1) * (M - 1) + effective_range = 2 * doa_range - max_offset + + if effective_range <= 0: + raise ValueError(f"Cannot fit {M} sources with {gap}° separation in ±{doa_range}° range") + + # Sample positions + valid_range = np.arange(0, effective_range, doa_resolution) + sampled_values = sorted(sample(valid_range.tolist(), M)) + + # Convert to actual angles - put a gap of at least `gap` degrees between sources + # and shift the angles to start from -doa_range + DOA = [(gap - 1) * i + x - doa_range for i, x in enumerate(sampled_values)] + + return np.round(DOA, 3) + + def samples_creation(self, + noise_mean: float = 0, + noise_variance: float = 1, + signal_mean: float = 0, + signal_variance: float = 1): + """ + Create observation samples: X = A @ S + N + The noise variance is as desribed above and the signal variance is multiplied by the SNR in linear scale. + Args: - ----- - noise_mean (float, optional): Mean of the noise. Defaults to 0. - noise_variance (float, optional): Variance of the noise. Defaults to 1. - signal_mean (float, optional): Mean of the signal. Defaults to 0. - signal_variance (float, optional): Variance of the signal. Defaults to 1. - + noise_mean: Mean of noise (typically 0) + noise_variance: Variance of noise (typically 1) + signal_mean: Mean of signal (typically 0) + signal_variance: Variance of signal (typically 1) + Returns: - -------- - tuple: Tuple containing the created samples, signal, steering vectors, and noise. - - Raises: - ------- - Exception: If the signal_type is not defined. - + tuple: (samples, signal, steering_matrix, noise) + - samples: Complex observations (N, T) + - signal: Complex source signals (M, T) + - steering_matrix: Complex steering matrix (N, M) + - noise: Complex noise (N, T) """ - # Generate signal matrix - signal = self.signal_creation(signal_mean, signal_variance, source_number=source_number) - signal = torch.from_numpy(signal) - # Generate noise matrix - noise = self.noise_creation(noise_mean, noise_variance) - noise = torch.from_numpy(noise) - if self.params.signal_type.startswith("broadband"): - raise Exception("Samples.samples_creation: Broadband signal type is not defined for far field") - if self.params.field_type.startswith("far"): - A = self.steering_vec(self.angles, f_c=self.f_rng[self.params.signal_type]) - samples = (A @ signal) + noise - elif self.params.field_type.startswith("near"): - A = self.steering_vec(angles=self.angles, ranges=self.distances, nominal=False, generate_search_grid=False, - f_c=self.f_rng[self.params.signal_type]) - samples = (A @ signal) + noise - elif self.params.field_type.startswith("full"): - A = self.steering_vec_full_model(angles=self.angles, - ranges=self.distances) - samples = (A @ signal) + noise - else: - raise Exception(f"Samples.params.field_type: Field type {self.params.field_type} is not defined") + # Generate source signals (non-coherent) + signal = self.signal_creation(signal_mean, signal_variance) + signal = torch.from_numpy(signal).to(torch.complex64) + + # Generate noise + noise = self.noise_creation(noise_mean, noise_variance) + noise = torch.from_numpy(noise).to(torch.complex64) + + # Compute steering matrix for current angles + A = self.steering_vec(self.angles) + + # Create observations: X = A @ S + N + samples = (A @ signal) + noise + return samples, signal, A, noise - - def noise_creation(self, noise_mean, noise_variance): - """Creates noise based on the specified mean and variance. - + + def signal_creation(self, signal_mean: float = 0, signal_variance: float = 1): + """ + Generate non-coherent source signals + Args: - ----- - noise_mean (float): Mean of the noise. - noise_variance (float): Variance of the noise. - + signal_mean: Mean of signals + signal_variance: Variance of signals + Returns: - -------- - np.ndarray: Generated noise. - + Complex signal matrix (M, T) """ - # for NarrowBand signal_type Noise represented in the time domain - noise = ( - np.sqrt(noise_variance) - * (np.sqrt(2) / 2) - * ( - np.random.randn(self.params.N, self.params.T) - + 1j * np.random.randn(self.params.N, self.params.T) - ) - + noise_mean - ) - return noise - - def signal_creation(self, signal_mean: float = 0, signal_variance: float = 1, source_number: int = None): + M, T = self.params.M, self.params.T + + # Convert SNR from dB to linear scale + amplitude = np.sqrt(10 ** (self.params.snr / 10)) #NOTE: the SNR is the power ration between the variances, the sqrt is needed than to convert to std. + + # Generate M independent complex Gaussian signals + signals = (amplitude * (np.sqrt(2) / 2) * np.sqrt(signal_variance) * + (np.random.randn(M, T) + 1j * np.random.randn(M, T)) + signal_mean) + + return signals + + def noise_creation(self, noise_mean: float = 0, noise_variance: float = 1): """ - Creates signals based on the specified signal nature and parameters. - + Generate complex white Gaussian noise + Args: - ----- - signal_mean (float, optional): Mean of the signal. Defaults to 0. - signal_variance (float, optional): Variance of the signal. Defaults to 1. - + noise_mean: Mean of noise + noise_variance: Variance of noise + Returns: - -------- - np.ndarray: Created signals. - - Raises: - ------- - Exception: If the signal type is not defined. - Exception: If the signal nature is not defined. + Complex noise matrix (N, T) """ - M = source_number - if self.params.snr is None: - snr = np.random.uniform(-5, 5) - else: - snr = self.params.snr - amplitude = 10 ** (snr / 10) - # NarrowBand signal creation - if self.params.signal_type == "narrowband": - if self.params.signal_nature == "non-coherent": - # create M non-coherent signals - return ( - amplitude - * (np.sqrt(2) / 2) - * np.sqrt(signal_variance) - * ( - np.random.randn(M, self.params.T) - + 1j * np.random.randn(M, self.params.T) - ) - + signal_mean - ) + N, T = self.params.N, self.params.T + + # Generate complex white Gaussian noise + noise = (np.sqrt(noise_variance) * (np.sqrt(2) / 2) * + (np.random.randn(N, T) + 1j * np.random.randn(N, T)) + noise_mean) + + return noise + + +def create_single_sample(system_model_params: SystemModelParams, + true_doa: list = None): + """ + Simplified version of create_dataset for single sample generation + + Args: + system_model_params: System model parameters + true_doa: Predefined angles in degrees, or None for random + + Returns: + tuple: (observations, labels, samples_model) + - observations: Complex matrix (N, T) + - labels: Source angles in radians + - samples_model: The Samples object used + """ + # Create samples model + samples_model = Samples(system_model_params) + + # Set source angles + samples_model.set_labels(true_doa) + + # Generate observations + X, signal, A, noise = samples_model.samples_creation( + noise_mean=0, noise_variance=1, + signal_mean=0, signal_variance=1 + ) + + # Get ground truth labels + Y = samples_model.get_labels() + + return X, Y, samples_model - elif self.params.signal_nature == "coherent": - # Coherent signals: same amplitude and phase for all signals - sig = ( - amplitude - * (np.sqrt(2) / 2) - * np.sqrt(signal_variance) - * ( - np.random.randn(1, self.params.T) - + 1j * np.random.randn(1, self.params.T) - ) - + signal_mean - ) - return np.repeat(sig, M, axis=0) - else: - raise Exception(f"signal type {self.params.signal_type} is not defined") diff --git a/src/subspace_method.py b/src/subspace_method.py new file mode 100644 index 0000000..8b58d23 --- /dev/null +++ b/src/subspace_method.py @@ -0,0 +1,244 @@ +import torch +import torch.nn as nn +import matplotlib.pyplot as plt +import wandb +import numpy as np +import warnings + + +from src.utils import sample_covariance, spatial_smoothing_covariance + + + +class SubspaceMethod(nn.Module): + """ + Basic methods for all subspace methods. + """ + + def __init__(self, system_model, model_order_estimation: str = None, + physical_array: torch.Tensor = None, physical_gains: torch.Tensor = None): + super(SubspaceMethod, self).__init__() + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.system_model = system_model + self.eigen_threshold = nn.Parameter(torch.tensor(.5, requires_grad=False)) + self.normalized_eigenvals = None + self.normalized_eigenvals_mean = None + self.model_order_estimation = model_order_estimation + self.physical_array = physical_array + self.physical_gains = physical_gains + + def subspace_separation(self, + covariance: torch.Tensor, + number_of_sources: torch.tensor = None) \ + -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.tensor]: + """ + + Args: + covariance: from size (B, N, N) where B is the batch size and N is the number of antennas. + number_of_sources: if None it estimates the number of sources using the model_order_estimation method. + + Returns: + the signal and noise subspaces, both as torch.Tensor(), number of sources estimation, + and the regularization term for the eigenvalues (not used for now). + """ + eigenvalues, eigenvectors = torch.linalg.eigh(covariance) + sorted_idx = torch.argsort(torch.abs(eigenvalues), descending=True) + sorted_eigvectors = torch.gather(eigenvectors, 2, + sorted_idx.unsqueeze(-1).expand(-1, -1, covariance.shape[-1]).transpose(1, 2)) + # number of sources estimation + source_estimation, l_eig = self.estimate_number_of_sources(eigenvalues, + number_of_sources=number_of_sources) + if number_of_sources is None: + warnings.warn("Number of sources is not defined, using the number of sources estimation.") + # if source_estimation == sorted_eigvectors.shape[2]: + # source_estimation -= 1 + signal_subspace = sorted_eigvectors[:, :, :source_estimation] + noise_subspace = sorted_eigvectors[:, :, source_estimation:] + else: + signal_subspace = sorted_eigvectors[:, :, :number_of_sources] + noise_subspace = sorted_eigvectors[:, :, number_of_sources:] + + return signal_subspace.to(self.device), noise_subspace.to(self.device), source_estimation, l_eig + + def estimate_number_of_sources(self, eigenvalues, number_of_sources: int = None): + """ + + Args: + eigenvalues: + + Returns: + + """ + sorted_eigenvals = torch.sort(torch.real(eigenvalues), descending=True, dim=1).values + # try: + # if self.normalized_eigenvals_mean is None: + # self.normalized_eigenvals_mean = torch.mean(sorted_eigenvals, dim=0) + # else: + # self.normalized_eigenvals_mean = 0.9 * self.normalized_eigenvals_mean + 0.1 * torch.mean(sorted_eigenvals, dim=0) + # wandb.config.update({"eigenvalues": wandb.Histogram(self.normalized_eigenvals_mean.cpu().detach().numpy())}) + # except Exception: + # pass + l_eig = None + if self.model_order_estimation is None: + return None, None + elif self.model_order_estimation.lower().startswith("threshold"): + self.normalized_eigenvals = sorted_eigenvals / sorted_eigenvals[:, 0][:, None] + source_estimation = torch.linalg.norm( + nn.functional.relu( + self.normalized_eigenvals - self.__get_eigen_threshold() * torch.ones_like(self.normalized_eigenvals)), + dim=1, ord=0).to(torch.int) + # return regularization term if training + if self.training: + l_eig = self.eigen_regularization(number_of_sources) + elif self.model_order_estimation.lower() in ["mdl", "aic"]: + # mdl -> calculate the value of the mdl test for each number of sources + # and choose the number of sources that minimizes the mdl test + optimal_test = torch.ones(eigenvalues.shape[0], device=self.device) * float("inf") + optimal_m = torch.zeros(eigenvalues.shape[0], device=self.device) + for m in range(1, eigenvalues.shape[1]): + m = torch.tensor(m, device=self.device) + # calculate the test + test = self.hypothesis_testing(sorted_eigenvals, m) + # update the optimal number of sources by masking the current number of sources + optimal_m = torch.where(test < optimal_test, m, optimal_m) + # update the optimal mdl value + optimal_test = torch.where(test < optimal_test, test, optimal_test) + if self.training and m == number_of_sources: + # l_eig = torch.sum(test) + l_eig = test + + source_estimation = optimal_m + + else: + raise ValueError(f"SubspaceMethod.estimate_number_of_sources: method {self.model_order_estimation.lower()} is not recognized.") + return source_estimation, l_eig + + def hypothesis_testing(self, eigenvalues, number_of_sources): + # extract the number of snapshots and the number of antennas + T = self.system_model.params.T + N = self.system_model.params.N + M = number_of_sources + # calculate the number of degrees of freedom + # dof = (2 * M) * (N - M) + dof = (2 * N * M - M ** 2 + 1) / 2 + if self.model_order_estimation.lower().startswith("mdl"): + penalty = dof * np.log(T) + # penalty = dof * np.log(T) + else: # self.model_order_estimation.lower().startswith("aic"): + penalty = dof * 2 + ll = self.get_ll(eigenvalues, M) + mdl = ll + penalty + return mdl + + def snr_estimation(self, eigenvalues, M): + snr = 10 * torch.log10(torch.mean(eigenvalues[:, :M], dim=1) / torch.mean(eigenvalues[:, M:], dim=1)) + return snr + + def get_ll(self, eigenvalues, M): + T = self.system_model.params.T + N = self.system_model.params.N + ll = -T * torch.sum(torch.log(eigenvalues[:, M:]), dim=1) + T * (N - M) * torch.log(torch.mean(eigenvalues[:, M:], dim=1)) + return ll + + def get_noise_subspace(self, covariance: torch.Tensor, number_of_sources: int): + """ + + Args: + covariance: + number_of_sources: + + Returns: + + """ + _, noise_subspace, _, _ = self.subspace_separation(covariance, number_of_sources) + return noise_subspace + + def get_signal_subspace(self, covariance: torch.Tensor, number_of_sources: int): + """ + + Args: + covariance: + number_of_sources: + + Returns: + + """ + signal_subspace, _, _, _ = self.subspace_separation(covariance, number_of_sources) + return signal_subspace + + def eigen_regularization(self, number_of_sources: int): + """ + + Args: + normalized_eigenvalues: + number_of_sources: + + Returns: + + """ + l_eig = (self.normalized_eigenvals[:, number_of_sources - 1] - self.__get_eigen_threshold(level="high")) * \ + (self.normalized_eigenvals[:, number_of_sources] - self.__get_eigen_threshold(level="low")) + # l_eig = -(self.normalized_eigen[:, number_of_sources - 1] - self.__get_eigen_threshold(level="high")) + \ + # (self.normalized_eigen[:, number_of_sources] - self.__get_eigen_threshold(level="low")) + # l_eig = torch.sum(l_eig) + # eigen_regularization = nn.functional.elu(eigen_regularization, alpha=1.0) + return l_eig + + def test_step(self, batch, batch_idx): + raise NotImplementedError + + def __init_criteria(self): + raise NotImplementedError + + def __get_eigen_threshold(self, level: str = None): + # if self.training: + # if level is None: + # return self.eigen_threshold + # elif level == "high": + # return self.eigen_threshold + 0.0 + # elif level == "low": + # return self.eigen_threshold - 0.0 + # else: + # if self.system_model.params.M is not None: + # return self.eigen_threshold - self.system_model.params.M / self.system_model.params.N + # else: + # return self.eigen_threshold - 0.1 + return self.eigen_threshold + + def pre_processing(self, x: torch.Tensor, mode: str = "sample"): + if mode == "sample": + Rx = sample_covariance(x) + elif mode == "sps": + Rx = spatial_smoothing_covariance(x) + else: + raise ValueError( + f"SubspaceMethod.pre_processing: method {mode} is not recognized for covariance calculation.") + + return Rx + + + def plot_eigen_spectrum(self, batch_idx: int=0, normelized_eign: torch.Tensor = None): + """ + Plot the eigenvalues spectrum. + + Args: + ----- + batch_idx (int): Index of the batch to plot. + """ + if normelized_eign is None: + normelized_eign = self.normalized_eigenvals + plt.figure() + plt.stem(normelized_eign[batch_idx].cpu().detach().numpy(), label="Normalized Eigenvalues") + # ADD threshold line + plt.axhline(y=self.__get_eigen_threshold().cpu().detach().numpy(), color='r', linestyle='--', label="Threshold") + plt.title("Eigenvalues Spectrum") + plt.xlabel("Eigenvalue Index") + plt.ylabel("Eigenvalue") + plt.legend() + plt.grid() + plt.show() + + def source_estimation_accuracy(self, sources_num, source_estimation): + if sources_num is None or source_estimation is None: + return 0 + return torch.sum(source_estimation == sources_num * torch.ones_like(source_estimation).float()).item() \ No newline at end of file diff --git a/src/utils.py b/src/utils.py index 463ffcb..22d3ab9 100644 --- a/src/utils.py +++ b/src/utils.py @@ -1,25 +1,3 @@ -"""Subspace-Net -Details ----------- -Name: utils.py -Authors: D. H. Shmuel -Created: 01/10/21 -Edited: 17/03/23 - -Purpose: --------- -This script defines some helpful functions: - * sum_of_diag: returns the some of each diagonal in a given matrix. - * sum_of_diag_torch: returns the some of each diagonal in a given matrix, Pytorch oriented. - * find_roots: solves polynomial equation defines by polynomial coefficients. - * find_roots_torch: solves polynomial equation defines by polynomial coefficients, Pytorch oriented.. - * set_unified_seed: Sets unified seed for all random attributed in the simulation. - * get_k_angles: Retrieves the top-k angles from a prediction tensor. - * get_k_peaks: Retrieves the top-k peaks (angles) from a prediction tensor using peak finding. - * gram_diagonal_overload(self, Kx: torch.Tensor, eps: float): generates Hermitian and PSD (Positive Semi-Definite) matrix, - using gram operation and diagonal loading. -""" - # Imports import numpy as np import torch @@ -28,63 +6,116 @@ import random import scipy import warnings +from pathlib import Path +import pickle from pathlib import Path -from src.config import device import matplotlib.pyplot as plt import torch.nn as nn +from datetime import datetime + + + +def plot_spectrums(angles_grid: torch.Tensor, + spectrums_dict : dict, + true_angles: torch.Tensor = None, + save: bool = False, + title: str = "MUSIC spectrums" + ): + + """ Plots the spectrums for given angles and spectrums dictionary. + Args: + angles_grid (torch.Tensor): A tensor containing the angles in radians. + spectrums_dict (dict): A dictionary where keys are labels and values are tensors of spectrum values. + true_angles (torch.Tensor, optional): A tensor containing the true angles in radians to highlight on the plot. + save (bool, optional): If True, saves the plot as a PDF file. Defaults to False. + title (str, optional): Title of the plot. Defaults to "MUSIC spectrums"." + + """ + + angles_deg = torch.rad2deg(angles_grid).cpu().numpy() + plt.figure(figsize=(10, 6)) + for key, value in spectrums_dict.items(): + if not isinstance(value, torch.Tensor): + raise TypeError(f"plot_spectrums: Expected torch.Tensor, got {type(value)}") + elif len(angles_grid) != len(value): + raise ValueError(f"plot_spectrums: Length of angles_grid ({len(angles_grid)}) " + f"does not match length of spectrum ({len(value)}).") + plt.plot(angles_deg, value.cpu().detach().numpy(), 'r--', linewidth=1, label= f"{key}") + + + if true_angles is not None: + highlight_deg = torch.rad2deg(true_angles).cpu().numpy() + for i, angle in enumerate(highlight_deg): + plt.axvline(x=angle, color='r', linestyle='--', alpha=0.7, + label='True DoA' if i == 0 else "") + + plt.xlabel('Angle [degrees]') + plt.ylabel('Spectrum Power') + plt.title(title) + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + + if save: + plt.savefig('diffmusic_spectrum.pdf') + plt.show() + +def plot_parameters(results_dict : dict): + """Plots the array and gains for each method in the results dictionary. + Need to also implement for refrence the nominal array and gains.""" + #TODO: implement as stated above. + pass + + +def save_data_to_file(data_path: Path, *args): + """ + Saves the provided data to a file in the specified data path. -# Constants -R2D = 180 / np.pi -D2R = 1 / R2D -plot_styles = { - 'CCRB': {'color': 'r', 'linestyle': '-', 'marker': 'o', "markersize": 8}, - 'Beamformer': {'color': 'r', 'linestyle': '--', 'marker': 's', "markersize": 8}, - 'DCD-MUSIC': {'color': 'g', 'linestyle': '-', 'marker': 'D', "markersize": 8}, - 'DCD-MUSIC_V2': {'color': 'g', 'linestyle': '--', 'marker': 'd', "markersize": 8}, - 'TransMUSIC': {'color': 'm', 'linestyle': '-.', 'marker': 'P', "markersize": 8}, - '2D-MUSIC': {'color': 'c', 'linestyle': ':', 'marker': '^', "markersize": 8}, - '2D-MUSIC(SPS)': {'color': 'c', 'linestyle': '--', 'marker': 'v', "markersize": 8}, - 'SubspaceNet': {'color': 'k', 'linestyle': '-', 'marker': 'X', "markersize": 8}, - 'NFSubspaceNet': {'color': 'k', 'linestyle': '--', 'marker': 'p', "markersize": 8}, - 'NFSubspaceNet_V2': {'color': 'b', 'linestyle': '-.', 'marker': 'h', "markersize": 8}, - 'ESPRIT': {'color': 'r', 'linestyle': '-', 'marker': 'v', "markersize": 8}, - 'esprit(SPS)': {'color': 'r', 'linestyle': '--', 'marker': 'v', "markersize": 8}, - '1D-MUSIC': {'color': 'y', 'linestyle': '-.', 'marker': 's', "markersize": 8}, - 'music(SPS)': {'color': 'y', 'linestyle': ':', 'marker': 's', "markersize": 8}, -} - -def validate_constant_sources_number(number_of_sources: torch.tensor): - """ - Validate that the number of sources in the batch is equal for all samples. Args: - number_of_sources: The number of sources in the batch. + data_path (Path): The path where the data will be saved. + *args: Data to be saved. Returns: None - Raises: - ValueError: If the number of sources in the batch is not equal for all samples - """ - if (number_of_sources != number_of_sources[0]).any(): - raise ValueError(f"validate_constant_sources_number: " - f"Number of sources in the batch is not equal for all samples.") + with open(data_path / "data.pkl", "wb") as f: + pickle.dump(args, f) -def initialize_data_paths(path: Path): - datasets_path = path / "datasets" - simulations_path = path / "simulations" - saving_path = path / "weights" +def load_data_from_file(data_path: str): + """ + Loads data from a file in the specified data path. + Args: + data_path (Path): The path from where the data will be loaded. + Returns: + tuple: A tuple containing the loaded data. + Raises: + FileNotFoundError: If the specified data file does not exist. + Examples: + >>> measurements, signals, steering_mat, noise, true_angles, array, gain_impairments = load_data_from_file(Path("data")) + """ + data_file = Path(data_path) + with open(data_file, "rb") as f: + data = pickle.load(f) + if not data: + raise FileNotFoundError(f"load_data_from_file: No data found in {data_file}.") + return data + +def initialize_paths(main_path: Path, system_model_params) -> tuple: + dt_string_for_save = system_model_params.dt_string_for_save + indicating_str = (f"N:{system_model_params.N}_M:{system_model_params.M}_T:{system_model_params.T}_" + + f"snr:{system_model_params.snr}_location_pert_boundary:{system_model_params.location_perturbation}_" + + f"gain_perturbation_var:{system_model_params.gain_perturbation_var}_" + + f"seed:{system_model_params.seed}") + datasets_path = main_path / "datasets" / indicating_str / dt_string_for_save + results_path = main_path / "results" / indicating_str / dt_string_for_save # create folders if not exists datasets_path.mkdir(parents=True, exist_ok=True) - (datasets_path / "train").mkdir(parents=True, exist_ok=True) - (datasets_path / "test").mkdir(parents=True, exist_ok=True) - simulations_path.mkdir(parents=True, exist_ok=True) - saving_path.mkdir(parents=True, exist_ok=True) - (saving_path / "final_models").mkdir(parents=True, exist_ok=True) + results_path.mkdir(parents=True, exist_ok=True) - return datasets_path, simulations_path, saving_path + return datasets_path, results_path def sample_covariance(x: torch.Tensor) -> torch.Tensor: """ @@ -99,7 +130,7 @@ def sample_covariance(x: torch.Tensor) -> torch.Tensor: covariance_mat (np.ndarray): Covariance matrix. """ if x.dim() == 2: - x = x[None, :, :] + x = x.unsqueeze(0) # Add batch dimension if not present batch_size, sensor_number, samples_number = x.shape Rx = torch.einsum("bmt, btl -> bml", x, torch.conj(x).transpose(1, 2)) / samples_number return Rx @@ -132,177 +163,6 @@ def spatial_smoothing_covariance(x: torch.Tensor): # Divide overall matrix by the number of sources return Rx_smoothed -def tops_covariance(x: torch.Tensor, number_of_bins: int=1): - """ - Tops algorithm uses K bins to calculate the covariance by using STFT. - Args: - x: - number_of_bins: - - - Returns: - - """ - Rx = torch.zeros(x.shape[0], number_of_bins,x.shape[1], x.shape[1], dtype=torch.complex128, device=device) - bin_size = x.shape[2] // number_of_bins - for i in range(number_of_bins): - x_bin = x[:, :, i*bin_size:(i+1)*bin_size] - Rx[:, i, :, :] = sample_covariance(x_bin) - return Rx - - -def keep_far_enough_points(tensor, M, D): - # # Calculate pairwise distances between columns - # distances = cdist(tensor.T, tensor.T, metric="euclidean") - # - # # Keep the first M columns as far enough points - # selected_cols = [] - # for i in range(tensor.shape[1]): - # if len(selected_cols) >= M: - # break - # if all(distances[i, col] >= D for col in selected_cols): - # selected_cols.append(i) - # - # # Remove columns that are less than distance D from each other - # filtered_tensor = tensor[:, selected_cols] - # retrun filtered_tensor - ############################################## - # Extract x_coords (first dimension) - x_coords = tensor[0, :] - - # Keep the first M columns that are far enough apart in x_coords - selected_cols = [] - for i in range(tensor.shape[1]): - if len(selected_cols) >= M: - break - if i == 0: - selected_cols.append(i) - continue - if all(abs(x_coords[i] - x_coords[col]) >= D for col in selected_cols): - selected_cols.append(i) - - # Select the columns that meet the distance criterion - filtered_tensor = tensor[:, selected_cols] - - return filtered_tensor - -# Functions -# def sum_of_diag(matrix: np.ndarray) -> list: -def sum_of_diag(matrix: np.ndarray): - """Calculates the sum of diagonals in a square matrix. - - Args: - matrix (np.ndarray): Square matrix for which diagonals need to be summed. - - Returns: - list: A list containing the sums of all diagonals in the matrix, from left to right. - - Raises: - None - - Examples: - >>> matrix = np.array([[1, 2, 3], - [4, 5, 6], - [7, 8, 9]]) - >>> sum_of_diag(matrix) - [7, 12, 15, 8, 3] - - """ - diag_sum = [] - diag_index = np.linspace( - -matrix.shape[0] + 1, - matrix.shape[0] + 1, - 2 * matrix.shape[0] - 1, - endpoint=False, - dtype=int, - ) - for idx in diag_index: - diag_sum.append(np.sum(matrix.diagonal(idx))) - return diag_sum - - -def sum_of_diags_torch(matrix: torch.Tensor): - """Calculates the sum of diagonals in a square matrix. - equivalent sum_of_diag, but support Pytorch. - - Args: - matrix (torch.Tensor): Square matrix for which diagonals need to be summed. - - Returns: - torch.Tensor: A list containing the sums of all diagonals in the matrix, from left to right. - - Raises: - None - - Examples: - >>> matrix = torch.tensor([[1, 2, 3], - [4, 5, 6], - [7, 8, 9]]) - >>> sum_of_diag(matrix) - torch.tensor([7, 12, 15, 8, 3]) - """ - diag_sum = [] - diag_index = torch.linspace( - -matrix.shape[0] + 1, matrix.shape[0] - 1, 2 * matrix.shape[0] - 1, dtype=int - ) - for idx in diag_index: - diag_sum.append(torch.sum(torch.diagonal(matrix, idx))) - return torch.stack(diag_sum, dim=0) - - -# def find_roots(coefficients: list) -> np.ndarray: -def find_roots(coefficients: list): - """Finds the roots of a polynomial defined by its coefficients. - - Args: - coefficients (list): List of polynomial coefficients in descending order of powers. - - Returns: - np.ndarray: An array containing the roots of the polynomial. - - Raises: - None - - Examples: - >>> coefficients = [1, -5, 6] # x^2 - 5x + 6 - >>> find_roots(coefficients) - array([3., 2.]) - - """ - coefficients = np.array(coefficients) - A = np.diag(np.ones((len(coefficients) - 2,), coefficients.dtype), -1) - if np.abs(coefficients[0]) == 0: - A[0, :] = -coefficients[1:] / (coefficients[0] + 1e-9) - else: - A[0, :] = -coefficients[1:] / coefficients[0] - roots = np.array(np.linalg.eigvals(A)) - return roots - - -def find_roots_torch(coefficients: torch.Tensor): - """Finds the roots of a polynomial defined by its coefficients. - equivalent to src.utils.find_roots, but support Pytorch. - - Args: - coefficients (torch.Tensor): List of polynomial coefficients in descending order of powers. - - Returns: - torch.Tensor: An array containing the roots of the polynomial. - - Raises: - None - - Examples: - >>> coefficients = torch.tensor([1, -5, 6]) # x^2 - 5x + 6 - >>> find_roots(coefficients) - tensor([3., 2.]) - - """ - A = torch.diag(torch.ones(len(coefficients) - 2, dtype=coefficients.dtype), -1) - A[0, :] = -coefficients[1:] / coefficients[0] - roots = torch.linalg.eigvals(A) - return roots - def set_unified_seed(seed: int = 42): """ @@ -314,17 +174,14 @@ def set_unified_seed(seed: int = 42): Returns: None - Raises: - None - Examples: >>> set_unified_seed(42) """ random.seed(seed) np.random.seed(seed) - torch.manual_seed(0) - torch.cuda.manual_seed_all(0) + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False if torch.cuda.is_available(): @@ -333,237 +190,14 @@ def set_unified_seed(seed: int = 42): torch.use_deterministic_algorithms(True) -# def get_k_angles(grid_size: float, k: int, prediction: torch.Tensor) -> torch.Tensor: -def get_k_angles(grid_size: float, k: int, prediction: torch.Tensor): - """ - Retrieves the top-k angles from a prediction tensor. +def save_pickle(data, filepath): + """Save data to pickle file""" + with open(filepath, 'wb') as f: + pickle.dump(data, f) - Args: - grid_size (float): The size of the angle grid (range) in degrees. - k (int): The number of top angles to retrieve. - prediction (torch.Tensor): The prediction tensor containing angle probabilities, sizeof equal to grid_size . - Returns: - torch.Tensor: A tensor containing the top-k angles in degrees. - - Raises: - None +def load_pickle(filepath): + """Load data from pickle file""" + with open(filepath, 'rb') as f: + return pickle.load(f) - Examples: - >>> grid_size = 6 - >>> k = 3 - >>> prediction = torch.tensor([0.1, 0.3, 0.5, 0.2, 0.4, 0.6]) - >>> get_k_angles(grid_size, k, prediction) - tensor([ 90., -18., 54.]) - - """ - angles_grid = torch.linspace(-90, 90, grid_size) - doa_prediction = angles_grid[torch.topk(prediction.flatten(), k).indices] - return doa_prediction - - -# def get_k_peaks(grid_size, k: int, prediction) -> torch.Tensor: -def get_k_peaks(grid_size: int, k: int, prediction: torch.Tensor): - """ - Retrieves the top-k peaks (angles) from a prediction tensor using peak finding. - - Args: - grid_size (int): The size of the angle grid (range) in degrees. - k (int): The number of top peaks (angles) to retrieve. - prediction (torch.Tensor): The prediction tensor containing the peak values. - - Returns: - torch.Tensor: A tensor containing the top-k angles in degrees. - - Raises: - None - - Examples: - >>> grid_size = 6 - >>> k = 3 - >>> prediction = torch.tensor([0.1, 0.3, 0.5, 0.2, 0.4, 0.6]) - >>> get_k_angles(grid_size, k, prediction) - tensor([ 90., -18., 54.]) - - """ - angels_grid = torch.linspace(-90, 90, grid_size) - peaks, peaks_data = scipy.signal.find_peaks( - prediction.detach().numpy().flatten(), prominence=0.05, height=0.01 - ) - peaks = peaks[np.argsort(peaks_data["peak_heights"])[::-1]] - doa_prediction = angels_grid[peaks] - while doa_prediction.shape[0] < k: - doa_prediction = torch.cat( - ( - doa_prediction, - torch.Tensor(np.round(np.random.rand(1) * 180, decimals=2) - 90.00), - ), - 0, - ) - - return doa_prediction[:k] - - -# def gram_diagonal_overload(Kx: torch.Tensor, eps: float) -> torch.Tensor: -def gram_diagonal_overload(Kx: torch.Tensor, eps: float): - """Multiply a matrix Kx with its Hermitian conjecture (gram matrix), - and adds eps to the diagonal values of the matrix, - ensuring a Hermitian and PSD (Positive Semi-Definite) matrix. - - Args: - ----- - Kx (torch.Tensor): Complex matrix with shape [BS, N, N], - where BS is the batch size and N is the matrix size. - eps (float): Constant added to each diagonal element. - - Returns: - -------- - torch.Tensor: Hermitian and PSD matrix with shape [BS, N, N]. - - """ - # Insuring Tensor input - if not isinstance(Kx, torch.Tensor): - Kx = torch.tensor(Kx) - Kx = Kx.to(device) - - # Kx_garm = torch.matmul(torch.transpose(Kx.conj(), 1, 2).to("cpu"), Kx.to("cpu")).to(device) - Kx_garm = torch.bmm(Kx.conj().transpose(1, 2), Kx) - eps_addition = (eps * torch.diag(torch.ones(Kx_garm.shape[-1]))).to(device) - Kx_Out = Kx_garm + eps_addition - - # check if the matrix is Hermitian - A^H = A - mask = (torch.abs(Kx_Out - Kx_Out.conj().transpose(1, 2)) > 1e-6) - if mask.any(): - batch_mask = mask.any(dim=(1,2)) - warnings.warn(f"gram_diagonal_overload: {batch_mask.sum()} matrices in the batch aren't hermitian, taking the average of R and R^H.") - Kx_Out[batch_mask] = 0.5 * (Kx_Out[batch_mask] + Kx_Out[batch_mask].conj().transpose(1, 2)) - - return Kx_Out - - -# def _spatial_smoothing_covariance(sampels: torch.Tensor): -# """ -# Calculates the covariance matrix using spatial smoothing technique. -# -# Args: -# ----- -# X (np.ndarray): Input samples matrix. -# -# Returns: -# -------- -# covariance_mat (np.ndarray): Covariance matrix. -# """ -# -# X = sampels.squeeze() -# N = X.shape[0] -# # Define the sub-arrays size -# sub_array_size = int(N / 2) + 1 -# # Define the number of sub-arrays -# number_of_sub_arrays = N - sub_array_size + 1 -# # Initialize covariance matrix -# covariance_mat = torch.zeros((sub_array_size, sub_array_size), dtype=torch.complex128) -# -# for j in range(number_of_sub_arrays): -# # Run over all sub-arrays -# x_sub = X[j: j + sub_array_size, :] -# # Calculate sample covariance matrix for each sub-array -# sub_covariance = torch.cov(x_sub) -# # Aggregate sub-arrays covariances -# covariance_mat += sub_covariance / number_of_sub_arrays -# # Divide overall matrix by the number of sources -# return covariance_mat - - -def parse_loss_results_for_plotting(loss_results: dict, tested_param: str): - plt_res = {} - plt_acc = False - for test, results in loss_results.items(): - for method, loss_ in results.items(): - if plt_res.get(method) is None: - plt_res[method] = {tested_param: []} - try: - plt_res[method][tested_param].append(loss_[tested_param]) - except KeyError: - plt_res[method][tested_param].append(loss_["Overall"]) - if loss_.get("Accuracy") is not None: - if "Accuracy" not in plt_res[method].keys(): - plt_res[method]["Accuracy"] = [] - plt_acc = True - plt_res[method]["Accuracy"].append(loss_["Accuracy"]) - return plt_res, plt_acc - - -def print_loss_results_from_simulation(loss_results: dict): - """ - Print the loss results from the simulation. - """ - for test, value_dict in loss_results.items(): - print("#" * 10 + f"{test} TEST RESULTS" + "#" * 10) - for test_value, results in value_dict.items(): - if test == "SNR": - print(f"{test} = {test_value} [dB]: ") - else: - print(f"{test} = {test_value}: ") - for method, loss in results.items(): - txt = f"\t{method.upper(): <30}: " - for key, value in loss.items(): - if value is not None: - if key == "Accuracy": - txt += f"{key}: {value * 100:.2f} %|" - else: - txt += f"{key}: {value:.6e} |" - print(txt) - print("\n") - print("\n") - -class AntiRectifier(nn.Module): - def __init__(self, relu_inplace=False): - super(AntiRectifier, self).__init__() - self.relu = nn.ReLU(inplace=relu_inplace) - - def forward(self, x): - return torch.cat((self.relu(x), self.relu(-x)), 1) - -class L2NormLayer(nn.Module): - def __init__(self, dim=(1, 2), eps=1e-6): - super(L2NormLayer, self).__init__() - self.dim = dim - self.eps = eps - - def forward(self, x): - return torch.nn.functional.normalize(x, p=2, dim=self.dim, eps=self.eps) + self.eps * torch.diag(torch.ones(x.shape[-1], device=x.device)) - -class TraceNorm(nn.Module): - def __init__(self, eps=1e-8): - super().__init__() - self.eps = eps - - def forward(self, Rz): - trace = torch.real(Rz.diagonal(dim1=-2, dim2=-1).sum(-1)).clamp(min=self.eps) # shape [B] - trace = trace.view(-1, 1, 1) - return Rz / trace - - -if __name__ == "__main__": - # sum_of_diag example - matrix = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) - sum_of_diag(matrix) - - matrix = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) - sum_of_diags_torch(matrix) - - # find_roots example - coefficients = [1, -5, 6] - find_roots(coefficients) - - # get_k_angles example - grid_size = 6 - k = 3 - prediction = torch.tensor([0.1, 0.3, 0.5, 0.2, 0.4, 0.6]) - get_k_angles(grid_size, k, prediction) - - # get_k_peaks example - grid_size = 6 - k = 3 - prediction = torch.tensor([0.1, 0.3, 0.5, 0.2, 0.4, 0.6]) - get_k_peaks(grid_size, k, prediction)