test.py

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 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 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.models import get_model

from src.data.strata import MolClassStrataNPClassifier
from src.data.strata import MolClassStrataClassyFire

from src.data.weighting import IgnoreConstantClass

from massspecgym.models.base import Stage


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

    if params['model'] != "MolClass":
        raise Exception('model: %s not in "{MolClass}"'%(params['model']))

    # Load dataset
    if params['unsupervised_mol_path'] is not None:
        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'] != None:
        unsupervised_strata = MolClassStrataClassyFire(params['formula_candidate_pth'], classes_used)
        if params['strata_mode'] == 'smiles':
            unsupervised_strata = SmilesStrata(unsupervised_strata)
    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)

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

    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

    # Init data module
    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)
    else:
        collate_fn = partial(MolClassDataset.collate_fn, num_classes=len(class_featurizer.get_all_classes()), spec_enc=params['spec_enc'], spectra_view=spec_view)
    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,
        unsupervised_fold='test',
        use_class_balance=params['use_class_bal'],
        class_weight_dict=class_weight_dict
    )

    model = get_model(params['model'], params, all_classes=class_featurizer.get_all_classes())
    
    # Init trainer
    trainer = Trainer(
        accelerator=params['accelerator'],
        devices=params['devices'],
        default_root_dir=params['experiment_dir'],
    )

    # Prepare data module to test
    data_module.prepare_data()
    data_module.setup()
    data_module.setup(stage="test")

    if params['model'] == "MolClass":
        if params['use_class_bal']:
            mol_dataset = data_module.train_mol_class_dataset

            ignore_class_bal = IgnoreConstantClass(mol_dataset.class_weight_dict, mol_dataset)
            ignore_pos_weights, ignore_neg_weights = ignore_class_bal()

            model.set_weights(Stage.TRAIN, weight_by_label=True, pos_weights=ignore_pos_weights, neg_weights=ignore_neg_weights)

        
    # Test
    trainer.test(model, datamodule=data_module)


def setup_parser(args, params):
    # Experiment directory
    exp_dir = ''
    run_name = params['run_name']

    if not exp_dir:
        now = datetime.datetime.now().strftime("%Y%m%d")
        exp_dir = str(TEST_RESULTS_DIR / f"{now}_{params['run_name']}")
    params['experiment_dir'] = exp_dir
    
    # Checkpoint path
    if args is not None and args.checkpoint_pth:
        params['checkpoint_pth'] = args.checkpoint_pth
    
    if not params['checkpoint_pth']:
        print("No checkpoint provided. Using the checkpoint in the experiment directory")
        for f in os.listdir(exp_dir):
            if f.endswith("ckpt") and f.startswith("epoch") and args is not None and args.checkpoint_choice in f:
                checkpoint_path = os.path.join(exp_dir, f)
                params['checkpoint_pth'] = checkpoint_path
                break
    assert(params['checkpoint_pth'] != '')

    if params['df_test_path'] is None:
        now = datetime.datetime.now().strftime("%Y%m%d")
        if params['unsupervised_mol_path'] is not None:
            params['df_test_path'] = str(TEST_RESULTS_DIR / f"{now}_{params['run_name']}"/ f"preds.tsv")
        else:
            params['df_test_path'] = str(TEST_RESULTS_DIR / f"{now}_{params['run_name']}"/ f"preds.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_test.yaml")
    parser.add_argument('--checkpoint_pth', type=str, default='')
    parser.add_argument('--checkpoint_choice', type=str, default='train', choices=['train', 'val'])
    parser.add_argument('--df_test_pth', type=str, help='result file name')

    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)