From 9ab68b51053a0deeeb271d7608354dc6e4db8886 Mon Sep 17 00:00:00 2001 From: oritalp Date: Tue, 3 Jun 2025 15:42:43 +0300 Subject: [PATCH 01/13] Building Data Creation --- .gitignore | 3 + README.md | 32 +- archive/main.py | 290 ++++++ archive/run_simulation.py | 286 ++++++ {src => archive/src}/__init__.py | 0 {src => archive/src}/config.py | 0 {src => archive/src}/data_handler.py | 0 {src => archive/src}/evaluation.py | 0 {src => archive/src}/methods_pack/.gitignore | 0 {src => archive/src}/methods_pack/__init__.py | 0 .../src}/methods_pack/beamformer.py | 0 .../src}/methods_pack/csestimator.py | 0 {src => archive/src}/methods_pack/esprit.py | 0 {src => archive/src}/methods_pack/mle.py | 0 {src => archive/src}/methods_pack/music.py | 0 .../src}/methods_pack/root_music.py | 0 .../src}/methods_pack/subspace_method.py | 0 {src => archive/src}/metrics/__init__.py | 0 .../src}/metrics/beamformer_loss.py | 0 .../src}/metrics/cartesian_loss.py | 0 .../src}/metrics/music_spec_loss.py | 0 {src => archive/src}/metrics/rmse_loss.py | 0 {src => archive/src}/metrics/rmspe_loss.py | 0 {src => archive/src}/models.py | 0 {src => archive/src}/models_pack/.gitignore | 0 {src => archive/src}/models_pack/__init__.py | 0 {src => archive/src}/models_pack/dcd_music.py | 0 .../src}/models_pack/deep_augmented_music.py | 0 {src => archive/src}/models_pack/deep_cnn.py | 0 .../src}/models_pack/deep_root_music.py | 0 .../src}/models_pack/parent_model.py | 0 .../src}/models_pack/subspacenet.py | 0 .../src}/models_pack/trans_music.py | 0 {src => archive/src}/plotting.py | 0 archive/src/signal_creation.py | 315 ++++++ {src => archive/src}/system_model.py | 0 {src => archive/src}/training.py | 0 archive/src/utils.py | 566 +++++++++++ train_dcd.py => archive/train_dcd.py | 0 check.py | 3 + full_environment.yml | 2 +- main.py | 197 ++-- run_simulation.py | 320 ++---- src/signal_creation.py | 570 +++++------ src/utils.py | 922 +++++++++--------- 45 files changed, 2380 insertions(+), 1126 deletions(-) create mode 100644 archive/main.py create mode 100644 archive/run_simulation.py rename {src => archive/src}/__init__.py (100%) rename {src => archive/src}/config.py (100%) rename {src => archive/src}/data_handler.py (100%) rename {src => archive/src}/evaluation.py (100%) rename {src => archive/src}/methods_pack/.gitignore (100%) rename {src => archive/src}/methods_pack/__init__.py (100%) rename {src => archive/src}/methods_pack/beamformer.py (100%) rename {src => archive/src}/methods_pack/csestimator.py (100%) rename {src => archive/src}/methods_pack/esprit.py (100%) rename {src => archive/src}/methods_pack/mle.py (100%) rename {src => archive/src}/methods_pack/music.py (100%) rename {src => archive/src}/methods_pack/root_music.py (100%) rename {src => archive/src}/methods_pack/subspace_method.py (100%) rename {src => archive/src}/metrics/__init__.py (100%) rename {src => archive/src}/metrics/beamformer_loss.py (100%) rename {src => archive/src}/metrics/cartesian_loss.py (100%) rename {src => archive/src}/metrics/music_spec_loss.py (100%) rename {src => archive/src}/metrics/rmse_loss.py (100%) rename {src => archive/src}/metrics/rmspe_loss.py (100%) rename {src => archive/src}/models.py (100%) rename {src => archive/src}/models_pack/.gitignore (100%) rename {src => archive/src}/models_pack/__init__.py (100%) rename {src => archive/src}/models_pack/dcd_music.py (100%) rename {src => archive/src}/models_pack/deep_augmented_music.py (100%) rename {src => archive/src}/models_pack/deep_cnn.py (100%) rename {src => archive/src}/models_pack/deep_root_music.py (100%) rename {src => archive/src}/models_pack/parent_model.py (100%) rename {src => archive/src}/models_pack/subspacenet.py (100%) rename {src => archive/src}/models_pack/trans_music.py (100%) rename {src => archive/src}/plotting.py (100%) create mode 100644 archive/src/signal_creation.py rename {src => archive/src}/system_model.py (100%) rename {src => archive/src}/training.py (100%) create mode 100644 archive/src/utils.py rename train_dcd.py => archive/train_dcd.py (100%) create mode 100644 check.py diff --git a/.gitignore b/.gitignore index 6a17814..bad7ae0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ .idea/ *.ini __pycache__/ +**/*.pyc .venv*/ data/ src/__pycache__/ @@ -8,3 +9,5 @@ weights/ wandb/ *.png *.pdf +*.out +*.sh \ 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 100% rename from src/methods_pack/subspace_method.py rename to archive/src/methods_pack/subspace_method.py 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..d411ca8 --- /dev/null +++ b/check.py @@ -0,0 +1,3 @@ +import torch + +print("aabcd".find("b")) \ No newline at end of file diff --git a/full_environment.yml b/full_environment.yml index c51ba88..412b4b7 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 diff --git a/main.py b/main.py index 079933b..fa2f120 100644 --- a/main.py +++ b/main.py @@ -26,6 +26,11 @@ import matplotlib.pyplot as plt from run_simulation import run_simulation import argparse +import torch + +#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") @@ -39,134 +44,54 @@ } 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, + "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 + "wavelength": 1, # The carrier wavelength of the signal in meters, 1 can be fine for reserch, + # 0.06 is for wifi 5 GHz for example. + + "location_perturbation": "wavelength/4", # The boundaries of the location perturbation in meters, + # insert any vlaid float between 0 and wavelength/4 or "wavelength/n" to use with refrence 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 + } model_config = { - "model_type": "SubspaceNet", # SubspaceNet, DCD-MUSIC, DeepCNN, TransMUSIC, DR_MUSIC + "model_type": "diffMUSIC", # diffMUSIC "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", - ] + "use_wandb": False } +training_params["device"] = torch.device("cuda" if torch.cuda.is_available() else "cpu") def parse_arguments(): @@ -179,15 +104,14 @@ 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; diffMUSIC, SubspaceNet, DCD-MUSIC, DeepCNN, TransMUSIC, DR_MUSIC', 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) @@ -195,15 +119,10 @@ def parse_arguments(): 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('-w', '--wandb', action="store_true", help='Use wandb') - 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('-c', '--create', action="store_true", help='create a new dataset') return parser.parse_args() @@ -234,6 +153,18 @@ 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 + 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 @@ -242,15 +173,6 @@ def parse_arguments(): 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: @@ -265,22 +187,21 @@ 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 + 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 oreder switching between neigboring 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) print("Total time: ", time.time() - start) diff --git a/run_simulation.py b/run_simulation.py index 24d5c81..f75097d 100644 --- a/run_simulation.py +++ b/run_simulation.py @@ -11,14 +11,12 @@ """ # 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 def __run_simulation(**kwargs): @@ -26,20 +24,11 @@ def __run_simulation(**kwargs): 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 + 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" @@ -48,177 +37,33 @@ def __run_simulation(**kwargs): # 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") + utils.set_unified_seed(SYSTEM_MODEL_PARAMS["seed"]) - # 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 + system_model_params = SystemModelParams(**SYSTEM_MODEL_PARAMS) + # Initialize paths + data_saving_path, results_path = utils.initialize_paths(Path(__file__).parent, system_model_params, + dt_string_for_save) + data_loading_path = SIMULATION_COMMANDS["data_loading_path"] #ONLY USED if CREATE_DATA is False! + # Create system 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() + array = signals_creator.get_array() + gain_impairments = signals_creator.get_gain_perturbations() + gain_impairments_norm = torch.linalg.norm(gain_impairments, ord=2, dim=0) #just for testing purposes + # Save the created data under the data_path + utils.save_data_to_file(data_saving_path, measurements, signals, steering_mat, + noise, true_angles, array, gain_impairments) + else: + # Load data from file + measurements, signals, steering_mat, noise, true_angles, array, gain_impairments = \ + utils.load_data_from_file(data_loading_path) + return None def run_simulation(**kwargs): @@ -229,56 +74,63 @@ 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 """ + #TODO: check if anything is missing for the scenario_dict option once needed. 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 + + # from this on the option of multiple scenarios is used. this is activated when we specify the sceario_dict + # in main.py + # TODO: Later adjust it to our way. + + # 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__": diff --git a/src/signal_creation.py b/src/signal_creation.py index 381fefe..4de373e 100644 --- a/src/signal_creation.py +++ b/src/signal_creation.py @@ -1,315 +1,327 @@ -"""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: + def __init__(self, **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 """ - Class used for defining and creating signals and observations. - Inherits from SystemModel class. - ... + def __init__(self, system_model_params: SystemModelParams): + 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() - 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. + 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: + """ + 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_gain_perturbations(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): + """ + 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): """ - 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/utils.py b/src/utils.py index 463ffcb..14eaef2 100644 --- a/src/utils.py +++ b/src/utils.py @@ -28,280 +28,315 @@ 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 + +# # 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 -# 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): +# 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 save_data_to_file(data_path: Path, *args): """ - Validate that the number of sources in the batch is equal for all samples. + Saves the provided data to a file in the specified data path. + 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, dt_string_for_save: str) -> tuple: + 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: - """ - Calculates the sample covariance matrix for each element in the batch. +# 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. +# 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 +# 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. +# 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. +# Args: +# ----- +# X (np.ndarray): Input samples matrix. - Returns: - -------- - covariance_mat (np.ndarray): Covariance 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: +# 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. +# Returns: - Args: - matrix (np.ndarray): Square matrix for which diagonals need to be summed. +# """ +# 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. - Returns: - list: A list containing the sums of all diagonals in the matrix, from left to right. +# Args: +# matrix (np.ndarray): Square matrix for which diagonals need to be summed. - Raises: - None +# Returns: +# list: A list containing the sums of all diagonals in the matrix, from left to right. - Examples: - >>> matrix = np.array([[1, 2, 3], - [4, 5, 6], - [7, 8, 9]]) - >>> sum_of_diag(matrix) - [7, 12, 15, 8, 3] +# Raises: +# None - """ - 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. +# Examples: +# >>> matrix = np.array([[1, 2, 3], +# [4, 5, 6], +# [7, 8, 9]]) +# >>> sum_of_diag(matrix) +# [7, 12, 15, 8, 3] - Args: - matrix (torch.Tensor): Square matrix for which diagonals need to be summed. +# """ +# 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. - Returns: - torch.Tensor: A list containing the sums of all diagonals in the matrix, from left to right. +# Args: +# matrix (torch.Tensor): Square matrix for which diagonals need to be summed. - Raises: - None +# Returns: +# torch.Tensor: A list containing the sums of all diagonals in the matrix, from left to right. - 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) +# 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. +# # def find_roots(coefficients: list) -> np.ndarray: +# def find_roots(coefficients: list): +# """Finds the roots of a polynomial defined by its coefficients. - Returns: - np.ndarray: An array containing the roots of the polynomial. +# Args: +# coefficients (list): List of polynomial coefficients in descending order of powers. - Raises: - None +# Returns: +# np.ndarray: An array containing the roots of the polynomial. - Examples: - >>> coefficients = [1, -5, 6] # x^2 - 5x + 6 - >>> find_roots(coefficients) - array([3., 2.]) +# Raises: +# None - """ - 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 +# 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. +# 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. - Returns: - torch.Tensor: An array containing the roots of the polynomial. +# Args: +# coefficients (torch.Tensor): List of polynomial coefficients in descending order of powers. - Raises: - None +# Returns: +# torch.Tensor: An array containing the roots of the polynomial. - Examples: - >>> coefficients = torch.tensor([1, -5, 6]) # x^2 - 5x + 6 - >>> find_roots(coefficients) - tensor([3., 2.]) +# Raises: +# None - """ - 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 +# 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 +349,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 +365,237 @@ 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. - - 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. +# # 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. - Returns: - torch.Tensor: A tensor containing the top-k angles in degrees. +# 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 . - Raises: - None +# Returns: +# torch.Tensor: A tensor containing the top-k angles in degrees. - 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.]) +# Raises: +# None - """ - 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. +# 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.]) - 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. +# """ +# angles_grid = torch.linspace(-90, 90, grid_size) +# doa_prediction = angles_grid[torch.topk(prediction.flatten(), k).indices] +# return doa_prediction - 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) +# # 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. - # 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 +# 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. - # 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)) +# Returns: +# torch.Tensor: A tensor containing the top-k angles in degrees. - return Kx_Out +# 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.]) -# def _spatial_smoothing_covariance(sampels: torch.Tensor): # """ -# Calculates the covariance matrix using spatial smoothing technique. -# +# 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: # ----- -# X (np.ndarray): Input samples matrix. -# +# 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: # -------- -# covariance_mat (np.ndarray): Covariance matrix. +# torch.Tensor: Hermitian and PSD matrix with shape [BS, N, N]. + # """ -# -# 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) +# # 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) From 0029cceade850a9b81456ca844faa45d7f372000 Mon Sep 17 00:00:00 2001 From: oritalp Date: Sun, 8 Jun 2025 16:22:10 +0300 Subject: [PATCH 02/13] Started working on diffMUSIC --- run_simulation.py | 4 +- src/diffMUSIC.py | 594 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 596 insertions(+), 2 deletions(-) create mode 100644 src/diffMUSIC.py diff --git a/run_simulation.py b/run_simulation.py index f75097d..bdf0d97 100644 --- a/run_simulation.py +++ b/run_simulation.py @@ -58,10 +58,10 @@ def __run_simulation(**kwargs): gain_impairments_norm = torch.linalg.norm(gain_impairments, ord=2, dim=0) #just for testing purposes # Save the created data under the data_path utils.save_data_to_file(data_saving_path, measurements, signals, steering_mat, - noise, true_angles, array, gain_impairments) + noise, true_angles, array, gain_impairments, system_model_params) else: # Load data from file - measurements, signals, steering_mat, noise, true_angles, array, gain_impairments = \ + measurements, signals, steering_mat, noise, true_angles, array, gain_impairments, system_model_params = \ utils.load_data_from_file(data_loading_path) return None diff --git a/src/diffMUSIC.py b/src/diffMUSIC.py new file mode 100644 index 0000000..90f121f --- /dev/null +++ b/src/diffMUSIC.py @@ -0,0 +1,594 @@ +import warnings +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import matplotlib.pyplot as plt +import scipy as sc + +from src.signal_creation import SystemModel, SystemModelParams +from src.utils import set_unified_seed +from scipy.ndimage import maximum_filter +from scipy.ndimage import label +from scipy.ndimage import find_objects + +#TODO: Haven't checked this method yet. +def find_k_highest_peaks(matrix, k): + """ + Find the k highest peaks in a 2D matrix using SciPy tools. A peak is defined as a + local maximum surrounded by smaller values. + """ + # Apply maximum filter to find local maxima + neighborhood = maximum_filter(matrix, size=21, mode='constant', cval=-np.inf) + local_max = (matrix == neighborhood) + + # Label the connected components of local maxima + labeled, num_features = label(local_max) + slices = find_objects(labeled) + + # Extract peak positions and values + peaks = [] + try: + for sl in slices: + row = int((sl[0].start + sl[0].stop - 1) / 2) + col = int((sl[1].start + sl[1].stop - 1) / 2) + value = matrix[row, col] + peaks.append((row, col, value)) + except Exception as e: + pass + + # Sort peaks by value (descending) and select the top k + peaks = sorted(peaks, key=lambda x: x[2], reverse=True)[:k] + if len(peaks) < k: + warnings.warn(f"find_k_highest_peaks: Less than {k} peaks found.") + # add random peaks + x_random = np.random.randint(0, matrix.shape[0], (k - len(peaks),)) + y_random = np.random.randint(0, matrix.shape[1], (k - len(peaks),)) + for i in range(k - len(peaks)): + peaks.append((x_random[i], y_random[i], matrix[x_random[i], y_random[i]])) + + return peaks + + +class DiffMUSIC(nn.Module): + """ + Differentiable MUSIC implementation with learnable antenna gains and positions. + Integrates with the existing SystemModel architecture for far-field, non-coherent, narrowband scenarios. + """ + + def __init__(self, + system_model_params: SystemModelParams, + init_antenna_positions: torch.Tensor = None, + init_antenna_gains: torch.Tensor = None, + gain_constraint: str = "positive", # "positive", "normalized", "none" + temperature: float = 1.0, + cell_size_coeff: float = 0.2 # between 0 and 1, determines the size of tthe window for softmax peak + # finding, the size is computed as len(grid) * cell_size_coeff + ): + """ + Initialize DiffMUSIC with learnable antenna parameters. + + Args: + system_model_params: SystemModelParams instance + init_antenna_positions: Initial antenna positions [N] for ULA case + init_antenna_gains: Initial antenna gains [N] + gain_constraint: Type of gain constraint + temperature: Temperature for soft peak finding + cell_size_coeff: Coefficient for cell size in soft peak finding + """ + super().__init__() + + self.params = system_model_params + self.N = self.params.N # Number of antennas + self.M = self.params.M # Number of sources + self.gain_constraint = gain_constraint + self.temperature = temperature + self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + # Initialize antenna positions as learnable parameters + if init_antenna_positions is None: + init_antenna_positions = torch.arange(self.N, dtype=torch.float64) * (self.params.wavelength / 2) + + self.antenna_positions = nn.Parameter(init_antenna_positions.clone()) + assert init_antenna_positions.requires_grad is True, "Initial antenna positions must be a differentiable tensor." + + # Initialize antenna gains as learnable parameters + if init_antenna_gains is None: + init_antenna_gains = torch.ones(self.N, dtype=torch.complex64) + self.antenna_gains = nn.Parameter(init_antenna_gains.clone()) + + # Initialize angle grid for far-field DOA estimation + self._init_angle_grid() + + # Initialize cell size for soft peak finding + self.cell_size = int(self.angles_dict.shape[0] * cell_size_coeff) + if self.cell_size % 2 == 0: + self.cell_size += 1 + + # Store current MUSIC spectrum for analysis + self.music_spectrum = None + + 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_dict = torch.arange(-angle_range, angle_range + angle_resolution, + angle_resolution, dtype=torch.float64) + self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) + + def _apply_position_constraints(self): + """Reorders the antenna positions, ensuring we don't get stucked if some replace positions. + I don't think we need this because it is mathematically fine that thy will replace order. + In fact, this replacement may occur some incontinuity in the loss. + We do need to keep it in mind though. + """ + # with torch.no_grad(): + # if self.position_constraint == "ula": + # sorted_positions, _ = torch.sort(self.antenna_positions) + # self.antenna_positions.data = sorted_positions + pass + + def _apply_gain_constraints(self): + """Prevent extrme gains from future random walk. For now i leave commented.""" + # with torch.no_grad(): + # if self.gain_constraint == "positive": + # # Just clamp to reasonable ranges to prevent numerical issues + # real_part = torch.clamp(self.antenna_gains.real, min=0.1, max=10.0) + # imag_part = torch.clamp(self.antenna_gains.imag, min=-2.0, max=2.0) + # self.antenna_gains.data = torch.complex(real_part, imag_part, dtype=torch.complex64) + pass + + def get_constrained_parameters(self): + """Get antenna parameters without any sorting - order doesn't matter mathematically""" + positions = self.antenna_positions + gains = self.antenna_gains + return positions, gains + +#NOTE: ORI - I read up to this point. + + def compute_steering_matrix(self, angles: torch.Tensor): + """ + Compute steering matrix for far-field sources using current antenna configuration. + Follows the pattern from your SystemModel.steering_vec_far_field method. + + Args: + angles: DOA angles in radians [num_angles] + + Returns: + Complex steering matrix [N, num_angles] + """ + positions, gains = self.get_constrained_parameters() + + # Convert angles to proper shape for broadcasting + if angles.dim() == 1: + angles = angles.unsqueeze(0) # [1, num_angles] + + # Reshape positions for broadcasting: [N, 1] + positions = positions.view(-1, 1) + + # Compute phase delays: [N, 1] * sin([1, num_angles]) -> [N, num_angles] + time_delay = positions @ torch.sin(angles) + + # Compute steering vectors (following your formula) + steering_matrix = torch.exp(-2j * torch.pi * time_delay / self.params.wavelength) + steering_matrix = steering_matrix.to(torch.complex64) + + # Apply antenna gains + gains_diag = torch.diag(gains) + + # Normalize gains (following your normalization pattern) + normalization_factor = 1 / torch.linalg.norm(gains, ord=2) + steering_matrix = normalization_factor * gains_diag @ steering_matrix + + return steering_matrix + + def sample_covariance(self, x: torch.Tensor) -> torch.Tensor: + """ + Compute sample covariance matrix (following your pattern). + + Args: + x: Input samples [batch_size, N, T] or [N, T] + + Returns: + Covariance matrices [batch_size, N, N] + """ + if x.dim() == 2: + x = x.unsqueeze(0) + 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 subspace_separation(self, cov: torch.Tensor, number_of_sources: int): + """ + Perform eigendecomposition and separate signal/noise subspaces. + + Args: + cov: Covariance matrix [batch_size, N, N] + number_of_sources: Number of sources + + Returns: + signal_subspace, noise_subspace, source_estimation, eigen_regularization + """ + # Eigendecomposition + eigenvalues, eigenvectors = torch.linalg.eigh(cov) + + # Sort in descending order + sorted_indices = torch.argsort(eigenvalues, dim=-1, descending=True) + eigenvalues = torch.gather(eigenvalues, -1, sorted_indices) + eigenvectors = torch.gather(eigenvectors, -1, sorted_indices.unsqueeze(-2).expand_as(eigenvectors)) + + # Separate subspaces + signal_subspace = eigenvectors[:, :, :number_of_sources] + noise_subspace = eigenvectors[:, :, number_of_sources:] + + # Source estimation (simplified) + source_estimation = number_of_sources + + # Eigen regularization + eigen_regularization = torch.mean(eigenvalues[:, number_of_sources:]) + + return signal_subspace, noise_subspace, source_estimation, eigen_regularization + + def get_inverse_spectrum(self, noise_subspace: torch.Tensor): + """ + Compute inverse MUSIC spectrum using current antenna configuration. + + Args: + noise_subspace: Noise subspace [batch_size, N, N-M] + + Returns: + Inverse spectrum [batch_size, num_angles] + """ + # Compute steering matrix for all angles + steering_dict = self.compute_steering_matrix(self.angles_dict) + steering_dict = steering_dict[:noise_subspace.shape[1]].to(self.device) + + # Compute projection onto noise subspace + var1 = torch.einsum("an, bnm -> bam", + steering_dict.conj().transpose(0, 1)[:, :noise_subspace.shape[1]], + noise_subspace) + inverse_spectrum = torch.norm(var1, dim=2) ** 2 + + return inverse_spectrum + + def soft_peak_finding_1d(self, spectrum: torch.Tensor, search_space: torch.Tensor, num_sources: int): + """ + Differentiable 1D peak finding using temperature-scaled softmax. + Adapted from your __maskpeak_1d method. + + Args: + spectrum: MUSIC spectrum [batch_size, num_points] + search_space: Search grid values [num_points] + num_sources: Number of sources to find + + Returns: + Estimated parameters [batch_size, num_sources] + """ + batch_size = spectrum.shape[0] + + # Find hard peaks for initialization + peaks = torch.zeros(batch_size, num_sources, dtype=torch.int64, device=self.device) + for batch in range(batch_size): + music_spectrum_np = spectrum[batch].cpu().detach().numpy().squeeze() + # Find spectrum peaks + peaks_tmp = sc.signal.find_peaks(music_spectrum_np, threshold=0.0)[0] + if len(peaks_tmp) < num_sources: + # Take top values if not enough peaks + random_peaks = torch.topk(torch.from_numpy(music_spectrum_np), + num_sources - peaks_tmp.shape[0], largest=True).indices.cpu().detach().numpy() + peaks_tmp = np.concatenate((peaks_tmp, random_peaks)) + # Sort by amplitude + sorted_peaks = peaks_tmp[np.argsort(music_spectrum_np[peaks_tmp])[::-1]] + peaks[batch] = torch.from_numpy(sorted_peaks[0:num_sources]).to(self.device) + + # Soft peak finding (differentiable) + soft_decision = torch.zeros(batch_size, num_sources, dtype=torch.float64, device=self.device) + top_indxs = peaks.to(self.device) + + for source in range(num_sources): + # Create cell around each peak + cell_idx = (top_indxs[:, source][:, None] + - self.cell_size + + torch.arange(2 * self.cell_size + 1, dtype=torch.long, device=self.device)) + + # Handle boundaries + out_of_bounds_mask = (cell_idx < 0) | (cell_idx >= spectrum.shape[1]) + cell_idx[out_of_bounds_mask] = top_indxs[:, source].unsqueeze(1).expand_as(cell_idx)[out_of_bounds_mask] + cell_idx = cell_idx.reshape(batch_size, -1, 1) + + # Extract spectrum values in cell + metrix_thr = torch.gather(spectrum.unsqueeze(-1).expand(-1, -1, cell_idx.size(-1)), 1, + cell_idx).requires_grad_(True) + + # Apply temperature-scaled softmax + soft_max = torch.softmax(metrix_thr / self.temperature, dim=1) + + # Compute weighted average + soft_decision[:, source] = torch.einsum("bms, bms -> bs", + search_space[cell_idx.cpu()].to(self.device), + soft_max).squeeze() + + return soft_decision + + def forward(self, x: torch.Tensor, number_of_sources: int = None): + """ + Forward pass of DiffMUSIC. + + Args: + x: Input samples [batch_size, N, T] or [N, T] + number_of_sources: Number of sources (defaults to self.M) + + Returns: + Estimated DOA angles, source estimation, eigen regularization + """ + if number_of_sources is None: + number_of_sources = self.M + + # Apply constraints (for monitoring during training) + if self.training: + self._apply_position_constraints() + self._apply_gain_constraints() + + # Compute covariance matrix + cov = self.sample_covariance(x) + + # Subspace separation + _, noise_subspace, source_estimation, eigen_regularization = self.subspace_separation(cov, number_of_sources) + + # Compute inverse spectrum + inverse_spectrum = self.get_inverse_spectrum(noise_subspace) + + # Compute MUSIC spectrum + self.music_spectrum = 1 / (inverse_spectrum + 1e-10) + + # Differentiable peak finding + params = self.soft_peak_finding_1d(self.music_spectrum, self.angles_dict, number_of_sources) + + return params, source_estimation, eigen_regularization + + def plot_spectrum(self, batch: int = 0, true_angles: torch.Tensor = None, save: bool = False): + """Plot MUSIC spectrum""" + if self.music_spectrum is None: + print("No spectrum available. Run forward pass first.") + return + + x = np.rad2deg(self.angles_dict.detach().cpu().numpy()) + y = self.music_spectrum[batch].detach().cpu().numpy() + + plt.figure(figsize=(10, 6)) + plt.plot(x, y, label="DiffMUSIC Spectrum", linewidth=2) + + if true_angles is not None: + true_angles_deg = np.rad2deg(true_angles.detach().cpu().numpy()) + for i, angle in enumerate(true_angles_deg): + plt.axvline(angle, color='red', linestyle='--', alpha=0.7, + label=f'True DOA {i+1}' if i == 0 else "") + + plt.xlabel('Angle [degrees]') + plt.ylabel('Spectrum Power') + plt.title('DiffMUSIC Spectrum') + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + + if save: + plt.savefig('diffmusic_spectrum.pdf') + plt.show() + + def plot_array_geometry(self, save: bool = False): + """Plot current antenna array geometry""" + positions, gains = self.get_constrained_parameters() + positions_np = positions.detach().cpu().numpy() + gains_np = torch.abs(gains).detach().cpu().numpy() + + plt.figure(figsize=(12, 4)) + + # Plot antenna positions + plt.scatter(positions_np, np.zeros_like(positions_np), + c=gains_np, s=100, cmap='viridis', + edgecolors='black', linewidth=1, marker='s') + plt.colorbar(label='Antenna Gain Magnitude') + + # Add antenna numbers + for i, pos in enumerate(positions_np): + plt.annotate(f'{i}', (pos, 0.01), + xytext=(0, 10), textcoords='offset points', + ha='center', va='bottom') + + plt.xlabel('Position [meters]') + plt.ylabel('') + plt.title('DiffMUSIC Antenna Array Geometry') + plt.grid(True, alpha=0.3) + plt.ylim(-0.1, 0.1) + + # Add wavelength reference + plt.axhline(0, color='black', linewidth=0.5) + spacing_text = f'λ/2 = {self.params.wavelength/2:.3f}m' + plt.text(0.02, 0.98, spacing_text, transform=plt.gca().transAxes, + bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8)) + + plt.tight_layout() + + if save: + plt.savefig('diffmusic_array_geometry.pdf') + plt.show() + + def get_antenna_info(self): + """Get current antenna configuration""" + positions, gains = self.get_constrained_parameters() + return { + 'positions': positions.detach().cpu().numpy(), + 'gains': gains.detach().cpu().numpy(), + 'position_constraint': self.position_constraint, + 'gain_constraint': self.gain_constraint, + 'N': self.N, + 'wavelength': self.params.wavelength + } + + +class DiffMUSICTrainer: + """Training wrapper for DiffMUSIC integrated with your simulation framework""" + + def __init__(self, diffmusic_model: DiffMUSIC, training_params: dict): + self.model = diffmusic_model + self.training_params = training_params + + # Initialize optimizer + if training_params["optimizer"] == "Adam": + self.optimizer = torch.optim.Adam( + self.model.parameters(), + lr=training_params["learning_rate"], + weight_decay=training_params.get("weight_decay", 0) + ) + elif training_params["optimizer"] == "SGD": + self.optimizer = torch.optim.SGD( + self.model.parameters(), + lr=training_params["learning_rate"], + weight_decay=training_params.get("weight_decay", 0) + ) + + # Initialize scheduler + if training_params["scheduler"] == "ReduceLROnPlateau": + self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + self.optimizer, mode='min', factor=0.5, patience=10 + ) + elif training_params["scheduler"] == "StepLR": + self.scheduler = torch.optim.lr_scheduler.StepLR( + self.optimizer, + step_size=training_params.get("step_size", 50), + gamma=0.5 + ) + + def train_step(self, samples_batch: torch.Tensor, targets_batch: torch.Tensor, num_sources: int = None): + """Single training step""" + self.model.train() + self.optimizer.zero_grad() + + # Forward pass + predictions, _, eigen_reg = self.model(samples_batch, num_sources) + + # Compute RMSE loss + loss = torch.sqrt(torch.mean((predictions - targets_batch) ** 2)) + + # Add regularization terms + reg_loss = 0.01 * eigen_reg # Eigenvalue regularization + + # Add antenna position regularization (encourage smooth spacing) + positions, _ = self.model.get_constrained_parameters() + if len(positions) > 1: + pos_diff = positions[1:] - positions[:-1] + target_spacing = self.model.params.wavelength / 2 + spacing_reg = 0.001 * torch.mean((pos_diff - target_spacing) ** 2) + else: + spacing_reg = 0.0 + + total_loss = loss + reg_loss + spacing_reg + + # Backward pass + total_loss.backward() + + # Gradient clipping + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) + + self.optimizer.step() + + return total_loss.item(), loss.item(), reg_loss.item(), spacing_reg + + def validate(self, val_samples: torch.Tensor, val_targets: torch.Tensor, num_sources: int = None): + """Validation step""" + self.model.eval() + + with torch.no_grad(): + predictions, _, _ = self.model(val_samples, num_sources) + loss = torch.sqrt(torch.mean((predictions - val_targets) ** 2)) + + if hasattr(self.scheduler, 'step'): + if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + self.scheduler.step(loss) + else: + self.scheduler.step() + + return loss.item() + + +# Integration function for your simulation framework +def create_diffmusic_model(system_model_params: SystemModelParams, model_config: dict): + """ + Create DiffMUSIC model for integration with your simulation framework. + + Args: + system_model_params: Your SystemModelParams instance + model_config: Model configuration dictionary + + Returns: + DiffMUSIC model instance + """ + model_params = model_config.get("model_params", {}) + + diffmusic = DiffMUSIC( + system_model_params=system_model_params, + position_constraint=model_params.get("position_constraint", "ula"), + gain_constraint=model_params.get("gain_constraint", "positive"), + temperature=model_params.get("temperature", 0.1), + cell_size_coeff=model_params.get("cell_size_coeff", 0.2) + ) + + return diffmusic + + +# Example usage function compatible with your framework +def run_diffmusic_simulation(system_model_params: SystemModelParams, + training_params: dict, + samples: torch.Tensor, + true_angles: torch.Tensor): + """ + Example function showing how to use DiffMUSIC with your simulation framework. + + Args: + system_model_params: Your SystemModelParams instance + training_params: Training parameters dictionary + samples: Signal samples [batch_size, N, T] or [N, T] + true_angles: True DOA angles [batch_size, M] or [M] + + Returns: + Trained DiffMUSIC model and final loss + """ + # Create model + model_config = { + "model_type": "diffMUSIC", + "model_params": { + "position_constraint": "ula", + "gain_constraint": "positive", + "temperature": 0.1 + } + } + + diffmusic = create_diffmusic_model(system_model_params, model_config) + trainer = DiffMUSICTrainer(diffmusic, training_params) + + # Move to device + device = training_params["device"] + diffmusic = diffmusic.to(device) + samples = samples.to(device) + true_angles = true_angles.to(device) + + # Training loop + num_epochs = training_params["epochs"] + batch_size = training_params["batch_size"] + + # Simple training (you can adapt this to your batch loading pattern) + for epoch in range(num_epochs): + total_loss, data_loss, reg_loss, spacing_reg = trainer.train_step( + samples, true_angles, system_model_params.M + ) + + if epoch % 10 == 0: + print(f"Epoch {epoch}: Total Loss = {total_loss:.6f}, Data Loss = {data_loss:.6f}") + + # Validation + val_loss = trainer.validate(samples, true_angles, system_model_params.M) + print(f"Final Validation Loss: {val_loss:.6f}") + + return diffmusic, val_loss \ No newline at end of file From 9677c49958d2fb8a8ed26b5dcbd58010b34066f0 Mon Sep 17 00:00:00 2001 From: oritalp Date: Tue, 10 Jun 2025 17:10:16 +0300 Subject: [PATCH 03/13] Keep rearanging the diffmusic --- archive/src/methods_pack/subspace_method.py | 3 +- main.py | 7 +- src/diffmusic.py | 632 ++++++++++++++++++ src/music.py | 698 ++++++++++++++++++++ src/{ => old}/diffMUSIC.py | 112 +--- src/old/diff_subspace_method | 140 ++++ src/old/subspace_method.py | 242 +++++++ src/signal_creation.py | 6 +- src/subspace_method.py | 241 +++++++ src/utils.py | 78 +-- 10 files changed, 2007 insertions(+), 152 deletions(-) create mode 100644 src/diffmusic.py create mode 100644 src/music.py rename src/{ => old}/diffMUSIC.py (79%) create mode 100644 src/old/diff_subspace_method create mode 100644 src/old/subspace_method.py create mode 100644 src/subspace_method.py diff --git a/archive/src/methods_pack/subspace_method.py b/archive/src/methods_pack/subspace_method.py index dc39054..058ab35 100644 --- a/archive/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/main.py b/main.py index fa2f120..37232d2 100644 --- a/main.py +++ b/main.py @@ -75,9 +75,12 @@ "signal_nature": "non-coherent" # if defined, values in scenario_dict will be ignored } + +system_model_params["device"] = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model_config = { "model_type": "diffMUSIC", # diffMUSIC - "model_params": {} + "model_params": {"window_size": 21} } @@ -91,7 +94,7 @@ "use_wandb": False } -training_params["device"] = torch.device("cuda" if torch.cuda.is_available() else "cpu") + def parse_arguments(): diff --git a/src/diffmusic.py b/src/diffmusic.py new file mode 100644 index 0000000..7279a61 --- /dev/null +++ b/src/diffmusic.py @@ -0,0 +1,632 @@ +""" +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 RMSPELoss + + +class DiffMUSIC(SubspaceMethod): + """ + Differentiable MUSIC (diffMUSIC) implementation for DoA estimation with hardware impairment learning. + + 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, + window_size: int = 21, model_order_estimation: str = None): + """ + Initialize diffMUSIC + + Args: + system_model: System model object (kept for compatibility) + N: Number of antennas + wavelength: Signal wavelength + window_size: Size of angular window for softmax peak finding + model_order_estimation: Model order estimation method + """ + system_model = SystemModel(system_model_params) + super().__init__(system_model, model_order_estimation=model_order_estimation) + + self.params = system_model.params + self.N = N + self.wavelength = self.params.wavelength + self.window_size = window_size # For now it is inserted directly + + # 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) + self.gains_real = nn.Parameter(gains_real) + self.gains_imag = nn.Parameter(gains_imag) + self.complex_gain = torch.complex(gains_real, gains_imag, dtype=torch.complex64) + + 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_dict = torch.arange(-angle_range, angle_range + angle_resolution, + angle_resolution, dtype=torch.float64) + self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) + + 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 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) + + # Get complex gains + complex_gains = self.complex_gain + + # 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 = self.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) + +#TODO: Read up to the peak_finder inside the forward-pass. Looks fine, need to understand the peak_finder. +# If the hard decision is fine, we can use this also for MUSIC. Keep reading and notice the first dimension +# of cov is the batch_dim. Another thing is to change system_model_params to also include the training_params +# (all params in general). For now, the gains real and imag parts are parameters, change just +# the whole complex gain to be a parameter, and use it in the steering matrix. + + 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) + 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 = self._peak_finder(number_of_sources) + + return estimated_angles, 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) # (batch_size, num_angles, N-M) + # 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) + + def _hard_peak_finder(self, number_of_sources: int) -> torch.Tensor: + """ + Non-differentiable peak finding for inference + + Args: + number_of_sources: Number of peaks to find + + 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("diffMUSIC: Not enough peaks found, using highest values") + # 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) + + # 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 + + 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) + + for batch in range(batch_size): + spectrum = self.music_spectrum[batch] + + # Find initial peaks (non-differentiable, but gradients will flow through softmax) + peaks_indices = self._find_initial_peaks(spectrum, number_of_sources) + + # For each peak, apply differentiable refinement + for source_idx in range(number_of_sources): + peak_idx = peaks_indices[source_idx] + + # Create angular mask around peak + mask_indices = self._create_angular_mask(peak_idx, spectrum.shape[0]) + + # 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) + + return estimated_angles + + def _find_initial_peaks(self, spectrum: torch.Tensor, number_of_sources: int) -> torch.Tensor: + """Find initial peak locations (non-differentiable but provides starting points)""" + spectrum_np = spectrum.cpu().detach().numpy() + peaks_indices = sc.signal.find_peaks(spectrum_np, threshold=0.0)[0] + + if len(peaks_indices) < number_of_sources: + # Use highest values if not enough peaks + additional_peaks = torch.topk(spectrum, + number_of_sources - len(peaks_indices), + largest=True).indices.cpu().numpy() + peaks_indices = np.concatenate([peaks_indices, additional_peaks]) + + # Sort by amplitude and take top peaks + sorted_peaks = peaks_indices[np.argsort(spectrum_np[peaks_indices])[::-1]] + return torch.from_numpy(sorted_peaks[:number_of_sources]).to(self.device) + + 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 get_learned_parameters(self): + """ + Get the learned array parameters + + Returns: + dict: Dictionary containing learned positions and gains + """ + return { + 'antenna_positions': self.antenna_positions.detach().cpu().numpy(), + 'complex_gains': self.get_complex_gains().detach().cpu().numpy(), + 'gains_magnitude': torch.abs(self.get_complex_gains()).detach().cpu().numpy(), + 'gains_phase': torch.angle(self.get_complex_gains()).detach().cpu().numpy() + } + + 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, batch_idx: int = 0, highlight_angles: torch.Tensor = None, + save: bool = False, title: str = "diffMUSIC Spectrum"): + """ + 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(title) + plt.grid(True, alpha=0.3) + plt.legend() + plt.tight_layout() + + if save: + plt.savefig('diffmusic_spectrum.pdf') + plt.show() + + def plot_learned_array(self, save: bool = False): + """ + Plot the learned antenna array geometry + + Args: + save: Whether to save the plot + """ + learned_params = self.get_learned_parameters() + positions = learned_params['antenna_positions'] + gains_mag = learned_params['gains_magnitude'] + gains_phase = learned_params['gains_phase'] + + # Nominal positions for comparison + nominal_positions = np.arange(self.N) * (self.wavelength / 2) + + fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 10)) + + # Plot 1: Antenna positions + ax1.scatter(nominal_positions, np.zeros_like(nominal_positions), + marker='o', s=100, alpha=0.5, label='Nominal', color='blue') + ax1.scatter(positions, np.zeros_like(positions), + marker='s', s=100, label='Learned', color='red') + ax1.set_xlabel('Position [wavelengths]') + ax1.set_title('Antenna Positions') + ax1.legend() + ax1.grid(True, alpha=0.3) + + # Plot 2: Gain magnitudes + antenna_indices = np.arange(self.N) + ax2.bar(antenna_indices - 0.2, np.ones(self.N), width=0.4, + alpha=0.7, label='Nominal', color='blue') + ax2.bar(antenna_indices + 0.2, gains_mag, width=0.4, + alpha=0.7, label='Learned', color='red') + ax2.set_xlabel('Antenna Index') + ax2.set_ylabel('Gain Magnitude') + ax2.set_title('Complex Gain Magnitudes') + ax2.legend() + ax2.grid(True, alpha=0.3) + + # Plot 3: Gain phases + ax3.bar(antenna_indices, gains_phase, width=0.6, alpha=0.7, color='green') + ax3.set_xlabel('Antenna Index') + ax3.set_ylabel('Gain Phase [radians]') + ax3.set_title('Complex Gain Phases') + ax3.grid(True, alpha=0.3) + + plt.tight_layout() + + if save: + plt.savefig('diffmusic_learned_array.pdf') + plt.show() + + def test_step(self, batch, batch_idx, model: nn.Module = None): + """ + Test step compatible with the existing framework + + Args: + batch: Test batch (x, sources_num, label) + batch_idx: Batch index + model: Model (not used, kept for compatibility) + + Returns: + tuple: (rmspe, accuracy, test_length) + """ + x, sources_num, label = batch + if x.dim() == 2: + x = x.unsqueeze(0) + + test_length = x.shape[0] + x = x.to(self.device) + angles = label.to(self.device) + + # Check if sources number is consistent + if (sources_num != sources_num[0]).any(): + raise Exception("diffMUSIC test_step: Inconsistent number of sources in batch") + + sources_num = sources_num[0] + + # Compute covariance + if self.system_model.params.signal_nature == "non-coherent": + Rx = self.pre_processing(x, mode="sample") + else: + Rx = self.pre_processing(x, mode="sample") + + # Run diffMUSIC + predictions, sources_num_estimation, _ = self(Rx, number_of_sources=sources_num) + + # Compute RMSPE + criterion = RMSPELoss(balance_factor=1.0) + rmspe = criterion(predictions, angles).sum().item() + + # Compute accuracy + acc = self.source_estimation_accuracy(sources_num, sources_num_estimation) + + return rmspe, acc, test_length + + def __str__(self): + return "diffMUSIC" + + def _get_name(self): + return "diffMUSIC" + + +# class DiffMUSICLoss(nn.Module): +# """ +# Loss functions for diffMUSIC training +# Implements both supervised learning strategies from the paper: +# - LSL,θ: RMSPE on estimated DoAs +# - LSL,P: Maximize spectrum amplitude at true DoA locations +# """ + +# def __init__(self, loss_type: str = "rmspe"): +# """ +# Args: +# loss_type: "rmspe" for LSL,θ or "spectrum" for LSL,P +# """ +# super().__init__() +# self.loss_type = loss_type +# self.rmspe_loss = RMSPELoss(balance_factor=1.0) + +# def forward(self, predictions, targets, spectrum=None, angles_grid=None): +# """ +# Compute loss based on specified type + +# Args: +# predictions: Predicted DoAs (for RMSPE loss) +# targets: True DoAs +# spectrum: MUSIC spectrum (for spectrum loss) +# angles_grid: Angular grid (for spectrum loss) + +# Returns: +# Loss value +# """ +# if self.loss_type == "rmspe": +# return self.rmspe_loss(predictions, targets) + +# elif self.loss_type == "spectrum": +# if spectrum is None or angles_grid is None: +# raise ValueError("Spectrum and angles_grid required for spectrum loss") + +# return self._spectrum_loss(targets, spectrum, angles_grid) + +# else: +# raise ValueError(f"Unknown loss type: {self.loss_type}") + +# def _spectrum_loss(self, true_angles, spectrum, angles_grid): +# """ +# Spectrum-based loss (LSL,P from paper) +# Maximizes spectrum amplitude at true DoA locations +# """ +# 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] + +# return total_loss / (batch_size * true_angles.shape[1]) + + +# class UnsupervisedDiffMUSIC(DiffMUSIC): +# """ +# Unsupervised diffMUSIC using Jain's Index (LUL from paper) +# Maximizes spectrum sharpness without requiring true DoA labels +# """ + +# def __init__(self, *args, **kwargs): +# super().__init__(*args, **kwargs) +# self.jains_index_loss = JainsIndexLoss() + +# def compute_unsupervised_loss(self): +# """ +# Compute unsupervised loss using Jain's Index on spectrum peaks + +# Returns: +# Loss value encouraging sharp peaks +# """ +# if self.music_spectrum is None: +# raise ValueError("No spectrum computed. Run forward pass first.") + +# total_loss = 0 +# batch_size = self.music_spectrum.shape[0] + +# for batch in range(batch_size): +# spectrum = self.music_spectrum[batch] + +# # Find initial peaks to create masks +# peaks_indices = self._find_initial_peaks(spectrum, self.system_model.params.M) + +# # Apply Jain's index to each peak region +# for peak_idx in peaks_indices: +# mask_indices = self._create_angular_mask(peak_idx, spectrum.shape[0]) +# masked_spectrum = spectrum[mask_indices] + +# # Jain's index encourages sharp peaks +# jains_loss = self.jains_index_loss(masked_spectrum) +# total_loss += jains_loss + +# 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().__init__() + +# def forward(self, x): +# """ +# 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 \ No newline at end of file diff --git a/src/music.py b/src/music.py new file mode 100644 index 0000000..3a2d728 --- /dev/null +++ b/src/music.py @@ -0,0 +1,698 @@ +import warnings + +import numpy as np +import torch +import torch.nn as nn +import matplotlib.pyplot as plt +import scipy as sc + +from src.signal_creation import SystemModel +from src.subspace_method import SubspaceMethod +from src.utils import * +# from src.metrics import RMSPELoss, CartesianLoss + +from scipy.ndimage import maximum_filter +from scipy.ndimage import label +from scipy.ndimage import find_objects + +#TODO: HAVENT CHECKED THE FILE YET, this is claude's suggestion for our new case + +def find_k_highest_peaks(matrix, k): + """ + Find the k highest peaks in a 2D matrix using SciPy tools. A peak is defined as a + local maximum surrounded by smaller values. + + Parameters: + - matrix (2D array-like): Input matrix. + - k (int): Number of highest peaks to extract. + + Returns: + - peaks (list): List of tuples (row, col, value) representing the positions and values of the k highest peaks. + """ + # Apply maximum filter to find local maxima + neighborhood = maximum_filter(matrix, size=21, mode='constant', cval=-np.inf) + local_max = (matrix == neighborhood) + + # Label the connected components of local maxima + labeled, num_features = label(local_max) + slices = find_objects(labeled) + + # Extract peak positions and values + peaks = [] + try: + for sl in slices: + row = int((sl[0].start + sl[0].stop - 1) / 2) + col = int((sl[1].start + sl[1].stop - 1) / 2) + value = matrix[row, col] + peaks.append((row, col, value)) + except Exception as e: + pass + + # Sort peaks by value (descending) and select the top k + peaks = sorted(peaks, key=lambda x: x[2], reverse=True)[:k] + if len(peaks) < k: + warnings.warn(f"find_k_highest_peaks: Less than {k} peaks found.") + # add random peaks + x_random = np.random.randint(0, matrix.shape[0], (k - len(peaks),)) + y_random = np.random.randint(0, matrix.shape[1], (k - len(peaks),)) + for i in range(k - len(peaks)): + peaks.append((x_random[i], y_random[i], matrix[x_random[i], y_random[i]])) + + return peaks + + +class MUSIC(SubspaceMethod): + """ + This is implementation of the MUSIC method for localization in Far and Near field environments. + For Far field - only "angle" can be estimated + For Near field - "angle", "range" and "angle, range" are the possible options. + """ + + def __init__(self, system_model: SystemModel, estimation_parameter: str, model_order_estimation: str = None): + """ + + Args: + system_model: + estimation_parameter: + """ + super().__init__(system_model, model_order_estimation=model_order_estimation) + self.estimation_params = estimation_parameter + self.angles_dict = None + self.ranges_dict = None + self.steering_dict = None + self.music_spectrum = None + self.cell_size = None + self.cell_size_angle = None + self.cell_size_range = None + self.noise_subspace = None + self.criterion = None + self.separated_criterion = None + + self.__init_grid_params() + self.__init_cells(0.2) + self.__init_criteria() + self.__init_search_grid() + + def forward(self, cov: torch.Tensor, number_of_sources: int, known_angles=None, known_distances=None): + """ + + Args: + cov: The covariance matrix of the input signal. + number_of_sources: the number of sources in the signal. Needed in case the dataset comprises mix number of sources. + known_angles: in case we are dealing with the Near field, the known angles should be passed. + known_distances: in case we are dealing with the Near field, the known distances should be passed. + + Returns: + tuple: the predicted parameters, the source estimation and the eigen regularization value. + """ + # single param estimation: the search grid should be updated for each batch, else, it's the same search grid. + if self.system_model.params.field_type in ["near", "full"] and self.estimation_params in ["range"]: + if known_angles.shape[-1] == 1: + self.set_search_grid(known_angles=known_angles, known_distances=known_distances) + else: + params = torch.zeros((cov.shape[0], number_of_sources), dtype=torch.float64, device=self.device) + for source in range(number_of_sources): + params_source, _, _ = self.forward(cov, number_of_sources=number_of_sources, + known_angles=known_angles[:, source][:, None]) + params[:, source] = params_source.squeeze() + return params + _, noise_subspace, source_estimation, eigen_regularization = self.subspace_separation(cov.to(torch.complex128), number_of_sources) + inverse_spectrum = self.get_inverse_spectrum(noise_subspace.to(self.device)).to(self.device) + if self._get_name() == "TOPS": + self.music_spectrum = torch.sum(1 / (inverse_spectrum + 1e-10), dim=-1) + else: + self.music_spectrum = 1 / (inverse_spectrum + 1e-10) + params = self.peak_finder(number_of_sources) + return params, source_estimation, eigen_regularization + + def get_music_spectrum_from_noise_subspace(self, noise_subspace: torch.Tensor) -> torch.Tensor: + inverse_spectrum = self.get_inverse_spectrum(noise_subspace.to(torch.complex128)) + self.music_spectrum = 1 / inverse_spectrum + return self.music_spectrum + + def update_number_of_sensors(self, number_of_sensors: int): + self.system_model.create_array(number_of_sensors) + self.set_search_grid() + + def adjust_cell_size(self): + if self.estimation_params == "range": + if self.cell_size > 1: + self.cell_size = int(0.8 * self.cell_size) + if self.cell_size % 2 == 0: + self.cell_size -= 1 + elif self.estimation_params == "angle, range": + if self.cell_size_angle > 1: + self.cell_size_angle = int(0.95 * self.cell_size_angle) + if self.cell_size_angle % 2 == 0: + self.cell_size_angle -= 1 + if self.cell_size_range > 1: + self.cell_size_range = int(0.95 * self.cell_size_range) + if self.cell_size_range % 2 == 0: + self.cell_size_range -= 1 + elif self.estimation_params == "angle": + if self.cell_size > 1: + self.cell_size = int(0.95 * self.cell_size) + if self.cell_size % 2 == 0: + self.cell_size -= 1 + + def get_inverse_spectrum(self, noise_subspace: torch.Tensor): + """ + + Parameters + ---------- + noise_subspace - the noise related subspace vectors of size BatchSizex#SENSORSx(#SENSORS-#SOURCES) + + Returns + ------- + in all cases it will return the inverse spectrum, + in case of single param estimation it will be 1D inverse spectrum: BatchSizex(length_search_grid) + in case of dual param estimation it will be 2D inverse spectrum: + BatchSizex(length_search_grid_angle)x(length_search_grid_distance) + """ + # steering_dict = self.steering_dict.to(device) + if self.system_model.params.field_type.startswith("far"): + steering_dict = self.steering_dict[:noise_subspace.shape[1]].to(self.device) + var1 = torch.einsum("an, bnm -> bam", steering_dict.conj().transpose(0, 1)[:, :noise_subspace.shape[1]], + noise_subspace) + inverse_spectrum = torch.norm(var1, dim=2) ** 2 + else: + if self.estimation_params.startswith("angle, range"): + steering_dict = self.steering_dict[:noise_subspace.shape[1]].conj().transpose(0, 2).transpose(0, 1).to(self.device) + try: + var1 = torch.einsum("adk, bkl -> badl", + steering_dict, + noise_subspace) + # get the norm value for each element in the batch. + inverse_spectrum = torch.norm(var1, dim=-1) ** 2 + except RuntimeError: + warnings.warn("MUSIC.get_inverse_spectrum: Out of memory error, trying to free some memory and convert the batch operation to for loop.") + torch.cuda.empty_cache() + inverse_spectrum = torch.zeros((noise_subspace.shape[0], self.angles_dict.shape[0], self.ranges_dict.shape[0]), dtype=torch.float64, device=self.device) + for batch in range(noise_subspace.shape[0]): + var1 = torch.einsum("adk, kl -> adl", + steering_dict, + noise_subspace[batch]) + inverse_spectrum[batch] = torch.norm(var1, dim=-1) ** 2 + + del var1 + + elif self.estimation_params.endswith("angle"): + steering_dict = self.steering_dict[:noise_subspace.shape[1]].to(self.device) + var1 = torch.einsum("an, nbm -> abm", steering_dict.conj().transpose(0, 1), + noise_subspace.transpose(0, 1)) + inverse_spectrum = torch.norm(var1, dim=-1).T ** 2 + elif self.estimation_params.startswith("range"): + steering_dict = self.steering_dict[:noise_subspace.shape[1]].to(self.device) + var1 = torch.bmm(steering_dict.conj().transpose(0, 2).transpose(0, 1), noise_subspace) + inverse_spectrum = torch.norm(var1, dim=-1) ** 2 + if torch.isnan(inverse_spectrum).any(): + raise ValueError("Nan values in inverse spectrum") + else: + raise ValueError(f"MUSIC.get_inverse_spectrum: unknown estimation param {self.estimation_params}") + del steering_dict + try: + torch.cuda.empty_cache() + except AttributeError: + pass + return inverse_spectrum + + def peak_finder(self, source_number: int): + """ + + Parameters + ---------- + is_soft: this boolean paramter will determine wether to use derivative approxamtion of the peak_finder for + the training stage. + + Returns + ------- + the predicted param(torch.Tensor) or params(tuple) + """ + if self.system_model.params.field_type.lower().startswith("far"): + return self._peak_finder_1d(self.angles_dict, source_number) + else: + if self.estimation_params.startswith("angle, range"): + return self._peak_finder_2d(source_number) + elif self.estimation_params.endswith("angle"): + return self._peak_finder_1d(self.angles_dict, source_number) + elif self.estimation_params.startswith("range"): + return self._peak_finder_1d(self.ranges_dict, source_number) + + def set_search_grid(self, known_angles: torch.Tensor = None, known_distances: torch.Tensor = None): + if self.system_model.params.field_type.startswith("far"): + self.__set_search_grid_far_field() + elif self.system_model.params.field_type in ["near", "full"]: + self.__set_search_grid_near_field(known_angles=known_angles, known_distances=known_distances) + else: + raise ValueError(f"MUSIC.set_search_grid: Unrecognized field type: {self.system_model.params.field_type}") + + def plot_spectrum(self, highlight_corrdinates=None, batch: int = 0, method: str = "heatmap", music_spectrum = None, add_title: bool = False, save: bool = False): + if self.estimation_params == "angle, range": + self._plot_3d_spectrum(highlight_corrdinates, batch, method, music_spectrum=music_spectrum, add_title=add_title, save=save) + else: + self._plot_1d_spectrum(highlight_corrdinates, batch, add_title=add_title, save=save) + + def test_step(self, batch, batch_idx, model: nn.Module=None): + x, sources_num, label = batch + if x.dim() == 2: + x = x.unsqueeze(0) + test_length = x.shape[0] + x = x.to(self.device) + if self.estimation_params == "angle, range": + angles, ranges = torch.split(label, max(sources_num), dim=1) + angles = angles.to(self.device) + ranges = ranges.to(self.device) + else: + angles = label.to(self.device) # only angles + # Check if the sources number is the same for all samples in the batch + if (sources_num != sources_num[0]).any(): + # in this case, the sources number is not the same for all samples in the batch + raise Exception(f"train_model:" + f" The sources number is not the same for all samples in the batch.") + else: + sources_num = sources_num[0] + if model is not None: + try: + Rx = model.get_surrogate_covariance(x) + except NotImplementedError as e: + raise e + else: + if self.system_model.params.signal_nature == "non-coherent": + Rx = self.pre_processing(x, mode="sample") + else: + # Rx = self.pre_processing(x, mode="sps") + Rx = self.pre_processing(x, mode="sample") + predictions, sources_num_estimation, _ = self(Rx, number_of_sources=sources_num) + if self.estimation_params == "angle, range": + angles_prediction, ranges_prediction = predictions + rmspe = self.criterion(angles_prediction, angles, ranges_prediction, ranges).sum(-1).item() + _, rmspe_angle, rmspe_range = self.separated_criterion(angles_prediction, angles, ranges_prediction, ranges) + rmspe = (rmspe, rmspe_angle.sum(-1).item(), rmspe_range.sum(-1).item()) + else: + rmspe = self.criterion(predictions, angles).sum().item() + + acc = self.source_estimation_accuracy(sources_num, sources_num_estimation) + + return rmspe, acc, test_length + + def _peak_finder_1d(self, search_space, source_number: int): + if self.estimation_params == "range": + source_number = 1 # for the range estimation, only one source is expected. + + batch_size = self.music_spectrum.shape[0] + + peaks = torch.zeros(batch_size, source_number, dtype=torch.int64, device=self.device) + for batch in range(batch_size): + music_spectrum = self.music_spectrum[batch].cpu().detach().numpy().squeeze() + # Find spectrum peaks + peaks_tmp = sc.signal.find_peaks(music_spectrum, threshold=0.0)[0] + if len(peaks_tmp) < source_number: + warnings.warn(f"MUSIC._peak_finder_1d: No peaks were found! taking max values instead.") + # random_peaks = np.random.randint(0, search_space.shape[0], (source_number - peaks_tmp.shape[0],)) + random_peaks = torch.topk(torch.from_numpy(music_spectrum), source_number - peaks_tmp.shape[0], + largest=True).indices.cpu().detach().numpy() + peaks_tmp = np.concatenate((peaks_tmp, random_peaks)) + # Sort the peak by their amplitude + sorted_peaks = peaks_tmp[np.argsort(music_spectrum[peaks_tmp])[::-1]] + peaks[batch] = torch.from_numpy(sorted_peaks[0:source_number]).to(self.device) + if not self.training: + # if the model is not in training mode, return the peaks + if peaks.dim() == 1: + return search_space[peaks] + else: + labels = torch.gather(search_space.unsqueeze(1).repeat(1, source_number).to(self.device), 0, peaks) + return labels + else: + return self.__maskpeak_1d(peaks, search_space, source_number) + + def _peak_finder_2d(self, source_number: int): + batch_size = self.music_spectrum.shape[0] + + max_row = torch.zeros((batch_size, source_number) + , dtype=torch.int64, device=self.device) + max_col = torch.zeros((batch_size, source_number) + , dtype=torch.int64, device=self.device) + for batch in range(batch_size): + music_spectrum = self.music_spectrum[batch].detach().cpu().numpy().squeeze() + peaks = find_k_highest_peaks(music_spectrum, source_number) + original_idx = torch.from_numpy(np.array(peaks)[:, :2]).T + max_row[batch] = original_idx[0][0: source_number] + max_col[batch] = original_idx[1][0: source_number] + if not self.training: + # if the model is not in training mode, return the peaks. + angle_dict = self.angles_dict.to(self.device) + range_dict = self.ranges_dict.to(self.device) + angles_pred = angle_dict[max_row] + distances_pred = range_dict[max_col] + del angle_dict, range_dict + try: + torch.cuda.empty_cache() + except AttributeError: + pass + return angles_pred, distances_pred + else: + return self.__maskpeak_2d(max_row, max_col, source_number) + + def __maskpeak_1d(self, peaks, search_space, source_number: int = None): + + batch_size = self.music_spectrum.shape[0] + soft_decision = torch.zeros(batch_size, source_number, dtype=torch.float64, device=self.device) + top_indxs = peaks.to(self.device) + + for source in range(source_number): + cell_idx = (top_indxs[:, source][:, None] + - self.cell_size + + torch.arange(2 * self.cell_size + 1, dtype=torch.long, device=self.device)) + # cell_idx %= self.music_spectrum.shape[1] + out_of_bounds_mask = (cell_idx < 0) | (cell_idx >= self.music_spectrum.shape[1]) + cell_idx[out_of_bounds_mask] = top_indxs[:, source].unsqueeze(1).expand_as(cell_idx)[out_of_bounds_mask] + cell_idx = cell_idx.reshape(batch_size, -1, 1) + metrix_thr = torch.gather(self.music_spectrum.unsqueeze(-1).expand(-1, -1, cell_idx.size(-1)), 1, + cell_idx).requires_grad_(True) + soft_max = torch.softmax(metrix_thr, dim=1) + soft_decision[:, source][:, None] = torch.einsum("bms, bms -> bs", search_space[cell_idx.cpu()].to(self.device), soft_max).to( + self.device) + + return soft_decision + + def __maskpeak_2d(self, peaks_r, peaks_c, source_number): + batch_size = self.music_spectrum.shape[0] + soft_row = torch.zeros((batch_size, source_number), device=self.device) + soft_col = torch.zeros((batch_size, source_number), device=self.device) + + for source in range(source_number): + max_row_cell_idx = (peaks_r[:, source][:, None] + - self.cell_size_angle + + torch.arange(2 * self.cell_size_angle + 1, dtype=torch.int32, device=self.device)) + max_row_cell_idx %= self.music_spectrum.shape[1] + max_row_cell_idx = max_row_cell_idx.reshape(batch_size, -1, 1) + + max_col_cell_idx = (peaks_c[:, source][:, None] + - self.cell_size_range + + torch.arange(2 * self.cell_size_range + 1, dtype=torch.int32, device=self.device)) + max_col_cell_idx %= self.music_spectrum.shape[2] + max_col_cell_idx = max_col_cell_idx.reshape(batch_size, 1, -1) + + metrix_thr = self.music_spectrum.gather(1, + max_row_cell_idx.expand(-1, -1, self.music_spectrum.shape[2])) + metrix_thr = metrix_thr.gather(2, max_col_cell_idx.repeat(1, max_row_cell_idx.shape[-2], 1)) + soft_max = torch.softmax(metrix_thr.view(batch_size, -1), dim=1).reshape(metrix_thr.shape) + soft_row[:, source][:, None] = torch.einsum("bla, bad -> bl", + self.angles_dict[max_row_cell_idx].transpose(1, 2), + torch.sum(soft_max, dim=2).unsqueeze(-1)) + soft_col[:, source][:, None] = torch.einsum("bmc, bcm -> bm", + self.ranges_dict[max_col_cell_idx], + torch.sum(soft_max, dim=1).unsqueeze(-1)) + + return soft_row, soft_col + + def _init_spectrum(self, batch_size): + if self.system_model.params.field_type == "Far": + self.music_spectrum = torch.zeros(batch_size, len(self.angles_dict)) + else: + if self.estimation_params.startswith("angle, range"): + self.music_spectrum = torch.zeros(batch_size, len(self.angles_dict), len(self.ranges_dict)) + elif self.estimation_params.endswith("angle"): + self.music_spectrum = torch.zeros(batch_size, len(self.angles_dict)) + elif self.estimation_params.startswith("range"): + self.music_spectrum = torch.zeros(batch_size, len(self.ranges_dict)) + + def __init_grid_params(self): + angle_range = np.deg2rad(self.system_model.params.doa_range) + angle_resolution = np.deg2rad(self.system_model.params.doa_resolution / 2) + angle_decimals = int(np.ceil(np.log10(1 / angle_resolution))) + + if self.system_model.params.field_type.startswith("far"): + # if it's the Far field case, need to init angles range. + self.angles_dict = torch.arange(-angle_range, angle_range + angle_resolution, angle_resolution, + dtype=torch.float64).to(torch.float64) + self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) + elif self.system_model.params.field_type in ["near", "full"]: + # if it's the Near field, there are 3 possabilities. + fresnel = self.system_model.fresnel + fraunhofer = self.system_model.fraunhofer + if self.estimation_params.startswith("angle"): + self.angles_dict = torch.arange(-angle_range, angle_range + angle_resolution, angle_resolution, + dtype=torch.float64).to(torch.float64) + # self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) + + + if self.estimation_params.endswith("range"): + fraunhofer_ratio = self.system_model.params.max_range_ratio_to_limit + distance_resolution = self.system_model.params.range_resolution / 2 + max_distance = min(self.system_model.fraunhofer, fraunhofer * fraunhofer_ratio + distance_resolution) + self.ranges_dict = torch.arange(np.ceil(fresnel), + max_distance, + distance_resolution, dtype=torch.float64) + else: + raise ValueError(f"MUSIC.__define_grid_params: Unrecognized field type for MUSIC class init stage," + f" got {self.system_model.params.field_type} but only Far and Near are allowed.") + + def __init_search_grid(self): + # if this is the music 2D case, the search grid is constant and can be calculated once. + if self.system_model.params.field_type in ["near", "full"]: + if self.angles_dict is not None and self.ranges_dict is not None: + self.set_search_grid() + elif self.angles_dict is not None: # Near field case with Far field inference + self.__set_search_grid_far_field() + else: + self.set_search_grid() + + def __init_cells(self, coeff: float = 0.1): + + if self.estimation_params == "range": + self.cell_size = int(self.ranges_dict.shape[0] * coeff) + elif self.estimation_params == "angle": + self.cell_size = int(self.angles_dict.shape[0] * coeff) + elif self.estimation_params == "angle, range": + self.cell_size_angle = int(self.angles_dict.shape[0] * coeff) + self.cell_size_range = int(self.ranges_dict.shape[0] * coeff) + + if self.cell_size is not None: + if self.cell_size % 2 == 0: + self.cell_size += 1 + if self.cell_size_angle is not None: + if self.cell_size_angle % 2 == 0: + self.cell_size_angle += 1 + if self.cell_size_range is not None: + if self.cell_size_range % 2 == 0: + self.cell_size_range += 1 + + def init_cells(self, coeff: float = 0.2): + self.__init_cells(coeff) + + def _plot_1d_spectrum(self, highlight_corrdinates, batch, add_title: bool = False, save: bool = False): + if self.estimation_params == "angle": + x = np.rad2deg(self.angles_dict.detach().cpu().numpy()) + x_label = "angle [deg]" + elif self.estimation_params == "range": + x = self.ranges_dict.detach().cpu().numpy() + x_label = "distance [m]" + else: + raise ValueError(f"MUSIC._plot_1d_spectrum: No such option for param estimation.") + y = self.music_spectrum[batch].detach().cpu().numpy() + plt.figure() + plt.plot(x, y.T, label="Music Spectrum") + if highlight_corrdinates is not None: + for idx, dot in enumerate(highlight_corrdinates): + plt.vlines(dot, np.min(y), np.max(y), colors='r', linestyles='dashed', label=f"Ground Truth") + if add_title: + plt.title("MUSIC SPECTRUM") + plt.grid() + plt.ylabel("Spectrum power") + plt.xlabel(x_label) + plt.legend() + plt.tight_layout() + if save: + plt.savefig("1d_music_spectrum.pdf") + plt.show() + + def _plot_3d_spectrum(self, highlight_coordinates, batch, method, music_spectrum=None, add_title: bool = False, save: bool = False): + """ + Plot the MUSIC 2D spectrum. + + """ + if method == "3D": + # Creating figure + distances = self.ranges_dict.detach().cpu().numpy() + angles = self.angles_dict.detach().cpu().numpy() + if music_spectrum is None: + spectrum = self.music_spectrum[batch].detach().cpu().numpy() + else: + spectrum = music_spectrum[batch].detach().cpu().numpy() + x, y = np.meshgrid(distances, np.rad2deg(angles)) + # Plotting the 3D surface + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + ax.plot_surface(x, y, 10 * np.log10(spectrum), cmap='viridis') + + if highlight_coordinates: + highlight_coordinates = np.array(highlight_coordinates) + ax.scatter( + highlight_coordinates[:, 0], + np.rad2deg(highlight_coordinates[:, 1]), + np.log1p(highlight_coordinates[:, 2]), + color='red', + s=50, + label='Ground Truth', + marker="x" + ) + if add_title: + ax.set_title('MUSIC spectrum') + ax.set_xlim(distances[0], distances[-1]) + ax.set_ylim(np.rad2deg(angles[0]), np.rad2deg(angles[-1])) + # Adding labels + ax.set_ylabel('Theta [deg]') + ax.set_xlabel('Radius [m]') + ax.set_zlabel('Power [dB]') + plt.colorbar(ax.plot_surface(x, y, 10 * np.log10(spectrum), cmap='viridis'), shrink=0.5, aspect=5) + + if highlight_coordinates: + ax.legend() # Adding a legend + + # Display the plot + plt.tight_layout() + if save: + plt.savefig("3d_music_spectrum.pdf") + plt.show() + elif method == "heatmap": + xmin, xmax = np.min(self.ranges_dict.cpu().detach().numpy()), np.max(self.ranges_dict.cpu().detach().numpy()) + ymin, ymax = np.min(self.angles_dict.cpu().detach().numpy()), np.max(self.angles_dict.cpu().detach().numpy()) + if music_spectrum is None: + spectrum = self.music_spectrum[batch].cpu().detach().numpy() + else: + spectrum = music_spectrum[batch].cpu().detach().numpy() + plt.imshow(spectrum, cmap="hot", + extent=[xmin, xmax, np.rad2deg(ymin), np.rad2deg(ymax)], origin='lower', aspect="auto") + if highlight_coordinates is not None: + for idx, dot in enumerate(highlight_coordinates): + x = self.ranges_dict.cpu().detach().numpy()[dot[1]] + y = np.rad2deg(self.angles_dict.cpu().detach().numpy()[dot[0]]) + plt.plot(x, y, label=f"{x:.1f} [m], {y:.1f} [deg]", marker='o', markerfacecolor='none', + markeredgecolor='white', linestyle='-', color='white', markersize=10) + # plt.plot(x, y, marker='x', linestyle='', color='green', markersize=8) + plt.legend() + plt.colorbar() + if add_title: + plt.title("MUSIC Spectrum heatmap") + plt.xlabel("Distances [m]") + plt.ylabel("Angles [deg]") + + plt.figaspect(2) + plt.tight_layout() + if save: + plt.savefig("heatmap_music_spectrum.pdf") + plt.show() + elif method == "slice": + x = self.ranges_dict.detach().cpu().numpy() + x_label = "distance [m]" + y = self.music_spectrum[batch].detach().cpu().numpy()[highlight_coordinates[0]] + plt.figure() + plt.plot(x, y.T, label="Music Spectrum") + if highlight_coordinates is not None: + for idx, dot in enumerate(highlight_coordinates[1:]): + plt.vlines(dot, np.min(y), np.max(y), colors='r', linestyles='dashed', label=f"Ground Truth") + if add_title: + plt.title(f"MUSIC SPECTRUM Slice at {torch.round(torch.rad2deg(self.angles_dict[highlight_coordinates[0]]))}") + plt.grid() + plt.ylabel("Spectrum power") + plt.xlabel(x_label) + plt.legend() + plt.tight_layout() + if save: + plt.savefig("slice_music_spectrum.pdf") + plt.show() + + def __set_search_grid_far_field(self): + self.steering_dict = self.system_model.steering_vec_far_field(self.angles_dict, f_c=None, nominal=True, fix_sv_noise=True).squeeze(-1) + + def __set_search_grid_near_field(self, known_angles: torch.Tensor = None, known_distances: torch.Tensor = None): + """ + + Returns: + + """ + if known_angles is None: + known_angles = self.angles_dict + if known_distances is None: + known_distances = self.ranges_dict + self.steering_dict = self.system_model.steering_vec_near_field(angles=known_angles, ranges=known_distances, + generate_search_grid=True, nominal=True, + f_c=None).squeeze(-1).cpu() + if torch.isnan(self.steering_dict).any(): + raise ValueError("Nan values in steering matrix") + + def __str__(self): + if self.estimation_params == "angle": + return "music_angle" + elif self.estimation_params == "range": + return "music_range" + elif self.estimation_params == "angle, range": + return "2d_music" + + def __init_criteria(self): + if self.estimation_params == "angle": + self.criterion = RMSPELoss(balance_factor=1.0) + elif self.estimation_params == "range": + self.criterion = RMSPELoss(balance_factor=0.0) + elif self.estimation_params == "angle, range": + self.criterion = CartesianLoss() + self.separated_criterion = RMSPELoss(1.0) + else: + raise ValueError(f"MUSIC.__init_criteria: Unrecognized estimation param {self.estimation_params}") + + def _get_name(self): + return "MUSIC" + + +# Import the diffMUSIC implementation +from src.methods_pack.diffmusic import DiffMUSIC + + +class Filter(nn.Module): + def __init__(self, min_cell_size, max_cell_size, number_of_filter=10): + super(Filter, self).__init__() + self.number_of_filters = number_of_filter + self.cell_sizes = torch.linspace(min_cell_size, max_cell_size, number_of_filter).to(torch.int32).to(self.device) + self.cell_bank = {} + for cell_size in enumerate(self.cell_sizes.data): + cell_size = cell_size[1] + self.cell_bank[cell_size] = torch.arange(-cell_size, cell_size, 1, dtype=torch.long, device=self.device) + self.fc = nn.Linear(self.number_of_filters, 1) + self.fc.weight.data = torch.randn(1, number_of_filter) / 100 + (1 / number_of_filter) + self.fc.weight.data = self.fc.weight.data.to(torch.float64) + self.fc.bias.data = torch.Tensor([0]) + self.fc.bias.data = self.fc.bias.data.to(torch.float64) + self.fc.bias.requires_grad_(False) + self.relu = nn.ReLU() + + def forward(self, input, search_space): + peaks = torch.zeros(input.shape[0], 1).to(torch.int64) + for batch in range(peaks.shape[0]): + music_spectrum = input[batch].cpu().detach().numpy().squeeze() + # Find spectrum peaks + peaks_tmp = list(sc.signal.find_peaks(music_spectrum)[0]) + # Sort the peak by their amplitude + peaks_tmp.sort(key=lambda x: music_spectrum[x], reverse=True) + if len(peaks_tmp) == 0: + peaks_tmp = torch.randint(search_space.shape[0], (1,)) + else: + peaks_tmp = peaks_tmp[0] + peaks[batch] = peaks_tmp + top_1 = peaks + output = torch.zeros(input.shape[0], self.number_of_filters).to(self.device).to(torch.float64) + for idx, cell in enumerate(self.cell_bank.values()): + tmp_cell = top_1 + cell + tmp_cell %= input.shape[1] + tmp_cell = tmp_cell.unsqueeze(-1) + metrix_thr = torch.gather(input.unsqueeze(-1).expand(-1, -1, tmp_cell.size(-1)), 1, tmp_cell) + soft_max = torch.softmax(metrix_thr, dim=1) + output[:, idx] = torch.einsum("bkm, bkm -> bm", search_space[tmp_cell], soft_max).squeeze() + output = self.fc(output) + output = self.relu(output) + self.clip_weights_values() + return output + + def clip_weights_values(self): + self.fc.weight.data = torch.clip(self.fc.weight.data, 0.1, 1) + self.fc.weight.data /= torch.sum(self.fc.weight.data) \ No newline at end of file diff --git a/src/diffMUSIC.py b/src/old/diffMUSIC.py similarity index 79% rename from src/diffMUSIC.py rename to src/old/diffMUSIC.py index 90f121f..496e0df 100644 --- a/src/diffMUSIC.py +++ b/src/old/diffMUSIC.py @@ -76,113 +76,13 @@ def __init__(self, temperature: Temperature for soft peak finding cell_size_coeff: Coefficient for cell size in soft peak finding """ - super().__init__() + super().__init__(system_model_params, + init_antenna_positions, + init_antenna_gains, + gain_constraint, + temperature, + cell_size_coeff) - self.params = system_model_params - self.N = self.params.N # Number of antennas - self.M = self.params.M # Number of sources - self.gain_constraint = gain_constraint - self.temperature = temperature - self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - - # Initialize antenna positions as learnable parameters - if init_antenna_positions is None: - init_antenna_positions = torch.arange(self.N, dtype=torch.float64) * (self.params.wavelength / 2) - - self.antenna_positions = nn.Parameter(init_antenna_positions.clone()) - assert init_antenna_positions.requires_grad is True, "Initial antenna positions must be a differentiable tensor." - - # Initialize antenna gains as learnable parameters - if init_antenna_gains is None: - init_antenna_gains = torch.ones(self.N, dtype=torch.complex64) - self.antenna_gains = nn.Parameter(init_antenna_gains.clone()) - - # Initialize angle grid for far-field DOA estimation - self._init_angle_grid() - - # Initialize cell size for soft peak finding - self.cell_size = int(self.angles_dict.shape[0] * cell_size_coeff) - if self.cell_size % 2 == 0: - self.cell_size += 1 - - # Store current MUSIC spectrum for analysis - self.music_spectrum = None - - 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_dict = torch.arange(-angle_range, angle_range + angle_resolution, - angle_resolution, dtype=torch.float64) - self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) - - def _apply_position_constraints(self): - """Reorders the antenna positions, ensuring we don't get stucked if some replace positions. - I don't think we need this because it is mathematically fine that thy will replace order. - In fact, this replacement may occur some incontinuity in the loss. - We do need to keep it in mind though. - """ - # with torch.no_grad(): - # if self.position_constraint == "ula": - # sorted_positions, _ = torch.sort(self.antenna_positions) - # self.antenna_positions.data = sorted_positions - pass - - def _apply_gain_constraints(self): - """Prevent extrme gains from future random walk. For now i leave commented.""" - # with torch.no_grad(): - # if self.gain_constraint == "positive": - # # Just clamp to reasonable ranges to prevent numerical issues - # real_part = torch.clamp(self.antenna_gains.real, min=0.1, max=10.0) - # imag_part = torch.clamp(self.antenna_gains.imag, min=-2.0, max=2.0) - # self.antenna_gains.data = torch.complex(real_part, imag_part, dtype=torch.complex64) - pass - - def get_constrained_parameters(self): - """Get antenna parameters without any sorting - order doesn't matter mathematically""" - positions = self.antenna_positions - gains = self.antenna_gains - return positions, gains - -#NOTE: ORI - I read up to this point. - - def compute_steering_matrix(self, angles: torch.Tensor): - """ - Compute steering matrix for far-field sources using current antenna configuration. - Follows the pattern from your SystemModel.steering_vec_far_field method. - - Args: - angles: DOA angles in radians [num_angles] - - Returns: - Complex steering matrix [N, num_angles] - """ - positions, gains = self.get_constrained_parameters() - - # Convert angles to proper shape for broadcasting - if angles.dim() == 1: - angles = angles.unsqueeze(0) # [1, num_angles] - - # Reshape positions for broadcasting: [N, 1] - positions = positions.view(-1, 1) - - # Compute phase delays: [N, 1] * sin([1, num_angles]) -> [N, num_angles] - time_delay = positions @ torch.sin(angles) - - # Compute steering vectors (following your formula) - steering_matrix = torch.exp(-2j * torch.pi * time_delay / self.params.wavelength) - steering_matrix = steering_matrix.to(torch.complex64) - - # Apply antenna gains - gains_diag = torch.diag(gains) - - # Normalize gains (following your normalization pattern) - normalization_factor = 1 / torch.linalg.norm(gains, ord=2) - steering_matrix = normalization_factor * gains_diag @ steering_matrix - - return steering_matrix def sample_covariance(self, x: torch.Tensor) -> torch.Tensor: """ diff --git a/src/old/diff_subspace_method b/src/old/diff_subspace_method new file mode 100644 index 0000000..0655b4e --- /dev/null +++ b/src/old/diff_subspace_method @@ -0,0 +1,140 @@ +import warnings +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import matplotlib.pyplot as plt +import scipy as sc + +from src.signal_creation import SystemModel, SystemModelParams +from src.utils import set_unified_seed +from scipy.ndimage import maximum_filter +from scipy.ndimage import label +from scipy.ndimage import find_objects + +from src.subspace_method import SubspaceMethod + + +class DiffSubspaceMethod(SubspaceMethod): + """ + Differentiable extension of SubspaceMethod with learnable antenna parameters + """ + def __init__(self, + system_model_params: SystemModelParams, + init_antenna_positions: torch.Tensor = None, + init_antenna_gains: torch.Tensor = None, + gain_constraint: str = "none", + temperature: float = 1.0, + cell_size_coeff: float = 0.2, + model_order_estimation: str = None): + + # Create system model for parent class + system_model = SystemModel(system_model_params) + super().__init__(system_model, model_order_estimation) + + # Store parameters + self.params = system_model_params + self.N = system_model_params.N + self.gain_constraint = gain_constraint + self.temperature = temperature + + # Initialize antenna positions as learnable parameters + if init_antenna_positions is None: + init_antenna_positions = torch.arange(self.N, dtype=torch.float64) * (self.params.wavelength / 2) + + self.antenna_positions = nn.Parameter(init_antenna_positions.clone()) + assert init_antenna_positions.requires_grad is True, "Initial antenna positions must be a differentiable tensor." + + # Initialize antenna gains as learnable parameters + if init_antenna_gains is None: + init_antenna_gains = torch.ones(self.N, dtype=torch.complex64) + self.antenna_gains = nn.Parameter(init_antenna_gains.clone()) + + # Initialize angle grid for far-field DOA estimation + self._init_angle_grid() + + # Initialize cell size for soft peak finding + self.cell_size = int(self.angles_dict.shape[0] * cell_size_coeff) + if self.cell_size % 2 == 0: + self.cell_size += 1 + + # Store current MUSIC spectrum for analysis + self.music_spectrum = None + + 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_dict = torch.arange(-angle_range, angle_range + angle_resolution, + angle_resolution, dtype=torch.float64) + self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) + + def _apply_position_constraints(self): + """Reorders the antenna positions, ensuring we don't get stucked if some replace positions. + I don't think we need this because it is mathematically fine that thy will replace order. + In fact, this replacement may occur some incontinuity in the loss. + We do need to keep it in mind though. + """ + # with torch.no_grad(): + # if self.position_constraint == "ula": + # sorted_positions, _ = torch.sort(self.antenna_positions) + # self.antenna_positions.data = sorted_positions + pass + + def _apply_gain_constraints(self): + """Prevent extrme gains from future random walk. For now i leave commented.""" + # with torch.no_grad(): + # if self.gain_constraint == "positive": + # # Just clamp to reasonable ranges to prevent numerical issues + # real_part = torch.clamp(self.antenna_gains.real, min=0.1, max=10.0) + # imag_part = torch.clamp(self.antenna_gains.imag, min=-2.0, max=2.0) + # self.antenna_gains.data = torch.complex(real_part, imag_part, dtype=torch.complex64) + pass + + def get_constrained_parameters(self): + """Get antenna parameters without any sorting - order doesn't matter mathematically""" + positions = self.antenna_positions + gains = self.antenna_gains + return positions, gains + + #NOTE: I read up to this point + + def compute_steering_matrix(self, angles: torch.Tensor): + """ + Compute steering matrix for far-field sources using current antenna configuration. + Follows the pattern from your SystemModel.steering_vec_far_field method. + + Args: + angles: DOA angles in radians [num_angles] + + Returns: + Complex steering matrix [N, num_angles] + """ + positions, gains = self.get_constrained_parameters() + + # Convert angles to proper shape for broadcasting + if angles.dim() == 1: + angles = angles.unsqueeze(0) # [1, num_angles] + + # Reshape positions for broadcasting: [N, 1] + positions = positions.view(-1, 1) + + # Compute phase delays: [N, 1] * sin([1, num_angles]) -> [N, num_angles] + time_delay = positions @ torch.sin(angles) + + # Compute steering vectors (following your formula) + steering_matrix = torch.exp(-2j * torch.pi * time_delay / self.params.wavelength) + steering_matrix = steering_matrix.to(torch.complex64) + + # Apply antenna gains + gains_diag = torch.diag(gains) + + # Normalize gains (following your normalization pattern) + normalization_factor = 1 / torch.linalg.norm(gains, ord=2) + steering_matrix = normalization_factor * gains_diag @ steering_matrix + + return steering_matrix + + \ No newline at end of file diff --git a/src/old/subspace_method.py b/src/old/subspace_method.py new file mode 100644 index 0000000..00eb552 --- /dev/null +++ b/src/old/subspace_method.py @@ -0,0 +1,242 @@ +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 +from src.signal_creation import SystemModel + + + +#TODO: Haven't really checked the code here yet, just inheritrs for now. +class SubspaceMethod(nn.Module): + """ + Basic methods for all subspace methods. + """ + + def __init__(self, system_model: SystemModel, model_order_estimation: str = 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 + + def subspace_separation(self, + covariance: torch.Tensor, + number_of_sources: torch.tensor = None) \ + -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.tensor]: + """ + + Args: + covariance: + number_of_sources: + + Returns: + the signal ana noise subspaces, both as torch.Tensor(). + """ + 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/signal_creation.py b/src/signal_creation.py index 4de373e..2d56d8a 100644 --- a/src/signal_creation.py +++ b/src/signal_creation.py @@ -78,7 +78,7 @@ def steering_vec(self, angles: np.ndarray) -> torch.Tensor: #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: + def steering_vec_far_field(self, angles: np.ndarray | torch.tensor) -> torch.Tensor: """ Compute far-field steering vectors @@ -154,7 +154,7 @@ def __init__(self, system_model_params: SystemModelParams): super().__init__(system_model_params) self.angles = None - def set_labels(self, angles: list = None): + def set_labels(self, angles: list | None = None): """ Set the angles for the sources @@ -167,7 +167,7 @@ def get_labels(self): """Get the current source angles as tensor""" return torch.tensor(self.angles, dtype=torch.float32) - def set_angles(self, doa: list = None): + def set_angles(self, doa: list | None = None): """ Set direction of arrival angles diff --git a/src/subspace_method.py b/src/subspace_method.py new file mode 100644 index 0000000..aeca4f1 --- /dev/null +++ b/src/subspace_method.py @@ -0,0 +1,241 @@ +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): + 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 + + 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 14eaef2..c82c2a6 100644 --- a/src/utils.py +++ b/src/utils.py @@ -121,51 +121,51 @@ def initialize_paths(main_path: Path, system_model_params, dt_string_for_save: s return datasets_path, results_path -# def sample_covariance(x: torch.Tensor) -> torch.Tensor: -# """ -# Calculates the sample covariance matrix for each element in the batch. +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. + 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 + Returns: + -------- + covariance_mat (np.ndarray): Covariance matrix. + """ + if x.dim() == 2: + 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 -# def spatial_smoothing_covariance(x: torch.Tensor): -# """ -# Calculates the covariance matrix using spatial smoothing technique for each element in the batch. +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. + Args: + ----- + X (np.ndarray): Input samples matrix. -# Returns: -# -------- -# covariance_mat (np.ndarray): Covariance 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 + 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): # """ From d3f5f08644cd8a2eacd9650005ad9eaee3e4c920 Mon Sep 17 00:00:00 2001 From: oritalp Date: Wed, 18 Jun 2025 13:39:47 +0300 Subject: [PATCH 04/13] DiffMUSIC implementation almost done --- .gitignore | 3 +- check.py | 29 +- doa_runner.py | 606 ++++++++++++++++++++++++++++++ full_environment.yml | 1 + main.py | 123 +++--- run_simulation.py | 92 ++++- src/diffmusic.py | 410 ++++++-------------- src/metrics.py | 179 +++++++++ src/music.py | 698 ----------------------------------- src/old/diffMUSIC.py | 494 ------------------------- src/old/diff_subspace_method | 140 ------- src/old/subspace_method.py | 242 ------------ src/signal_creation.py | 9 +- src/subspace_method.py | 5 +- src/utils.py | 508 +++---------------------- 15 files changed, 1142 insertions(+), 2397 deletions(-) create mode 100644 doa_runner.py create mode 100644 src/metrics.py delete mode 100644 src/music.py delete mode 100644 src/old/diffMUSIC.py delete mode 100644 src/old/diff_subspace_method delete mode 100644 src/old/subspace_method.py diff --git a/.gitignore b/.gitignore index bad7ae0..5f9a405 100644 --- a/.gitignore +++ b/.gitignore @@ -10,4 +10,5 @@ wandb/ *.png *.pdf *.out -*.sh \ No newline at end of file +*.sh +./datasets/ \ No newline at end of file diff --git a/check.py b/check.py index d411ca8..5bd6f41 100644 --- a/check.py +++ b/check.py @@ -1,3 +1,30 @@ import torch -print("aabcd".find("b")) \ No newline at end of file + +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..26e8b32 --- /dev/null +++ b/doa_runner.py @@ -0,0 +1,606 @@ +""" +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 + +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}_ws_{window_size}_{self.system_params.loss_type}" + else: + run_name = f"ws_{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 = results + return results + + 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 + """ + import time + + best_rmspe = float('inf') + best_results = None + best_window_size = None + + total_start_time = time.time() + 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': [] + } + + 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) + + # Store training history + self.training_history['train_loss'].append(epoch_loss) + 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, + "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}") + + 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 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_coplex_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 + } + + 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 + ) + return None \ No newline at end of file diff --git a/full_environment.yml b/full_environment.yml index 412b4b7..f9311e4 100644 --- a/full_environment.yml +++ b/full_environment.yml @@ -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 37232d2..d6e5178 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,26 +10,35 @@ --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: 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. +#TODO: check wandb (short single case) +# run multiple window sizes cases. +# pass over the running once again +# defined missions for rearanging the results and where they are saved, split it into single_run_ressults and multi_run_results. +# define graphs plotting class for the multi-case results. + +#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. # Initialization os.system("cls||clear") @@ -44,14 +52,14 @@ } simulation_commands = { - "CREATE_DATA": True, + "create_data": True, + "save_data": False, # Save data after creation + "plot_results": True, # Plot data after creation "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": 16, # number of antennas "M": 5, # number of sources @@ -61,42 +69,51 @@ "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 - "wavelength": 1, # The carrier wavelength of the signal in meters, 1 can be fine for reserch, + "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 vlaid float between 0 and wavelength/4 or "wavelength/n" to use with refrence to the wavelength + # 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################################## + ############################### 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 - } -system_model_params["device"] = torch.device("cuda" if torch.cuda.is_available() else "cpu") - -model_config = { - "model_type": "diffMUSIC", # diffMUSIC - "model_params": {"window_size": 21} +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.2, 0.4, 0.05), # 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 } training_params = { - "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, "step_size": 50, + "weight_decay": 0.0, "use_wandb": False } - def parse_arguments(): parser = argparse.ArgumentParser(description="Run simulation with optional parameters.") parser.add_argument('-s', "--snr" ,type=int, help='SNR value', default=None) @@ -111,11 +128,10 @@ def parse_arguments(): 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; diffMUSIC, MUSIC', default=None) + parser.add_argument('-lt', '--loss_type', type=str, help='Loss type; rmspe, spectrum, unsupervised', default=None) - parser.add_argument('-mt', '--model_type', type=str, help='Model type; diffMUSIC, SubspaceNet, DCD-MUSIC, DeepCNN, TransMUSIC, DR_MUSIC', 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) @@ -124,16 +140,15 @@ def parse_arguments(): parser.add_argument('-step', '--step_size', type=int, help='Step size', default=None) 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: @@ -168,14 +183,13 @@ def parse_arguments(): 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 + # 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: @@ -194,17 +208,34 @@ def parse_arguments(): 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["load_data"] = not simulation_commands["create_data"] + # 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 oreder switching between neigboring array sensors.") + "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, - 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 isinstance(results, dict) and 'rmspe' in results: + print(f"\nFinal Results:") + print(f"RMSPE: {results['rmspe']:.6f} degrees") + 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 and "learned_antenna_gains" in results: + if model_config["model_type"].lower() == "music": + print("These shoulf be the tandard one:") + print(f"Learned antenna positions: {results['learned_antenna_positions']}") + print(f"Learned antenna gains: {results['learned_antenna_gains']}") \ No newline at end of file diff --git a/run_simulation.py b/run_simulation.py index bdf0d97..c816682 100644 --- a/run_simulation.py +++ b/run_simulation.py @@ -17,6 +17,7 @@ from datetime import datetime import torch import numpy as np +from doa_runner import DoARunner def __run_simulation(**kwargs): @@ -24,46 +25,99 @@ def __run_simulation(**kwargs): SYSTEM_MODEL_PARAMS = kwargs["system_model_params"] MODEL_CONFIG = kwargs["model_config"] TRAINING_PARAMS = kwargs["training_params"] - create_data = SIMULATION_COMMANDS["CREATE_DATA"] # Creating new dataset - load_data = SIMULATION_COMMANDS["LOAD_DATA"] # Load specific model for training + 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...") - 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 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 - # Define system model parameters - system_model_params = SystemModelParams(**SYSTEM_MODEL_PARAMS) # Initialize paths + #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, dt_string_for_save) data_loading_path = SIMULATION_COMMANDS["data_loading_path"] #ONLY USED if CREATE_DATA is False! - # Create system model + + # Prepare data dictionary + data_dict = {} 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() - array = signals_creator.get_array() - gain_impairments = signals_creator.get_gain_perturbations() - gain_impairments_norm = torch.linalg.norm(gain_impairments, ord=2, dim=0) #just for testing purposes + physical_array = signals_creator.get_array() + physical_antennas_gains = signals_creator.get_antenna_gains() + antennas_gains_norm = torch.linalg.norm(physical_antennas_gains, ord=2, dim=0) #just for testing purposes + + # 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 - utils.save_data_to_file(data_saving_path, measurements, signals, steering_mat, - noise, true_angles, array, gain_impairments, system_model_params) + 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 - measurements, signals, steering_mat, noise, true_angles, array, gain_impairments, system_model_params = \ - utils.load_data_from_file(data_loading_path) - return None + 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 + } + + # 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_simulation(**kwargs): @@ -78,8 +132,8 @@ def run_simulation(**kwargs): """ #TODO: check if anything is missing for the scenario_dict option once needed. if kwargs["scenario_dict"] == {}: - loss = __run_simulation(**kwargs) - return loss + results = __run_simulation(**kwargs) + return results # from this on the option of multiple scenarios is used. this is activated when we specify the sceario_dict # in main.py @@ -134,4 +188,4 @@ def run_simulation(**kwargs): 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 index 7279a61..7b372da 100644 --- a/src/diffmusic.py +++ b/src/diffmusic.py @@ -21,12 +21,14 @@ 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) implementation for DoA estimation with hardware impairment learning. + 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 @@ -40,8 +42,8 @@ class DiffMUSIC(SubspaceMethod): - Softmax-based peak finding for differentiability """ - def __init__(self, system_model_params, N: int, - window_size: int = 21, model_order_estimation: str = None): + 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 @@ -49,16 +51,16 @@ def __init__(self, system_model_params, N: int, system_model: System model object (kept for compatibility) N: Number of antennas wavelength: Signal wavelength - window_size: Size of angular window for softmax peak finding model_order_estimation: Model order estimation method """ system_model = SystemModel(system_model_params) - super().__init__(system_model, model_order_estimation=model_order_estimation) + 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 = window_size # For now it is inserted directly + self.window_size = self.params.softmax_window_size # Initialize learnable parameters self._init_learnable_parameters() @@ -82,9 +84,8 @@ def _init_learnable_parameters(self): # 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) - self.gains_real = nn.Parameter(gains_real) - self.gains_imag = nn.Parameter(gains_imag) - self.complex_gain = torch.complex(gains_real, gains_imag, dtype=torch.complex64) + 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""" @@ -92,16 +93,47 @@ def _init_angle_grid(self): 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_dict = torch.arange(-angle_range, angle_range + angle_resolution, + self.angles_grid = torch.arange(-angle_range, angle_range + angle_resolution, angle_resolution, dtype=torch.float64) - self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) + 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: """ @@ -117,13 +149,14 @@ def compute_steering_matrix(self, angles: torch.Tensor) -> torch.Tensor: angles = angles.unsqueeze(0) # Get complex gains - complex_gains = self.complex_gain + 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 = self.antenna_positions.unsqueeze(1) @ torch.sin(angles).unsqueeze(0) + 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) @@ -137,18 +170,14 @@ def compute_steering_matrix(self, angles: torch.Tensor) -> torch.Tensor: return steering_matrix.to(torch.complex64) -#TODO: Read up to the peak_finder inside the forward-pass. Looks fine, need to understand the peak_finder. -# If the hard decision is fine, we can use this also for MUSIC. Keep reading and notice the first dimension -# of cov is the batch_dim. Another thing is to change system_model_params to also include the training_params -# (all params in general). For now, the gains real and imag parts are parameters, change just -# the whole complex gain to be a parameter, and use it in the steering matrix. +#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) + 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) @@ -168,9 +197,9 @@ def forward(self, cov: torch.Tensor, number_of_sources: int, known_angles=None): self.music_spectrum = 1 / (inverse_spectrum + 1e-10) # Peak finding (differentiable during training, hard during inference) - estimated_angles = self._peak_finder(number_of_sources) + estimated_angles, peaks_masks = self._peak_finder(number_of_sources) - return estimated_angles, source_estimation, eigen_regularization #eigen_regularization is not used in this implementation + 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: """ @@ -189,7 +218,7 @@ def _compute_inverse_spectrum(self, noise_subspace: torch.Tensor) -> torch.Tenso # steering_grid: (N, num_angles), noise_subspace: (batch_size, N, N-M) var1 = torch.einsum("na, bnm -> bam", steering_grid.conj(), - noise_subspace) # (batch_size, num_angles, N-M) + 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) @@ -209,14 +238,15 @@ def _peak_finder(self, number_of_sources: int) -> torch.Tensor: 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 self._hard_peak_finder(number_of_sources, return_angles=True) - def _hard_peak_finder(self, number_of_sources: int) -> torch.Tensor: + 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) @@ -231,7 +261,7 @@ def _hard_peak_finder(self, number_of_sources: int) -> torch.Tensor: peaks_indices = sc.signal.find_peaks(spectrum, threshold=0.0)[0] if len(peaks_indices) < number_of_sources: - warnings.warn("diffMUSIC: Not enough peaks found, using highest values") + 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), @@ -242,13 +272,17 @@ def _hard_peak_finder(self, number_of_sources: int) -> torch.Tensor: sorted_peaks = peaks_indices[np.argsort(spectrum[peaks_indices])[::-1]] peaks[batch] = torch.from_numpy(sorted_peaks[:number_of_sources]).to(self.device) - # 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 + 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: """ @@ -264,18 +298,23 @@ def _differentiable_peak_finder(self, number_of_sources: int) -> torch.Tensor: 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 = self._find_initial_peaks(spectrum, number_of_sources) + 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[source_idx] + 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] @@ -286,24 +325,11 @@ def _differentiable_peak_finder(self, number_of_sources: int) -> torch.Tensor: # 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 + return estimated_angles, peaks_masks - def _find_initial_peaks(self, spectrum: torch.Tensor, number_of_sources: int) -> torch.Tensor: - """Find initial peak locations (non-differentiable but provides starting points)""" - spectrum_np = spectrum.cpu().detach().numpy() - peaks_indices = sc.signal.find_peaks(spectrum_np, threshold=0.0)[0] - - if len(peaks_indices) < number_of_sources: - # Use highest values if not enough peaks - additional_peaks = torch.topk(spectrum, - number_of_sources - len(peaks_indices), - largest=True).indices.cpu().numpy() - peaks_indices = np.concatenate([peaks_indices, additional_peaks]) - - # Sort by amplitude and take top peaks - sorted_peaks = peaks_indices[np.argsort(spectrum_np[peaks_indices])[::-1]] - return torch.from_numpy(sorted_peaks[:number_of_sources]).to(self.device) def _create_angular_mask(self, center_idx: int, spectrum_length: int) -> torch.Tensor: """ @@ -326,19 +352,6 @@ def _create_angular_mask(self, center_idx: int, spectrum_length: int) -> torch.T return mask_indices - def get_learned_parameters(self): - """ - Get the learned array parameters - - Returns: - dict: Dictionary containing learned positions and gains - """ - return { - 'antenna_positions': self.antenna_positions.detach().cpu().numpy(), - 'complex_gains': self.get_complex_gains().detach().cpu().numpy(), - 'gains_magnitude': torch.abs(self.get_complex_gains()).detach().cpu().numpy(), - 'gains_phase': torch.angle(self.get_complex_gains()).detach().cpu().numpy() - } def set_nominal_parameters(self, positions: torch.Tensor = None, gains: torch.Tensor = None): """ @@ -361,8 +374,8 @@ def set_nominal_parameters(self, positions: torch.Tensor = None, gains: torch.Te self.gains_real.copy_(gains.real.to(torch.float64)) self.gains_imag.copy_(gains.imag.to(torch.float64)) - def plot_spectrum(self, batch_idx: int = 0, highlight_angles: torch.Tensor = None, - save: bool = False, title: str = "diffMUSIC Spectrum"): + def plot_spectrum(self, highlight_angles: torch.Tensor = None, + path_to_save: str | None = False, batch_idx: int = 0): """ Plot the MUSIC spectrum @@ -390,243 +403,52 @@ def plot_spectrum(self, batch_idx: int = 0, highlight_angles: torch.Tensor = Non plt.xlabel('Angle [degrees]') plt.ylabel('Spectrum Power') - plt.title(title) + plt.title(("diffMUSIC" if self.training else "MUSIC") + " Spectrum") plt.grid(True, alpha=0.3) plt.legend() plt.tight_layout() - if save: - plt.savefig('diffmusic_spectrum.pdf') + if path_to_save is not None: + plt.savefig(path_to_save / f"Spectrum.png", dpi=300) plt.show() - def plot_learned_array(self, save: bool = False): - """ - Plot the learned antenna array geometry - - Args: - save: Whether to save the plot - """ - learned_params = self.get_learned_parameters() - positions = learned_params['antenna_positions'] - gains_mag = learned_params['gains_magnitude'] - gains_phase = learned_params['gains_phase'] - - # Nominal positions for comparison - nominal_positions = np.arange(self.N) * (self.wavelength / 2) - - fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 10)) - - # Plot 1: Antenna positions - ax1.scatter(nominal_positions, np.zeros_like(nominal_positions), - marker='o', s=100, alpha=0.5, label='Nominal', color='blue') - ax1.scatter(positions, np.zeros_like(positions), - marker='s', s=100, label='Learned', color='red') - ax1.set_xlabel('Position [wavelengths]') - ax1.set_title('Antenna Positions') - ax1.legend() - ax1.grid(True, alpha=0.3) - - # Plot 2: Gain magnitudes - antenna_indices = np.arange(self.N) - ax2.bar(antenna_indices - 0.2, np.ones(self.N), width=0.4, - alpha=0.7, label='Nominal', color='blue') - ax2.bar(antenna_indices + 0.2, gains_mag, width=0.4, - alpha=0.7, label='Learned', color='red') - ax2.set_xlabel('Antenna Index') - ax2.set_ylabel('Gain Magnitude') - ax2.set_title('Complex Gain Magnitudes') - ax2.legend() - ax2.grid(True, alpha=0.3) - - # Plot 3: Gain phases - ax3.bar(antenna_indices, gains_phase, width=0.6, alpha=0.7, color='green') - ax3.set_xlabel('Antenna Index') - ax3.set_ylabel('Gain Phase [radians]') - ax3.set_title('Complex Gain Phases') - ax3.grid(True, alpha=0.3) - - plt.tight_layout() - - if save: - plt.savefig('diffmusic_learned_array.pdf') - plt.show() - def test_step(self, batch, batch_idx, model: nn.Module = None): +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): """ - Test step compatible with the existing framework - Args: - batch: Test batch (x, sources_num, label) - batch_idx: Batch index - model: Model (not used, kept for compatibility) - - Returns: - tuple: (rmspe, accuracy, test_length) + loss_type: "rmspe" for LSL,θ, "spectrum" for LSL,P, or "unsupervised" for LUL + **kwargs: Additional arguments for specific loss functions """ - x, sources_num, label = batch - if x.dim() == 2: - x = x.unsqueeze(0) - - test_length = x.shape[0] - x = x.to(self.device) - angles = label.to(self.device) - - # Check if sources number is consistent - if (sources_num != sources_num[0]).any(): - raise Exception("diffMUSIC test_step: Inconsistent number of sources in batch") - - sources_num = sources_num[0] - - # Compute covariance - if self.system_model.params.signal_nature == "non-coherent": - Rx = self.pre_processing(x, mode="sample") + 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: - Rx = self.pre_processing(x, mode="sample") - - # Run diffMUSIC - predictions, sources_num_estimation, _ = self(Rx, number_of_sources=sources_num) - - # Compute RMSPE - criterion = RMSPELoss(balance_factor=1.0) - rmspe = criterion(predictions, angles).sum().item() - - # Compute accuracy - acc = self.source_estimation_accuracy(sources_num, sources_num_estimation) - - return rmspe, acc, test_length - - def __str__(self): - return "diffMUSIC" - - def _get_name(self): - return "diffMUSIC" - - -# class DiffMUSICLoss(nn.Module): -# """ -# Loss functions for diffMUSIC training -# Implements both supervised learning strategies from the paper: -# - LSL,θ: RMSPE on estimated DoAs -# - LSL,P: Maximize spectrum amplitude at true DoA locations -# """ - -# def __init__(self, loss_type: str = "rmspe"): -# """ -# Args: -# loss_type: "rmspe" for LSL,θ or "spectrum" for LSL,P -# """ -# super().__init__() -# self.loss_type = loss_type -# self.rmspe_loss = RMSPELoss(balance_factor=1.0) - -# def forward(self, predictions, targets, spectrum=None, angles_grid=None): -# """ -# Compute loss based on specified type - -# Args: -# predictions: Predicted DoAs (for RMSPE loss) -# targets: True DoAs -# spectrum: MUSIC spectrum (for spectrum loss) -# angles_grid: Angular grid (for spectrum loss) - -# Returns: -# Loss value -# """ -# if self.loss_type == "rmspe": -# return self.rmspe_loss(predictions, targets) - -# elif self.loss_type == "spectrum": -# if spectrum is None or angles_grid is None: -# raise ValueError("Spectrum and angles_grid required for spectrum loss") - -# return self._spectrum_loss(targets, spectrum, angles_grid) - -# else: -# raise ValueError(f"Unknown loss type: {self.loss_type}") - -# def _spectrum_loss(self, true_angles, spectrum, angles_grid): -# """ -# Spectrum-based loss (LSL,P from paper) -# Maximizes spectrum amplitude at true DoA locations -# """ -# 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] - -# return total_loss / (batch_size * true_angles.shape[1]) - - -# class UnsupervisedDiffMUSIC(DiffMUSIC): -# """ -# Unsupervised diffMUSIC using Jain's Index (LUL from paper) -# Maximizes spectrum sharpness without requiring true DoA labels -# """ + raise ValueError(f"Unknown loss type: {loss_type}") -# def __init__(self, *args, **kwargs): -# super().__init__(*args, **kwargs) -# self.jains_index_loss = JainsIndexLoss() - -# def compute_unsupervised_loss(self): -# """ -# Compute unsupervised loss using Jain's Index on spectrum peaks - -# Returns: -# Loss value encouraging sharp peaks -# """ -# if self.music_spectrum is None: -# raise ValueError("No spectrum computed. Run forward pass first.") - -# total_loss = 0 -# batch_size = self.music_spectrum.shape[0] - -# for batch in range(batch_size): -# spectrum = self.music_spectrum[batch] - -# # Find initial peaks to create masks -# peaks_indices = self._find_initial_peaks(spectrum, self.system_model.params.M) - -# # Apply Jain's index to each peak region -# for peak_idx in peaks_indices: -# mask_indices = self._create_angular_mask(peak_idx, spectrum.shape[0]) -# masked_spectrum = spectrum[mask_indices] - -# # Jain's index encourages sharp peaks -# jains_loss = self.jains_index_loss(masked_spectrum) -# total_loss += jains_loss - -# 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().__init__() - -# def forward(self, x): -# """ -# 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 \ No newline at end of file + 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/music.py b/src/music.py deleted file mode 100644 index 3a2d728..0000000 --- a/src/music.py +++ /dev/null @@ -1,698 +0,0 @@ -import warnings - -import numpy as np -import torch -import torch.nn as nn -import matplotlib.pyplot as plt -import scipy as sc - -from src.signal_creation import SystemModel -from src.subspace_method import SubspaceMethod -from src.utils import * -# from src.metrics import RMSPELoss, CartesianLoss - -from scipy.ndimage import maximum_filter -from scipy.ndimage import label -from scipy.ndimage import find_objects - -#TODO: HAVENT CHECKED THE FILE YET, this is claude's suggestion for our new case - -def find_k_highest_peaks(matrix, k): - """ - Find the k highest peaks in a 2D matrix using SciPy tools. A peak is defined as a - local maximum surrounded by smaller values. - - Parameters: - - matrix (2D array-like): Input matrix. - - k (int): Number of highest peaks to extract. - - Returns: - - peaks (list): List of tuples (row, col, value) representing the positions and values of the k highest peaks. - """ - # Apply maximum filter to find local maxima - neighborhood = maximum_filter(matrix, size=21, mode='constant', cval=-np.inf) - local_max = (matrix == neighborhood) - - # Label the connected components of local maxima - labeled, num_features = label(local_max) - slices = find_objects(labeled) - - # Extract peak positions and values - peaks = [] - try: - for sl in slices: - row = int((sl[0].start + sl[0].stop - 1) / 2) - col = int((sl[1].start + sl[1].stop - 1) / 2) - value = matrix[row, col] - peaks.append((row, col, value)) - except Exception as e: - pass - - # Sort peaks by value (descending) and select the top k - peaks = sorted(peaks, key=lambda x: x[2], reverse=True)[:k] - if len(peaks) < k: - warnings.warn(f"find_k_highest_peaks: Less than {k} peaks found.") - # add random peaks - x_random = np.random.randint(0, matrix.shape[0], (k - len(peaks),)) - y_random = np.random.randint(0, matrix.shape[1], (k - len(peaks),)) - for i in range(k - len(peaks)): - peaks.append((x_random[i], y_random[i], matrix[x_random[i], y_random[i]])) - - return peaks - - -class MUSIC(SubspaceMethod): - """ - This is implementation of the MUSIC method for localization in Far and Near field environments. - For Far field - only "angle" can be estimated - For Near field - "angle", "range" and "angle, range" are the possible options. - """ - - def __init__(self, system_model: SystemModel, estimation_parameter: str, model_order_estimation: str = None): - """ - - Args: - system_model: - estimation_parameter: - """ - super().__init__(system_model, model_order_estimation=model_order_estimation) - self.estimation_params = estimation_parameter - self.angles_dict = None - self.ranges_dict = None - self.steering_dict = None - self.music_spectrum = None - self.cell_size = None - self.cell_size_angle = None - self.cell_size_range = None - self.noise_subspace = None - self.criterion = None - self.separated_criterion = None - - self.__init_grid_params() - self.__init_cells(0.2) - self.__init_criteria() - self.__init_search_grid() - - def forward(self, cov: torch.Tensor, number_of_sources: int, known_angles=None, known_distances=None): - """ - - Args: - cov: The covariance matrix of the input signal. - number_of_sources: the number of sources in the signal. Needed in case the dataset comprises mix number of sources. - known_angles: in case we are dealing with the Near field, the known angles should be passed. - known_distances: in case we are dealing with the Near field, the known distances should be passed. - - Returns: - tuple: the predicted parameters, the source estimation and the eigen regularization value. - """ - # single param estimation: the search grid should be updated for each batch, else, it's the same search grid. - if self.system_model.params.field_type in ["near", "full"] and self.estimation_params in ["range"]: - if known_angles.shape[-1] == 1: - self.set_search_grid(known_angles=known_angles, known_distances=known_distances) - else: - params = torch.zeros((cov.shape[0], number_of_sources), dtype=torch.float64, device=self.device) - for source in range(number_of_sources): - params_source, _, _ = self.forward(cov, number_of_sources=number_of_sources, - known_angles=known_angles[:, source][:, None]) - params[:, source] = params_source.squeeze() - return params - _, noise_subspace, source_estimation, eigen_regularization = self.subspace_separation(cov.to(torch.complex128), number_of_sources) - inverse_spectrum = self.get_inverse_spectrum(noise_subspace.to(self.device)).to(self.device) - if self._get_name() == "TOPS": - self.music_spectrum = torch.sum(1 / (inverse_spectrum + 1e-10), dim=-1) - else: - self.music_spectrum = 1 / (inverse_spectrum + 1e-10) - params = self.peak_finder(number_of_sources) - return params, source_estimation, eigen_regularization - - def get_music_spectrum_from_noise_subspace(self, noise_subspace: torch.Tensor) -> torch.Tensor: - inverse_spectrum = self.get_inverse_spectrum(noise_subspace.to(torch.complex128)) - self.music_spectrum = 1 / inverse_spectrum - return self.music_spectrum - - def update_number_of_sensors(self, number_of_sensors: int): - self.system_model.create_array(number_of_sensors) - self.set_search_grid() - - def adjust_cell_size(self): - if self.estimation_params == "range": - if self.cell_size > 1: - self.cell_size = int(0.8 * self.cell_size) - if self.cell_size % 2 == 0: - self.cell_size -= 1 - elif self.estimation_params == "angle, range": - if self.cell_size_angle > 1: - self.cell_size_angle = int(0.95 * self.cell_size_angle) - if self.cell_size_angle % 2 == 0: - self.cell_size_angle -= 1 - if self.cell_size_range > 1: - self.cell_size_range = int(0.95 * self.cell_size_range) - if self.cell_size_range % 2 == 0: - self.cell_size_range -= 1 - elif self.estimation_params == "angle": - if self.cell_size > 1: - self.cell_size = int(0.95 * self.cell_size) - if self.cell_size % 2 == 0: - self.cell_size -= 1 - - def get_inverse_spectrum(self, noise_subspace: torch.Tensor): - """ - - Parameters - ---------- - noise_subspace - the noise related subspace vectors of size BatchSizex#SENSORSx(#SENSORS-#SOURCES) - - Returns - ------- - in all cases it will return the inverse spectrum, - in case of single param estimation it will be 1D inverse spectrum: BatchSizex(length_search_grid) - in case of dual param estimation it will be 2D inverse spectrum: - BatchSizex(length_search_grid_angle)x(length_search_grid_distance) - """ - # steering_dict = self.steering_dict.to(device) - if self.system_model.params.field_type.startswith("far"): - steering_dict = self.steering_dict[:noise_subspace.shape[1]].to(self.device) - var1 = torch.einsum("an, bnm -> bam", steering_dict.conj().transpose(0, 1)[:, :noise_subspace.shape[1]], - noise_subspace) - inverse_spectrum = torch.norm(var1, dim=2) ** 2 - else: - if self.estimation_params.startswith("angle, range"): - steering_dict = self.steering_dict[:noise_subspace.shape[1]].conj().transpose(0, 2).transpose(0, 1).to(self.device) - try: - var1 = torch.einsum("adk, bkl -> badl", - steering_dict, - noise_subspace) - # get the norm value for each element in the batch. - inverse_spectrum = torch.norm(var1, dim=-1) ** 2 - except RuntimeError: - warnings.warn("MUSIC.get_inverse_spectrum: Out of memory error, trying to free some memory and convert the batch operation to for loop.") - torch.cuda.empty_cache() - inverse_spectrum = torch.zeros((noise_subspace.shape[0], self.angles_dict.shape[0], self.ranges_dict.shape[0]), dtype=torch.float64, device=self.device) - for batch in range(noise_subspace.shape[0]): - var1 = torch.einsum("adk, kl -> adl", - steering_dict, - noise_subspace[batch]) - inverse_spectrum[batch] = torch.norm(var1, dim=-1) ** 2 - - del var1 - - elif self.estimation_params.endswith("angle"): - steering_dict = self.steering_dict[:noise_subspace.shape[1]].to(self.device) - var1 = torch.einsum("an, nbm -> abm", steering_dict.conj().transpose(0, 1), - noise_subspace.transpose(0, 1)) - inverse_spectrum = torch.norm(var1, dim=-1).T ** 2 - elif self.estimation_params.startswith("range"): - steering_dict = self.steering_dict[:noise_subspace.shape[1]].to(self.device) - var1 = torch.bmm(steering_dict.conj().transpose(0, 2).transpose(0, 1), noise_subspace) - inverse_spectrum = torch.norm(var1, dim=-1) ** 2 - if torch.isnan(inverse_spectrum).any(): - raise ValueError("Nan values in inverse spectrum") - else: - raise ValueError(f"MUSIC.get_inverse_spectrum: unknown estimation param {self.estimation_params}") - del steering_dict - try: - torch.cuda.empty_cache() - except AttributeError: - pass - return inverse_spectrum - - def peak_finder(self, source_number: int): - """ - - Parameters - ---------- - is_soft: this boolean paramter will determine wether to use derivative approxamtion of the peak_finder for - the training stage. - - Returns - ------- - the predicted param(torch.Tensor) or params(tuple) - """ - if self.system_model.params.field_type.lower().startswith("far"): - return self._peak_finder_1d(self.angles_dict, source_number) - else: - if self.estimation_params.startswith("angle, range"): - return self._peak_finder_2d(source_number) - elif self.estimation_params.endswith("angle"): - return self._peak_finder_1d(self.angles_dict, source_number) - elif self.estimation_params.startswith("range"): - return self._peak_finder_1d(self.ranges_dict, source_number) - - def set_search_grid(self, known_angles: torch.Tensor = None, known_distances: torch.Tensor = None): - if self.system_model.params.field_type.startswith("far"): - self.__set_search_grid_far_field() - elif self.system_model.params.field_type in ["near", "full"]: - self.__set_search_grid_near_field(known_angles=known_angles, known_distances=known_distances) - else: - raise ValueError(f"MUSIC.set_search_grid: Unrecognized field type: {self.system_model.params.field_type}") - - def plot_spectrum(self, highlight_corrdinates=None, batch: int = 0, method: str = "heatmap", music_spectrum = None, add_title: bool = False, save: bool = False): - if self.estimation_params == "angle, range": - self._plot_3d_spectrum(highlight_corrdinates, batch, method, music_spectrum=music_spectrum, add_title=add_title, save=save) - else: - self._plot_1d_spectrum(highlight_corrdinates, batch, add_title=add_title, save=save) - - def test_step(self, batch, batch_idx, model: nn.Module=None): - x, sources_num, label = batch - if x.dim() == 2: - x = x.unsqueeze(0) - test_length = x.shape[0] - x = x.to(self.device) - if self.estimation_params == "angle, range": - angles, ranges = torch.split(label, max(sources_num), dim=1) - angles = angles.to(self.device) - ranges = ranges.to(self.device) - else: - angles = label.to(self.device) # only angles - # Check if the sources number is the same for all samples in the batch - if (sources_num != sources_num[0]).any(): - # in this case, the sources number is not the same for all samples in the batch - raise Exception(f"train_model:" - f" The sources number is not the same for all samples in the batch.") - else: - sources_num = sources_num[0] - if model is not None: - try: - Rx = model.get_surrogate_covariance(x) - except NotImplementedError as e: - raise e - else: - if self.system_model.params.signal_nature == "non-coherent": - Rx = self.pre_processing(x, mode="sample") - else: - # Rx = self.pre_processing(x, mode="sps") - Rx = self.pre_processing(x, mode="sample") - predictions, sources_num_estimation, _ = self(Rx, number_of_sources=sources_num) - if self.estimation_params == "angle, range": - angles_prediction, ranges_prediction = predictions - rmspe = self.criterion(angles_prediction, angles, ranges_prediction, ranges).sum(-1).item() - _, rmspe_angle, rmspe_range = self.separated_criterion(angles_prediction, angles, ranges_prediction, ranges) - rmspe = (rmspe, rmspe_angle.sum(-1).item(), rmspe_range.sum(-1).item()) - else: - rmspe = self.criterion(predictions, angles).sum().item() - - acc = self.source_estimation_accuracy(sources_num, sources_num_estimation) - - return rmspe, acc, test_length - - def _peak_finder_1d(self, search_space, source_number: int): - if self.estimation_params == "range": - source_number = 1 # for the range estimation, only one source is expected. - - batch_size = self.music_spectrum.shape[0] - - peaks = torch.zeros(batch_size, source_number, dtype=torch.int64, device=self.device) - for batch in range(batch_size): - music_spectrum = self.music_spectrum[batch].cpu().detach().numpy().squeeze() - # Find spectrum peaks - peaks_tmp = sc.signal.find_peaks(music_spectrum, threshold=0.0)[0] - if len(peaks_tmp) < source_number: - warnings.warn(f"MUSIC._peak_finder_1d: No peaks were found! taking max values instead.") - # random_peaks = np.random.randint(0, search_space.shape[0], (source_number - peaks_tmp.shape[0],)) - random_peaks = torch.topk(torch.from_numpy(music_spectrum), source_number - peaks_tmp.shape[0], - largest=True).indices.cpu().detach().numpy() - peaks_tmp = np.concatenate((peaks_tmp, random_peaks)) - # Sort the peak by their amplitude - sorted_peaks = peaks_tmp[np.argsort(music_spectrum[peaks_tmp])[::-1]] - peaks[batch] = torch.from_numpy(sorted_peaks[0:source_number]).to(self.device) - if not self.training: - # if the model is not in training mode, return the peaks - if peaks.dim() == 1: - return search_space[peaks] - else: - labels = torch.gather(search_space.unsqueeze(1).repeat(1, source_number).to(self.device), 0, peaks) - return labels - else: - return self.__maskpeak_1d(peaks, search_space, source_number) - - def _peak_finder_2d(self, source_number: int): - batch_size = self.music_spectrum.shape[0] - - max_row = torch.zeros((batch_size, source_number) - , dtype=torch.int64, device=self.device) - max_col = torch.zeros((batch_size, source_number) - , dtype=torch.int64, device=self.device) - for batch in range(batch_size): - music_spectrum = self.music_spectrum[batch].detach().cpu().numpy().squeeze() - peaks = find_k_highest_peaks(music_spectrum, source_number) - original_idx = torch.from_numpy(np.array(peaks)[:, :2]).T - max_row[batch] = original_idx[0][0: source_number] - max_col[batch] = original_idx[1][0: source_number] - if not self.training: - # if the model is not in training mode, return the peaks. - angle_dict = self.angles_dict.to(self.device) - range_dict = self.ranges_dict.to(self.device) - angles_pred = angle_dict[max_row] - distances_pred = range_dict[max_col] - del angle_dict, range_dict - try: - torch.cuda.empty_cache() - except AttributeError: - pass - return angles_pred, distances_pred - else: - return self.__maskpeak_2d(max_row, max_col, source_number) - - def __maskpeak_1d(self, peaks, search_space, source_number: int = None): - - batch_size = self.music_spectrum.shape[0] - soft_decision = torch.zeros(batch_size, source_number, dtype=torch.float64, device=self.device) - top_indxs = peaks.to(self.device) - - for source in range(source_number): - cell_idx = (top_indxs[:, source][:, None] - - self.cell_size - + torch.arange(2 * self.cell_size + 1, dtype=torch.long, device=self.device)) - # cell_idx %= self.music_spectrum.shape[1] - out_of_bounds_mask = (cell_idx < 0) | (cell_idx >= self.music_spectrum.shape[1]) - cell_idx[out_of_bounds_mask] = top_indxs[:, source].unsqueeze(1).expand_as(cell_idx)[out_of_bounds_mask] - cell_idx = cell_idx.reshape(batch_size, -1, 1) - metrix_thr = torch.gather(self.music_spectrum.unsqueeze(-1).expand(-1, -1, cell_idx.size(-1)), 1, - cell_idx).requires_grad_(True) - soft_max = torch.softmax(metrix_thr, dim=1) - soft_decision[:, source][:, None] = torch.einsum("bms, bms -> bs", search_space[cell_idx.cpu()].to(self.device), soft_max).to( - self.device) - - return soft_decision - - def __maskpeak_2d(self, peaks_r, peaks_c, source_number): - batch_size = self.music_spectrum.shape[0] - soft_row = torch.zeros((batch_size, source_number), device=self.device) - soft_col = torch.zeros((batch_size, source_number), device=self.device) - - for source in range(source_number): - max_row_cell_idx = (peaks_r[:, source][:, None] - - self.cell_size_angle - + torch.arange(2 * self.cell_size_angle + 1, dtype=torch.int32, device=self.device)) - max_row_cell_idx %= self.music_spectrum.shape[1] - max_row_cell_idx = max_row_cell_idx.reshape(batch_size, -1, 1) - - max_col_cell_idx = (peaks_c[:, source][:, None] - - self.cell_size_range - + torch.arange(2 * self.cell_size_range + 1, dtype=torch.int32, device=self.device)) - max_col_cell_idx %= self.music_spectrum.shape[2] - max_col_cell_idx = max_col_cell_idx.reshape(batch_size, 1, -1) - - metrix_thr = self.music_spectrum.gather(1, - max_row_cell_idx.expand(-1, -1, self.music_spectrum.shape[2])) - metrix_thr = metrix_thr.gather(2, max_col_cell_idx.repeat(1, max_row_cell_idx.shape[-2], 1)) - soft_max = torch.softmax(metrix_thr.view(batch_size, -1), dim=1).reshape(metrix_thr.shape) - soft_row[:, source][:, None] = torch.einsum("bla, bad -> bl", - self.angles_dict[max_row_cell_idx].transpose(1, 2), - torch.sum(soft_max, dim=2).unsqueeze(-1)) - soft_col[:, source][:, None] = torch.einsum("bmc, bcm -> bm", - self.ranges_dict[max_col_cell_idx], - torch.sum(soft_max, dim=1).unsqueeze(-1)) - - return soft_row, soft_col - - def _init_spectrum(self, batch_size): - if self.system_model.params.field_type == "Far": - self.music_spectrum = torch.zeros(batch_size, len(self.angles_dict)) - else: - if self.estimation_params.startswith("angle, range"): - self.music_spectrum = torch.zeros(batch_size, len(self.angles_dict), len(self.ranges_dict)) - elif self.estimation_params.endswith("angle"): - self.music_spectrum = torch.zeros(batch_size, len(self.angles_dict)) - elif self.estimation_params.startswith("range"): - self.music_spectrum = torch.zeros(batch_size, len(self.ranges_dict)) - - def __init_grid_params(self): - angle_range = np.deg2rad(self.system_model.params.doa_range) - angle_resolution = np.deg2rad(self.system_model.params.doa_resolution / 2) - angle_decimals = int(np.ceil(np.log10(1 / angle_resolution))) - - if self.system_model.params.field_type.startswith("far"): - # if it's the Far field case, need to init angles range. - self.angles_dict = torch.arange(-angle_range, angle_range + angle_resolution, angle_resolution, - dtype=torch.float64).to(torch.float64) - self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) - elif self.system_model.params.field_type in ["near", "full"]: - # if it's the Near field, there are 3 possabilities. - fresnel = self.system_model.fresnel - fraunhofer = self.system_model.fraunhofer - if self.estimation_params.startswith("angle"): - self.angles_dict = torch.arange(-angle_range, angle_range + angle_resolution, angle_resolution, - dtype=torch.float64).to(torch.float64) - # self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) - - - if self.estimation_params.endswith("range"): - fraunhofer_ratio = self.system_model.params.max_range_ratio_to_limit - distance_resolution = self.system_model.params.range_resolution / 2 - max_distance = min(self.system_model.fraunhofer, fraunhofer * fraunhofer_ratio + distance_resolution) - self.ranges_dict = torch.arange(np.ceil(fresnel), - max_distance, - distance_resolution, dtype=torch.float64) - else: - raise ValueError(f"MUSIC.__define_grid_params: Unrecognized field type for MUSIC class init stage," - f" got {self.system_model.params.field_type} but only Far and Near are allowed.") - - def __init_search_grid(self): - # if this is the music 2D case, the search grid is constant and can be calculated once. - if self.system_model.params.field_type in ["near", "full"]: - if self.angles_dict is not None and self.ranges_dict is not None: - self.set_search_grid() - elif self.angles_dict is not None: # Near field case with Far field inference - self.__set_search_grid_far_field() - else: - self.set_search_grid() - - def __init_cells(self, coeff: float = 0.1): - - if self.estimation_params == "range": - self.cell_size = int(self.ranges_dict.shape[0] * coeff) - elif self.estimation_params == "angle": - self.cell_size = int(self.angles_dict.shape[0] * coeff) - elif self.estimation_params == "angle, range": - self.cell_size_angle = int(self.angles_dict.shape[0] * coeff) - self.cell_size_range = int(self.ranges_dict.shape[0] * coeff) - - if self.cell_size is not None: - if self.cell_size % 2 == 0: - self.cell_size += 1 - if self.cell_size_angle is not None: - if self.cell_size_angle % 2 == 0: - self.cell_size_angle += 1 - if self.cell_size_range is not None: - if self.cell_size_range % 2 == 0: - self.cell_size_range += 1 - - def init_cells(self, coeff: float = 0.2): - self.__init_cells(coeff) - - def _plot_1d_spectrum(self, highlight_corrdinates, batch, add_title: bool = False, save: bool = False): - if self.estimation_params == "angle": - x = np.rad2deg(self.angles_dict.detach().cpu().numpy()) - x_label = "angle [deg]" - elif self.estimation_params == "range": - x = self.ranges_dict.detach().cpu().numpy() - x_label = "distance [m]" - else: - raise ValueError(f"MUSIC._plot_1d_spectrum: No such option for param estimation.") - y = self.music_spectrum[batch].detach().cpu().numpy() - plt.figure() - plt.plot(x, y.T, label="Music Spectrum") - if highlight_corrdinates is not None: - for idx, dot in enumerate(highlight_corrdinates): - plt.vlines(dot, np.min(y), np.max(y), colors='r', linestyles='dashed', label=f"Ground Truth") - if add_title: - plt.title("MUSIC SPECTRUM") - plt.grid() - plt.ylabel("Spectrum power") - plt.xlabel(x_label) - plt.legend() - plt.tight_layout() - if save: - plt.savefig("1d_music_spectrum.pdf") - plt.show() - - def _plot_3d_spectrum(self, highlight_coordinates, batch, method, music_spectrum=None, add_title: bool = False, save: bool = False): - """ - Plot the MUSIC 2D spectrum. - - """ - if method == "3D": - # Creating figure - distances = self.ranges_dict.detach().cpu().numpy() - angles = self.angles_dict.detach().cpu().numpy() - if music_spectrum is None: - spectrum = self.music_spectrum[batch].detach().cpu().numpy() - else: - spectrum = music_spectrum[batch].detach().cpu().numpy() - x, y = np.meshgrid(distances, np.rad2deg(angles)) - # Plotting the 3D surface - fig = plt.figure() - ax = fig.add_subplot(111, projection='3d') - ax.plot_surface(x, y, 10 * np.log10(spectrum), cmap='viridis') - - if highlight_coordinates: - highlight_coordinates = np.array(highlight_coordinates) - ax.scatter( - highlight_coordinates[:, 0], - np.rad2deg(highlight_coordinates[:, 1]), - np.log1p(highlight_coordinates[:, 2]), - color='red', - s=50, - label='Ground Truth', - marker="x" - ) - if add_title: - ax.set_title('MUSIC spectrum') - ax.set_xlim(distances[0], distances[-1]) - ax.set_ylim(np.rad2deg(angles[0]), np.rad2deg(angles[-1])) - # Adding labels - ax.set_ylabel('Theta [deg]') - ax.set_xlabel('Radius [m]') - ax.set_zlabel('Power [dB]') - plt.colorbar(ax.plot_surface(x, y, 10 * np.log10(spectrum), cmap='viridis'), shrink=0.5, aspect=5) - - if highlight_coordinates: - ax.legend() # Adding a legend - - # Display the plot - plt.tight_layout() - if save: - plt.savefig("3d_music_spectrum.pdf") - plt.show() - elif method == "heatmap": - xmin, xmax = np.min(self.ranges_dict.cpu().detach().numpy()), np.max(self.ranges_dict.cpu().detach().numpy()) - ymin, ymax = np.min(self.angles_dict.cpu().detach().numpy()), np.max(self.angles_dict.cpu().detach().numpy()) - if music_spectrum is None: - spectrum = self.music_spectrum[batch].cpu().detach().numpy() - else: - spectrum = music_spectrum[batch].cpu().detach().numpy() - plt.imshow(spectrum, cmap="hot", - extent=[xmin, xmax, np.rad2deg(ymin), np.rad2deg(ymax)], origin='lower', aspect="auto") - if highlight_coordinates is not None: - for idx, dot in enumerate(highlight_coordinates): - x = self.ranges_dict.cpu().detach().numpy()[dot[1]] - y = np.rad2deg(self.angles_dict.cpu().detach().numpy()[dot[0]]) - plt.plot(x, y, label=f"{x:.1f} [m], {y:.1f} [deg]", marker='o', markerfacecolor='none', - markeredgecolor='white', linestyle='-', color='white', markersize=10) - # plt.plot(x, y, marker='x', linestyle='', color='green', markersize=8) - plt.legend() - plt.colorbar() - if add_title: - plt.title("MUSIC Spectrum heatmap") - plt.xlabel("Distances [m]") - plt.ylabel("Angles [deg]") - - plt.figaspect(2) - plt.tight_layout() - if save: - plt.savefig("heatmap_music_spectrum.pdf") - plt.show() - elif method == "slice": - x = self.ranges_dict.detach().cpu().numpy() - x_label = "distance [m]" - y = self.music_spectrum[batch].detach().cpu().numpy()[highlight_coordinates[0]] - plt.figure() - plt.plot(x, y.T, label="Music Spectrum") - if highlight_coordinates is not None: - for idx, dot in enumerate(highlight_coordinates[1:]): - plt.vlines(dot, np.min(y), np.max(y), colors='r', linestyles='dashed', label=f"Ground Truth") - if add_title: - plt.title(f"MUSIC SPECTRUM Slice at {torch.round(torch.rad2deg(self.angles_dict[highlight_coordinates[0]]))}") - plt.grid() - plt.ylabel("Spectrum power") - plt.xlabel(x_label) - plt.legend() - plt.tight_layout() - if save: - plt.savefig("slice_music_spectrum.pdf") - plt.show() - - def __set_search_grid_far_field(self): - self.steering_dict = self.system_model.steering_vec_far_field(self.angles_dict, f_c=None, nominal=True, fix_sv_noise=True).squeeze(-1) - - def __set_search_grid_near_field(self, known_angles: torch.Tensor = None, known_distances: torch.Tensor = None): - """ - - Returns: - - """ - if known_angles is None: - known_angles = self.angles_dict - if known_distances is None: - known_distances = self.ranges_dict - self.steering_dict = self.system_model.steering_vec_near_field(angles=known_angles, ranges=known_distances, - generate_search_grid=True, nominal=True, - f_c=None).squeeze(-1).cpu() - if torch.isnan(self.steering_dict).any(): - raise ValueError("Nan values in steering matrix") - - def __str__(self): - if self.estimation_params == "angle": - return "music_angle" - elif self.estimation_params == "range": - return "music_range" - elif self.estimation_params == "angle, range": - return "2d_music" - - def __init_criteria(self): - if self.estimation_params == "angle": - self.criterion = RMSPELoss(balance_factor=1.0) - elif self.estimation_params == "range": - self.criterion = RMSPELoss(balance_factor=0.0) - elif self.estimation_params == "angle, range": - self.criterion = CartesianLoss() - self.separated_criterion = RMSPELoss(1.0) - else: - raise ValueError(f"MUSIC.__init_criteria: Unrecognized estimation param {self.estimation_params}") - - def _get_name(self): - return "MUSIC" - - -# Import the diffMUSIC implementation -from src.methods_pack.diffmusic import DiffMUSIC - - -class Filter(nn.Module): - def __init__(self, min_cell_size, max_cell_size, number_of_filter=10): - super(Filter, self).__init__() - self.number_of_filters = number_of_filter - self.cell_sizes = torch.linspace(min_cell_size, max_cell_size, number_of_filter).to(torch.int32).to(self.device) - self.cell_bank = {} - for cell_size in enumerate(self.cell_sizes.data): - cell_size = cell_size[1] - self.cell_bank[cell_size] = torch.arange(-cell_size, cell_size, 1, dtype=torch.long, device=self.device) - self.fc = nn.Linear(self.number_of_filters, 1) - self.fc.weight.data = torch.randn(1, number_of_filter) / 100 + (1 / number_of_filter) - self.fc.weight.data = self.fc.weight.data.to(torch.float64) - self.fc.bias.data = torch.Tensor([0]) - self.fc.bias.data = self.fc.bias.data.to(torch.float64) - self.fc.bias.requires_grad_(False) - self.relu = nn.ReLU() - - def forward(self, input, search_space): - peaks = torch.zeros(input.shape[0], 1).to(torch.int64) - for batch in range(peaks.shape[0]): - music_spectrum = input[batch].cpu().detach().numpy().squeeze() - # Find spectrum peaks - peaks_tmp = list(sc.signal.find_peaks(music_spectrum)[0]) - # Sort the peak by their amplitude - peaks_tmp.sort(key=lambda x: music_spectrum[x], reverse=True) - if len(peaks_tmp) == 0: - peaks_tmp = torch.randint(search_space.shape[0], (1,)) - else: - peaks_tmp = peaks_tmp[0] - peaks[batch] = peaks_tmp - top_1 = peaks - output = torch.zeros(input.shape[0], self.number_of_filters).to(self.device).to(torch.float64) - for idx, cell in enumerate(self.cell_bank.values()): - tmp_cell = top_1 + cell - tmp_cell %= input.shape[1] - tmp_cell = tmp_cell.unsqueeze(-1) - metrix_thr = torch.gather(input.unsqueeze(-1).expand(-1, -1, tmp_cell.size(-1)), 1, tmp_cell) - soft_max = torch.softmax(metrix_thr, dim=1) - output[:, idx] = torch.einsum("bkm, bkm -> bm", search_space[tmp_cell], soft_max).squeeze() - output = self.fc(output) - output = self.relu(output) - self.clip_weights_values() - return output - - def clip_weights_values(self): - self.fc.weight.data = torch.clip(self.fc.weight.data, 0.1, 1) - self.fc.weight.data /= torch.sum(self.fc.weight.data) \ No newline at end of file diff --git a/src/old/diffMUSIC.py b/src/old/diffMUSIC.py deleted file mode 100644 index 496e0df..0000000 --- a/src/old/diffMUSIC.py +++ /dev/null @@ -1,494 +0,0 @@ -import warnings -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -import matplotlib.pyplot as plt -import scipy as sc - -from src.signal_creation import SystemModel, SystemModelParams -from src.utils import set_unified_seed -from scipy.ndimage import maximum_filter -from scipy.ndimage import label -from scipy.ndimage import find_objects - -#TODO: Haven't checked this method yet. -def find_k_highest_peaks(matrix, k): - """ - Find the k highest peaks in a 2D matrix using SciPy tools. A peak is defined as a - local maximum surrounded by smaller values. - """ - # Apply maximum filter to find local maxima - neighborhood = maximum_filter(matrix, size=21, mode='constant', cval=-np.inf) - local_max = (matrix == neighborhood) - - # Label the connected components of local maxima - labeled, num_features = label(local_max) - slices = find_objects(labeled) - - # Extract peak positions and values - peaks = [] - try: - for sl in slices: - row = int((sl[0].start + sl[0].stop - 1) / 2) - col = int((sl[1].start + sl[1].stop - 1) / 2) - value = matrix[row, col] - peaks.append((row, col, value)) - except Exception as e: - pass - - # Sort peaks by value (descending) and select the top k - peaks = sorted(peaks, key=lambda x: x[2], reverse=True)[:k] - if len(peaks) < k: - warnings.warn(f"find_k_highest_peaks: Less than {k} peaks found.") - # add random peaks - x_random = np.random.randint(0, matrix.shape[0], (k - len(peaks),)) - y_random = np.random.randint(0, matrix.shape[1], (k - len(peaks),)) - for i in range(k - len(peaks)): - peaks.append((x_random[i], y_random[i], matrix[x_random[i], y_random[i]])) - - return peaks - - -class DiffMUSIC(nn.Module): - """ - Differentiable MUSIC implementation with learnable antenna gains and positions. - Integrates with the existing SystemModel architecture for far-field, non-coherent, narrowband scenarios. - """ - - def __init__(self, - system_model_params: SystemModelParams, - init_antenna_positions: torch.Tensor = None, - init_antenna_gains: torch.Tensor = None, - gain_constraint: str = "positive", # "positive", "normalized", "none" - temperature: float = 1.0, - cell_size_coeff: float = 0.2 # between 0 and 1, determines the size of tthe window for softmax peak - # finding, the size is computed as len(grid) * cell_size_coeff - ): - """ - Initialize DiffMUSIC with learnable antenna parameters. - - Args: - system_model_params: SystemModelParams instance - init_antenna_positions: Initial antenna positions [N] for ULA case - init_antenna_gains: Initial antenna gains [N] - gain_constraint: Type of gain constraint - temperature: Temperature for soft peak finding - cell_size_coeff: Coefficient for cell size in soft peak finding - """ - super().__init__(system_model_params, - init_antenna_positions, - init_antenna_gains, - gain_constraint, - temperature, - cell_size_coeff) - - - def sample_covariance(self, x: torch.Tensor) -> torch.Tensor: - """ - Compute sample covariance matrix (following your pattern). - - Args: - x: Input samples [batch_size, N, T] or [N, T] - - Returns: - Covariance matrices [batch_size, N, N] - """ - if x.dim() == 2: - x = x.unsqueeze(0) - 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 subspace_separation(self, cov: torch.Tensor, number_of_sources: int): - """ - Perform eigendecomposition and separate signal/noise subspaces. - - Args: - cov: Covariance matrix [batch_size, N, N] - number_of_sources: Number of sources - - Returns: - signal_subspace, noise_subspace, source_estimation, eigen_regularization - """ - # Eigendecomposition - eigenvalues, eigenvectors = torch.linalg.eigh(cov) - - # Sort in descending order - sorted_indices = torch.argsort(eigenvalues, dim=-1, descending=True) - eigenvalues = torch.gather(eigenvalues, -1, sorted_indices) - eigenvectors = torch.gather(eigenvectors, -1, sorted_indices.unsqueeze(-2).expand_as(eigenvectors)) - - # Separate subspaces - signal_subspace = eigenvectors[:, :, :number_of_sources] - noise_subspace = eigenvectors[:, :, number_of_sources:] - - # Source estimation (simplified) - source_estimation = number_of_sources - - # Eigen regularization - eigen_regularization = torch.mean(eigenvalues[:, number_of_sources:]) - - return signal_subspace, noise_subspace, source_estimation, eigen_regularization - - def get_inverse_spectrum(self, noise_subspace: torch.Tensor): - """ - Compute inverse MUSIC spectrum using current antenna configuration. - - Args: - noise_subspace: Noise subspace [batch_size, N, N-M] - - Returns: - Inverse spectrum [batch_size, num_angles] - """ - # Compute steering matrix for all angles - steering_dict = self.compute_steering_matrix(self.angles_dict) - steering_dict = steering_dict[:noise_subspace.shape[1]].to(self.device) - - # Compute projection onto noise subspace - var1 = torch.einsum("an, bnm -> bam", - steering_dict.conj().transpose(0, 1)[:, :noise_subspace.shape[1]], - noise_subspace) - inverse_spectrum = torch.norm(var1, dim=2) ** 2 - - return inverse_spectrum - - def soft_peak_finding_1d(self, spectrum: torch.Tensor, search_space: torch.Tensor, num_sources: int): - """ - Differentiable 1D peak finding using temperature-scaled softmax. - Adapted from your __maskpeak_1d method. - - Args: - spectrum: MUSIC spectrum [batch_size, num_points] - search_space: Search grid values [num_points] - num_sources: Number of sources to find - - Returns: - Estimated parameters [batch_size, num_sources] - """ - batch_size = spectrum.shape[0] - - # Find hard peaks for initialization - peaks = torch.zeros(batch_size, num_sources, dtype=torch.int64, device=self.device) - for batch in range(batch_size): - music_spectrum_np = spectrum[batch].cpu().detach().numpy().squeeze() - # Find spectrum peaks - peaks_tmp = sc.signal.find_peaks(music_spectrum_np, threshold=0.0)[0] - if len(peaks_tmp) < num_sources: - # Take top values if not enough peaks - random_peaks = torch.topk(torch.from_numpy(music_spectrum_np), - num_sources - peaks_tmp.shape[0], largest=True).indices.cpu().detach().numpy() - peaks_tmp = np.concatenate((peaks_tmp, random_peaks)) - # Sort by amplitude - sorted_peaks = peaks_tmp[np.argsort(music_spectrum_np[peaks_tmp])[::-1]] - peaks[batch] = torch.from_numpy(sorted_peaks[0:num_sources]).to(self.device) - - # Soft peak finding (differentiable) - soft_decision = torch.zeros(batch_size, num_sources, dtype=torch.float64, device=self.device) - top_indxs = peaks.to(self.device) - - for source in range(num_sources): - # Create cell around each peak - cell_idx = (top_indxs[:, source][:, None] - - self.cell_size - + torch.arange(2 * self.cell_size + 1, dtype=torch.long, device=self.device)) - - # Handle boundaries - out_of_bounds_mask = (cell_idx < 0) | (cell_idx >= spectrum.shape[1]) - cell_idx[out_of_bounds_mask] = top_indxs[:, source].unsqueeze(1).expand_as(cell_idx)[out_of_bounds_mask] - cell_idx = cell_idx.reshape(batch_size, -1, 1) - - # Extract spectrum values in cell - metrix_thr = torch.gather(spectrum.unsqueeze(-1).expand(-1, -1, cell_idx.size(-1)), 1, - cell_idx).requires_grad_(True) - - # Apply temperature-scaled softmax - soft_max = torch.softmax(metrix_thr / self.temperature, dim=1) - - # Compute weighted average - soft_decision[:, source] = torch.einsum("bms, bms -> bs", - search_space[cell_idx.cpu()].to(self.device), - soft_max).squeeze() - - return soft_decision - - def forward(self, x: torch.Tensor, number_of_sources: int = None): - """ - Forward pass of DiffMUSIC. - - Args: - x: Input samples [batch_size, N, T] or [N, T] - number_of_sources: Number of sources (defaults to self.M) - - Returns: - Estimated DOA angles, source estimation, eigen regularization - """ - if number_of_sources is None: - number_of_sources = self.M - - # Apply constraints (for monitoring during training) - if self.training: - self._apply_position_constraints() - self._apply_gain_constraints() - - # Compute covariance matrix - cov = self.sample_covariance(x) - - # Subspace separation - _, noise_subspace, source_estimation, eigen_regularization = self.subspace_separation(cov, number_of_sources) - - # Compute inverse spectrum - inverse_spectrum = self.get_inverse_spectrum(noise_subspace) - - # Compute MUSIC spectrum - self.music_spectrum = 1 / (inverse_spectrum + 1e-10) - - # Differentiable peak finding - params = self.soft_peak_finding_1d(self.music_spectrum, self.angles_dict, number_of_sources) - - return params, source_estimation, eigen_regularization - - def plot_spectrum(self, batch: int = 0, true_angles: torch.Tensor = None, save: bool = False): - """Plot MUSIC spectrum""" - if self.music_spectrum is None: - print("No spectrum available. Run forward pass first.") - return - - x = np.rad2deg(self.angles_dict.detach().cpu().numpy()) - y = self.music_spectrum[batch].detach().cpu().numpy() - - plt.figure(figsize=(10, 6)) - plt.plot(x, y, label="DiffMUSIC Spectrum", linewidth=2) - - if true_angles is not None: - true_angles_deg = np.rad2deg(true_angles.detach().cpu().numpy()) - for i, angle in enumerate(true_angles_deg): - plt.axvline(angle, color='red', linestyle='--', alpha=0.7, - label=f'True DOA {i+1}' if i == 0 else "") - - plt.xlabel('Angle [degrees]') - plt.ylabel('Spectrum Power') - plt.title('DiffMUSIC Spectrum') - plt.grid(True, alpha=0.3) - plt.legend() - plt.tight_layout() - - if save: - plt.savefig('diffmusic_spectrum.pdf') - plt.show() - - def plot_array_geometry(self, save: bool = False): - """Plot current antenna array geometry""" - positions, gains = self.get_constrained_parameters() - positions_np = positions.detach().cpu().numpy() - gains_np = torch.abs(gains).detach().cpu().numpy() - - plt.figure(figsize=(12, 4)) - - # Plot antenna positions - plt.scatter(positions_np, np.zeros_like(positions_np), - c=gains_np, s=100, cmap='viridis', - edgecolors='black', linewidth=1, marker='s') - plt.colorbar(label='Antenna Gain Magnitude') - - # Add antenna numbers - for i, pos in enumerate(positions_np): - plt.annotate(f'{i}', (pos, 0.01), - xytext=(0, 10), textcoords='offset points', - ha='center', va='bottom') - - plt.xlabel('Position [meters]') - plt.ylabel('') - plt.title('DiffMUSIC Antenna Array Geometry') - plt.grid(True, alpha=0.3) - plt.ylim(-0.1, 0.1) - - # Add wavelength reference - plt.axhline(0, color='black', linewidth=0.5) - spacing_text = f'λ/2 = {self.params.wavelength/2:.3f}m' - plt.text(0.02, 0.98, spacing_text, transform=plt.gca().transAxes, - bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8)) - - plt.tight_layout() - - if save: - plt.savefig('diffmusic_array_geometry.pdf') - plt.show() - - def get_antenna_info(self): - """Get current antenna configuration""" - positions, gains = self.get_constrained_parameters() - return { - 'positions': positions.detach().cpu().numpy(), - 'gains': gains.detach().cpu().numpy(), - 'position_constraint': self.position_constraint, - 'gain_constraint': self.gain_constraint, - 'N': self.N, - 'wavelength': self.params.wavelength - } - - -class DiffMUSICTrainer: - """Training wrapper for DiffMUSIC integrated with your simulation framework""" - - def __init__(self, diffmusic_model: DiffMUSIC, training_params: dict): - self.model = diffmusic_model - self.training_params = training_params - - # Initialize optimizer - if training_params["optimizer"] == "Adam": - self.optimizer = torch.optim.Adam( - self.model.parameters(), - lr=training_params["learning_rate"], - weight_decay=training_params.get("weight_decay", 0) - ) - elif training_params["optimizer"] == "SGD": - self.optimizer = torch.optim.SGD( - self.model.parameters(), - lr=training_params["learning_rate"], - weight_decay=training_params.get("weight_decay", 0) - ) - - # Initialize scheduler - if training_params["scheduler"] == "ReduceLROnPlateau": - self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( - self.optimizer, mode='min', factor=0.5, patience=10 - ) - elif training_params["scheduler"] == "StepLR": - self.scheduler = torch.optim.lr_scheduler.StepLR( - self.optimizer, - step_size=training_params.get("step_size", 50), - gamma=0.5 - ) - - def train_step(self, samples_batch: torch.Tensor, targets_batch: torch.Tensor, num_sources: int = None): - """Single training step""" - self.model.train() - self.optimizer.zero_grad() - - # Forward pass - predictions, _, eigen_reg = self.model(samples_batch, num_sources) - - # Compute RMSE loss - loss = torch.sqrt(torch.mean((predictions - targets_batch) ** 2)) - - # Add regularization terms - reg_loss = 0.01 * eigen_reg # Eigenvalue regularization - - # Add antenna position regularization (encourage smooth spacing) - positions, _ = self.model.get_constrained_parameters() - if len(positions) > 1: - pos_diff = positions[1:] - positions[:-1] - target_spacing = self.model.params.wavelength / 2 - spacing_reg = 0.001 * torch.mean((pos_diff - target_spacing) ** 2) - else: - spacing_reg = 0.0 - - total_loss = loss + reg_loss + spacing_reg - - # Backward pass - total_loss.backward() - - # Gradient clipping - torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) - - self.optimizer.step() - - return total_loss.item(), loss.item(), reg_loss.item(), spacing_reg - - def validate(self, val_samples: torch.Tensor, val_targets: torch.Tensor, num_sources: int = None): - """Validation step""" - self.model.eval() - - with torch.no_grad(): - predictions, _, _ = self.model(val_samples, num_sources) - loss = torch.sqrt(torch.mean((predictions - val_targets) ** 2)) - - if hasattr(self.scheduler, 'step'): - if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): - self.scheduler.step(loss) - else: - self.scheduler.step() - - return loss.item() - - -# Integration function for your simulation framework -def create_diffmusic_model(system_model_params: SystemModelParams, model_config: dict): - """ - Create DiffMUSIC model for integration with your simulation framework. - - Args: - system_model_params: Your SystemModelParams instance - model_config: Model configuration dictionary - - Returns: - DiffMUSIC model instance - """ - model_params = model_config.get("model_params", {}) - - diffmusic = DiffMUSIC( - system_model_params=system_model_params, - position_constraint=model_params.get("position_constraint", "ula"), - gain_constraint=model_params.get("gain_constraint", "positive"), - temperature=model_params.get("temperature", 0.1), - cell_size_coeff=model_params.get("cell_size_coeff", 0.2) - ) - - return diffmusic - - -# Example usage function compatible with your framework -def run_diffmusic_simulation(system_model_params: SystemModelParams, - training_params: dict, - samples: torch.Tensor, - true_angles: torch.Tensor): - """ - Example function showing how to use DiffMUSIC with your simulation framework. - - Args: - system_model_params: Your SystemModelParams instance - training_params: Training parameters dictionary - samples: Signal samples [batch_size, N, T] or [N, T] - true_angles: True DOA angles [batch_size, M] or [M] - - Returns: - Trained DiffMUSIC model and final loss - """ - # Create model - model_config = { - "model_type": "diffMUSIC", - "model_params": { - "position_constraint": "ula", - "gain_constraint": "positive", - "temperature": 0.1 - } - } - - diffmusic = create_diffmusic_model(system_model_params, model_config) - trainer = DiffMUSICTrainer(diffmusic, training_params) - - # Move to device - device = training_params["device"] - diffmusic = diffmusic.to(device) - samples = samples.to(device) - true_angles = true_angles.to(device) - - # Training loop - num_epochs = training_params["epochs"] - batch_size = training_params["batch_size"] - - # Simple training (you can adapt this to your batch loading pattern) - for epoch in range(num_epochs): - total_loss, data_loss, reg_loss, spacing_reg = trainer.train_step( - samples, true_angles, system_model_params.M - ) - - if epoch % 10 == 0: - print(f"Epoch {epoch}: Total Loss = {total_loss:.6f}, Data Loss = {data_loss:.6f}") - - # Validation - val_loss = trainer.validate(samples, true_angles, system_model_params.M) - print(f"Final Validation Loss: {val_loss:.6f}") - - return diffmusic, val_loss \ No newline at end of file diff --git a/src/old/diff_subspace_method b/src/old/diff_subspace_method deleted file mode 100644 index 0655b4e..0000000 --- a/src/old/diff_subspace_method +++ /dev/null @@ -1,140 +0,0 @@ -import warnings -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -import matplotlib.pyplot as plt -import scipy as sc - -from src.signal_creation import SystemModel, SystemModelParams -from src.utils import set_unified_seed -from scipy.ndimage import maximum_filter -from scipy.ndimage import label -from scipy.ndimage import find_objects - -from src.subspace_method import SubspaceMethod - - -class DiffSubspaceMethod(SubspaceMethod): - """ - Differentiable extension of SubspaceMethod with learnable antenna parameters - """ - def __init__(self, - system_model_params: SystemModelParams, - init_antenna_positions: torch.Tensor = None, - init_antenna_gains: torch.Tensor = None, - gain_constraint: str = "none", - temperature: float = 1.0, - cell_size_coeff: float = 0.2, - model_order_estimation: str = None): - - # Create system model for parent class - system_model = SystemModel(system_model_params) - super().__init__(system_model, model_order_estimation) - - # Store parameters - self.params = system_model_params - self.N = system_model_params.N - self.gain_constraint = gain_constraint - self.temperature = temperature - - # Initialize antenna positions as learnable parameters - if init_antenna_positions is None: - init_antenna_positions = torch.arange(self.N, dtype=torch.float64) * (self.params.wavelength / 2) - - self.antenna_positions = nn.Parameter(init_antenna_positions.clone()) - assert init_antenna_positions.requires_grad is True, "Initial antenna positions must be a differentiable tensor." - - # Initialize antenna gains as learnable parameters - if init_antenna_gains is None: - init_antenna_gains = torch.ones(self.N, dtype=torch.complex64) - self.antenna_gains = nn.Parameter(init_antenna_gains.clone()) - - # Initialize angle grid for far-field DOA estimation - self._init_angle_grid() - - # Initialize cell size for soft peak finding - self.cell_size = int(self.angles_dict.shape[0] * cell_size_coeff) - if self.cell_size % 2 == 0: - self.cell_size += 1 - - # Store current MUSIC spectrum for analysis - self.music_spectrum = None - - 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_dict = torch.arange(-angle_range, angle_range + angle_resolution, - angle_resolution, dtype=torch.float64) - self.angles_dict = torch.round(self.angles_dict, decimals=angle_decimals) - - def _apply_position_constraints(self): - """Reorders the antenna positions, ensuring we don't get stucked if some replace positions. - I don't think we need this because it is mathematically fine that thy will replace order. - In fact, this replacement may occur some incontinuity in the loss. - We do need to keep it in mind though. - """ - # with torch.no_grad(): - # if self.position_constraint == "ula": - # sorted_positions, _ = torch.sort(self.antenna_positions) - # self.antenna_positions.data = sorted_positions - pass - - def _apply_gain_constraints(self): - """Prevent extrme gains from future random walk. For now i leave commented.""" - # with torch.no_grad(): - # if self.gain_constraint == "positive": - # # Just clamp to reasonable ranges to prevent numerical issues - # real_part = torch.clamp(self.antenna_gains.real, min=0.1, max=10.0) - # imag_part = torch.clamp(self.antenna_gains.imag, min=-2.0, max=2.0) - # self.antenna_gains.data = torch.complex(real_part, imag_part, dtype=torch.complex64) - pass - - def get_constrained_parameters(self): - """Get antenna parameters without any sorting - order doesn't matter mathematically""" - positions = self.antenna_positions - gains = self.antenna_gains - return positions, gains - - #NOTE: I read up to this point - - def compute_steering_matrix(self, angles: torch.Tensor): - """ - Compute steering matrix for far-field sources using current antenna configuration. - Follows the pattern from your SystemModel.steering_vec_far_field method. - - Args: - angles: DOA angles in radians [num_angles] - - Returns: - Complex steering matrix [N, num_angles] - """ - positions, gains = self.get_constrained_parameters() - - # Convert angles to proper shape for broadcasting - if angles.dim() == 1: - angles = angles.unsqueeze(0) # [1, num_angles] - - # Reshape positions for broadcasting: [N, 1] - positions = positions.view(-1, 1) - - # Compute phase delays: [N, 1] * sin([1, num_angles]) -> [N, num_angles] - time_delay = positions @ torch.sin(angles) - - # Compute steering vectors (following your formula) - steering_matrix = torch.exp(-2j * torch.pi * time_delay / self.params.wavelength) - steering_matrix = steering_matrix.to(torch.complex64) - - # Apply antenna gains - gains_diag = torch.diag(gains) - - # Normalize gains (following your normalization pattern) - normalization_factor = 1 / torch.linalg.norm(gains, ord=2) - steering_matrix = normalization_factor * gains_diag @ steering_matrix - - return steering_matrix - - \ No newline at end of file diff --git a/src/old/subspace_method.py b/src/old/subspace_method.py deleted file mode 100644 index 00eb552..0000000 --- a/src/old/subspace_method.py +++ /dev/null @@ -1,242 +0,0 @@ -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 -from src.signal_creation import SystemModel - - - -#TODO: Haven't really checked the code here yet, just inheritrs for now. -class SubspaceMethod(nn.Module): - """ - Basic methods for all subspace methods. - """ - - def __init__(self, system_model: SystemModel, model_order_estimation: str = 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 - - def subspace_separation(self, - covariance: torch.Tensor, - number_of_sources: torch.tensor = None) \ - -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.tensor]: - """ - - Args: - covariance: - number_of_sources: - - Returns: - the signal ana noise subspaces, both as torch.Tensor(). - """ - 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/signal_creation.py b/src/signal_creation.py index 2d56d8a..a1beb75 100644 --- a/src/signal_creation.py +++ b/src/signal_creation.py @@ -3,7 +3,12 @@ from random import sample class SystemModelParams: + """ + Parameters for the for the simulation at one claass. + """ + def __init__(self, **kwargs): + kwargs = {**kwargs} for key, value in kwargs.items(): if isinstance(value, str): value = value.lower() @@ -78,7 +83,7 @@ def steering_vec(self, angles: np.ndarray) -> torch.Tensor: #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: + def steering_vec_far_field(self, angles: np.ndarray | torch.Tensor) -> torch.Tensor: """ Compute far-field steering vectors @@ -129,7 +134,7 @@ def get_array(self): """ return self.array - def get_gain_perturbations(self): + def get_antenna_gains(self): """ Get the gain perturbation applied to the steering vector diff --git a/src/subspace_method.py b/src/subspace_method.py index aeca4f1..8b58d23 100644 --- a/src/subspace_method.py +++ b/src/subspace_method.py @@ -15,7 +15,8 @@ class SubspaceMethod(nn.Module): Basic methods for all subspace methods. """ - def __init__(self, system_model, model_order_estimation: str = None): + 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 @@ -23,6 +24,8 @@ def __init__(self, system_model, model_order_estimation: str = None): 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, diff --git a/src/utils.py b/src/utils.py index c82c2a6..79c07ac 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 @@ -36,42 +14,59 @@ import torch.nn as nn from datetime import datetime -# # 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 +def plot_spectrums(angles_grid: torch.Tensor, + spectrums_dict : dict, + true_angles: torch.Tensor = None, + save: bool = False, + title: str = "MUSIC spectrums" + ): -# Raises: -# ValueError: If the number of sources in the batch is not equal for all samples + """ 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(arg)}") + 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(arg)}).") + 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 -# """ -# 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 save_data_to_file(data_path: Path, *args): """ @@ -167,177 +162,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): """ @@ -365,237 +189,3 @@ 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. - -# 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) From 2afb26b50f8449aa4c8c0c088d58795d07b01705 Mon Sep 17 00:00:00 2001 From: oritalp Date: Wed, 18 Jun 2025 17:42:47 +0300 Subject: [PATCH 05/13] minor changes --- .gitignore | 3 ++- doa_runner.py | 17 ++++++++++------- main.py | 35 ++++++++++++++++++++++++----------- run_simulation.py | 12 +++++------- src/utils.py | 7 ++++--- 5 files changed, 45 insertions(+), 29 deletions(-) diff --git a/.gitignore b/.gitignore index 5f9a405..fc94012 100644 --- a/.gitignore +++ b/.gitignore @@ -10,5 +10,6 @@ wandb/ *.png *.pdf *.out -*.sh +interactive.sh +*.pth ./datasets/ \ No newline at end of file diff --git a/doa_runner.py b/doa_runner.py index 26e8b32..d9e3eb4 100644 --- a/doa_runner.py +++ b/doa_runner.py @@ -12,6 +12,7 @@ from typing import Dict, Any, Optional, Tuple from tqdm import tqdm from datetime import datetime +import time from src.utils import sample_covariance @@ -244,9 +245,9 @@ def _setup_wandb(self, window_size, trial_idx=None): # Create run name if trial_idx is not None: - run_name = f"trial_{trial_idx+1}_ws_{window_size}_{self.system_params.loss_type}" + run_name = f"trial_{trial_idx+1}_window_size_{window_size}_{self.system_params.loss_type}" else: - run_name = f"ws_{window_size}_{self.system_params.loss_type}" + run_name = f"winsow_size_{window_size}_{self.system_params.loss_type}" # Initialize wandb wandb.init( @@ -305,9 +306,11 @@ def run(self) -> Dict[str, Any]: 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 - self.results = results - return results + + return self.results.copy() # Return a copy to avoid external modifications def _create_fresh_algorithm(self, window_size): """ @@ -345,13 +348,13 @@ def _run_with_window_optimization(self, window_sizes, optimization_mode=True) -> Returns: Results from best configuration (or single configuration) plus stats """ - import time + best_rmspe = float('inf') best_results = None best_window_size = None - total_start_time = time.time() + trial_times = [] for i, window_size in enumerate(window_sizes): @@ -563,7 +566,7 @@ def _run_evaluation(self) -> Dict[str, Any]: 'true_angles': true_angles.cpu().numpy(), 'rmspe': rmspe.item(), "learned_antenna_positions": learned_antenna_positions, - "learned_coplex_gains": learned_coplex_gains, + "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 } diff --git a/main.py b/main.py index d6e5178..5c48d7e 100644 --- a/main.py +++ b/main.py @@ -30,16 +30,25 @@ import argparse import torch -#TODO: check wandb (short single case) -# run multiple window sizes cases. -# pass over the running once again -# defined missions for rearanging the results and where they are saved, split it into single_run_ressults and multi_run_results. -# define graphs plotting class for the multi-case results. +#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: -#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. + # Initialization os.system("cls||clear") plt.close("all") @@ -94,7 +103,7 @@ "softmax_window_size": 0.1, # % of angle grid length # Case 3: Window size optimization (array/list of values) - # "softmax_window_size": np.arange(0.2, 0.4, 0.05), # Relative sizes: [0.2, 0.25, 0.3, 0.35] + # "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 } @@ -234,8 +243,12 @@ def parse_arguments(): 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 and "learned_antenna_gains" in results: + if "learned_antenna_positions" in results.keys() and "learned_antennas_gains" in results.keys(): if model_config["model_type"].lower() == "music": - print("These shoulf be the tandard one:") - print(f"Learned antenna positions: {results['learned_antenna_positions']}") - print(f"Learned antenna gains: {results['learned_antenna_gains']}") \ No newline at end of file + 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("Compared to the physical antennas gains:") + print(f"Physical antennas gains: {results['physical_antennas_gains']}") \ No newline at end of file diff --git a/run_simulation.py b/run_simulation.py index c816682..5f3cf9f 100644 --- a/run_simulation.py +++ b/run_simulation.py @@ -48,8 +48,7 @@ def __run_simulation(**kwargs): # Initialize paths #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, - dt_string_for_save) + 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 @@ -62,7 +61,6 @@ def __run_simulation(**kwargs): true_angles = signals_creator.get_labels() physical_array = signals_creator.get_array() physical_antennas_gains = signals_creator.get_antenna_gains() - antennas_gains_norm = torch.linalg.norm(physical_antennas_gains, ord=2, dim=0) #just for testing purposes # Pack data into dictionary data_dict = { @@ -109,10 +107,10 @@ def __run_simulation(**kwargs): 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) + # # 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}") diff --git a/src/utils.py b/src/utils.py index 79c07ac..96d450c 100644 --- a/src/utils.py +++ b/src/utils.py @@ -37,10 +37,10 @@ def plot_spectrums(angles_grid: torch.Tensor, 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(arg)}") + 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(arg)}).") + f"does not match length of spectrum ({len(value)}).") plt.plot(angles_deg, value.cpu().detach().numpy(), 'r--', linewidth=1, label= f"{key}") @@ -102,7 +102,8 @@ def load_data_from_file(data_path: str): raise FileNotFoundError(f"load_data_from_file: No data found in {data_file}.") return data -def initialize_paths(main_path: Path, system_model_params, dt_string_for_save: str) -> tuple: +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}_" + From 236994834d569b5cd43128b9165f7ca6d3e3451c Mon Sep 17 00:00:00 2001 From: oritalp Date: Wed, 18 Jun 2025 18:57:41 +0300 Subject: [PATCH 06/13] minor note added --- main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/main.py b/main.py index 5c48d7e..8f4b7aa 100644 --- a/main.py +++ b/main.py @@ -48,6 +48,7 @@ # 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") From edfe2bb8bb37a8b4674805cf62e9d9ef1b6b6c09 Mon Sep 17 00:00:00 2001 From: Sheli Date: Sat, 21 Jun 2025 13:01:46 +0300 Subject: [PATCH 07/13] Adding support for Multi-Experiment --- main.py | 26 ++- run_simulation.py | 103 +++++----- src/multi_experiment_runner.py | 335 +++++++++++++++++++++++++++++++++ src/utils.py | 11 ++ 4 files changed, 416 insertions(+), 59 deletions(-) create mode 100644 src/multi_experiment_runner.py diff --git a/main.py b/main.py index 8f4b7aa..6b92e14 100644 --- a/main.py +++ b/main.py @@ -55,10 +55,28 @@ 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], +# "SNR_sweep": { +# "parameter": "snr", +# "values": [-10, -5, 0, 5, 10, 15, 20, 25, 30], +# "fixed_params": {"T": 100}, +# "plot_config": { +# "title": "RMSPE vs SNR (T=100)", +# "x_label": "SNR (dB)", +# "y_label": "RMSPE (degrees)", +# "save_name": "rmspe_vs_snr" +# } +# }, +# "T_sweep": { +# "parameter": "T", +# "values": [10, 20, 30, 50, 70, 100, 150, 200], +# "fixed_params": {"snr": 30}, +# "plot_config": { +# "title": "RMSPE vs T (SNR=30dB)", +# "x_label": "T (snapshots)", +# "y_label": "RMSPE (degrees)", +# "save_name": "rmspe_vs_T" +# } +# } } simulation_commands = { diff --git a/run_simulation.py b/run_simulation.py index 5f3cf9f..ed4ebef 100644 --- a/run_simulation.py +++ b/run_simulation.py @@ -20,7 +20,7 @@ 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"] @@ -117,6 +117,45 @@ def __run_simulation(**kwargs): 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_simulation(**kwargs): """ @@ -129,60 +168,14 @@ def run_simulation(**kwargs): - M: a list of number of sources to be tested """ #TODO: check if anything is missing for the scenario_dict option once needed. - if kwargs["scenario_dict"] == {}: - results = __run_simulation(**kwargs) - return results - - # from this on the option of multiple scenarios is used. this is activated when we specify the sceario_dict - # in main.py - # TODO: Later adjust it to our way. - - # 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 + scenario_dict = kwargs["scenario_dict"] + + if not scenario_dict: # Single experiment (current behavior) + return run_single_simulation(**kwargs) + else: # Multi-experiment sweep + return __run_parameter_sweeps(**kwargs) + + if __name__ == "__main__": diff --git a/src/multi_experiment_runner.py b/src/multi_experiment_runner.py new file mode 100644 index 0000000..6b0e83a --- /dev/null +++ b/src/multi_experiment_runner.py @@ -0,0 +1,335 @@ +""" +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", {}) + + # 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 + + Returns: + Aggregated sweep results + """ + 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 + self._save_sweep_results(sweep_results) + + # Generate plots + 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 _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}" + 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 _prepare_experiment_kwargs(self, param_value) -> Dict[str, Any]: + """ + Prepare kwargs for a single experiment with the specified parameter value + + Args: + param_value: Value for the sweep parameter + + 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 + + # 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 _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\n") + + f.write(f"Results:\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 the sweep results""" + plotter = SweepPlotter(sweep_results, self.sweep_results_path / "plots") + plotter.plot_rmspe_vs_parameter() + + +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 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): + """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}') + 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/utils.py b/src/utils.py index 96d450c..22d3ab9 100644 --- a/src/utils.py +++ b/src/utils.py @@ -190,3 +190,14 @@ def set_unified_seed(seed: int = 42): torch.use_deterministic_algorithms(True) +def save_pickle(data, filepath): + """Save data to pickle file""" + with open(filepath, 'wb') as f: + pickle.dump(data, f) + + +def load_pickle(filepath): + """Load data from pickle file""" + with open(filepath, 'rb') as f: + return pickle.load(f) + From 5f961d2e2d35231e07f354dd8510c920e78342a2 Mon Sep 17 00:00:00 2001 From: Sheli Date: Sat, 21 Jun 2025 13:36:23 +0300 Subject: [PATCH 08/13] Support for multi loss --- main.py | 62 +++--- src/multi_experiment_runner.py | 345 +++++++++++++++++++++++++++++---- 2 files changed, 352 insertions(+), 55 deletions(-) diff --git a/main.py b/main.py index 6b92e14..32f81e7 100644 --- a/main.py +++ b/main.py @@ -55,28 +55,46 @@ plt.close("all") scenario_dict = { -# "SNR_sweep": { -# "parameter": "snr", -# "values": [-10, -5, 0, 5, 10, 15, 20, 25, 30], -# "fixed_params": {"T": 100}, -# "plot_config": { -# "title": "RMSPE vs SNR (T=100)", -# "x_label": "SNR (dB)", -# "y_label": "RMSPE (degrees)", -# "save_name": "rmspe_vs_snr" -# } -# }, -# "T_sweep": { -# "parameter": "T", -# "values": [10, 20, 30, 50, 70, 100, 150, 200], -# "fixed_params": {"snr": 30}, -# "plot_config": { -# "title": "RMSPE vs T (SNR=30dB)", -# "x_label": "T (snapshots)", -# "y_label": "RMSPE (degrees)", -# "save_name": "rmspe_vs_T" -# } -# } + # 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 = { diff --git a/src/multi_experiment_runner.py b/src/multi_experiment_runner.py index 6b0e83a..0885b55 100644 --- a/src/multi_experiment_runner.py +++ b/src/multi_experiment_runner.py @@ -35,6 +35,10 @@ def __init__(self, sweep_config: Dict[str, Any], base_kwargs: Dict[str, Any]): 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) @@ -45,10 +49,18 @@ def __init__(self, sweep_config: Dict[str, Any], base_kwargs: Dict[str, Any]): 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}") @@ -92,10 +104,8 @@ def run_sweep(self) -> Dict[str, Any]: total_time = time.time() - sweep_start_time sweep_results = self.results_manager.finalize_results(total_time, individual_times) - # Save results + # Save results and generate plots self._save_sweep_results(sweep_results) - - # Generate plots self._generate_plots(sweep_results) print(f"\nSweep completed in {total_time:.2f}s") @@ -103,34 +113,82 @@ def run_sweep(self) -> Dict[str, Any]: return sweep_results - def _setup_results_directory(self): - """Setup directory structure for sweep results""" - from datetime import datetime + 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}") - # Create timestamp - dt_string = datetime.now().strftime("%d_%m_%Y_%H_%M") + # Setup results directory + self._setup_results_directory() - # 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}" - sweep_name += f"_{dt_string}" + # Initialize multi-loss results manager + multi_loss_results = MultiLossResultsManager(self.sweep_config) - # Setup paths - self.sweep_results_path = self.base_path / "results" / "parameter_sweeps" / sweep_name - self.sweep_results_path.mkdir(parents=True, exist_ok=True) + total_experiments = len(self.sweep_values) * len(self.loss_functions) + experiment_count = 0 - # Create subdirectories - (self.sweep_results_path / "plots").mkdir(exist_ok=True) - (self.sweep_results_path / "individual_experiments").mkdir(exist_ok=True) + 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) -> Dict[str, Any]: + 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 @@ -147,6 +205,10 @@ def _prepare_experiment_kwargs(self, param_value) -> Dict[str, Any]: 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 @@ -154,6 +216,34 @@ def _prepare_experiment_kwargs(self, param_value) -> Dict[str, Any]: 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 @@ -185,25 +275,137 @@ def _save_summary_text(self, sweep_results: Dict[str, Any], summary_file: Path): 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\n") + f.write(f"Fixed Parameters: {sweep_results['fixed_params']}\n") - f.write(f"Results:\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") + # 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"\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 the sweep results""" + """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 @@ -277,6 +479,83 @@ def finalize_results(self, total_time: float, individual_times: List[float]) -> 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 @@ -287,7 +566,7 @@ def __init__(self, sweep_results: Dict[str, Any], plots_path: Path): self.plots_path = plots_path self.plots_path.mkdir(parents=True, exist_ok=True) - def plot_rmspe_vs_parameter(self): + def plot_rmspe_vs_parameter(self, suffix=''): """Plot RMSPE vs the sweep parameter""" import matplotlib.pyplot as plt @@ -324,7 +603,7 @@ def plot_rmspe_vs_parameter(self): bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8)) # Save plot - save_name = plot_config.get('save_name', f'rmspe_vs_{param_name}') + 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') From 1a26b04f185bbec60f0aea724a53543f6e6cf35f Mon Sep 17 00:00:00 2001 From: Sheli Date: Mon, 23 Jun 2025 12:10:30 +0300 Subject: [PATCH 09/13] adding plot_learned_parameters --- doa_runner.py | 23 ++++++++++ main.py | 80 ++++++++++++++++---------------- src/diffmusic.py | 116 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 180 insertions(+), 39 deletions(-) diff --git a/doa_runner.py b/doa_runner.py index d9e3eb4..117eba8 100644 --- a/doa_runner.py +++ b/doa_runner.py @@ -606,4 +606,27 @@ def plot_graphs(self, path: Path): 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 \ No newline at end of file diff --git a/main.py b/main.py index 32f81e7..427d7c6 100644 --- a/main.py +++ b/main.py @@ -55,46 +55,46 @@ plt.close("all") scenario_dict = { - # 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 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 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" - } - } + # # 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 = { @@ -287,5 +287,7 @@ def parse_arguments(): 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']}") \ No newline at end of file + 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/src/diffmusic.py b/src/diffmusic.py index 7b372da..48879f1 100644 --- a/src/diffmusic.py +++ b/src/diffmusic.py @@ -412,6 +412,122 @@ def plot_spectrum(self, highlight_angles: torch.Tensor = 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): """ From 61b750f61ea2c647869c96e27b17f7a9452a65b6 Mon Sep 17 00:00:00 2001 From: Sheli Date: Mon, 23 Jun 2025 13:49:06 +0300 Subject: [PATCH 10/13] adding multi_loss_comparison --- doa_runner.py | 421 +++++++++++++++++++++++++++++++++++++++++++++- main.py | 36 ++-- run_simulation.py | 137 ++++++++++++++- 3 files changed, 577 insertions(+), 17 deletions(-) diff --git a/doa_runner.py b/doa_runner.py index 117eba8..9cc4a4a 100644 --- a/doa_runner.py +++ b/doa_runner.py @@ -13,6 +13,7 @@ from tqdm import tqdm from datetime import datetime import time +from typing import Dict, Any, Optional, Tuple, List from src.utils import sample_covariance @@ -629,4 +630,422 @@ def plot_graphs(self, path: Path): path_to_save=path, title=title ) - return None \ No newline at end of file + 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 + colors = {'rmspe': 'blue', 'spectrum': 'red', 'unsupervised': 'green'} + line_styles = {'rmspe': '-', 'spectrum': '--', 'unsupervised': '-.'} + default_colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown'] + default_styles = ['-', '--', '-.', ':', '-', '--'] + + # 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 + 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") + 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() \ No newline at end of file diff --git a/main.py b/main.py index 427d7c6..70303c2 100644 --- a/main.py +++ b/main.py @@ -82,25 +82,27 @@ # } # }, - # # 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" - # } - # } + # 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 = { "create_data": True, "save_data": False, # Save data after creation - "plot_results": True, # Plot data after creation + "plot_results": False, # Plot data after creation + "multi_loss_comparison": False, # Enable multi-loss spectrum and learned parameters comparison + "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. @@ -272,7 +274,11 @@ def parse_arguments(): print("Total time: ", time.time() - start) # Print final results summary - if isinstance(results, dict) and 'rmspe' in results: + 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 'estimated_angles' in results and 'true_angles' in results: diff --git a/run_simulation.py b/run_simulation.py index ed4ebef..d9252c3 100644 --- a/run_simulation.py +++ b/run_simulation.py @@ -155,6 +155,137 @@ def __run_parameter_sweeps(**kwargs): 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}") + + 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"]) + + # 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"] + + # Prepare data dictionary + data_dict = {} + + 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 + } + + # 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')) + print(f"{loss_func.upper()}: RMSPE = {rmspe:.6f}°") + else: + 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 + def run_simulation(**kwargs): @@ -169,8 +300,12 @@ def run_simulation(**kwargs): """ #TODO: check if anything is missing for the scenario_dict option once needed. scenario_dict = kwargs["scenario_dict"] + simulation_commands = kwargs["simulation_commands"] + - if not scenario_dict: # Single experiment (current behavior) + 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) From 9dad0c7c640ca62fe4a5c87070920a03f2e5f3b6 Mon Sep 17 00:00:00 2001 From: Sheli Date: Mon, 23 Jun 2025 16:49:10 +0300 Subject: [PATCH 11/13] adding print of steering_matrix_mse --- doa_runner.py | 45 ++++++++++++++++++++++++++++++++++++++++++--- main.py | 8 +++++--- run_simulation.py | 4 ++-- src/diffmusic.py | 3 +++ 4 files changed, 52 insertions(+), 8 deletions(-) diff --git a/doa_runner.py b/doa_runner.py index 9cc4a4a..74e067a 100644 --- a/doa_runner.py +++ b/doa_runner.py @@ -554,6 +554,8 @@ def _run_evaluation(self) -> Dict[str, Any]: 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 @@ -569,7 +571,8 @@ def _run_evaluation(self) -> Dict[str, Any]: "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 + '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): @@ -844,10 +847,14 @@ def plot_multi_loss_spectrum_comparison(self, multi_loss_results: Dict[str, Any] 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") + 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)) @@ -1048,4 +1055,36 @@ def plot_multi_loss_learned_parameters(self, multi_loss_results: Dict[str, Any], plt.savefig(path / "multi_loss_learned_parameters.pdf", bbox_inches='tight') print(f"Multi-loss learned parameters plot saved to: {save_path}") - plt.show() \ No newline at end of file + 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() \ No newline at end of file diff --git a/main.py b/main.py index 70303c2..0652069 100644 --- a/main.py +++ b/main.py @@ -54,7 +54,7 @@ os.system("cls||clear") plt.close("all") -scenario_dict = { +scenario_dict = { # Fill in to do multiple experiments # # Example 1: Single loss function sweep (original behavior) # "SNR_sweep_rmspe": { # "parameter": "snr", @@ -100,8 +100,8 @@ simulation_commands = { "create_data": True, "save_data": False, # Save data after creation - "plot_results": False, # Plot data after creation - "multi_loss_comparison": False, # Enable multi-loss spectrum and learned parameters comparison + "plot_results": True, # Plot data after creation + "multi_loss_comparison": False, # 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! @@ -281,6 +281,8 @@ def parse_arguments(): 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'])}") diff --git a/run_simulation.py b/run_simulation.py index d9252c3..2a19026 100644 --- a/run_simulation.py +++ b/run_simulation.py @@ -114,7 +114,6 @@ def run_single_simulation(**kwargs): print(f"Results saved to: {results_path}") print(f"RMSPE: {results.get('rmspe', 'N/A'):.6f}") - return results def __run_parameter_sweeps(**kwargs): @@ -278,7 +277,8 @@ def run_multi_loss_comparison(**kwargs): result = multi_loss_results['results'].get(loss_func, {}) if 'error' not in result: rmspe = result.get('rmspe', float('inf')) - print(f"{loss_func.upper()}: RMSPE = {rmspe:.6f}°") + steering_mse = result.get('steering_matrix_mse', float('nan')) + print(f"{loss_func.upper()}: RMSPE = {rmspe:.6f}°, Steering MSE = {steering_mse:.6f}") else: print(f"{loss_func.upper()}: ERROR - {result.get('error', 'Unknown error')}") diff --git a/src/diffmusic.py b/src/diffmusic.py index 48879f1..f6b92db 100644 --- a/src/diffmusic.py +++ b/src/diffmusic.py @@ -148,6 +148,9 @@ def compute_steering_matrix(self, angles: torch.Tensor) -> torch.Tensor: 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 From c9e6d480ae7d1802faedb1e84be2e91d3b8c91ff Mon Sep 17 00:00:00 2001 From: Sheli Date: Mon, 23 Jun 2025 17:33:01 +0300 Subject: [PATCH 12/13] Adding a spectrum with the real parameters --- doa_runner.py | 77 ++++++++++++++++++++++++++++++++++++++++++++++++--- main.py | 24 ++++++++-------- 2 files changed, 85 insertions(+), 16 deletions(-) diff --git a/doa_runner.py b/doa_runner.py index 74e067a..647b439 100644 --- a/doa_runner.py +++ b/doa_runner.py @@ -785,12 +785,19 @@ def plot_multi_loss_spectrum_comparison(self, multi_loss_results: Dict[str, Any] # Create figure plt.figure(figsize=(14, 8)) - # Define colors and line styles for different loss functions - colors = {'rmspe': 'blue', 'spectrum': 'red', 'unsupervised': 'green'} - line_styles = {'rmspe': '-', 'spectrum': '--', 'unsupervised': '-.'} + # 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) @@ -1087,4 +1094,66 @@ def _compute_steering_matrix_mse(self, true_angles: torch.Tensor) -> float: # Compute MSE mse = torch.mean(torch.abs(real_steering_matrix - estimated_steering_matrix) ** 2) - return mse.item() \ No newline at end of file + 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/main.py b/main.py index 0652069..f86fef3 100644 --- a/main.py +++ b/main.py @@ -83,18 +83,18 @@ # }, # 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" - } - } + # "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 = { From 1330341573394aca10b3837d28aa7c41091efddb Mon Sep 17 00:00:00 2001 From: Sheli Date: Sat, 5 Jul 2025 14:23:00 +0300 Subject: [PATCH 13/13] printing every epoch Steering MSE --- doa_runner.py | 11 +++++++++-- main.py | 4 ++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/doa_runner.py b/doa_runner.py index 647b439..04342a6 100644 --- a/doa_runner.py +++ b/doa_runner.py @@ -439,7 +439,8 @@ def _run_training(self) -> Dict[str, Any]: self.training_history = { 'train_loss': [], 'val_loss': [], - 'learning_rates': [] + 'learning_rates': [], + 'steering_mse': [] } print(f"Training for {self.system_params.epochs} epochs...") @@ -447,8 +448,13 @@ def _run_training(self) -> Dict[str, Any]: 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) @@ -460,6 +466,7 @@ def _run_training(self) -> Dict[str, Any]: wandb.log({ "epoch": epoch, "train_loss": epoch_loss, + "steering_mse": steering_mse, "learning_rate": current_lr if self.optimizer else 0 }) except: @@ -474,7 +481,7 @@ def _run_training(self) -> Dict[str, Any]: # Print progress if (epoch + 1) % 10 == 0 or epoch == 0: - print(f"Epoch {epoch+1}/{self.system_params.epochs}, Loss: {epoch_loss:.6f}") + print(f"Epoch {epoch+1}/{self.system_params.epochs}, Loss: {epoch_loss:.6f}, Steering MSE: {steering_mse:.6f}") return { 'training_history': self.training_history, diff --git a/main.py b/main.py index f86fef3..1384a40 100644 --- a/main.py +++ b/main.py @@ -82,7 +82,7 @@ # } # }, - # Example 3: T sweep with multi-loss + # # Example 3: T sweep with multi-loss # "T_sweep_multi_loss": { # "parameter": "T", # "values": [20, 30, 50, 70, 100, 150, 200], @@ -101,7 +101,7 @@ "create_data": True, "save_data": False, # Save data after creation "plot_results": True, # Plot data after creation - "multi_loss_comparison": False, # Enable multi-loss spectrum and learned parameters comparison, if multiple experiments then need to be False + "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!