Source code for schrodinger.application.vss.config

import itertools
import os

from ruamel import yaml
from voluptuous import All
from voluptuous import Range
from voluptuous import Required
from voluptuous import Schema

from schrodinger.application.vss import database
from schrodinger.utils import fileutils

CONFIG_FILE_SCHEMA = Schema({
    'databases': [str],
    Required('glide_inp'): str,
    'ligprep_args': [str],
    Required('decoys_pool'): str,
    'glide_al_train_host': str,
    'glide_al_eval_host': str,
    'glide_al_driver_host': str,
    'glide_al_njobs': All(int, Range(min=1)),
    'glide_al_block_size': All(int, Range(min=1)),
    'glide_al_num_rescore_ligand': All(int, Range(min=1)),
    'glide_al_glide_subjob_size': All(int, Range(min=1)),
    Required('glide_al_train_size', default=20): All(int, Range(min=1)),
    Required('glide_al_num_iter', default=2): All(int, Range(min=1)),
    'ligand_ml_train_host': str,
    'ligand_ml_eval_host': str,
    'ligand_ml_driver_host': str,
    'ligand_ml_njobs': All(int, Range(min=1)),
    'glide_host': str,
    'glide_driver_host': str,
    'shape_host': str,
    'shape_driver_host': str,
})


class Config:
    '''
    Runner configuration data.
    '''

    def __init__(self, filename):

        with open(filename, 'r') as fp:
            data = CONFIG_FILE_SCHEMA(yaml.safe_load(fp))

        keywords = [str(k) for k in CONFIG_FILE_SCHEMA.schema]
        for keyword in keywords:
            setattr(self, keyword, data.get(keyword))

        with fileutils.chdir(os.path.dirname(filename) or os.curdir):
            self.filename = os.path.abspath(filename)
            self.decoys_pool = os.path.abspath(data['decoys_pool'])
            self.databases = [
                database.Database(p) for p in data.get('databases', [])
            ]

    def get_smiles_csv(self, dbnames):
        return itertools.chain.from_iterable(
            db.get_smiles_csv() for db in self.databases if db.name in dbnames)

    def get_shape_data(self, dbnames, *, shape_type):
        shape_data = []
        shape_data_is_local = False
        for db in self.databases:
            if db.name in dbnames:
                db_data = list(db.get_shape_data(shape_type))
                if db_data:
                    shape_data.extend(db_data)
                    if db.local:
                        shape_data_is_local = True
        return shape_data, shape_data_is_local