train_model.py

import argparse
import pickle
import os
import sys
import gzip
import torch
import pyrootutils

ROOT_DIR = pyrootutils.setup_root(__file__, indicator="README.md", pythonpath=True)
os.chdir(ROOT_DIR)
RAMIFICATIONS_DIR = os.path.join(ROOT_DIR, "ramifications")
if not os.path.exists(RAMIFICATIONS_DIR):
    print(f"Ramifications datasets in {RAMIFICATIONS_DIR} does not exist. Please run the scripts to generate them.")
    sys.exit(1)
MODELS_DIR = os.path.join(ROOT_DIR, "trained_models")

def main():

    parser = argparse.ArgumentParser(description='Train surrogate model on an ABM')
    parser.add_argument('--abm_model', type=str, default='predatorprey', help='choose between predatorprey and schelling (default: predatorprey)')
    parser.add_argument('--model_type', type=str, default='GDN', help='choose between GDN, diffusion-only and gnn-only (default: GDN)')
    parser.add_argument('--parameter', type=str, default='psi1', help='choose parameter between: xi1, xi2, xi3 for schelling; and psi1, psi2, psi3, psi4 for predator-prey (default: psi1)')
    parser.add_argument('--learning_rate', type=float, default=1e-5, help='learning rate for the model (default: 1e-5)')
    parser.add_argument('--T_diffusion', type=int, default=100, help='number of diffusion steps (default: 100)')
    parser.add_argument('--n_epochs', type=int, default=100, help='number of epochs for training (default: 100)')
    args = parser.parse_args()

    # Number of ramifications to use for training (500 adviced, but can be reduced for testing)
    N_RAMIFICATIONS = 500

    # Load the configuration
    abm_model = args.abm_model
    model_type = args.model_type
    parameter = args.parameter
    learning_rate = args.learning_rate
    T_diffusion = args.T_diffusion
    n_epochs = args.n_epochs


    # Load the training data
    with gzip.open(f'{RAMIFICATIONS_DIR}/{abm_model}/ramification_training_{parameter}.pickle.gz', 'rb') as f:
        ramification_training = pickle.load(f)


    # Initialize the model
    if abm_model == 'predatorprey':
        from models.abm.predpreyfeaturizer import PredPreyFeaturizer
        featurizer = PredPreyFeaturizer()
        if model_type == 'GDN':
            from models.surrogate.nnmodel import NNModel
            model = NNModel(n_features = 6, learning_rate=learning_rate, abm_featurizer = featurizer,diffusion_timesteps=T_diffusion,aggregation='add')
        elif model_type == 'diffusion-only':
            from models.surrogate.nnmodel_diffusion_only import NNModel_diffusion_only
            model = NNModel_diffusion_only(n_features = 6, learning_rate=learning_rate, abm_featurizer = featurizer, diffusion_timesteps=T_diffusion,
                                        domain_dim=featurizer.scale_abm_state(ramification_training[0][0]).flatten().shape[0])
        elif model_type == 'gnn-only':
            from models.surrogate.nnmodel_gnn_only import GNNModel
            model = GNNModel(n_features = 6, abm_featurizer = featurizer, aggregation='add')
        else:
            raise ValueError(
                f"model_type '{model_type}' not recognized. Choose among "
                "['GDN', 'diffusion-only', 'gnn-only']."
            )
    
    elif abm_model == 'schelling':
        from models.abm.schellingfeaturizer import SchellingFeaturizer
        featurizer = SchellingFeaturizer()
        if model_type == 'GDN':
            from models.surrogate.nnmodel import NNModel
            model = NNModel(n_features = 2, learning_rate=learning_rate, abm_featurizer = featurizer, diffusion_timesteps=T_diffusion, aggregation='mean')
        elif model_type == 'diffusion-only':
            from models.surrogate.nnmodel_diffusion_only import NNModel_diffusion_only
            model = NNModel_diffusion_only(n_features = 2, learning_rate=learning_rate, abm_featurizer = featurizer, diffusion_timesteps=T_diffusion,
                                        domain_dim=featurizer.scale_abm_state(ramification_training[0][0]).flatten().shape[0])
        elif model_type == 'gnn-only':
            from models.surrogate.nnmodel_gnn_only import GNNModel
            model = GNNModel(n_features = 2, abm_featurizer = featurizer, aggregation='mean')
        else:
            raise ValueError(
                f"model_type '{model_type}' not recognized. Choose among "
                "['GDN', 'diffusion-only', 'gnn-only']."
            )

    model.train(ramifications = ramification_training, n_ramifications= 500, n_epochs=n_epochs)

    # Save the model
    save_dir = f'{MODELS_DIR}/{abm_model}/{model_type}/'
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)   

    if model_type == 'GDN':
        save_path = os.path.join(save_dir, f"model_GDN_{parameter}.pth")
        torch.save({'ld_model_state_dict': model.ld_model.state_dict(),
                    'graph_model_state_dict': model.graph_model.state_dict(),
                    'diffusion_timesteps': model.diffusion_timesteps,
                    'aggregation': model.aggregation,
                    'learning_rate': model.lr_ld,
                    'losses': model.losses
                    }, save_path)
    
    elif model_type == 'diffusion-only':
        save_path = os.path.join(save_dir, f"model_diffusion-only_{parameter}.pth")
        torch.save({'ld_model_state_dict': model.ld_model.state_dict(),
                    'diffusion_timesteps': model.diffusion_timesteps,
                    'learning_rate': model.lr_ld,
                    'losses': model.losses
                    }, save_path)
    
    elif model_type == 'gnn-only':
        save_path = os.path.join(save_dir, f"model_gnn-only_{parameter}.pth")
        torch.save({'graph_model_state_dict': model.graph_model.state_dict(),
                    'learning_rate': model.lr_gnn,
                    'aggregation': model.aggregation,
                    'losses': model.losses
                    }, save_path)

    print(f'Model saved to {save_path}')

if __name__ == "__main__":
    main()