Create hydranet main

common_utils/artifacts_utils.py

@@ -0,0 +1,65 @@
+import os
+
+
+def get_model_files(path, run_type):
+    """
+    Retrieve model files from a directory that match the given run type and common extensions.
+
+    Args:
+        path (str): The directory path where model files are stored.
+        run_type (str): The type of run (e.g., calibration, testing).
+
+    Returns:
+        list: List of matching model file paths.
+    """
+    # Define the common model file extensions - more can be added as needed
+    common_extensions = ['.pt', '.pth', '.h5', '.hdf5', '.pkl', '.json', '.bst', '.txt', '.bin', '.cbm', '.onnx']
+
+    # Retrieve files that start with run_type and end with any of the common extensions
+    model_files = [f for f in os.listdir(path) if f.startswith(f"{run_type}_model_") and any(f.endswith(ext) for ext in common_extensions)]
+
+    return model_files
+
+
+def get_latest_model_artifact(path, run_type):
+    """
+    Retrieve the latest model artifact for a given run type based on the modification time.
+
+    Args:
+        path (str): The model specifc directory path where artifacts are stored. 
+        Where PATH_ARTIFACTS = setup_artifacts_paths(PATH) executed in the model specifc main.py script.
+        and PATH = Path(__file__)
+
+        run_type (str): The type of run (e.g., calibration, testing, forecasting).
+
+    Returns:
+        str: The path to the latest model artifact given the run type.
+
+    Raises:
+        FileNotFoundError: If no model artifacts are found for the given run type.
+    """
+
+    # List all model files for the given specific run_type with the expected filename pattern
+    model_files = get_model_files(path, run_type) #[f for f in os.listdir(path) if f.startswith(f"{run_type}_model_") and f.endswith('.pt')]
+
+    if not model_files:
+        raise FileNotFoundError(f"No model artifacts found for run type '{run_type}' in path '{path}'")
+
+    # Sort the files based on the timestamp embedded in the filename. With format %Y%m%d_%H%M%S For example, '20210831_123456.pt'
+    model_files.sort(reverse=True)
+
+    #print statements for debugging
+    print(f"artifacts availible: {model_files}")
+    print(f"artifact used: {model_files[0]}")
+
+    # Return the latest model file
+    return os.path.join(path, model_files[0])
+
+    # notes on stepshifted models:
+    # There will be some thinking here in regards to how we store, denote (naming convention), and retrieve the model artifacts from stepshifted models.
+    # It is not a big issue, but it is something to consider os we don't do something headless. 
+    # A possible format could be: <run_type>_model_s<step>_<timestamp>.pt example: calibration_model_s00_20210831_123456.pt, calibration_model_s01_20210831_123456.pt, etc.
+    # And the rest of the code maded in a way to handle this naming convention without any issues. Could be a simple fix.
+    # Alternatively, we could store the model artifacts in a subfolder for each stepshifted model. This would make it easier to handle the artifacts, but it would also make it harder to retrieve the latest artifact for a given run type.
+    # Lastly, the solution Xiaolong is working on might allow us the store multiple models (steps) in one artifact, which would make this whole discussion obsolete and be the best solution.
+

common_utils/cli_parser_utils.py

@@ -0,0 +1,76 @@
+import sys
+import argparse
+
+def parse_args():
+
+    """
+    CLI parser for model specific main.py scripts.
+    """
+
+    parser = argparse.ArgumentParser(description='Run model pipeline with specified run type.')
+
+    parser.add_argument('-r', '--run_type',
+                        choices=['calibration', 'testing', 'forecasting'],
+                        type=str,
+                        default='calibration',
+                        help='Choose the run type for the model: calibration, testing, or forecasting. Default is calibration. '
+                             'Note: If --sweep is flagged, --run_type must be calibration.')
+
+    parser.add_argument('-s', '--sweep',
+                        action='store_true',
+                        help='Set flag to run the model pipeline as part of a sweep. No explicit flag means no sweep.'
+                             'Note: If --sweep is flagged, --run_type must be calibration, and both training and evaluation is automatically implied.')
+
+    parser.add_argument('-t', '--train',
+                        action='store_true',
+                        help='Flag to indicate if a new model should be trained. '
+                             'Note: If --sweep is flagged, --train will also automatically be flagged.')
+
+    parser.add_argument('-e', '--evaluate',
+                        action='store_true',
+                        help='Flag to indicate if the model should be evaluated. '
+                             'Note: If --sweep is specified, --evaluate will also automatically be flagged. '
+                             'Cannot be used with --run_type forecasting.')
+
+    parser.add_argument('-a', '--artifact_name',
+                        type=str,
+                        help='Specify the name of the model artifact to be used for evaluation. '
+                             'The file extension will be added in the main and fit with the specific model algorithm.'
+                             'The artifact name should be in the format: <run_type>_model_<timestamp>.pt.'
+                             'where <run_type> is calibration, testing, or forecasting, and <timestamp> is in the format YMD_HMS.'
+                             'If not provided, the latest artifact will be used by default.')
+
+    return parser.parse_args()
+
+def validate_arguments(args):
+    if args.sweep:
+        if args.run_type != 'calibration':
+            print("Error: Sweep runs must have --run_type set to 'calibration'. Exiting.")
+            print("To fix: Use --run_type calibration when --sweep is flagged.")
+            sys.exit(1)
+
+    if args.run_type in ['testing', 'forecasting'] and args.sweep:
+        print("Error: Sweep cannot be performed with testing or forecasting run types. Exiting.")
+        print("To fix: Remove --sweep flag or set --run_type to 'calibration'.")
+        sys.exit(1)
+
+    if args.run_type == 'forecasting' and args.evaluate:
+        print("Error: Forecasting runs cannot evaluate. Exiting.")
+        print("To fix: Remove --evaluate flag when --run_type is 'forecasting'.")
+        sys.exit(1)
+
+    if args.run_type in ['calibration', 'testing'] and not args.train and not args.evaluate and not args.sweep:
+        print(f"Error: Run type is {args.run_type} but neither --train, --evaluate, nor --sweep flag is set. Nothing to do... Exiting.")
+        print("To fix: Add --train and/or --evaluate flag. Or use --sweep to run both training and evaluation in a WadnB sweep loop.")
+        sys.exit(1)
+
+
+    # notes on stepshifted models:
+    # There will be some thinking here in regards to how we store, denote (naming convention), and retrieve the model artifacts from stepshifted models.
+    # It is not a big issue, but it is something to consider os we don't do something headless. 
+    # A possible format could be: <run_type>_model_s<step>_<timestamp>.pt example: calibration_model_s00_20210831_123456.pt, calibration_model_s01_20210831_123456.pt, etc.
+    # And the rest of the code maded in a way to handle this naming convention without any issues. Could be a simple fix.
+    # Alternatively, we could store the model artifacts in a subfolder for each stepshifted model. This would make it easier to handle the artifacts, but it would also make it harder to retrieve the latest artifact for a given run type.
+    # Lastly, the solution Xiaolong is working on might allow us the store multiple models (steps) in one artifact, which would make this whole discussion obsolete and be the best solution.
+
+

models/purple_alien/configs/config_hyperparameters.py

@@ -8,7 +8,7 @@ def get_hp_config():
     'scheduler' : 'WarmupDecay', #  'CosineAnnealingLR' 'OneCycleLR'
     'total_hidden_channels' : 32,
     'min_events' : 5,
-    'samples': 600, # 10 just for debug
+    'samples': 600, # 600 for actual trainnig, 10 for debug
     'batch_size': 3, 
     'dropout_rate' : 0.125,
     'learning_rate' :  0.001,
@@ -24,7 +24,7 @@ def get_hp_config():
     'loss_reg': 'b',
     'loss_reg_a' : 258, 
     'loss_reg_c' :  0.001, # 0.05 works...
-    'test_samples': 128,
+    'test_samples': 128, # 128 for actual testing, 10 for debug
     'np_seed' : 4,
     'torch_seed' : 4,
     'window_dim' : 32,

models/purple_alien/configs/config_sweep.py

@@ -17,7 +17,7 @@ def get_swep_config():
         'scheduler' : {'value': 'WarmupDecay'}, #CosineAnnealingLR004  'CosineAnnealingLR' 'OneCycleLR'
         'total_hidden_channels': {'value': 32}, # you like need 32, it seems from qualitative results
         'min_events': {'value': 5},
-        'samples': {'value': 600}, # should be a function of batches becaus batch 3 and sample 1000 = 3000....
+        'samples': {'value': 600}, # 600 for run 10 for debug. should be a function of batches becaus batch 3 and sample 1000 = 3000....
         'batch_size': {'value':  3}, # just speed running here..
         "dropout_rate" : {'value' : 0.125},
         'learning_rate': {'value' :  0.001}, #0.001 default, but 0.005 might be better
@@ -33,7 +33,7 @@ def get_swep_config():
         'loss_reg' : { 'value' :  'b'},
         'loss_reg_a' : { 'value' : 256},
         'loss_reg_c' : { 'value' : 0.001},
-        'test_samples': { 'value' : 128},
+        'test_samples': { 'value' :128}, # 128 for actual testing, 10 for debug
         'np_seed' : {'values' : [4,8]},
         'torch_seed' : {'values' : [4,8]},
         'window_dim' : {'value' : 32},
@@ -43,6 +43,7 @@ def get_swep_config():
         'first_feature_idx' : {'value' : 5},
         'norm_target' : {'value' : False},
         'freeze_h' : {'value' : "hl"},
+        'time_steps' : {'value' : 36}
         }
 
     sweep_config['parameters'] = parameters_dict

models/purple_alien/main.py

@@ -0,0 +1,57 @@
+import time
+
+import wandb
+
+import sys
+from pathlib import Path
+
+PATH = Path(__file__)
+sys.path.insert(0, str(Path(*[i for i in PATH.parts[:PATH.parts.index("views_pipeline")+1]]) / "common_utils")) # PATH_COMMON_UTILS  
+from set_path import setup_project_paths, setup_artifacts_paths
+setup_project_paths(PATH)
+
+from cli_parser_utils import parse_args, validate_arguments
+#from artifacts_utils import get_latest_model_artifact
+
+from model_run_handlers import handle_sweep_run, handle_single_run
+
+#from mode_run_manager import model_run_manager
+
+if __name__ == "__main__":
+
+    # new argpars solution.
+    args = parse_args()
+    #print(args)
+
+    # Validate the parsed arguments to ensure they conform to the required logic and combinations.
+    validate_arguments(args)
+
+    # wandb login
+    wandb.login()
+
+    start_t = time.time()
+
+    # Test if and why a model_metadata_dict.py was saved in the artifacts folder..
+
+    # first you need to check if you are running a sweep or not, because the sweep will overwrite the train and evaluate flags
+    if args.sweep == True:
+
+        handle_sweep_run(args)
+
+    elif args.sweep == False:
+
+        handle_single_run(args)
+
+    end_t = time.time()
+    minutes = (end_t - start_t)/60
+    print(f'Done. Runtime: {minutes:.3f} minutes')
+
+    # notes on stepshifted models:
+    # There will be some thinking here in regards to how we store, denote (naming convention), and retrieve the model artifacts from stepshifted models.
+    # It is not a big issue, but it is something to consider os we don't do something headless. 
+    # A possible format could be: <run_type>_model_s<step>_<timestamp>.pt example: calibration_model_s00_20210831_123456.pt, calibration_model_s01_20210831_123456.pt, etc.
+    # And the rest of the code maded in a way to handle this naming convention without any issues. Could be a simple fix.
+    # Alternatively, we could store the model artifacts in a subfolder for each stepshifted model. This would make it easier to handle the artifacts, but it would also make it harder to retrieve the latest artifact for a given run type.
+    # Lastly, the solution Xiaolong is working on might allow us the store multiple models (steps) in one artifact, which would make this whole discussion obsolete and be the best solution.
+
+

models/purple_alien/src/forecasting/generate_forecast.py

@@ -0,0 +1,149 @@
+import os
+
+import numpy as np
+import pickle
+import time
+import functools
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import wandb
+
+import sys
+from pathlib import Path
+
+PATH = Path(__file__)
+sys.path.insert(0, str(Path(*[i for i in PATH.parts[:PATH.parts.index("views_pipeline")+1]]) / "common_utils")) # PATH_COMMON_UTILS  
+from set_path import setup_project_paths, setup_data_paths
+setup_project_paths(PATH)
+
+
+from utils import choose_model, choose_loss, choose_sheduler, get_train_tensors, get_full_tensor, apply_dropout, execute_freeze_h_option, get_log_dict, train_log, init_weights, get_data
+from utils_prediction import predict, sample_posterior
+from config_hyperparameters import get_hp_config
+
+
+def generate_forecast(model, views_vol, config, device, PATH):
+    """
+    Function to generate forecast using the provided model and views_vol.
+    It saves the generated posterior distributions and out-of-sample volumes.
+
+    Args:
+        model: The trained model used for forecasting.
+        views_vol: The input data tensor for forecasting.
+        config: Configuration object containing settings.
+        device: The device (CPU or GPU) to run the predictions on.
+        PATH: The base path where generated data will be saved.
+
+    Returns:
+        None
+    """
+    # Ensure the model is in evaluation mode
+    model.eval()
+    model.apply(apply_dropout)
+
+    # Generate posterior samples and out-of-sample volumes
+    posterior_list, posterior_list_class, out_of_sample_vol, _ = sample_posterior(model, views_vol, config, device) # the _ is the full tensor. 
+
+    # I suspect you'll need the out_of_sample_vol to create the df (it has pg and ocean info)
+    # However, I see in the test_prediction_store notebook in "conflictnet" repo that I load the "calibration_vol" from the pickle file.... Investigate... 
+
+
+    # Set up paths for storing generated data
+    _, _, PATH_GENERATED = setup_data_paths(PATH)
+
+    # Create the directory if it does not exist
+    os.makedirs(PATH_GENERATED, exist_ok=True)
+
+    # Print the path for debugging
+    print(f'PATH to generated data: {PATH_GENERATED}')
+
+    # Create a dictionary to store posterior data
+    posterior_dict = {
+        'posterior_list': posterior_list,
+        'posterior_list_class': posterior_list_class,
+        'out_of_sample_vol': out_of_sample_vol          # you might need this for the df creation before predstore. Experiments in notebook test_to_prediction_store.ipynb
+    }
+
+    # Save the posterior data to a pickle file
+    filename = f'posterior_dict_{config.time_steps}_{config.run_type}_{config.model_time_stamp}.pkl'
+    with open(os.path.join(PATH_GENERATED, filename), 'wb') as file:
+        pickle.dump(posterior_dict, file)
+
+    print('Posterior dict and test vol pickled and dumped!')
+
+
+def handle_forecasting(config, device, views_vol, PATH_ARTIFACTS, artifact_name=None):
+
+    # the thing above might work, but it needs to be tested thoroughly....
+    raise NotImplementedError('Forecasting not implemented yet')
+
+
+
+
+# Ensure utils_prediction.py and any other dependencies are imported correctly
+# from utils_prediction import sample_posterior, apply_dropout
+# from utils_data import setup_data_paths
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+## you always load an artifact for forecasting - like with the evaluate you take the latest artifact unless you specify another one
+## But that is done in main.py - just passed to here as an argument
+#
+## Then the load the offical forescasting partition
+## And the first steps must be usign the function from utils_prediction.py to get the predictions and the posetrior
+#
+## model, views_vol, config, device should be passed as arguments to this function
+#
+#def generate_forecast(model, views_vol, config, device):
+#
+#
+#    # THIS IS ALL PURE MESS RIGHT NOW!!! 
+#
+#
+#    posterior_list, posterior_list_class, out_of_sample_vol, full_tensor = sample_posterior(model, views_vol, config, device)
+#
+## then to prediction store I guess? Or perhaps just the generated data for now... 
+#
+#    _ , _, PATH_GENERATED = setup_data_paths(PATH)
+#
+#    # if the path does not exist, create it
+#
+#    if not os.path.exists(PATH_GENERATED):
+#
+#        os.makedirs(PATH_GENERATED)
+#
+#    # print for debugging
+#    print(f'PATH to generated data: {PATH_GENERATED}')
+#
+#    # pickle the posterior dict, metric dict, and test vol
+#
+#    # Should be time_steps and run_type in the name....
+#    posterior_dict = {'posterior_list' : posterior_list, 'posterior_list_class': posterior_list_class, 'out_of_sample_vol' : out_of_sample_vol}
+#
+#
+#    with open(f'{PATH_GENERATED}/posterior_dict_{config.time_steps}_{config.run_type}_{config.model_time_stamp}.pkl', 'wb') as file:
+#
+#        pickle.dump(posterior_dict, file)       
+#
+#
+#    print('Posterior dict, metric dict and test vol pickled and dumped!')
+#
+#