train.py

import warnings

def warn(*args, **kwargs):
    pass
warnings.warn = warn

import argparse
import datetime

import os
import pickle

from rdkit import RDLogger
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks.early_stopping import EarlyStopping

from massspecgym.models.base import Stage

from src.data.data_module import MolClassDataModule

from src.definitions import TEST_RESULTS_DIR
import yaml
from src.data.datasets import MolClassDataset
from functools import partial

from torch.utils.data import Subset
import torch

from src.utils.data import get_ms_dataset, get_spec_featurizer, get_mol_featurizer
from src.utils.data import get_class_featurizer
from src.utils.data import get_formula_featurizer
from src.utils.data import get_valence_featurizer
from src.utils.models import get_model

from src.data.strata import MolClassStrataClassyFire

from src.data.weighting import IgnoreConstantClass

import numpy as np


def main(params):
    # Seed everything
    pl.seed_everything(params['seed'])
    print("EXPERIMENT directory: ", params['experiment_dir'])

    # Load dataset
    if params['unsupervised_mol_path'] is not None:
        if params['mol_to_class'] is False:
            raise Exception('mol_to_class cannot be False, when using unsupervised_mol_path')
        spec_featurizer = None
    else:
        spec_featurizer = get_spec_featurizer(params['spectra_view'], params)

    mol_featurizer = get_mol_featurizer(params['molecule_view'], params)
    class_featurizer = get_class_featurizer(params)
    formula_featurizer = get_formula_featurizer(params)

    if params['use_valence']:
        valence_featurizer = get_valence_featurizer(params)
    else:
        valence_featurizer = None

    if params['path_classes'] is None:
        classes_used = class_featurizer.get_valid_classes()
    else:
        with open(params['path_classes'], 'rb') as f:
            classes_used = pickle.load(f)

    strata = MolClassStrataClassyFire([params['classyfire_pth'], params['canopus_unsupervised_structures']], classes_used, bal_mol=params['mol_to_class'])

    if params['unsupervised_mol_path'] is not None:
        unsupervised_strata = MolClassStrataClassyFire(params['canopus_unsupervised_structures'], classes_used)
    else:
        unsupervised_strata = None

    dataset = get_ms_dataset(params['spectra_view'], params['molecule_view'], spec_featurizer, mol_featurizer, class_featurizer, formula_featurizer, valence_featurizer, params)

    print('JESTR SIZE:', len(dataset))
    
    # Init data module
    if params['use_class_bal'] and params['load_class_weight']:
        with open(params['class_weight_path'], 'rb') as f:
            class_weight_dict = pickle.load(f)['class_weight']['train']
    else:
        class_weight_dict = None

    if params['unsupervised_mol_path'] is None:
        spec_view = params['spectra_view']
    else:
        spec_view = "None"

    print('spec_view:', spec_view)

    if params['path_classes'] is not None:
        collate_fn = partial(MolClassDataset.collate_fn, num_classes=len(classes_used), spec_enc=params['spec_enc'], spectra_view=spec_view, gym_batch_size=params['gym_batch_size'])
    else:
        collate_fn = partial(MolClassDataset.collate_fn, num_classes=len(class_featurizer.get_all_classes()), spec_enc=params['spec_enc'], spectra_view=spec_view, gym_batch_size=params['gym_batch_size'])

    data_module = MolClassDataModule(
        dataset=dataset,
        strata=strata,
        collate_fn=collate_fn,
        unique_mol=params['unique_mol'],
        split_pth=params['split_pth'],
        batch_size=params['batch_size'],
        num_workers=params['num_workers'],
        unsupervised_strata=unsupervised_strata,
        class_weight_dict=class_weight_dict,
        equal_iceberg=params['equal_iceberg'],
        gym_batch_size=params['gym_batch_size']
    )

    if params['path_classes'] is None:
        model = get_model(params['model'], params, all_classes=class_featurizer.get_all_classes())
    else:
        model = get_model(params['model'], params, all_classes = classes_used) 

    if params['freeze_enc'] or params['freeze_dec']:
        model.freeze_weight()

    # Init logger
    logger = None

    # Init callbacks for checkpointing and early stopping
    callbacks = [pl.callbacks.ModelCheckpoint(save_last=False) ]
    for i, monitor in enumerate(model.get_checkpoint_monitors()):
        monitor_name = monitor['monitor']
        
        checkpoint = pl.callbacks.ModelCheckpoint(
            monitor=monitor_name,
            save_top_k=1,
            mode=monitor['mode'],
            dirpath=params['experiment_dir'],
            filename=f'{{epoch}}-{{{monitor_name}:.2f}}',
            auto_insert_metric_name=True,
            save_last=(i == 0)
        )
        callbacks.append(checkpoint)
        
        if monitor.get('early_stopping', False):
            early_stopping = EarlyStopping(
                monitor=monitor_name,
                mode=monitor['mode'],
                verbose=True,
                patience=params['early_stopping_patience'],
            )

            callbacks.append(early_stopping)

    # Init trainer
    trainer = Trainer(
        accelerator=params['accelerator'],
        devices=params['devices'],
        max_epochs=params['max_epochs'],
        logger=logger,
        log_every_n_steps=params['log_every_n_steps'],
        val_check_interval=params['val_check_interval'],
        callbacks=callbacks,
        default_root_dir=params['experiment_dir'],
        gradient_clip_val=0.5, gradient_clip_algorithm="value"
    )

    # Prepare data module to validate or test before training
    data_module.prepare_data()
    data_module.setup()

    if params['save_class_weight']:
        with open(params['class_weight_path'], 'wb') as f:
            pickle.dump({
                'class_weight': {
                    "train": mol_dataset.class_weight_dict
                }   
            },f)

    with open(params['class_weight_path'], 'rb') as f:
        class_weight_dict = pickle.load(f)['class_weight']['train']


    class_count_arr = torch.zeros(len(class_weight_dict))
    for curr_class in class_weight_dict:
        class_i = class_featurizer.get_class_i(curr_class)
        class_count_arr[class_i] = class_weight_dict[curr_class]

    model.set_class_count_arr(class_count_arr, len(data_module.train_dataset))

    # Validate before training
    trainer.validate(model, datamodule=data_module)

    # Train
    trainer.fit(model, datamodule=data_module)
        

def setup_parser(args, params):
    # Get current time
    now = datetime.datetime.now()
    now_formatted = now.strftime("%Y%m%d")
    
    res_folder_name = f"{now_formatted}_{params['run_name']}"
    experiment_dir = str(TEST_RESULTS_DIR / res_folder_name)
    params['experiment_dir'] = experiment_dir

    if not params['df_test_path']:
        params['df_test_path'] = os.path.join(experiment_dir, "result.pkl")

    return params


if __name__ == "__main__":
    # Suppress RDKit warnings and errors
    lg = RDLogger.logger()
    lg.setLevel(RDLogger.CRITICAL)

    parser = argparse.ArgumentParser()
    parser.add_argument("--param_pth", type=str, default="configs/params_train.yaml")

    args = parser.parse_args([] if "__file__" not in globals() else None)

    # Load
    with open(args.param_pth) as f:
        params = yaml.load(f, Loader=yaml.FullLoader)

    params = setup_parser(args, params)

    main(params)